Ollama Model-Anti-staller

gistfile1.txt

#!/usr/bin/env node

/**
 * ollama-pull.js — Reliable Ollama model downloader with adaptive stall detection
 */

const yargs = require("yargs");
const { hideBin } = require("yargs/helpers");

const argv = yargs(hideBin(process.argv))
  .usage("Usage: $0 [model] [options]")
  .option("model", {
    alias: "m",
    type: "string",
    description: "Model to pull",
    default: "llama3.2:latest",
  })
  .option("host", {
    alias: "H",
    type: "string",
    description: "Ollama server URL",
    default: process.env.OLLAMA_HOST || "http://127.0.0.1:11434",
  })
  .option("band-size", {
    alias: "b",
    type: "number",
    description: "% band size to track (e.g. 10 = check every 10%)",
    default: 10,
  })
  .option("stall-factor", {
    alias: "f",
    type: "number",
    description: "How many times slower the current band can be vs the last before restarting",
    default: 10,
  })
  .option("first-band-timeout", {
    alias: "t",
    type: "number",
    description: "Seconds to wait for the first band crossing before restarting",
    default: 300,
  })
  .example("$0 qwen3:30b", "Pull qwen3:30b with default stall detection")
  .example("$0 -m mistral:7b -b 5 -f 5", "5% bands, restart if 5x slower than last band")
  .help()
  .alias("help", "h")
  .argv;

const model = argv._[0] || argv.model;
const OLLAMA_HOST = argv.host;
const BAND_SIZE = argv.bandSize;   // % per band
const STALL_FACTOR = argv.stallFactor;
const FIRST_BAND_TIMEOUT_MS = argv.firstBandTimeout * 1000;

function formatBytes(bytes) {
  if (!bytes) return "?";
  const gb = bytes / 1024 ** 3;
  if (gb >= 1) return `${gb.toFixed(2)} GB`;
  const mb = bytes / 1024 ** 2;
  if (mb >= 1) return `${mb.toFixed(1)} MB`;
  return `${(bytes / 1024).toFixed(1)} KB`;
}

function renderBar(completed, total, width = 30) {
  if (!total) return "[" + "░".repeat(width) + "]";
  const filled = Math.round((completed / total) * width);
  return "[" + "█".repeat(filled) + "░".repeat(width - filled) + "]";
}

function overallPct(layerProgress) {
  let totalBytes = 0, completedBytes = 0;
  for (const info of Object.values(layerProgress)) {
    if (info.total) totalBytes += info.total;
    if (info.completed) completedBytes += info.completed;
  }
  if (!totalBytes) return 0;
  return (completedBytes / totalBytes) * 100;
}

