fix(onboard): handle Ollama unified-memory probe failures on DGX Spark (#3251) - #3389
Conversation
…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>
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds 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 ChangesOllama unified-memory host validation
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
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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.
Built for teams:
One agent for your entire SDLC. Right inside Slack. Comment |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
src/lib/inference/local.test.ts (1)
488-520: ⚡ Quick winAssert that the retry actually switches to 300s.
These cases only prove that
validateOllamaModelcallscapturetwice. They would still pass if the second probe accidentally reused the 120s timeout. Recording the argv and asserting the retry contains--max-timewith300would 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
📒 Files selected for processing (2)
src/lib/inference/local.test.tssrc/lib/inference/local.ts
…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>
eefbbb4 to
7ddf332
Compare
There was a problem hiding this comment.
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
📒 Files selected for processing (3)
src/lib/inference/local.test.tssrc/lib/inference/local.tssrc/lib/runner.ts
🚧 Files skipped from review as they are similar to previous changes (1)
- src/lib/inference/local.ts
| 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/); | ||
| }); |
There was a problem hiding this comment.
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.
| 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>
bd3c776 to
5844542
Compare
There was a problem hiding this comment.
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
📒 Files selected for processing (4)
src/lib/inference/local.test.tssrc/lib/inference/local.tssrc/lib/runner.tstest/ollama-tools-capability.test.ts
🚧 Files skipped from review as they are similar to previous changes (1)
- src/lib/inference/local.test.ts
…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>
f14b274 to
f8e1047
Compare
|
✨ 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: |
Summary
Fixes two failure modes that block
nemoclaw onboardon DGX Spark (128 GB unified memory) when selectingnemotron-3-nano:30b. Both are caused byvalidateOllamaModelnot accounting for Spark's unified-memory architecture.Related Issue
Closes #3251
Changes
src/lib/inference/local.ts): When Ollama returns"requires more system memory", intercept the error, runfree -mto 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.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.src/lib/inference/local.test.ts): 5 new unit tests with mockedrunCapture, covering mode 1, mode 2, and the composite case (mode 2 timeout on first probe → mode 1 OOM error on retry).Type of Change
Verification
npx prek run --all-filespassesnpm testpasses (Test Files 1 passed, Tests 45 passed)Note on
npx prek run --all-files: 4 pre-existing test failures inblueprint/stateandonboard/configare present onmainand 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
Tests