Skip to content

fix(onboard): handle Ollama unified-memory probe failures on DGX Spark (#3251) - #3389

Merged
cv merged 6 commits into
NVIDIA:mainfrom
tiaz-hh:fix/spark-ollama-memory-probe-unified
May 15, 2026
Merged

fix(onboard): handle Ollama unified-memory probe failures on DGX Spark (#3251)#3389
cv merged 6 commits into
NVIDIA:mainfrom
tiaz-hh:fix/spark-ollama-memory-probe-unified

Conversation

@tiaz-hh

@tiaz-hh tiaz-hh commented May 12, 2026

Copy link
Copy Markdown
Contributor

Summary

Fixes two failure modes that block nemoclaw onboard on DGX Spark (128 GB unified memory) when selecting nemotron-3-nano:30b. Both are caused by validateOllamaModel not accounting for Spark's unified-memory architecture.

Related Issue

Closes #3251

Changes

  • Mode 1 fix (src/lib/inference/local.ts): When Ollama returns "requires more system memory", intercept the error, run free -m to check total RAM. If total covers the requirement, return { ok: true } — Ollama's available-RAM check is a false positive on unified-memory hardware where GPU and CPU share the same 128 GB pool.
  • Mode 2 fix (src/lib/inference/local.ts): When the first probe returns empty (120 s timeout exceeded), retry once with a 300 s timeout. Covers the case where loading a 22 GB model from disk into unified memory takes >2 min. Normal hosts that respond quickly are unaffected; truly unhealthy models fail after both attempts.
  • Tests (src/lib/inference/local.test.ts): 5 new unit tests with mocked runCapture, covering mode 1, mode 2, and the composite case (mode 2 timeout on first probe → mode 1 OOM error on retry).

Type of Change

  • Code change (feature, bug fix, or refactor)

Verification

  • npx prek run --all-files passes
  • npm test passes (Test Files 1 passed, Tests 45 passed)
  • Tests added or updated for new or changed behavior
  • No secrets, API keys, or credentials committed
  • Docs updated for user-facing behavior changes

Note on npx prek run --all-files: 4 pre-existing test failures in blueprint/state and onboard/config are present on main and unrelated to this change.

Note on end-to-end reproduction: Requires a DGX Spark with GNOME desktop running per NVBugs#6157916. Our QA Spark lacks a desktop session so available RAM stays above the trigger threshold; unit tests cover both failure paths via mocked runCapture.


Signed-off-by: Tian Zhang tiazhang@nvidia.com

Summary by CodeRabbit

  • Bug Fixes

    • Smarter model validation on unified‑memory (Spark/DGX) hosts: conditional longer probe retry for slow responses and treating certain probe OOM messages as non‑fatal when host total RAM meets the model requirement.
  • Tests

    • Expanded tests for memory-detection edge cases, probe timeout/retry logic, and mixed probe/OOM outcomes.

Review Change Stack

…mory hosts

On DGX Spark (128 GB unified memory), Ollama checks available RAM instead
of total RAM when loading a model. With GNOME + browser open, available RAM
drops to ~5 GB even though 128 GB total is present, causing a false OOM
rejection after the user has already downloaded the model.

When validateOllamaModel receives a 'requires more system memory' error,
fall back to checking total system RAM via `free -m`. If total RAM covers
the model's requirement, treat the probe as passing — Ollama's
available-RAM check is a false positive on unified-memory hardware.

Fixes: NVIDIA#3251
Signed-off-by: Tian Zhang <tiazhang@nvidia.com>
@coderabbitai

coderabbitai Bot commented May 12, 2026

Copy link
Copy Markdown
Contributor

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review
📝 Walkthrough

Walkthrough

Adds optional Spark detection to validateOllamaModel, retries a no-output local probe with a ~300s timeout only on Spark/unified-memory hosts, and treats specific Ollama OOM errors as non-fatal when total system RAM (from free -m) meets the required GiB.

Changes

Ollama unified-memory host validation

Layer / File(s) Summary
API surface & imports
src/lib/inference/local.ts
Import detectNvidiaPlatform, add RunCaptureExFn type, and extend validateOllamaModel signature to accept optional isSparkImpl and runCaptureExImpl overrides.
runCaptureEx implementation and export
src/lib/runner.ts
Add CaptureResult and runCaptureEx that returns { stdout, exitCode, timedOut } and export it for richer probe handling.
Probe retry logic for slow probes
src/lib/inference/local.ts, src/lib/inference/local.test.ts
When the first /api/generate probe returns empty, retry once with --max-time ≈300s only if host is Spark/unified-memory and the first probe timed out (or exit code 28); tests cover retry/no-retry, fast failures, and double-empty failure.
OOM parsing and total RAM check
src/lib/inference/local.ts, src/lib/inference/local.test.ts
If Ollama reports “requires more system memory … than is available”, parse required GiB from the error, run free -m, convert Mem: total to GiB and allow validation to pass when total GiB ≥ required GiB (Spark-gated); tests cover sufficient/insufficient totals and composite retry+OOM interactions.
Tests: captureEx updates
src/lib/inference/local.test.ts, test/ollama-tools-capability.test.ts
Existing tests updated to supply explicit captureEx results (stdout, exitCode, timedOut) and new cases added for Spark vs non‑Spark, retry behavior, and OOM/free permutations.

Sequence Diagram(s)

sequenceDiagram
  participant Validator as validateOllamaModel
  participant Detector as detectNvidiaPlatform
  participant Probe as /api/generate
  participant Shell as runCapture (shell)

  Validator->>Detector: determine isSpark (default detectNvidiaPlatform()==="spark")
  Validator->>Probe: probe (short timeout)
  Probe-->>Validator: (stdout | empty | OOM error | fast failure)

  alt empty && isSpark && timedOut
    Validator->>Probe: retry probe (--max-time ~300s)
    Probe-->>Validator: (stdout | empty | OOM error)
  end

  alt OOM error observed
    Validator->>Shell: run `free -m`
    Shell-->>Validator: `Mem:` total MB
    Note right of Validator: convert MB → GiB and compare to required GiB
    alt total GiB >= required
      Validator-->>Validator: return success
    else
      Validator-->>Validator: return validation failure
    end
  end
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Poem

A rabbit waits while models load slow,
It counts the RAM in rows aglow.
If free shows enough for Ollama's plea,
It hops ahead — validation: free! 🐰✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title 'fix(onboard): handle Ollama unified-memory probe failures on DGX Spark' directly and specifically describes the main changes: fixing Ollama probe failures on DGX Spark unified-memory hardware, which is the core objective of the PR.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Tip

💬 Introducing Slack Agent: The best way for teams to turn conversations into code.

Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.

  • Generate code and open pull requests
  • Plan features and break down work
  • Investigate incidents and troubleshoot customer tickets together
  • Automate recurring tasks and respond to alerts with triggers
  • Summarize progress and report instantly

Built for teams:

  • Shared memory across your entire org—no repeating context
  • Per-thread sandboxes to safely plan and execute work
  • Governance built-in—scoped access, auditability, and budget controls

One agent for your entire SDLC. Right inside Slack.

👉 Get started


Comment @coderabbitai help to get the list of available commands and usage tips.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

🧹 Nitpick comments (1)
src/lib/inference/local.test.ts (1)

488-520: ⚡ Quick win

Assert that the retry actually switches to 300s.

These cases only prove that validateOllamaModel calls capture twice. They would still pass if the second probe accidentally reused the 120s timeout. Recording the argv and asserting the retry contains --max-time with 300 would lock down the regression this PR is fixing.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/lib/inference/local.test.ts` around lines 488 - 520, Update the tests for
validateOllamaModel so they assert the retry probe uses a 300s timeout instead
of just counting calls: record the argv/command passed into the capture stub in
the two tests ("retries with extended timeout when first probe returns empty"
and "passes when first probe times out then retry returns OOM error but total
RAM is sufficient") and add an assertion that the second invocation's command
array or string includes the flag "--max-time" with the value "300" (or contains
"--max-time 300"), while keeping the existing assertions that callCount is 2 and
result.ok expectations; use the existing capture function signature to inspect
the passed command and validate the presence of the 300s timeout on the retry.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/lib/inference/local.ts`:
- Around line 562-575: The current total-RAM bypass around memMatch incorrectly
accepts models on any machine with sufficient total RAM; restrict this override
to only unified-memory hosts by adding a check (e.g., isUnifiedMemoryHost() or a
hasUnifiedMemory/Spark signal) before returning { ok: true }. Specifically, in
the block using memMatch, freeOut, memLine, totalMB/totalGiB and requiresGiB,
replace the unconditional return with a conditional that only returns { ok: true
} when the unified-memory detection helper (create or call a function like
isUnifiedMemoryHost or check the Spark/unified-memory flag) is true; otherwise
preserve Ollama's original available-memory error path. Ensure the new helper is
clearly named and used where memMatch is handled so the bypass cannot apply
globally.
- Around line 535-539: The current logic treats any falsy capture(probeCmd, {
ignoreError: true }) as a reason to retry with a 300s probe, which converts fast
failures into long stalls; change this to retry only on timeout-specific signals
by capturing the full result (stdout/stderr/exit code) from probeCmd instead of
collapsing errors to an empty string, then only call
getOllamaProbeCommand(model, 300) if the probe indicates a real timeout (e.g.,
curl exit code 28 or stderr contains "timed out"/"Operation timed out"); update
the code around capture(probeCmd, ...) and the retry branch to inspect those
timeout indicators and otherwise surface the original error immediately.

---

Nitpick comments:
In `@src/lib/inference/local.test.ts`:
- Around line 488-520: Update the tests for validateOllamaModel so they assert
the retry probe uses a 300s timeout instead of just counting calls: record the
argv/command passed into the capture stub in the two tests ("retries with
extended timeout when first probe returns empty" and "passes when first probe
times out then retry returns OOM error but total RAM is sufficient") and add an
assertion that the second invocation's command array or string includes the flag
"--max-time" with the value "300" (or contains "--max-time 300"), while keeping
the existing assertions that callCount is 2 and result.ok expectations; use the
existing capture function signature to inspect the passed command and validate
the presence of the 300s timeout on the retry.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: fd9dbc4a-9023-4fc8-96d6-eb46fab5651a

📥 Commits

Reviewing files that changed from the base of the PR and between 4b47ab4 and eefbbb4.

📒 Files selected for processing (2)
  • src/lib/inference/local.test.ts
  • src/lib/inference/local.ts

Comment thread src/lib/inference/local.ts Outdated
Comment thread src/lib/inference/local.ts Outdated
…g unified-memory hosts

On DGX Spark (128 GB unified memory), loading a 22 GB model from disk
into unified memory can take well over 2 minutes. The existing 120 s
probe timeout expires before the model is ready, causing:

  Selected Ollama model '...' did not answer the local probe in time.

Fix: when the first probe returns empty (timeout), retry once with a
300 s timeout. Normal hosts that respond quickly are unaffected.
Machines where the model is genuinely unhealthy or missing will still
fail after both attempts.

This is failure mode 2 of issue NVIDIA#3251 (failure mode 1 — Ollama
available-RAM OOM rejection — was fixed in the previous commit).

Signed-off-by: Tian Zhang <tiazhang@nvidia.com>
@tiaz-hh
tiaz-hh force-pushed the fix/spark-ollama-memory-probe-unified branch from eefbbb4 to 7ddf332 Compare May 12, 2026 07:24

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/lib/inference/local.test.ts`:
- Around line 535-543: The test for validateOllamaModel ("does not retry on
Spark when probe fails fast...") is asserting a timeout-style message even
though captureEx returns exitCode: 7 and timedOut: false (connection refused);
update the assertion on result.message to expect wording that reflects a
connection-refused/fast-failure condition (e.g., contains "connection refused",
"exit code 7", or similar probe failure text) instead of "did not answer the
local probe in time", keep the callCount and result.ok assertions as-is, and
ensure the captureEx/probe simulation remains the same so the test validates the
non-timeout failure path for validateOllamaModel.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c622171a-1a0b-4817-b2cf-0f07d1f23806

📥 Commits

Reviewing files that changed from the base of the PR and between f9fcfc6 and bd3c776.

📒 Files selected for processing (3)
  • src/lib/inference/local.test.ts
  • src/lib/inference/local.ts
  • src/lib/runner.ts
🚧 Files skipped from review as they are similar to previous changes (1)
  • src/lib/inference/local.ts

Comment on lines +535 to +543
it("does not retry on Spark when probe fails fast (connection refused, not a timeout)", () => {
// exit code 7 = curl connection refused — should surface immediately, not stall 300s.
let callCount = 0;
const captureEx = () => { callCount++; return { stdout: "", exitCode: 7, timedOut: false }; };
const result = validateOllamaModel("nemotron-3-nano:30b", () => "", () => true, captureEx);
expect(result.ok).toBe(false);
expect(callCount).toBe(1);
expect(result.message).toMatch(/did not answer the local probe in time/);
});

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Non-timeout path is asserting a timeout message.

The test says this is a fast connection-refused case (timedOut: false, exit 7), but it still expects timeout wording. That can lock in misleading diagnostics for users and hide message regressions on this path.

Suggested test assertion adjustment
-    expect(result.message).toMatch(/did not answer the local probe in time/);
+    expect(result.message).toMatch(/connection refused|curl failed|exit 7/i);
+    expect(result.message).not.toMatch(/in time/);
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
it("does not retry on Spark when probe fails fast (connection refused, not a timeout)", () => {
// exit code 7 = curl connection refused — should surface immediately, not stall 300s.
let callCount = 0;
const captureEx = () => { callCount++; return { stdout: "", exitCode: 7, timedOut: false }; };
const result = validateOllamaModel("nemotron-3-nano:30b", () => "", () => true, captureEx);
expect(result.ok).toBe(false);
expect(callCount).toBe(1);
expect(result.message).toMatch(/did not answer the local probe in time/);
});
it("does not retry on Spark when probe fails fast (connection refused, not a timeout)", () => {
// exit code 7 = curl connection refused — should surface immediately, not stall 300s.
let callCount = 0;
const captureEx = () => { callCount++; return { stdout: "", exitCode: 7, timedOut: false }; };
const result = validateOllamaModel("nemotron-3-nano:30b", () => "", () => true, captureEx);
expect(result.ok).toBe(false);
expect(callCount).toBe(1);
expect(result.message).toMatch(/connection refused|curl failed|exit 7/i);
expect(result.message).not.toMatch(/in time/);
});
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/lib/inference/local.test.ts` around lines 535 - 543, The test for
validateOllamaModel ("does not retry on Spark when probe fails fast...") is
asserting a timeout-style message even though captureEx returns exitCode: 7 and
timedOut: false (connection refused); update the assertion on result.message to
expect wording that reflects a connection-refused/fast-failure condition (e.g.,
contains "connection refused", "exit code 7", or similar probe failure text)
instead of "did not answer the local probe in time", keep the callCount and
result.ok assertions as-is, and ensure the captureEx/probe simulation remains
the same so the test validates the non-timeout failure path for
validateOllamaModel.

Distinguish curl timeout (exit 28 / ETIMEDOUT) from fast failures
(connection refused, Ollama not running) in the DGX Spark extended-probe
path. Previously any empty probe result triggered a 300 s retry, stalling
the wizard when Ollama was simply unreachable.

Introduces runCaptureEx in runner.ts to return structured
{stdout, exitCode, timedOut} without collapsing errors to ''.

Signed-off-by: Tian Zhang <tiazhang@nvidia.com>
@tiaz-hh
tiaz-hh force-pushed the fix/spark-ollama-memory-probe-unified branch from bd3c776 to 5844542 Compare May 12, 2026 09:49

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/lib/inference/local.ts`:
- Around line 548-550: The retry logic incorrectly requires !output, so when
curl times out but writes partial stdout the Spark retry is skipped; change the
condition to trigger on probeResult.timedOut alone (i.e., if (isSpark() &&
probeResult.timedOut) ) and replace the retry call to use
captureEx(getOllamaProbeCommand(model, 300)) (or the codebase's captureEx
wrapper) with ignoreError behavior so the retry runs even when partial output
exists; update references in the local probe block where getOllamaProbeCommand,
isSpark, and capture are used to use captureEx and timedOut-only gating.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: ca2df312-331e-445b-ada0-3f531c28d599

📥 Commits

Reviewing files that changed from the base of the PR and between bd3c776 and f14b274.

📒 Files selected for processing (4)
  • src/lib/inference/local.test.ts
  • src/lib/inference/local.ts
  • src/lib/runner.ts
  • test/ollama-tools-capability.test.ts
🚧 Files skipped from review as they are similar to previous changes (1)
  • src/lib/inference/local.test.ts

Comment thread src/lib/inference/local.ts Outdated
…pability tests

Fixes CI failure where validateOllamaModel was called without a captureExImpl,
causing the real runCaptureEx to run on macOS CI (no Ollama available) and return
an empty stdout, bypassing the tools-error-message path entirely.

Signed-off-by: Tian Zhang <tiazhang@nvidia.com>
@tiaz-hh
tiaz-hh force-pushed the fix/spark-ollama-memory-probe-unified branch from f14b274 to f8e1047 Compare May 12, 2026 12:05
@wscurran wscurran added Platform: DGX Spark provider: ollama Ollama local model provider behavior labels May 12, 2026
@wscurran

Copy link
Copy Markdown
Contributor

✨ Thanks for submitting this detailed PR to handle Ollama unified-memory probe failures on DGX Spark. This change aims to improve the onboard process by addressing false positives and timeouts when loading models into unified memory.


Related open issues:

@tiaz-hh
tiaz-hh requested a review from ericksoa May 13, 2026 02:05
@cv
cv enabled auto-merge (squash) May 15, 2026 15:54
@cv
cv merged commit 40a99e8 into NVIDIA:main May 15, 2026
18 checks passed
@miyoungc miyoungc mentioned this pull request May 16, 2026
12 tasks
@wscurran wscurran added area: local-models Local model providers, downloads, launch, or connectivity area: providers Inference provider integrations and provider behavior bug-fix PR fixes a bug or regression labels Jun 3, 2026
@wscurran wscurran added platform: dgx-spark Affects DGX Spark hardware or workflows and removed priority: high labels Jun 3, 2026
@wscurran wscurran added the NV QA Bugs found by the NVIDIA QA Team label Jun 26, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area: local-models Local model providers, downloads, launch, or connectivity area: providers Inference provider integrations and provider behavior bug-fix PR fixes a bug or regression NV QA Bugs found by the NVIDIA QA Team platform: dgx-spark Affects DGX Spark hardware or workflows provider: ollama Ollama local model provider behavior

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[DGX Spark][Onboard] Memory probe rejects Ollama model on Spark unified memory — checks available RAM instead of total

3 participants