Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/lib/onboard/model-router-process.ts
Original file line number Diff line number Diff line change
Expand Up @@ -41,8 +41,8 @@ const ROUTER_HEALTH_BODY_MAX_BYTES = 64 * 1024;

/**
* Fetch /health and keep the response body for diagnosis (#8962). Unlike
* `isRouterHealthy`, this waits for the body, so callers pass a longer
* timeout; `startModelRouter` uses 30 seconds. LiteLLM's /health probes
* `isRouterHealthy`, this waits for the body, so a caller that must read it
* budgets for it; the final startup snapshot uses 30 seconds. LiteLLM's /health probes
* every upstream endpoint per request and can answer well after the
* 3-second liveness budget. The timeout is a wall-clock deadline, not a
* socket idle timeout, so a responder that trickles bytes cannot hold the
Expand Down
89 changes: 89 additions & 0 deletions src/lib/onboard/model-router-reconcile.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

import { beforeEach, describe, expect, it, vi } from "vitest";

import { reconcileModelRouter } from "./model-router";

const RECORDED_ROUTER_PID = 4321;

const holder = vi.hoisted(() => ({
snapshotBody: null as string | null,
stopped: [] as Array<[number, number]>,
reachabilityProbes: 0,
}));

// `stopModelRouterProcess` throws a sentinel so each case ends at the
// reuse decision. Restarting the router is `startModelRouter`'s contract and
// is covered by `test/onboard-model-router.test.ts`.
vi.mock("./model-router-process", () => ({
ROUTER_HEALTH_TIMEOUT_MS: 3_000,
getRouterHealthSnapshot: vi.fn(async () => ({ healthy: true, body: holder.snapshotBody })),
isRouterHealthy: vi.fn(async () => true),
doesModelRouterProcessOwnPort: vi.fn(() => true),
inspectModelRouterProcessForPort: vi.fn(() => ({ status: "missing" as const })),
stopModelRouterProcess: vi.fn(async (pid: number, port: number) => {
holder.stopped.push([pid, port]);
throw new Error("router restart reached");
}),
}));

vi.mock("../credentials/store", () => ({
normalizeCredentialValue: (value: string) => value,
resolveProviderCredential: () => "",
saveCredential: vi.fn(),
}));

vi.mock("./credential-env", () => ({
hydrateCredentialEnv: () => "nvapi-TEST-NOT-A-REAL-ROUTER-KEY",
}));

vi.mock("../state/onboard-session", () => ({
loadSession: () => ({
routerPid: RECORDED_ROUTER_PID,
routerCredentialHash: "MATCHING-HASH",
}),
updateSession: vi.fn(),
}));

vi.mock("../security/credential-hash", () => ({ hashCredential: () => "MATCHING-HASH" }));

vi.mock("./host-service-reachability", () => ({
probeHostServiceSandboxReachability: vi.fn(async () => {
holder.reachabilityProbes += 1;
return { ok: true };
}),
formatHostServiceUnreachableMessage: () => "",
}));

describe("model router reconciliation", () => {
beforeEach(() => {
holder.snapshotBody = null;
holder.stopped = [];
holder.reachabilityProbes = 0;
});

it("reuses a recorded router whose health snapshot names a healthy endpoint", async () => {
holder.snapshotBody = JSON.stringify({
healthy_endpoints: [{ api_base: "https://integrate.api.nvidia.com/v1" }],
unhealthy_endpoints: [],
});

await reconcileModelRouter();

expect(holder.stopped).toEqual([]);
expect(holder.reachabilityProbes).toBe(1);
});

it("restarts a recorded router that answers 2xx with no healthy endpoint (#9437)", async () => {
holder.snapshotBody = JSON.stringify({
healthy_endpoints: [],
unhealthy_endpoints: [{ api_base: "https://integrate.api.nvidia.com/v1" }],
});

await expect(reconcileModelRouter()).rejects.toThrow("router restart reached");

expect(holder.stopped).toEqual([[RECORDED_ROUTER_PID, expect.any(Number)]]);
expect(holder.reachabilityProbes).toBe(0);
});
});
37 changes: 24 additions & 13 deletions src/lib/onboard/model-router.ts
Original file line number Diff line number Diff line change
Expand Up @@ -61,9 +61,9 @@ const ROUTER_HEALTH_INTERVAL_MS = 2000;
const ROUTER_STARTUP_TIMEOUT_MS = 10 * 60_000;
// LiteLLM's /health live-probes every upstream endpoint per request, so it
// can need far longer than the 3-second liveness budget to answer (#8962).
// The startup poll keeps the 3-second budget: the status-only liveness
// probe must not accept a fast 200 that names zero healthy endpoints, so
// recovery for a slow-but-healthy router runs through the body-checked
// The startup poll keeps the 3-second budget and reads the body, so it
// never accepts a fast 200 that names zero healthy endpoints; recovery for
// a router whose /health outruns that budget runs through the body-checked
// final snapshot after the poll exhausts its retries.
const ROUTER_FINAL_HEALTH_SNAPSHOT_TIMEOUT_MS = 30_000;
const ROUTER_LOG_TAIL_LINES = 20;
Expand Down Expand Up @@ -468,7 +468,8 @@ export async function startModelRouter(
Math.min(ROUTER_HEALTH_REQUEST_TIMEOUT_MS, Math.ceil(remainingMs)),
);
healthAttempts += 1;
const healthy = await deps.isRouterHealthy(port, healthTimeoutMs);
const pollSnapshot = await deps.getRouterHealthSnapshot(port, healthTimeoutMs);
const healthy = isRouterSnapshotReady(pollSnapshot);
Comment thread
coderabbitai[bot] marked this conversation as resolved.
const processAlive = deps.isProcessAlive(pid);
if (healthy && processAlive) return pid;
if (!processAlive) {
Expand All @@ -484,7 +485,7 @@ export async function startModelRouter(
const finalSnapshot: RouterHealthSnapshot = childExited
? { healthy: false, body: null }
: await deps.getRouterHealthSnapshot(port, ROUTER_FINAL_HEALTH_SNAPSHOT_TIMEOUT_MS);
if (finalSnapshot.healthy && hasHealthyEndpoint(finalSnapshot.body) && deps.isProcessAlive(pid)) {
if (isRouterSnapshotReady(finalSnapshot) && deps.isProcessAlive(pid)) {
return pid;
}
try {
Expand All @@ -503,11 +504,11 @@ export async function startModelRouter(
);
}

/** True when the parsed /health body names at least one healthy endpoint. */
function hasHealthyEndpoint(body: string | null): boolean {
if (!body) return false;
/** Router readiness: /health answered 2xx and names at least one healthy endpoint. */
function isRouterSnapshotReady(snapshot: RouterHealthSnapshot): boolean {
if (!snapshot.healthy || !snapshot.body) return false;
try {
const parsed = JSON.parse(body) as { healthy_endpoints?: readonly unknown[] };
const parsed = JSON.parse(snapshot.body) as { healthy_endpoints?: readonly unknown[] };
return Array.isArray(parsed?.healthy_endpoints) && parsed.healthy_endpoints.length > 0;
} catch {
return false;
Expand Down Expand Up @@ -595,7 +596,7 @@ const MODEL_ROUTER_SERVICE_LABEL = "Model Router";
/**
* Verify the host Model Router is reachable from the OpenShell Docker network.
*
* `isRouterHealthy()` only proves the router answers on the host loopback. On
* A healthy /health answer only proves the router responds on the host loopback. On
* Linux Docker-driver hosts with UFW default-deny, a sandbox container can
* still fail to reach `host.openshell.internal:<routerPort>` even though the
* host curl succeeds (#4564). This mirrors the Ollama auth-proxy probe: on a
Expand Down Expand Up @@ -636,19 +637,29 @@ export async function reconcileModelRouter(): Promise<void> {
const recordedPid = session?.routerPid ?? null;
const recordedCredentialHash = session?.routerCredentialHash ?? null;

if (await isRouterHealthy(routerPort)) {
// One snapshot answers both questions: `healthy` is the occupied-port check,
// and `isRouterSnapshotReady` is the single authority for declaring the
// router usable, exactly as it is for the startup poll. Budget the body read,
// because /health probes every upstream endpoint and can answer well after
// the 3-second liveness budget.
const snapshot = await getRouterHealthSnapshot(
routerPort,
ROUTER_FINAL_HEALTH_SNAPSHOT_TIMEOUT_MS,
);
if (snapshot.healthy) {
const recordedProcessOwnsRouter = doesModelRouterProcessOwnPort(recordedPid, routerPort);
if (
routerCredentialHash &&
recordedCredentialHash === routerCredentialHash &&
recordedProcessOwnsRouter
recordedProcessOwnsRouter &&
isRouterSnapshotReady(snapshot)
) {
console.log(` ✓ Model router is already healthy on port ${routerPort}`);
await verifyModelRouterSandboxReachability(routerPort);
return;
}
if (recordedProcessOwnsRouter) {
console.log(" Restarting model router with updated credentials...");
console.log(" Restarting model router...");
await stopModelRouterProcess(
requireValue(recordedPid, "Expected recorded router PID"),
routerPort,
Expand Down
Loading
Loading