async function pullModel(modelName) {
  let attempt = 0;
  let lastBandMs = null; // persists across attempts — only ever increases

  while (true) {
    attempt++;
    console.log(`\n🦙 Pulling model: ${modelName}  (attempt ${attempt})`);
    console.log(`   Host: ${OLLAMA_HOST}`);
    console.log(`   Stall detection: ${BAND_SIZE}% bands, restart if ${STALL_FACTOR}x slower than last, first-band timeout: ${argv.firstBandTimeout}s\n`);

    const abortController = new AbortController();
    let stalled = false;

    try {
      const res = await fetch(`${OLLAMA_HOST}/api/pull`, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ name: modelName, stream: true }),
        signal: abortController.signal,
      });

      if (!res.ok) {
        const text = await res.text();
        throw new Error(`Ollama API error ${res.status}: ${text}`);
      }

      const decoder = new TextDecoder();
      const reader = res.body.getReader();
      const layerProgress = {};
      let lastStatus = "";

      // Adaptive stall detection state
      // Initialise to wherever we already are so a resume doesn't look instant.
      const initialPct = overallPct(layerProgress);
      let currentBand = Math.floor(initialPct / BAND_SIZE);
      let bandStartTime = Date.now();
      let firstCrossing = true; // first band crossing after (re)start — discard, can't trust timing
      let firstCrossingSince = Date.now(); // when we started waiting for first crossing

      // Poll pct every second to detect band crossings
      const stallTimer = setInterval(() => {
        const pct = overallPct(layerProgress);
        const band = Math.floor(pct / BAND_SIZE);

        // First-band timeout: if we've never crossed a band, restart after timeout
        if (firstCrossing && (Date.now() - firstCrossingSince) > FIRST_BAND_TIMEOUT_MS) {
          process.stdout.write(
            `\n\n⚠️  First band never crossed after ${argv.firstBandTimeout}s at ${pct.toFixed(1)}%. Restarting...\n`
          );
          stalled = true;
          abortController.abort();
          clearInterval(stallTimer);
          return;
        }

        if (band > currentBand) {
          const elapsed = Date.now() - bandStartTime;
          const bandLabel = `${currentBand * BAND_SIZE}–${band * BAND_SIZE}%`;

          if (firstCrossing) {
            firstCrossing = false;
            if (initialPct > 0) {
              // Resuming mid-download — elapsed includes pre-downloaded bytes, can't trust it
              process.stdout.write(
                `\n  ⏭  Band ${bandLabel} skipped (resuming from prior progress)\n\n`
              );
              // keep lastBandMs from prior attempt as baseline
            } else {
              // Fresh start — this IS a real measurement, use it as baseline
              process.stdout.write(
                `\n  ✓ Band ${bandLabel} in ${(elapsed/1000).toFixed(1)}s (baseline set)\n\n`
              );
              lastBandMs = Math.max(lastBandMs ?? 0, elapsed);
            }
            bandStartTime = Date.now();
            currentBand = band;
            return;
          }

          if (lastBandMs !== null && elapsed > lastBandMs * STALL_FACTOR) {
            process.stdout.write(
              `\n\n⚠️  Stall detected at ${pct.toFixed(1)}%! Band ${bandLabel} took ${(elapsed/1000).toFixed(1)}s ` +
              `(last band: ${(lastBandMs/1000).toFixed(1)}s, factor: ${(elapsed/lastBandMs).toFixed(1)}x). Restarting...\n`
            );
            stalled = true;
            abortController.abort();
            clearInterval(stallTimer);
            return;
          }

          // Band crossed cleanly — log and advance
          if (lastBandMs !== null) {
            process.stdout.write(
              `\n  ✓ Band ${bandLabel} in ${(elapsed/1000).toFixed(1)}s (last: ${(lastBandMs/1000).toFixed(1)}s)\n\n`
            );
          } else {
            process.stdout.write(
              `\n  ✓ Band ${bandLabel} in ${(elapsed/1000).toFixed(1)}s (baseline set)\n\n`
            );
          }

          lastBandMs = Math.max(lastBandMs ?? 0, elapsed); // never let baseline get faster
          bandStartTime = Date.now();
          currentBand = band;
        }
      }, 1000);

      process.stdout.write("\n");

      try {
        while (true) {
          const { done, value } = await reader.read();
          if (done) break;

          const chunk = decoder.decode(value, { stream: true });
          const lines = chunk.split("\n").filter((l) => l.trim());

          for (const line of lines) {
            let msg;
            try { msg = JSON.parse(line); } catch { continue; }

            if (msg.error) {
              process.stdout.write("\n");
              throw new Error(msg.error);
            }

            const { status, digest, total, completed } = msg;

            if (digest) {
              const shortDigest = digest.slice(-12);
              layerProgress[shortDigest] = { status, total, completed };

              process.stdout.write("\x1b[2K\r");

              const activeLayers = Object.entries(layerProgress).filter(
                ([, v]) => v.status !== "already exists"
              );

              if (activeLayers.length > 0) {
                if (lastStatus !== "") {
                  process.stdout.write(`\x1b[${activeLayers.length}A`);
                }

                for (const [id, info] of activeLayers) {
                  const pct =
                    info.total && info.completed
                      ? ((info.completed / info.total) * 100).toFixed(1)
                      : "—";
                  const bar = renderBar(info.completed, info.total);
                  const size = info.total
                    ? `${formatBytes(info.completed)} / ${formatBytes(info.total)}`
                    : "";
                  process.stdout.write(`\x1b[2K  ${id}  ${bar}  ${String(pct).padStart(5)}%  ${size}\n`);
                }

                lastStatus = "layers";
              }
            } else if (status) {
              if (status !== lastStatus) {
                process.stdout.write(`\x1b[2K\r  ⏳ ${status}\n`);
                lastStatus = status;
              }
            }

            if (status === "success") {
              clearInterval(stallTimer);
              process.stdout.write(`\n✅ Successfully pulled ${modelName}\n\n`);
              return;
            }
          }
        }
      } finally {
        clearInterval(stallTimer);
      }

      if (!stalled) {
        const finalPct = overallPct(layerProgress);
        if (finalPct >= 99.9) {
          console.log(`\n✅ Done — ${modelName} is ready.\n`);
          return;
        } else {
          // Stream dropped mid-download (the classic Ollama stall)
          process.stdout.write(
            `\n\n⚠️  Stream dropped at ${finalPct.toFixed(1)}% with no success message. Restarting...\n`
          );
          await new Promise((r) => setTimeout(r, 2000));
          continue;
        }
      }

    } catch (err) {
      if (stalled || err.name === "AbortError") {
        await new Promise((r) => setTimeout(r, 2000));
        continue;
      }
      throw err;
    }
  }
}

pullModel(model).catch((err) => {
  console.error(`\n❌ Error: ${err.message}\n`);
  process.exit(1);
});