diff --git a/docs/inference/use-local-inference.mdx b/docs/inference/use-local-inference.mdx index 96341ac4865..7b36bd6c802 100644 --- a/docs/inference/use-local-inference.mdx +++ b/docs/inference/use-local-inference.mdx @@ -44,7 +44,8 @@ $ nemoclaw onboard Select **Local Ollama** from the provider list. NemoClaw lists installed models or offers starter models if none are installed. -On hosts with at least 32 GiB of detected GPU memory, the starter list includes `qwen3.6:35b` and selects it by default. +On hosts where the larger starter models fit the currently available GPU memory, the starter list includes `qwen3.6:35b` and selects it by default. +When another GPU workload is using most of the memory at onboard time, NemoClaw downgrades the menu to the largest model that still fits. It pulls the selected model, loads it into memory, and validates it before continuing. If the selected model declares that it does not support tool calling, onboarding stops with guidance to choose a model whose `ollama show ` capabilities include `tools`. The validation also requires structured chat-completions tool calls. @@ -136,6 +137,8 @@ $ NEMOCLAW_PROVIDER=ollama \ ``` If `NEMOCLAW_MODEL` is not set, NemoClaw selects a default model based on available memory. +If `NEMOCLAW_MODEL` names a known bootstrap model (for example `qwen3.6:35b`) that does not fit the host's currently available GPU memory, NemoClaw warns and falls back to the largest known model that does fit. +Unknown or custom tags (any value the bootstrap registry has not seen) are still passed through; the Ollama runner validates the choice itself. `--yes` (or `NEMOCLAW_YES=1`) authorises the Ollama model download without an interactive confirmation prompt. Under `--non-interactive`, `--yes` (or `NEMOCLAW_YES=1`) is required to authorise the download — onboard exits otherwise, since it cannot prompt. diff --git a/src/lib/inference/local.test.ts b/src/lib/inference/local.test.ts index de807baf60a..a8475a0b146 100644 --- a/src/lib/inference/local.test.ts +++ b/src/lib/inference/local.test.ts @@ -8,10 +8,17 @@ import os from "node:os"; import path from "node:path"; // Import from compiled dist/ for correct coverage attribution. +import { OLLAMA_MODEL_REGISTRY } from "../../../dist/lib/inference/ollama-model-registry"; + +// Derive the "large enough to fit every registry entry" memory threshold +// from the registry itself so adding or resizing a model in the registry +// does not require updating these tests. +const LARGE_OLLAMA_FIT_MEMORY_MB = Math.max( + ...OLLAMA_MODEL_REGISTRY.map((entry) => entry.requiredMemoryMB), +); import { CONTAINER_REACHABILITY_IMAGE, DEFAULT_OLLAMA_MODEL, - LARGE_OLLAMA_MIN_MEMORY_MB, LOCAL_INFERENCE_SANDBOX_HOST_URL_ENV, QWEN3_6_OLLAMA_MODEL, getOllamaContainerPort, @@ -554,39 +561,146 @@ describe("local inference helpers", () => { it("falls back to bootstrap model options when no Ollama models are installed", () => { expect(getBootstrapOllamaModelOptions(null)).toEqual(["qwen2.5:7b"]); + // Below every registry entry's required memory: small only. expect( getBootstrapOllamaModelOptions({ type: "nvidia", - totalMemoryMB: LARGE_OLLAMA_MIN_MEMORY_MB - 1, + totalMemoryMB: 10_000, }), ).toEqual(["qwen2.5:7b"]); + // Comfortably above every registry entry's required memory: all options. expect( getBootstrapOllamaModelOptions({ type: "nvidia", - totalMemoryMB: LARGE_OLLAMA_MIN_MEMORY_MB, + totalMemoryMB: LARGE_OLLAMA_FIT_MEMORY_MB, }), ).toEqual(["qwen2.5:7b", DEFAULT_OLLAMA_MODEL, QWEN3_6_OLLAMA_MODEL]); - expect(getDefaultOllamaModel({ type: "nvidia", totalMemoryMB: 16384 }, () => "")).toBe( + expect(getDefaultOllamaModel({ type: "nvidia", totalMemoryMB: 10_000 }, () => "")).toBe( "qwen2.5:7b", ); expect( getDefaultOllamaModel( - { type: "nvidia", totalMemoryMB: LARGE_OLLAMA_MIN_MEMORY_MB }, + { type: "nvidia", totalMemoryMB: LARGE_OLLAMA_FIT_MEMORY_MB }, () => "", ), ).toBe(QWEN3_6_OLLAMA_MODEL); }); + it("downgrades the bootstrap menu when currently available memory is low", () => { + // Unified-memory host (e.g. DGX Spark) with another GPU workload eating + // the system pool: 128 GiB total, ~12 GiB currently free. The 23 GiB + // qwen3.6:35b model would crash the runner mid-load, so the bootstrap + // menu must only offer the small model. + expect( + getBootstrapOllamaModelOptions({ + type: "nvidia", + totalMemoryMB: 131_072, + availableMemoryMB: 12_000, + }), + ).toEqual(["qwen2.5:7b"]); + expect( + getDefaultOllamaModel( + { type: "nvidia", totalMemoryMB: 131_072, availableMemoryMB: 12_000 }, + () => "", + ), + ).toBe("qwen2.5:7b"); + }); + + it("filters installed-model selection by memory fit", async () => { + const { getDefaultOllamaModel: gdom } = await import("../../../dist/lib/inference/local"); + // Even though nemotron-3-nano:30b is installed, it does not fit a host + // with only 12 GiB available — the selector must downgrade to a fitting + // installed model rather than blindly returning DEFAULT_OLLAMA_MODEL. + const installed = () => "qwen2.5:7b abc 4 GB now\nnemotron-3-nano:30b def 19 GB now"; + expect( + gdom({ type: "nvidia", totalMemoryMB: 131_072, availableMemoryMB: 12_000 }, installed), + ).toBe("qwen2.5:7b"); + }); + + it("resolveNonInteractiveOllamaModel respects unknown tags and downgrades known oversize ones", async () => { + const { resolveNonInteractiveOllamaModel } = await import( + "../../../dist/lib/inference/local" + ); + const messages: string[] = []; + const log = (m: string) => messages.push(m); + + // Known model that does not fit → fallback + warning. + expect( + resolveNonInteractiveOllamaModel( + "qwen3.6:35b", + null, + { type: "nvidia", totalMemoryMB: 131_072, availableMemoryMB: 12_000 }, + log, + ), + ).toBe("qwen2.5:7b"); + expect(messages.some((m) => m.includes("qwen3.6:35b"))).toBe(true); + + // Unknown tag → respected as-is. + messages.length = 0; + expect( + resolveNonInteractiveOllamaModel( + "some-custom:model", + null, + { type: "nvidia", totalMemoryMB: 131_072, availableMemoryMB: 12_000 }, + log, + ), + ).toBe("some-custom:model"); + expect(messages).toEqual([]); + + // No explicit choice → falls through to getDefaultOllamaModel. + expect( + resolveNonInteractiveOllamaModel( + null, + null, + { type: "nvidia", totalMemoryMB: 131_072, availableMemoryMB: 131_072 }, + log, + ), + ).toBe(QWEN3_6_OLLAMA_MODEL); + }); + + it("resolveNonInteractiveOllamaModel surfaces the no-fit warning when even the smallest model exceeds available memory", async () => { + const { resolveNonInteractiveOllamaModel } = await import( + "../../../dist/lib/inference/local" + ); + const messages: string[] = []; + const log = (m: string) => messages.push(m); + + // Explicit oversize tag AND host has less than the smallest registry + // entry needs. The explicit-downgrade warning fires *and* the no-fit + // warning fires so the user sees both signals. + const result = resolveNonInteractiveOllamaModel( + "qwen3.6:35b", + null, + { type: "nvidia", totalMemoryMB: 16_384, availableMemoryMB: 4_000 }, + log, + ); + expect(result).toBe("qwen2.5:7b"); + expect(messages.some((m) => m.includes("qwen3.6:35b"))).toBe(true); + expect(messages.some((m) => m.includes("No known Ollama bootstrap model fits"))).toBe(true); + + // No explicit choice + nothing fits: only the no-fit warning fires. + messages.length = 0; + expect( + resolveNonInteractiveOllamaModel( + null, + null, + { type: "nvidia", totalMemoryMB: 16_384, availableMemoryMB: 4_000 }, + log, + ), + ).toBe("qwen2.5:7b"); + expect(messages.some((m) => m.includes("No known Ollama bootstrap model fits"))).toBe(true); + }); + it("offers the large Ollama model on Apple Silicon with sufficient unified memory", () => { expect( getBootstrapOllamaModelOptions({ type: "apple", - totalMemoryMB: LARGE_OLLAMA_MIN_MEMORY_MB, + totalMemoryMB: LARGE_OLLAMA_FIT_MEMORY_MB, }), ).toEqual(["qwen2.5:7b", DEFAULT_OLLAMA_MODEL, QWEN3_6_OLLAMA_MODEL]); expect( getDefaultOllamaModel( - { type: "apple", totalMemoryMB: LARGE_OLLAMA_MIN_MEMORY_MB }, + { type: "apple", totalMemoryMB: LARGE_OLLAMA_FIT_MEMORY_MB }, () => "", ), ).toBe(QWEN3_6_OLLAMA_MODEL); @@ -599,20 +713,20 @@ describe("local inference helpers", () => { // where totalMemoryMB is set but the device type is "generic" or // unspecified. expect( - getBootstrapOllamaModelOptions({ totalMemoryMB: LARGE_OLLAMA_MIN_MEMORY_MB }), + getBootstrapOllamaModelOptions({ totalMemoryMB: LARGE_OLLAMA_FIT_MEMORY_MB }), ).toEqual(["qwen2.5:7b"]); expect( - getDefaultOllamaModel({ totalMemoryMB: LARGE_OLLAMA_MIN_MEMORY_MB }, () => ""), + getDefaultOllamaModel({ totalMemoryMB: LARGE_OLLAMA_FIT_MEMORY_MB }, () => ""), ).toBe("qwen2.5:7b"); expect( getBootstrapOllamaModelOptions({ type: "generic", - totalMemoryMB: LARGE_OLLAMA_MIN_MEMORY_MB * 4, + totalMemoryMB: LARGE_OLLAMA_FIT_MEMORY_MB * 4, }), ).toEqual(["qwen2.5:7b"]); expect( getDefaultOllamaModel( - { type: "generic", totalMemoryMB: LARGE_OLLAMA_MIN_MEMORY_MB * 4 }, + { type: "generic", totalMemoryMB: LARGE_OLLAMA_FIT_MEMORY_MB * 4 }, () => "", ), ).toBe("qwen2.5:7b"); diff --git a/src/lib/inference/local.ts b/src/lib/inference/local.ts index 9bd9e2f6a8b..56440fe5736 100644 --- a/src/lib/inference/local.ts +++ b/src/lib/inference/local.ts @@ -19,6 +19,15 @@ const { shellQuote, runCapture, runCaptureEx } = require("../runner"); import { OLLAMA_PORT, OLLAMA_PROXY_PORT, VLLM_PORT } from "../core/ports"; import { sleepSeconds } from "../core/wait"; +import { + anyRegistryModelFits, + effectiveGpuMemoryMB, + fittableOllamaModelTags, + largestFittableOllamaModelTag, + modelFitsAvailableMemory, + OLLAMA_MODEL_REGISTRY, + SMALLEST_OLLAMA_MODEL_TAG, +} from "./ollama-model-registry"; const { containerCanReachHostLoopback, inferContainerRuntime, isWsl } = require("../platform"); const { dockerInfo } = require("../adapters/docker/info"); @@ -48,10 +57,24 @@ export function resetOllamaContainerPortCache(): void { export const HOST_GATEWAY_URL = "http://host.openshell.internal"; export const LOCAL_INFERENCE_SANDBOX_HOST_URL_ENV = "NEMOCLAW_LOCAL_INFERENCE_SANDBOX_HOST_URL"; export const CONTAINER_REACHABILITY_IMAGE = "curlimages/curl:8.10.1"; -export const DEFAULT_OLLAMA_MODEL = "nemotron-3-nano:30b"; -export const QWEN3_6_OLLAMA_MODEL = "qwen3.6:35b"; -export const SMALL_OLLAMA_MODEL = "qwen2.5:7b"; -export const LARGE_OLLAMA_MIN_MEMORY_MB = 32768; +// These tags are convenience aliases for callers that want to refer to a +// specific bootstrap model by role rather than by string. The canonical +// metadata (memory requirements, download sizes) lives in +// `ollama-model-registry.ts`; the assertion below makes module load fail +// loudly if a registry edit drops a tag a caller still references by +// name, so the two stay in sync. +function assertRegistryTag(tag: string): string { + if (!OLLAMA_MODEL_REGISTRY.some((entry) => entry.tag === tag)) { + throw new Error( + `Tag '${tag}' is not in OLLAMA_MODEL_REGISTRY. Update the registry first.`, + ); + } + return tag; +} + +export const SMALL_OLLAMA_MODEL = SMALLEST_OLLAMA_MODEL_TAG; +export const DEFAULT_OLLAMA_MODEL = assertRegistryTag("nemotron-3-nano:30b"); +export const QWEN3_6_OLLAMA_MODEL = assertRegistryTag("qwen3.6:35b"); export type RunCaptureFn = (cmd: string | string[], opts?: { ignoreError?: boolean }) => string; @@ -125,6 +148,14 @@ export interface GpuInfo { // does not get sized as if it were confirmed NVIDIA / Apple Silicon // (#3510). type?: string; + // Currently free GPU memory at probe time. Populated by `detectGpu` from + // `nvidia-smi memory.free`, `MemAvailable` on unified-memory hosts, or + // `vm_stat` reclaimable pages on macOS. Used by the bootstrap-model + // selector so an idle 128 GiB Spark and a 128 GiB Spark with another + // GPU workload eating 116 GiB do not get the same model recommendation. + // Absent => the selector falls back to `totalMemoryMB`, preserving the + // previous behaviour. + availableMemoryMB?: number; } export interface ValidationResult { @@ -734,26 +765,59 @@ export function getOllamaModelOptions(runCaptureImpl?: RunCaptureFn): string[] { return parseOllamaList(listOutput); } -function isLargeOllamaCapableGpu(gpu: GpuInfo | null): boolean { - // Only confirmed-NVIDIA and Apple-Silicon devices get the large-model - // default. Other detection outcomes (null, missing `type`, or a partial - // result that fell through the NVIDIA path with type set to something - // else) fall back to the smaller model so we never download a 22 GB - // model onto a host whose acceleration is unconfirmed (#3510). - return ( - !!gpu && - (gpu.type === "nvidia" || gpu.type === "apple") && - gpu.totalMemoryMB >= LARGE_OLLAMA_MIN_MEMORY_MB - ); +export function getBootstrapOllamaModelOptions(gpu: GpuInfo | null): string[] { + // Delegate to the registry so the menu reflects what the host can + // actually load right now. Only confirmed-NVIDIA and Apple-Silicon + // devices get larger options; ambiguous device types fall back to the + // smallest model so a partial GPU detection cannot promote a host to a + // 22 GB model. + return fittableOllamaModelTags(gpu); } -export function getBootstrapOllamaModelOptions(gpu: GpuInfo | null): string[] { - const options = [SMALL_OLLAMA_MODEL]; - if (isLargeOllamaCapableGpu(gpu)) { - options.push(DEFAULT_OLLAMA_MODEL); - options.push(QWEN3_6_OLLAMA_MODEL); +/** + * Resolve the non-interactive Ollama model selection. When the caller has + * passed an explicit `NEMOCLAW_MODEL` / recovered-session model that the + * registry knows is too big for the host's currently available memory, + * log a warning and fall back to the largest fittable registry entry so + * onboarding does not pull a model the runner will crash on. Unknown + * model tags (user-supplied values the registry has never seen) are + * respected as-is — the runner's own validation surfaces the failure if + * the choice was wrong. + */ +export function resolveNonInteractiveOllamaModel( + requestedModel: string | null, + recoveredModel: string | null, + gpu: GpuInfo | null, + log: (message: string) => void = (m) => console.warn(m), +): string { + const explicit = requestedModel || recoveredModel; + if (explicit && !modelFitsAvailableMemory(explicit, gpu)) { + const fallback = largestFittableOllamaModelTag(gpu); + log( + ` ! Requested Ollama model '${explicit}' is unlikely to fit currently available GPU memory; ` + + `falling back to '${fallback}'. Override by freeing memory and re-running, or unset NEMOCLAW_MODEL.`, + ); + if (!anyRegistryModelFits(gpu)) { + warnNoBootstrapModelFits(gpu, log); + } + return fallback; } - return options; + if (!explicit && !anyRegistryModelFits(gpu)) { + warnNoBootstrapModelFits(gpu, log); + } + return explicit || getDefaultOllamaModel(gpu); +} + +function warnNoBootstrapModelFits( + gpu: GpuInfo | null, + log: (message: string) => void, +): void { + const memory = effectiveGpuMemoryMB(gpu); + log( + ` ! No known Ollama bootstrap model fits the host's currently available GPU memory` + + `${memory ? ` (~${memory} MB free)` : ""}. Proceeding with the smallest known model; ` + + "the runner may still reject the load — free memory and re-run if it does.", + ); } export function getDefaultOllamaModel( @@ -762,12 +826,23 @@ export function getDefaultOllamaModel( ): string { const models = getOllamaModelOptions(runCaptureImpl); if (models.length === 0) { - if (isLargeOllamaCapableGpu(gpu)) { - return QWEN3_6_OLLAMA_MODEL; - } - return SMALL_OLLAMA_MODEL; + // No installed models — pick the largest registry entry that fits the + // host's currently available memory. + return largestFittableOllamaModelTag(gpu); + } + // Filter the installed list to entries we either don't know (unmanaged + // user pulls — let the runner validate) or that fit the registry's + // memory requirement at probe time. If everything has been filtered out, + // fall back to the largest registry entry that fits so the wizard never + // suggests a model the host can't load. + const fittingInstalled = models.filter((tag) => modelFitsAvailableMemory(tag, gpu)); + const pool = fittingInstalled.length > 0 ? fittingInstalled : null; + if (pool === null) { + return largestFittableOllamaModelTag(gpu); } - return models.includes(DEFAULT_OLLAMA_MODEL) ? DEFAULT_OLLAMA_MODEL : models[0]; + return pool.includes(DEFAULT_OLLAMA_MODEL) && modelFitsAvailableMemory(DEFAULT_OLLAMA_MODEL, gpu) + ? DEFAULT_OLLAMA_MODEL + : pool[0]; } export function getOllamaWarmupCommand(model: string, keepAlive = "15m"): string[] { diff --git a/src/lib/inference/nim.test.ts b/src/lib/inference/nim.test.ts index 2ffb7ee7810..8cd25afbca2 100644 --- a/src/lib/inference/nim.test.ts +++ b/src/lib/inference/nim.test.ts @@ -196,16 +196,18 @@ describe("nim", () => { }); it("populates name and memory from primary nvidia-smi path", () => { - // Primary path returns name+memory.total in a single CSV line per GPU. - // Regression guard for #2669: the GB300 preflight line was missing the - // GPU model because only memory.total was being queried. + // Primary path returns name+memory.total+memory.free in a single CSV + // line per GPU. Regression guard for #2669: the GB300 preflight line + // was missing the GPU model because only memory.total was being + // queried. memory.free is also captured so the bootstrap-model + // selector can size against currently free memory, not just total. const runCapture = vi.fn((cmd: string | string[]) => { if (!Array.isArray(cmd)) throw new Error("expected argv array"); if ( cmd[0] === "nvidia-smi" && cmd.some((a: string) => a.includes("name,memory.total")) ) { - return "NVIDIA GB300, 284208\n"; + return "NVIDIA GB300, 284208, 280000\n"; } return ""; }); @@ -217,6 +219,7 @@ describe("nim", () => { name: "NVIDIA GB300", count: 1, totalMemoryMB: 284208, + availableMemoryMB: 280000, perGpuMB: 284208, }); } finally { @@ -231,7 +234,7 @@ describe("nim", () => { cmd[0] === "nvidia-smi" && cmd.some((a: string) => a.includes("name,memory.total")) ) { - return "NVIDIA H100 80GB HBM3, 81920\nNVIDIA H100 80GB HBM3, 81920\n"; + return "NVIDIA H100 80GB HBM3, 81920, 81000\nNVIDIA H100 80GB HBM3, 81920, 60000\n"; } return ""; }); @@ -265,7 +268,7 @@ describe("nim", () => { cmd[0] === "nvidia-smi" && cmd.some((a: string) => a.includes("name,memory.total")) ) { - return "NVIDIA RTX A,B, 81920\n"; + return "NVIDIA RTX A,B, 81920, 80000\n"; } return ""; }); @@ -295,7 +298,7 @@ describe("nim", () => { cmd[0] === "nvidia-smi" && cmd.some((a: string) => a.includes("name,memory.total")) ) { - return "NVIDIA RTX PRO 6000 Blackwell Max-Q, 97887\nNVIDIA GB300, 256703\n"; + return "NVIDIA RTX PRO 6000 Blackwell Max-Q, 97887, 90000\nNVIDIA GB300, 256703, 250000\n"; } return ""; }); @@ -329,7 +332,7 @@ describe("nim", () => { cmd[0] === "nvidia-smi" && cmd.some((a: string) => a.includes("name,memory.total")) ) { - return "JMJWOA-Generic-GPU, 65471\n"; + return "JMJWOA-Generic-GPU, 65471, 65000\n"; } return ""; }); @@ -359,7 +362,7 @@ describe("nim", () => { cmd[0] === "nvidia-smi" && cmd.some((a: string) => a.includes("name,memory.total")) ) { - return "NVIDIA JMJWOA-Generic-GPU, 65471\n"; + return "NVIDIA JMJWOA-Generic-GPU, 65471, 65000\n"; } return ""; }); @@ -385,7 +388,7 @@ describe("nim", () => { cmd[0] === "nvidia-smi" && cmd.some((a: string) => a.includes("name,memory.total")) ) { - return "JMJWOA-Generic-GPU, 131072\n"; + return "JMJWOA-Generic-GPU, 131072, 12000\n"; } return ""; }); @@ -422,6 +425,10 @@ describe("nim", () => { name: "NVIDIA GB10", count: 1, totalMemoryMB: 131072, + // MemAvailable from the stubbed `free -m` row is propagated so the + // bootstrap-model selector can size against currently free memory + // on unified-memory hosts. + availableMemoryMB: 119808, perGpuMB: 131072, nimCapable: true, unifiedMemory: true, @@ -466,6 +473,7 @@ describe("nim", () => { name: "NVIDIA JMJWOA-Generic-GPU", count: 1, totalMemoryMB: 122543, + availableMemoryMB: 111279, perGpuMB: 122543, unifiedMemory: true, spark: true, @@ -523,6 +531,7 @@ describe("nim", () => { name: "NVIDIA Jetson AGX Orin", count: 1, totalMemoryMB: 32768, + availableMemoryMB: 27136, perGpuMB: 32768, nimCapable: true, unifiedMemory: true, @@ -539,7 +548,7 @@ describe("nim", () => { if (!Array.isArray(cmd)) throw new Error("expected argv array"); if (cmd[0] === "nvidia-smi") return ""; if (cmd[0] === "free" && cmd[1] === "-m") { - return " total used free\nMem: 65536 4096 50000\nSwap: 0 0 0"; + return " total used free shared buff/cache available\nMem: 65536 4096 50000 512 10928 60000\nSwap: 0 0 0"; } return ""; }); @@ -562,6 +571,7 @@ describe("nim", () => { name: "NVIDIA Jetson AGX Orin", count: 1, totalMemoryMB: 65536, + availableMemoryMB: 60000, perGpuMB: 65536, nimCapable: true, unifiedMemory: true, @@ -576,6 +586,35 @@ describe("nim", () => { } }); + it("omits availableMemoryMB when memory.free fails to parse on the primary path", () => { + // nvidia-smi sometimes reports `[N/A]` or empty strings for memory.free + // (driver / virtualisation quirks). Total still parses, so we keep + // surfacing it; available is dropped so callers fall back to total. + const runCapture = vi.fn((cmd: string | string[]) => { + if (!Array.isArray(cmd)) throw new Error("expected argv array"); + if ( + cmd[0] === "nvidia-smi" && + cmd.some((a: string) => a.includes("name,memory.total")) + ) { + return "NVIDIA H100 80GB HBM3, 81920, [N/A]\n"; + } + return ""; + }); + const { nimModule, restore } = loadNimWithMockedRunner(runCapture); + + try { + const result = nimModule.detectGpu(); + expect(result).toMatchObject({ + type: "nvidia", + name: "NVIDIA H100 80GB HBM3", + totalMemoryMB: 81920, + }); + expect(result?.availableMemoryMB).toBeUndefined(); + } finally { + restore(); + } + }); + // Same invariant as the primary-path mixed-model test: don't pin a // single name on hosts with multiple distinct GPU models, even on the // unified-memory fallback. Hypothetical today (no shipping platform mixes @@ -632,6 +671,87 @@ describe("nim", () => { restore(); } }); + + it("populates availableMemoryMB on macOS Apple Silicon via vm_stat", () => { + const runCapture = vi.fn((cmd: string | string[]) => { + if (!Array.isArray(cmd)) throw new Error("expected argv array"); + if (cmd[0] === "system_profiler" && cmd[1] === "SPDisplaysDataType") { + return [ + " Chipset Model: Apple M3 Max", + " Total Number of Cores: 40", + " VRAM (Dynamic, Max): 49152 MB", + ].join("\n"); + } + if (cmd[0] === "sysctl" && cmd[1] === "-n" && cmd[2] === "hw.memsize") { + return String(64 * 1024 * 1024 * 1024); + } + if (cmd[0] === "vm_stat") { + // 16 KiB page size; 32 000 free + 60 000 inactive + 1 000 speculative + // pages → (93 000 × 16 384) bytes ≈ 1 488 MiB available. + return [ + "Mach Virtual Memory Statistics: (page size of 16384 bytes)", + "Pages free: 32000.", + "Pages active: 500000.", + "Pages inactive: 60000.", + "Pages speculative: 1000.", + ].join("\n"); + } + return ""; + }); + const { nimModule, restore } = loadNimWithMockedRunner(runCapture); + const originalPlatform = Object.getOwnPropertyDescriptor(process, "platform"); + Object.defineProperty(process, "platform", { value: "darwin", configurable: true }); + + try { + const expectedAvailableMB = Math.floor(((32_000 + 60_000 + 1_000) * 16_384) / 1024 / 1024); + expect(nimModule.detectGpu()).toMatchObject({ + type: "apple", + name: "Apple M3 Max", + totalMemoryMB: 49152, + availableMemoryMB: expectedAvailableMB, + cores: 40, + }); + } finally { + if (originalPlatform) { + Object.defineProperty(process, "platform", originalPlatform); + } + restore(); + } + }); + + it("omits availableMemoryMB on macOS when vm_stat fails to parse", () => { + const runCapture = vi.fn((cmd: string | string[]) => { + if (!Array.isArray(cmd)) throw new Error("expected argv array"); + if (cmd[0] === "system_profiler" && cmd[1] === "SPDisplaysDataType") { + return [ + " Chipset Model: Apple M2", + " Total Number of Cores: 10", + " VRAM (Dynamic, Max): 16384 MB", + ].join("\n"); + } + // vm_stat returns nothing → readMacOsAvailableMemoryMB() returns 0, + // so availableMemoryMB must be absent from the result. + return ""; + }); + const { nimModule, restore } = loadNimWithMockedRunner(runCapture); + const originalPlatform = Object.getOwnPropertyDescriptor(process, "platform"); + Object.defineProperty(process, "platform", { value: "darwin", configurable: true }); + + try { + const result = nimModule.detectGpu(); + expect(result).toMatchObject({ + type: "apple", + name: "Apple M2", + totalMemoryMB: 16384, + }); + expect(result?.availableMemoryMB).toBeUndefined(); + } finally { + if (originalPlatform) { + Object.defineProperty(process, "platform", originalPlatform); + } + restore(); + } + }); }); describe("groupGpusByName", () => { diff --git a/src/lib/inference/nim.ts b/src/lib/inference/nim.ts index 0ccd314b821..5eb8358f5d8 100644 --- a/src/lib/inference/nim.ts +++ b/src/lib/inference/nim.ts @@ -76,6 +76,13 @@ export interface GpuDetection { gpus?: NimGpu[]; count: number; totalMemoryMB: number; + // Currently free GPU memory at probe time. NVIDIA: summed from + // `nvidia-smi memory.free`. Unified-memory (Spark/Jetson): approximated + // from host `MemAvailable` since GPU memory is the system pool. macOS: + // approximated from `vm_stat` reclaimable pages. Absent when every + // probe was inconclusive; downstream callers fall back to + // `totalMemoryMB`. + availableMemoryMB?: number; perGpuMB: number; cores?: number | null; nimCapable: boolean; @@ -180,6 +187,54 @@ function readHostMemoryMB(): number { return 0; } +// macOS equivalent of `MemAvailable`: parse `vm_stat` output, sum the +// kernel-reclaimable page classes (free + inactive + speculative), and +// scale by the reported page size. The result is the same "could I load +// a 22 GB model right now?" signal the unified-memory Linux path uses. +// Returns 0 when any expected field is missing so the caller can treat +// the figure as "unknown" and fall back to total memory. +function readMacOsAvailableMemoryMB(): number { + try { + const out = runCapture(["vm_stat"], { ignoreError: true }); + if (!out) return 0; + const pageMatch = out.match(/page size of (\d+) bytes/); + if (!pageMatch) return 0; + const pageBytes = parseInt(pageMatch[1], 10); + if (!Number.isFinite(pageBytes) || pageBytes <= 0) return 0; + const grab = (label: string): number => { + const match = out.match(new RegExp(`Pages ${label}:\\s+(\\d+)\\.`)); + return match ? parseInt(match[1], 10) : 0; + }; + const pages = grab("free") + grab("inactive") + grab("speculative"); + if (pages <= 0) return 0; + return Math.floor((pages * pageBytes) / 1024 / 1024); + } catch { + return 0; + } +} + +// `free -m` columns: total used free shared buff/cache available. +// "available" (column 6) is the kernel's estimate of memory that can be +// reclaimed without swapping — the right signal for "is there room for a +// 22 GB Ollama load right now?" on unified-memory hosts. Returns 0 when +// the column cannot be parsed; the caller treats 0 as "unknown" and falls +// back to total memory. +function readHostAvailableMemoryMB(): number { + try { + const freeOut = runCapture(["free", "-m"], { ignoreError: true }); + if (freeOut) { + const memLine = freeOut.split("\n").find((l: string) => l.includes("Mem:")); + if (memLine) { + const parts = memLine.split(/\s+/); + return parseInt(parts[6], 10) || 0; + } + } + } catch { + /* ignored */ + } + return 0; +} + function hostPathExists(path: string): boolean { try { return fs.existsSync(path); @@ -279,26 +334,41 @@ export function canRunNimWithMemory(totalMemoryMB: number): boolean { } export function detectGpu(): GpuDetection | null { - // Try NVIDIA first — query name and VRAM in a single call so the preflight - // line can show the GPU model alongside the memory size. + // Try NVIDIA first — query name, total, and free VRAM in a single call so + // the preflight line can show the GPU model alongside the memory size and + // the bootstrap-model selector can pick a model that fits currently + // available memory, not just the headline total. try { const output = runCapture( - ["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader,nounits"], + [ + "nvidia-smi", + "--query-gpu=name,memory.total,memory.free", + "--format=csv,noheader,nounits", + ], { ignoreError: true }, ); if (output) { - type ParsedGpu = { name: string; memoryMB: number }; + type ParsedGpu = { name: string; memoryMB: number; freeMemoryMB: number }; const parsed: ParsedGpu[] = []; for (const raw of output.split("\n")) { const line = raw.trim(); if (!line) continue; - // Split on the LAST comma — GPU names can contain commas in rare cases. - const idx = line.lastIndexOf(","); - if (idx === -1) continue; - const name = line.slice(0, idx).trim(); - const memoryMB = parseInt(line.slice(idx + 1).trim(), 10); + // Split on commas from the RIGHT: free MB, then total MB; the + // remainder is the GPU name (which can itself contain commas). + const lastIdx = line.lastIndexOf(","); + if (lastIdx === -1) continue; + const freeMemoryMB = parseInt(line.slice(lastIdx + 1).trim(), 10); + const beforeFree = line.slice(0, lastIdx); + const totalIdx = beforeFree.lastIndexOf(","); + if (totalIdx === -1) continue; + const memoryMB = parseInt(beforeFree.slice(totalIdx + 1).trim(), 10); + const name = beforeFree.slice(0, totalIdx).trim(); if (isNaN(memoryMB)) continue; - parsed.push({ name, memoryMB }); + parsed.push({ + name, + memoryMB, + freeMemoryMB: isNaN(freeMemoryMB) ? 0 : freeMemoryMB, + }); } if (parsed.length > 0) { const platform = detectNvidiaPlatform(); @@ -319,6 +389,10 @@ export function detectGpu(): GpuDetection | null { (sum: number, p: ParsedGpu) => sum + p.memoryMB, 0, ); + const availableMemoryMB = trusted.reduce( + (sum: number, p: ParsedGpu) => sum + p.freeMemoryMB, + 0, + ); const firstName = trusted[0].name; // Only surface a single name when every GPU reports the same model; // a mixed-GPU host would otherwise be misreported as `Nx `. @@ -330,6 +404,7 @@ export function detectGpu(): GpuDetection | null { gpus: trusted.map((p) => ({ name: p.name, memoryMB: p.memoryMB })), count: trusted.length, totalMemoryMB, + ...(availableMemoryMB > 0 ? { availableMemoryMB } : {}), perGpuMB: trusted[0].memoryMB, nimCapable: canRunNimWithMemory(totalMemoryMB), platform, @@ -388,12 +463,17 @@ export function detectGpu(): GpuDetection | null { // Memory.total is not available on unified-memory devices, so we split // the host RAM evenly across the named GPUs for the per-GPU breakdown. // Approximation, but the only number nvidia-smi gives us in this path. + // `availableMemoryMB` mirrors that approximation using MemAvailable so + // the bootstrap-model selector reacts to concurrent GPU workloads + // eating into the shared system pool. + const availableMemoryMB = readHostAvailableMemoryMB(); return { type: "nvidia", ...(allUnifiedSameName ? { name: firstUnifiedName } : {}), gpus: unifiedGpuNames.map((name: string) => ({ name, memoryMB: perGpuMB })), count, totalMemoryMB, + ...(availableMemoryMB > 0 ? { availableMemoryMB } : {}), perGpuMB: perGpuMB || totalMemoryMB, nimCapable: canRunNimWithMemory(totalMemoryMB), unifiedMemory: true, @@ -410,12 +490,14 @@ export function detectGpu(): GpuDetection | null { const tegraGpu = detectTegraHostGpu(); if (tegraGpu) { const totalMemoryMB = readHostMemoryMB(); + const availableMemoryMB = readHostAvailableMemoryMB(); return { type: "nvidia", name: tegraGpu.name, gpus: [{ name: tegraGpu.name, memoryMB: totalMemoryMB }], count: 1, totalMemoryMB, + ...(availableMemoryMB > 0 ? { availableMemoryMB } : {}), perGpuMB: totalMemoryMB, nimCapable: canRunNimWithMemory(totalMemoryMB), unifiedMemory: true, @@ -451,12 +533,14 @@ export function detectGpu(): GpuDetection | null { } } + const availableMemoryMB = readMacOsAvailableMemoryMB(); return { type: "apple", name, count: 1, cores: coresMatch ? parseInt(coresMatch[1], 10) : null, totalMemoryMB: memoryMB, + ...(availableMemoryMB > 0 ? { availableMemoryMB } : {}), perGpuMB: memoryMB, nimCapable: false, }; diff --git a/src/lib/inference/ollama-model-registry.test.ts b/src/lib/inference/ollama-model-registry.test.ts new file mode 100644 index 00000000000..68ef0257e34 --- /dev/null +++ b/src/lib/inference/ollama-model-registry.test.ts @@ -0,0 +1,180 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { describe, expect, it } from "vitest"; + +import { + effectiveGpuMemoryMB, + findOllamaModelEntry, + fittableOllamaModelTags, + largestFittableOllamaModelTag, + modelFitsAvailableMemory, + OLLAMA_DOWNLOAD_SIZE_FALLBACK_BYTES, + OLLAMA_MODEL_REGISTRY, + SMALLEST_OLLAMA_MODEL_TAG, +} from "../../../dist/lib/inference/ollama-model-registry"; + +describe("OLLAMA_MODEL_REGISTRY", () => { + it("is ordered largest-first by requiredMemoryMB", () => { + for (let i = 0; i < OLLAMA_MODEL_REGISTRY.length - 1; i++) { + expect(OLLAMA_MODEL_REGISTRY[i].requiredMemoryMB).toBeGreaterThan( + OLLAMA_MODEL_REGISTRY[i + 1].requiredMemoryMB, + ); + } + }); + + it("exposes the smallest tag as SMALLEST_OLLAMA_MODEL_TAG", () => { + const lastEntry = OLLAMA_MODEL_REGISTRY[OLLAMA_MODEL_REGISTRY.length - 1]; + expect(SMALLEST_OLLAMA_MODEL_TAG).toBe(lastEntry.tag); + }); +}); + +describe("findOllamaModelEntry", () => { + it("returns the registry entry by tag", () => { + const entry = findOllamaModelEntry(SMALLEST_OLLAMA_MODEL_TAG); + expect(entry).not.toBeNull(); + expect(entry?.tag).toBe(SMALLEST_OLLAMA_MODEL_TAG); + }); + + it("returns null for unknown tags", () => { + expect(findOllamaModelEntry("definitely-not-a-real-model:99b")).toBeNull(); + }); +}); + +describe("effectiveGpuMemoryMB", () => { + it("returns null when gpu is null", () => { + expect(effectiveGpuMemoryMB(null)).toBeNull(); + }); + + it("prefers availableMemoryMB when set", () => { + expect( + effectiveGpuMemoryMB({ type: "nvidia", totalMemoryMB: 131_072, availableMemoryMB: 12_000 }), + ).toBe(12_000); + }); + + it("falls back to totalMemoryMB when availableMemoryMB is absent", () => { + expect(effectiveGpuMemoryMB({ type: "nvidia", totalMemoryMB: 32_768 })).toBe(32_768); + }); + + it("ignores zero or negative availableMemoryMB so the caller's totalMemoryMB still wins", () => { + expect( + effectiveGpuMemoryMB({ type: "nvidia", totalMemoryMB: 32_768, availableMemoryMB: 0 }), + ).toBe(32_768); + }); +}); + +describe("fittableOllamaModelTags", () => { + it("returns the smallest tag for null gpus and ambiguous device types", () => { + expect(fittableOllamaModelTags(null)).toEqual([SMALLEST_OLLAMA_MODEL_TAG]); + expect(fittableOllamaModelTags({ type: "generic", totalMemoryMB: 131_072 })).toEqual([ + SMALLEST_OLLAMA_MODEL_TAG, + ]); + }); + + it("includes every entry that fits the available-memory figure (smallest-first)", () => { + const tags = fittableOllamaModelTags({ + type: "nvidia", + totalMemoryMB: 131_072, + availableMemoryMB: 131_072, + }); + expect(tags[0]).toBe(SMALLEST_OLLAMA_MODEL_TAG); + expect(tags.length).toBe(OLLAMA_MODEL_REGISTRY.length); + // Smallest-first: each subsequent entry should require at least as much + // memory as the previous one. + for (let i = 0; i < tags.length - 1; i++) { + const a = OLLAMA_MODEL_REGISTRY.find((e) => e.tag === tags[i]); + const b = OLLAMA_MODEL_REGISTRY.find((e) => e.tag === tags[i + 1]); + expect(a && b && a.requiredMemoryMB <= b.requiredMemoryMB).toBe(true); + } + }); + + it("falls back to the smallest tag when nothing in the registry fits available memory", () => { + // Unified-memory host with another GPU workload eating the system + // pool: 128 GiB total, ~12 GiB currently available. Nothing in the + // registry requires <= 12 GiB except the smallest model. + expect( + fittableOllamaModelTags({ + type: "nvidia", + totalMemoryMB: 131_072, + availableMemoryMB: 12_000, + }), + ).toEqual([SMALLEST_OLLAMA_MODEL_TAG]); + }); + + it("uses totalMemoryMB when availableMemoryMB is absent so legacy detection still works", () => { + expect( + fittableOllamaModelTags({ type: "nvidia", totalMemoryMB: 131_072 }).length, + ).toBe(OLLAMA_MODEL_REGISTRY.length); + }); +}); + +describe("modelFitsAvailableMemory", () => { + it("returns true for unknown tags so user-supplied model names are respected", () => { + expect( + modelFitsAvailableMemory("definitely-not-a-real-model:99b", { + type: "nvidia", + totalMemoryMB: 16_384, + availableMemoryMB: 4_000, + }), + ).toBe(true); + }); + + it("returns true when GPU memory is unknown so capacity gating does not fire blind", () => { + expect(modelFitsAvailableMemory(OLLAMA_MODEL_REGISTRY[0].tag, null)).toBe(true); + }); + + it("returns false when a known model exceeds the host's currently available memory", () => { + expect( + modelFitsAvailableMemory("qwen3.6:35b", { + type: "nvidia", + totalMemoryMB: 131_072, + availableMemoryMB: 12_000, + }), + ).toBe(false); + }); + + it("returns true when a known model fits", () => { + expect( + modelFitsAvailableMemory("qwen2.5:7b", { + type: "nvidia", + totalMemoryMB: 131_072, + availableMemoryMB: 12_000, + }), + ).toBe(true); + }); +}); + +describe("OLLAMA_DOWNLOAD_SIZE_FALLBACK_BYTES", () => { + it("mirrors the registry's downloadSizeBytes for every entry", () => { + for (const entry of OLLAMA_MODEL_REGISTRY) { + expect(OLLAMA_DOWNLOAD_SIZE_FALLBACK_BYTES[entry.tag]).toBe(entry.downloadSizeBytes); + } + }); + + it("exposes the largest fittable tag via largestFittableOllamaModelTag", () => { + expect( + largestFittableOllamaModelTag({ + type: "nvidia", + totalMemoryMB: 131_072, + availableMemoryMB: 12_000, + }), + ).toBe(SMALLEST_OLLAMA_MODEL_TAG); + const allFit = largestFittableOllamaModelTag({ + type: "nvidia", + totalMemoryMB: 131_072, + availableMemoryMB: 131_072, + }); + expect(allFit).toBe(OLLAMA_MODEL_REGISTRY[0].tag); + }); + + it("treats apple silicon the same as nvidia when availableMemoryMB is supplied", () => { + // The registry filter is identical across confirmed types — given the + // same availableMemoryMB it returns the same set of fittable tags. The + // macOS detection path populates availableMemoryMB from `vm_stat` + // reclaimable pages; this test exercises the filter logic directly so + // it does not depend on the macOS-only probe. + expect( + fittableOllamaModelTags({ type: "apple", totalMemoryMB: 131_072, availableMemoryMB: 12_000 }), + ).toEqual([SMALLEST_OLLAMA_MODEL_TAG]); + }); +}); diff --git a/src/lib/inference/ollama-model-registry.ts b/src/lib/inference/ollama-model-registry.ts new file mode 100644 index 00000000000..dc07f682d4d --- /dev/null +++ b/src/lib/inference/ollama-model-registry.ts @@ -0,0 +1,145 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +/** + * Memory-aware Ollama bootstrap-model registry. + * + * Central metadata for the bootstrap-model list. Every onboard path that + * cares about a known Ollama model — the menu, the non-interactive + * default, the requested-model capacity guard, and the download-size + * fallback table — reads from this single source so model facts are not + * duplicated across the codebase. + * + * Each entry pairs a tag with: + * + * - `requiredMemoryMB`: the GPU memory the runner needs to load the model + * at default context, set slightly above the on-disk weight size to + * leave headroom for the KV cache + context tokens. + * - `downloadSizeBytes`: the approximate compressed tarball size, used as + * a fallback when the live Ollama registry manifest probe fails. + * + * New models go here, in descending size order; the selector walks the + * list top-down and keeps every entry whose `requiredMemoryMB` fits the + * host's currently available memory. + */ + +import type { GpuInfo } from "./local"; + +export interface OllamaModelEntry { + tag: string; + requiredMemoryMB: number; + downloadSizeBytes: number; +} + +// Largest first. The selector walks this list, filters by available memory, +// and reverses the result so menus render smallest-first. +export const OLLAMA_MODEL_REGISTRY: readonly OllamaModelEntry[] = [ + { tag: "qwen3.6:35b", requiredMemoryMB: 26_000, downloadSizeBytes: 24_000_000_000 }, + { tag: "nemotron-3-nano:30b", requiredMemoryMB: 22_000, downloadSizeBytes: 19_000_000_000 }, + { tag: "qwen2.5:7b", requiredMemoryMB: 8_000, downloadSizeBytes: 4_683_073_184 }, +]; + +export const SMALLEST_OLLAMA_MODEL_TAG = + OLLAMA_MODEL_REGISTRY[OLLAMA_MODEL_REGISTRY.length - 1].tag; + +export function findOllamaModelEntry(tag: string): OllamaModelEntry | null { + return OLLAMA_MODEL_REGISTRY.find((entry) => entry.tag === tag) ?? null; +} + +/** + * Effective GPU memory for capacity decisions: prefer the currently + * available figure (from `nvidia-smi memory.free` or `MemAvailable`) and + * fall back to total when the host could not produce a usable free-memory + * reading. Total is a worse signal — it ignores concurrent workload + * footprints — but keeps the pre-registry behaviour on hosts where + * `availableMemoryMB` is missing. + */ +export function effectiveGpuMemoryMB(gpu: GpuInfo | null): number | null { + if (!gpu) return null; + if (typeof gpu.availableMemoryMB === "number" && gpu.availableMemoryMB > 0) { + return gpu.availableMemoryMB; + } + if (typeof gpu.totalMemoryMB === "number" && gpu.totalMemoryMB > 0) { + return gpu.totalMemoryMB; + } + return null; +} + +/** + * `true` when the registered tag fits the host's currently available + * memory. Unknown tags (e.g. user-supplied `NEMOCLAW_MODEL` values that + * the registry has never seen) and unknown memory both return `true` so + * the caller does not refuse to proceed when we have nothing to compare + * against — the runner's own validation is the final authority in that + * case. + */ +export function modelFitsAvailableMemory(tag: string, gpu: GpuInfo | null): boolean { + const entry = findOllamaModelEntry(tag); + if (!entry) return true; + const memory = effectiveGpuMemoryMB(gpu); + if (memory == null) return true; + return entry.requiredMemoryMB <= memory; +} + +/** + * Bootstrap model tags the host can plausibly load right now. Always + * includes `SMALLEST_OLLAMA_MODEL_TAG` so the menu has at least one + * fallback even when capacity probing says nothing in the registry fits; + * use `anyRegistryModelFits` to detect that under-spec case explicitly + * and warn the user before we hand them a model the runner is likely to + * reject too. + * + * Output is smallest-first so menu indices stay stable as registry entries + * are added. Only confirmed-NVIDIA and Apple-Silicon devices are eligible + * for larger entries; ambiguous device types fall back to the smallest + * model so a partial detection does not promote a host to a 22 GB model. + */ +export function fittableOllamaModelTags(gpu: GpuInfo | null): string[] { + const fallback = [SMALLEST_OLLAMA_MODEL_TAG]; + if (!gpu || (gpu.type !== "nvidia" && gpu.type !== "apple")) { + return fallback; + } + const memory = effectiveGpuMemoryMB(gpu); + if (memory == null) return fallback; + const fitting = OLLAMA_MODEL_REGISTRY.filter( + (entry) => entry.requiredMemoryMB <= memory && entry.tag !== SMALLEST_OLLAMA_MODEL_TAG, + ); + if (fitting.length === 0) return fallback; + return [SMALLEST_OLLAMA_MODEL_TAG, ...fitting.map((entry) => entry.tag).reverse()]; +} + +/** + * `true` when at least one registry entry fits the host's currently + * available memory. Returns `true` when memory is unknown so callers do + * not warn blind. Confirmed-eligible device types (`nvidia`, `apple`) + * compare against the registry; ambiguous types fall through to `true` + * for the same reason as `fittableOllamaModelTags` — we cannot tell, so + * the runner is left to surface any real failure. + */ +export function anyRegistryModelFits(gpu: GpuInfo | null): boolean { + if (!gpu || (gpu.type !== "nvidia" && gpu.type !== "apple")) return true; + const memory = effectiveGpuMemoryMB(gpu); + if (memory == null) return true; + return OLLAMA_MODEL_REGISTRY.some((entry) => entry.requiredMemoryMB <= memory); +} + +/** + * Largest tag in the smallest-first `fittableOllamaModelTags` output. Used + * by callers that want a single recommended default rather than the + * whole menu. + */ +export function largestFittableOllamaModelTag(gpu: GpuInfo | null): string { + const tags = fittableOllamaModelTags(gpu); + return tags[tags.length - 1]; +} + +/** + * Registry-derived download-size fallback table. Used by `model-size.ts` + * when the live `https://registry.ollama.ai` manifest probe fails. + */ +export const OLLAMA_DOWNLOAD_SIZE_FALLBACK_BYTES: Readonly> = + Object.freeze( + Object.fromEntries( + OLLAMA_MODEL_REGISTRY.map((entry) => [entry.tag, entry.downloadSizeBytes]), + ), + ); diff --git a/src/lib/inference/ollama/model-size.ts b/src/lib/inference/ollama/model-size.ts index 6efd0975822..2b4c1d1bc20 100644 --- a/src/lib/inference/ollama/model-size.ts +++ b/src/lib/inference/ollama/model-size.ts @@ -2,17 +2,18 @@ // SPDX-License-Identifier: Apache-2.0 import { runCapture } from "../../runner"; +import { OLLAMA_DOWNLOAD_SIZE_FALLBACK_BYTES } from "../ollama-model-registry"; const MANIFEST_HOST = "https://registry.ollama.ai"; const PROBE_TIMEOUT_SECONDS = 3; const MANIFEST_ACCEPT_HEADER = "Accept: application/vnd.docker.distribution.manifest.v2+json, application/vnd.oci.image.manifest.v1+json"; -const FALLBACK_SIZE_BYTES: Readonly> = { - "qwen2.5:7b": 4_683_073_184, - "nemotron-3-nano:30b": 19_000_000_000, - "qwen3.6:35b": 24_000_000_000, -}; +// Single source of truth lives in `ollama-model-registry.ts`. The fallback +// table is aliased locally so existing call sites in this module still +// read from a `FALLBACK_SIZE_BYTES` reference; nothing outside this file +// reaches in for it. +const FALLBACK_SIZE_BYTES = OLLAMA_DOWNLOAD_SIZE_FALLBACK_BYTES; export type CaptureFn = (cmd: readonly string[], opts?: { ignoreError?: boolean }) => string; diff --git a/src/lib/inference/ollama/proxy.test.ts b/src/lib/inference/ollama/proxy.test.ts new file mode 100644 index 00000000000..09e916eef2f --- /dev/null +++ b/src/lib/inference/ollama/proxy.test.ts @@ -0,0 +1,100 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import { createRequire } from "node:module"; +import { afterEach, describe, expect, it } from "vitest"; + +const require = createRequire(import.meta.url); +const PROXY_DIST = require.resolve("../../../../dist/lib/inference/ollama/proxy"); +const LOCAL_DIST = require.resolve("../../../../dist/lib/inference/local"); +const CREDS_DIST = require.resolve("../../../../dist/lib/credentials/store"); + +interface MockSetup { + installed: string[]; + promptValues: string[]; +} + +function loadProxyWithMocks(setup: MockSetup): { + proxy: typeof import("../../../../dist/lib/inference/ollama/proxy"); + promptArgs: string[]; + restore: () => void; +} { + const local = require(LOCAL_DIST); + const creds = require(CREDS_DIST); + const originalGetOllamaModelOptions = local.getOllamaModelOptions; + const originalPrompt = creds.prompt; + const promptArgs: string[] = []; + let promptCallIndex = 0; + + local.getOllamaModelOptions = () => setup.installed; + creds.prompt = async (message: string) => { + promptArgs.push(message); + const value = setup.promptValues[promptCallIndex]; + promptCallIndex += 1; + return value ?? ""; + }; + + delete require.cache[PROXY_DIST]; + const proxy = require(PROXY_DIST); + return { + proxy, + promptArgs, + restore() { + delete require.cache[PROXY_DIST]; + local.getOllamaModelOptions = originalGetOllamaModelOptions; + creds.prompt = originalPrompt; + }, + }; +} + +describe("promptOllamaModel installed-model fit filter", () => { + let active: { restore: () => void } | null = null; + afterEach(() => { + active?.restore(); + active = null; + }); + + it("downgrades to a starter model when the only installed entry exceeds available memory", async () => { + const setup = loadProxyWithMocks({ + installed: ["qwen3.6:35b"], + // Enter on the rendered default. + promptValues: [""], + }); + active = setup; + const result = await setup.proxy.promptOllamaModel({ + type: "nvidia", + totalMemoryMB: 131_072, + availableMemoryMB: 12_000, + }); + expect(result).toBe("qwen2.5:7b"); + }); + + it("keeps a fitting installed model as the default", async () => { + const setup = loadProxyWithMocks({ + installed: ["qwen2.5:7b", "qwen3.6:35b"], + promptValues: [""], + }); + active = setup; + const result = await setup.proxy.promptOllamaModel({ + type: "nvidia", + totalMemoryMB: 131_072, + availableMemoryMB: 12_000, + }); + // Only qwen2.5:7b fits; the menu offers only it, Enter selects it. + expect(result).toBe("qwen2.5:7b"); + }); + + it("respects unknown installed tags (not in the registry) even when nothing else fits", async () => { + const setup = loadProxyWithMocks({ + installed: ["my-custom:model"], + promptValues: [""], + }); + active = setup; + const result = await setup.proxy.promptOllamaModel({ + type: "nvidia", + totalMemoryMB: 131_072, + availableMemoryMB: 12_000, + }); + expect(result).toBe("my-custom:model"); + }); +}); diff --git a/src/lib/inference/ollama/proxy.ts b/src/lib/inference/ollama/proxy.ts index 5eef2307520..0ef15457ff9 100644 --- a/src/lib/inference/ollama/proxy.ts +++ b/src/lib/inference/ollama/proxy.ts @@ -5,6 +5,8 @@ // Ollama auth-proxy lifecycle: token persistence, PID management, // proxy start/stop, model pull and validation. +import type { GpuInfo } from "../local"; + const path = require("path"); const { spawn, spawnSync } = require("child_process"); const http = require("http"); @@ -21,6 +23,7 @@ const { probeOllamaModelCapabilities, validateOllamaModel, } = require("../local"); +const { anyRegistryModelFits, modelFitsAvailableMemory } = require("../ollama-model-registry"); const { buildSubprocessEnv } = require("../../subprocess-env"); const { prompt } = require("../../credentials/store"); const { promptManualModelId } = require("../model-prompts"); @@ -373,21 +376,40 @@ function probeOllamaAuthProxyHealth(): { ok: boolean; endpoint: string; detail: }; } -async function promptOllamaModel(gpu = null) { +async function promptOllamaModel(gpu: GpuInfo | null = null) { const installed = getOllamaModelOptions(); - const options = installed.length > 0 ? installed : getBootstrapOllamaModelOptions(gpu); + // Filter installed entries by registry-known memory fit so a host that + // currently cannot load the only installed model still gets a usable + // default — without the filter, pressing Enter would re-select the + // oversized model the runner is about to crash on. Unknown tags (user- + // pulled models the registry has never seen) pass the filter so the + // user's prior selection is respected. + const installedFitting = installed.filter((tag: string) => modelFitsAvailableMemory(tag, gpu)); + const usingInstalled = installedFitting.length > 0; + const options = usingInstalled ? installedFitting : getBootstrapOllamaModelOptions(gpu); const defaultModel = getDefaultOllamaModel(gpu); const defaultIndex = Math.max(0, options.indexOf(defaultModel)); console.log(""); - console.log(installed.length > 0 ? " Ollama models:" : " Ollama starter models:"); + console.log(usingInstalled ? " Ollama models:" : " Ollama starter models:"); options.forEach((option, index) => { console.log(` ${index + 1}) ${option}`); }); console.log(` ${options.length + 1}) Other...`); - if (installed.length === 0) { + if (!usingInstalled) { console.log(""); - console.log(" No local Ollama models are installed yet. Choose one to pull and load now."); + if (installed.length === 0) { + console.log(" No local Ollama models are installed yet. Choose one to pull and load now."); + } else { + console.log( + " No installed Ollama model fits the host's currently available memory; showing starter models instead.", + ); + } + } + if (!usingInstalled && !anyRegistryModelFits(gpu)) { + console.log( + " ! Even the smallest known bootstrap model may not fit currently available GPU memory; free memory or expect the runner to reject the load.", + ); } console.log(""); diff --git a/src/lib/onboard.ts b/src/lib/onboard.ts index 67b359821c6..dbebf63c021 100644 --- a/src/lib/onboard.ts +++ b/src/lib/onboard.ts @@ -4169,7 +4169,7 @@ async function selectAndValidateOllamaModel( const installedModels = getOllamaModelOptions(); let model: string | typeof BACK_TO_SELECTION; if (isNonInteractive()) { - model = requestedModel || recoveredModel || getDefaultOllamaModel(gpu); + model = localInference.resolveNonInteractiveOllamaModel(requestedModel, recoveredModel, gpu); } else { model = await promptOllamaModel(gpu); }