Skip to content

Commit 853a314

Browse files
HOYALIMHOYALIM
committed
perf(onboard): calibrate provider validation timeouts
Signed-off-by: Ho Lim <subhoya@gmail.com>
1 parent 05ff005 commit 853a314

6 files changed

Lines changed: 329 additions & 20 deletions

‎src/lib/inference/onboard-probes.ts‎

Lines changed: 85 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ const {
2828
const { isWsl } = require("../platform");
2929
const httpProbe = require("../adapters/http/probe");
3030
const authConfigModule = require("../adapters/http/auth-config");
31+
const trace = require("../trace");
3132
const {
3233
getHostDockerInternalProbeFailure,
3334
isHijackedDockerInternalUrl,
@@ -44,6 +45,7 @@ const {
4445
} = require("./probe-retry");
4546
const { probeAnthropicEndpoint } = require("./probe-anthropic");
4647
const {
48+
buildValidationProbeTimingProfile,
4749
getValidationProbeCurlArgs,
4850
getDeepSeekV4ProValidationProbeCurlArgs,
4951
getKimiK26ValidationProbeCurlArgs,
@@ -232,6 +234,55 @@ function getProbeAuthMode(_provider) {
232234
return undefined;
233235
}
234236

237+
function getProbeTimingOptions(options = {}) {
238+
const timingOptions = {};
239+
if (typeof options.isWsl === "boolean") {
240+
timingOptions.isWsl = options.isWsl;
241+
}
242+
if (options.validationTiming) {
243+
timingOptions.validationTiming = options.validationTiming;
244+
}
245+
return Object.keys(timingOptions).length > 0 ? timingOptions : undefined;
246+
}
247+
248+
function calibrateOpenAiLikeValidationTiming(baseUrl, options = {}) {
249+
return trace.withTraceSpan("nemoclaw.inference.validation_timeout_calibration", {}, () => {
250+
const url = `${baseUrl}/models`;
251+
const args = [
252+
"-sS",
253+
...buildResolvePinArgs(url, options.pinnedAddresses),
254+
"--connect-timeout",
255+
"3",
256+
"--max-time",
257+
"5",
258+
url,
259+
];
260+
const startedAtMs = Date.now();
261+
const result = runCurlProbe(args, {
262+
timeoutMs: getProbeProcessTimeoutMs(args),
263+
pinnedAddresses: options.pinnedAddresses,
264+
});
265+
const durationMs = Date.now() - startedAtMs;
266+
const calibration =
267+
result.curlStatus === 0 && result.httpStatus > 0
268+
? { ok: true, durationMs }
269+
: { ok: false, reason: result.message };
270+
const profile = buildValidationProbeTimingProfile({
271+
...(typeof options.isWsl === "boolean" ? { isWsl: options.isWsl } : {}),
272+
calibration,
273+
});
274+
trace.addTraceEvent("validation_timeout_profile", {
275+
calibration_curl_status: result.curlStatus,
276+
calibration_http_status: result.httpStatus,
277+
connect_timeout_seconds: profile.connectTimeoutSeconds,
278+
max_time_seconds: profile.maxTimeSeconds,
279+
observed_ms: profile.observedMs ?? null,
280+
source: profile.source,
281+
});
282+
return profile;
283+
});
284+
}
285+
235286
// ── Responses API probe ──────────────────────────────────────────
236287

237288
function probeResponsesToolCalling(endpointUrl, model, apiKey, options = {}) {
@@ -243,7 +294,7 @@ function probeResponsesToolCalling(endpointUrl, model, apiKey, options = {}) {
243294
[
244295
"-sS",
245296
...buildResolvePinArgs(`${baseUrl}/responses`, options.pinnedAddresses),
246-
...getValidationProbeCurlArgs(),
297+
...getValidationProbeCurlArgs(getProbeTimingOptions(options)),
247298
"-H",
248299
"Content-Type: application/json",
249300
...authConfig.args,
@@ -302,7 +353,9 @@ function probeChatCompletionsToolCalling(endpointUrl, model, apiKey, options = {
302353
let authConfig;
303354
try {
304355
authConfig = buildOpenAiLikeAuthConfig(apiKey, options);
305-
const timingArgs = options.timingArgs ?? getChatCompletionsProbeTimingArgs(model);
356+
const timingArgs =
357+
options.timingArgs ??
358+
getChatCompletionsProbeTimingArgs(model, getProbeTimingOptions(options));
306359
const args = [
307360
"-sS",
308361
...buildResolvePinArgs(`${baseUrl}/chat/completions`, options.pinnedAddresses),
@@ -483,9 +536,21 @@ export function getChatCompletionsProbeCurlArgs(opts: {
483536
url: string;
484537
isWsl?: boolean;
485538
pinnedAddresses?: readonly string[];
539+
validationTiming?: unknown;
486540
}) {
487-
const { credentialArgs, authHeader, model, url, isWsl: isWslOverride, pinnedAddresses } = opts;
488-
const platformOptions = typeof isWslOverride === "boolean" ? { isWsl: isWslOverride } : undefined;
541+
const {
542+
credentialArgs,
543+
authHeader,
544+
model,
545+
url,
546+
isWsl: isWslOverride,
547+
pinnedAddresses,
548+
validationTiming,
549+
} = opts;
550+
const platformOptions = getProbeTimingOptions({
551+
...(typeof isWslOverride === "boolean" ? { isWsl: isWslOverride } : {}),
552+
...(validationTiming ? { validationTiming } : {}),
553+
});
489554
const timingArgs = getChatCompletionsProbeTimingArgs(model, platformOptions);
490555
const credSlice = credentialArgs ?? authHeader ?? [];
491556
return [
@@ -508,13 +573,15 @@ function runChatCompletionsProbe({
508573
isWsl: isWslOverride,
509574
trustedConfigFiles,
510575
pinnedAddresses,
576+
validationTiming,
511577
}) {
512578
const args = getChatCompletionsProbeCurlArgs({
513579
credentialArgs,
514580
model,
515581
url,
516582
isWsl: isWslOverride,
517583
pinnedAddresses,
584+
validationTiming,
518585
});
519586
const probeOpts = { timeoutMs: getProbeProcessTimeoutMs(args), pinnedAddresses };
520587
if (trustedConfigFiles && trustedConfigFiles.length > 0) {
@@ -538,7 +605,7 @@ function runDoubledTimeoutChatCompletionsRetry({
538605
baseUrl,
539606
authConfig,
540607
}) {
541-
const platformOptions = typeof options.isWsl === "boolean" ? { isWsl: options.isWsl } : undefined;
608+
const platformOptions = getProbeTimingOptions(options);
542609
const baseArgs = getChatCompletionsProbeTimingArgs(model, platformOptions);
543610
const doubledArgs = baseArgs.map((arg) => (/^\d+$/.test(arg) ? String(Number(arg) * 2) : arg));
544611
const buildRetryArgs = () => [
@@ -651,6 +718,14 @@ function probeOpenAiLikeEndpoint(endpointUrl, model, apiKey, options = {}) {
651718
}
652719

653720
const baseUrl = String(endpointUrl).replace(/\/+$/, "");
721+
const validationTiming =
722+
options.validationTiming ??
723+
(options.calibrateTimeouts === true
724+
? calibrateOpenAiLikeValidationTiming(baseUrl, options)
725+
: undefined);
726+
if (validationTiming) {
727+
options = { ...options, validationTiming };
728+
}
654729
// Pin every probe curl to the SSRF-preflight-validated address(es) the caller
655730
// captured, so a second DNS lookup here cannot rebind the hostname to a
656731
// private/internal address after the public preflight (TOCTOU — cv, #6293).
@@ -667,6 +742,7 @@ function probeOpenAiLikeEndpoint(endpointUrl, model, apiKey, options = {}) {
667742
probeResponsesToolCalling(endpointUrl, model, apiKey, {
668743
authMode: options.authMode,
669744
pinnedAddresses,
745+
validationTiming,
670746
}),
671747
}
672748
: {
@@ -677,7 +753,7 @@ function probeOpenAiLikeEndpoint(endpointUrl, model, apiKey, options = {}) {
677753
[
678754
"-sS",
679755
...buildResolvePinArgs(`${baseUrl}/responses`, pinnedAddresses),
680-
...getValidationProbeCurlArgs(),
756+
...getValidationProbeCurlArgs(getProbeTimingOptions(options)),
681757
"-H",
682758
"Content-Type: application/json",
683759
...authConfig.args,
@@ -700,6 +776,7 @@ function probeOpenAiLikeEndpoint(endpointUrl, model, apiKey, options = {}) {
700776
? probeChatCompletionsToolCalling(endpointUrl, model, apiKey, {
701777
authMode: options.authMode,
702778
pinnedAddresses,
779+
validationTiming,
703780
})
704781
: runChatCompletionsProbe({
705782
credentialArgs: authConfig.args,
@@ -708,6 +785,7 @@ function probeOpenAiLikeEndpoint(endpointUrl, model, apiKey, options = {}) {
708785
isWsl: options.isWsl,
709786
trustedConfigFiles: authConfig.trustedConfigFiles,
710787
pinnedAddresses,
788+
validationTiming,
711789
}),
712790
};
713791

@@ -741,7 +819,7 @@ function probeOpenAiLikeEndpoint(endpointUrl, model, apiKey, options = {}) {
741819
[
742820
"-sS",
743821
...buildResolvePinArgs(`${baseUrl}/responses`, pinnedAddresses),
744-
...getValidationProbeCurlArgs(),
822+
...getValidationProbeCurlArgs(getProbeTimingOptions(options)),
745823
"-H",
746824
"Content-Type: application/json",
747825
...authConfig.args,

‎src/lib/inference/probe-http-helpers.test.ts‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import { afterEach, describe, expect, it, vi } from "vitest";
55

66
const {
7+
buildValidationProbeTimingProfile,
78
getKimiK26ValidationProbeCurlArgs,
89
getValidationProbeCurlArgs,
910
} = require("./probe-http-helpers");
@@ -13,6 +14,49 @@ afterEach(() => {
1314
});
1415

1516
describe("validation probe curl timing helpers", () => {
17+
it("derives a tighter fast-network profile from calibration latency", () => {
18+
expect(
19+
buildValidationProbeTimingProfile({ calibration: { ok: true, durationMs: 180 } }),
20+
).toEqual({
21+
connectTimeoutSeconds: 5,
22+
maxTimeSeconds: 15,
23+
observedMs: 180,
24+
source: "calibrated",
25+
});
26+
expect(getValidationProbeCurlArgs({ calibration: { ok: true, durationMs: 180 } })).toEqual([
27+
"--connect-timeout",
28+
"5",
29+
"--max-time",
30+
"15",
31+
]);
32+
});
33+
34+
it("derives a slower non-WSL profile from calibration latency", () => {
35+
expect(
36+
buildValidationProbeTimingProfile({ calibration: { ok: true, durationMs: 6_400 } }),
37+
).toEqual({
38+
connectTimeoutSeconds: 28,
39+
maxTimeSeconds: 42,
40+
observedMs: 6400,
41+
source: "calibrated",
42+
});
43+
});
44+
45+
it("falls back to the safe widened budget when calibration fails", () => {
46+
expect(
47+
getValidationProbeCurlArgs({ calibration: { ok: false, reason: "curl timed out" } }),
48+
).toEqual(["--connect-timeout", "20", "--max-time", "30"]);
49+
});
50+
51+
it("keeps the existing WSL fallback when no calibration result is available", () => {
52+
expect(getValidationProbeCurlArgs({ isWsl: true })).toEqual([
53+
"--connect-timeout",
54+
"20",
55+
"--max-time",
56+
"30",
57+
]);
58+
});
59+
1660
it("allows onboard validation max-time to be raised from the environment", () => {
1761
vi.stubEnv("NEMOCLAW_ONBOARD_VALIDATION_TIMEOUT_SECONDS", "300");
1862
expect(getValidationProbeCurlArgs({ isWsl: false })).toEqual([

‎src/lib/inference/probe-http-helpers.ts‎

Lines changed: 95 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -10,35 +10,121 @@
1010

1111
const { isWsl } = require("../platform");
1212

13-
type WslProbeOptions = { isWsl?: boolean } | undefined;
13+
type ValidationProbeCalibration = { ok: true; durationMs: number } | { ok: false; reason?: string };
14+
15+
export type ValidationProbeTimingProfile = {
16+
connectTimeoutSeconds: number;
17+
maxTimeSeconds: number;
18+
observedMs?: number;
19+
reason?: string;
20+
source: "standard" | "wsl-fallback" | "calibrated" | "calibration-fallback";
21+
};
22+
23+
type ValidationProbeOptions =
24+
| {
25+
isWsl?: boolean;
26+
validationTiming?: ValidationProbeTimingProfile;
27+
calibration?: ValidationProbeCalibration;
28+
}
29+
| undefined;
1430

1531
const ONBOARD_VALIDATION_TIMEOUT_ENV = "NEMOCLAW_ONBOARD_VALIDATION_TIMEOUT_SECONDS";
32+
const STANDARD_VALIDATION_TIMING: ValidationProbeTimingProfile = {
33+
connectTimeoutSeconds: 10,
34+
maxTimeSeconds: 15,
35+
source: "standard",
36+
};
37+
const WSL_VALIDATION_TIMING: ValidationProbeTimingProfile = {
38+
connectTimeoutSeconds: 20,
39+
maxTimeSeconds: 30,
40+
source: "wsl-fallback",
41+
};
42+
const CALIBRATION_FALLBACK_VALIDATION_TIMING: ValidationProbeTimingProfile = {
43+
connectTimeoutSeconds: 20,
44+
maxTimeSeconds: 30,
45+
source: "calibration-fallback",
46+
};
47+
const CALIBRATED_CONNECT_MIN_SECONDS = 5;
48+
const CALIBRATED_CONNECT_MAX_SECONDS = 30;
49+
const CALIBRATED_MAX_TIME_MIN_SECONDS = STANDARD_VALIDATION_TIMING.maxTimeSeconds;
50+
const CALIBRATED_MAX_TIME_MAX_SECONDS = 60;
51+
52+
function clampSeconds(value: number, min: number, max: number): number {
53+
return Math.min(max, Math.max(min, Math.ceil(value)));
54+
}
55+
56+
function copyTimingProfile(profile: ValidationProbeTimingProfile): ValidationProbeTimingProfile {
57+
return { ...profile };
58+
}
59+
60+
export function buildValidationProbeTimingProfile(
61+
opts?: ValidationProbeOptions,
62+
): ValidationProbeTimingProfile {
63+
if (opts?.validationTiming) return copyTimingProfile(opts.validationTiming);
64+
if (!opts?.calibration) {
65+
return copyTimingProfile(isWsl(opts) ? WSL_VALIDATION_TIMING : STANDARD_VALIDATION_TIMING);
66+
}
67+
if (!opts.calibration.ok) {
68+
return {
69+
...CALIBRATION_FALLBACK_VALIDATION_TIMING,
70+
reason: opts.calibration.reason,
71+
};
72+
}
73+
74+
const observedSeconds = Math.max(1, Math.ceil(opts.calibration.durationMs / 1000));
75+
const connectTimeoutSeconds = clampSeconds(
76+
observedSeconds * 4,
77+
CALIBRATED_CONNECT_MIN_SECONDS,
78+
CALIBRATED_CONNECT_MAX_SECONDS,
79+
);
80+
const maxTimeSeconds = clampSeconds(
81+
Math.max(connectTimeoutSeconds + 5, observedSeconds * 6),
82+
CALIBRATED_MAX_TIME_MIN_SECONDS,
83+
CALIBRATED_MAX_TIME_MAX_SECONDS,
84+
);
85+
return {
86+
connectTimeoutSeconds,
87+
maxTimeSeconds,
88+
observedMs: Math.max(0, Math.round(opts.calibration.durationMs)),
89+
source: "calibrated",
90+
};
91+
}
92+
93+
function buildCurlTimingArgs(profile: ValidationProbeTimingProfile): string[] {
94+
return [
95+
"--connect-timeout",
96+
String(profile.connectTimeoutSeconds),
97+
"--max-time",
98+
String(profile.maxTimeSeconds),
99+
];
100+
}
16101

17102
// Per-validation-probe curl timing. Tighter than the default 60s in
18103
// getCurlTimingArgs() because validation must not hang the wizard for a
19104
// minute on a misbehaving model. See issue #1601 (Bug 3).
20-
export function getValidationProbeCurlArgs(opts?: WslProbeOptions): string[] {
21-
const args = isWsl(opts)
22-
? ["--connect-timeout", "20", "--max-time", "30"]
23-
: ["--connect-timeout", "10", "--max-time", "15"];
24-
return withValidationMaxTimeOverride(args);
105+
export function getValidationProbeCurlArgs(opts?: ValidationProbeOptions): string[] {
106+
return withValidationMaxTimeOverride(
107+
buildCurlTimingArgs(buildValidationProbeTimingProfile(opts)),
108+
);
25109
}
26110

27-
export function getDeepSeekV4ProValidationProbeCurlArgs(opts?: WslProbeOptions): string[] {
111+
export function getDeepSeekV4ProValidationProbeCurlArgs(opts?: ValidationProbeOptions): string[] {
28112
const args = isWsl(opts)
29113
? ["--connect-timeout", "30", "--max-time", "150"]
30114
: ["--connect-timeout", "20", "--max-time", "120"];
31115
return withValidationMaxTimeOverride(args);
32116
}
33117

34-
export function getKimiK26ValidationProbeCurlArgs(opts?: WslProbeOptions): string[] {
118+
export function getKimiK26ValidationProbeCurlArgs(opts?: ValidationProbeOptions): string[] {
35119
const args = isWsl(opts)
36120
? ["--connect-timeout", "20", "--max-time", "90"]
37121
: ["--connect-timeout", "10", "--max-time", "60"];
38122
return withValidationMaxTimeOverride(args);
39123
}
40124

41-
export function getExtendedNvidiaEndpointValidationProbeCurlArgs(opts?: WslProbeOptions): string[] {
125+
export function getExtendedNvidiaEndpointValidationProbeCurlArgs(
126+
opts?: ValidationProbeOptions,
127+
): string[] {
42128
const args = isWsl(opts)
43129
? ["--connect-timeout", "30", "--max-time", "300"]
44130
: ["--connect-timeout", "10", "--max-time", "300"];

0 commit comments

Comments
 (0)