Skip to content

fix(onboard): prepare Jetson images before GPU probes #11877

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
latenighthackathon wants to merge 1 commit into from
Closed

fix(onboard): prepare Jetson images before GPU probes #11877

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
118 changes: 83 additions & 35 deletions src/lib/onboard/managed-bootstrap/docker-runtime.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -386,6 +386,34 @@ describe("Docker managed-bootstrap GPU probe diagnostics", () => {
});

describe("Docker managed-bootstrap GPU probe image", () => {
it("pulls an uncached Jetson image before bounded GPU mode probes (#11677)", async () => {
const { dependencies, dockerRun } = gpuModeDependencies();
const seed = authority("openclaw");
const originalInput = compatibilityLifecycleInput(seed, dependencies);
const input = {
...originalInput,
sandboxGpuConfig: { ...originalInput.sandboxGpuConfig, hostGpuPlatform: "jetson" as const },
};
const sandboxImage = `${input.image.repository}@${input.image.manifestDigest}`;
dockerAdapterMocks.imageInspect.mockReturnValue({ status: 1 });
dockerRun.mockImplementation(() => {
// A cold image cannot finish its implicit pull within the short GPU probe.
expect(dockerAdapterMocks.pullWithProgressWatchdog).toHaveBeenCalledOnce();
return { status: 0, stdout: "probe-id" };
});

await runCompatibilityCreate(input, seed);

expect(dockerAdapterMocks.pullWithProgressWatchdog).toHaveBeenCalledWith(
sandboxImage,
expect.objectContaining({ maxTimeoutMs: 30 * 60 * 1000 }),
);
expect(dockerRun.mock.calls[0]?.[0]).toEqual(
expect.arrayContaining(["--runtime", "nvidia", "--pull", "never"]),
);
expect(dockerRun.mock.calls[0]?.[1]).toEqual(expect.objectContaining({ timeout: 30_000 }));
});

it("pulls an uncached WSL sandbox image before bounded GPU mode probes", async () => {
const { dependencies, dockerRun } = gpuModeDependencies();
const seed = authority("openclaw");
Expand Down Expand Up @@ -473,21 +501,31 @@ describe("Docker managed-bootstrap GPU probe image", () => {
);
});

it("skips the pull when the exact WSL sandbox image is already local", async () => {
const { dependencies, dockerRun } = gpuModeDependencies();
const seed = authority("openclaw");
const input = compatibilityLifecycleInput(seed, dependencies);
const sandboxImage = `${input.image.repository}@${input.image.manifestDigest}`;
sandboxCreateMocks.isDockerDesktopWslRuntime.mockReturnValue(true);
it.each([
["WSL", true, "n1x"],
["Jetson", false, "jetson"],
] as const)(
"skips the pull when the exact %s sandbox image is already local",
async (_host, wsl, platform) => {
const { dependencies, dockerRun } = gpuModeDependencies();
const seed = authority("openclaw");
const originalInput = compatibilityLifecycleInput(seed, dependencies);
const input = {
...originalInput,
sandboxGpuConfig: { ...originalInput.sandboxGpuConfig, hostGpuPlatform: platform },
};
const sandboxImage = `${input.image.repository}@${input.image.manifestDigest}`;
sandboxCreateMocks.isDockerDesktopWslRuntime.mockReturnValue(wsl);

await runCompatibilityCreate(input, seed);
await runCompatibilityCreate(input, seed);

expect(dockerAdapterMocks.imageInspect).toHaveBeenCalledOnce();
expect(dockerAdapterMocks.pullWithProgressWatchdog).not.toHaveBeenCalled();
const probeArgs = dockerRun.mock.calls[0]?.[0] ?? [];
expect(probeArgs).toEqual(expect.arrayContaining(["--pull", "never"]));
expect(probeArgs.slice(-2)).toEqual([sandboxImage, "true"]);
});
expect(dockerAdapterMocks.imageInspect).toHaveBeenCalledOnce();
expect(dockerAdapterMocks.pullWithProgressWatchdog).not.toHaveBeenCalled();
const probeArgs = dockerRun.mock.calls[0]?.[0] ?? [];
expect(probeArgs).toEqual(expect.arrayContaining(["--pull", "never"]));
expect(probeArgs.slice(-2)).toEqual([sandboxImage, "true"]);
},
);

it("retains implicit pull behaviour on other Docker hosts", async () => {
const { dependencies, dockerRun } = gpuModeDependencies();
Expand All @@ -505,29 +543,39 @@ describe("Docker managed-bootstrap GPU probe image", () => {
expect(probeArgs).not.toContain("--pull");
});

it("stops before GPU mode probing when the WSL image pull exceeds its limit", async () => {
const { dependencies, dockerRun } = gpuModeDependencies();
const seed = authority("openclaw");
const input = compatibilityLifecycleInput(seed, dependencies);
sandboxCreateMocks.isDockerDesktopWslRuntime.mockReturnValue(true);
dockerAdapterMocks.imageInspect.mockReturnValue({ status: 1 });
dockerAdapterMocks.pullWithProgressWatchdog.mockResolvedValue({
status: 124,
signal: "SIGTERM",
output: "pull stopped",
timedOut: true,
timeoutKind: "max",
});
const lifecycle = createDockerManagedBootstrapSurface().createLifecycle(input);
it.each([
["WSL", true, "n1x"],
["Jetson", false, "jetson"],
] as const)(
"stops before GPU mode probing when the %s image pull exceeds its limit",
async (_host, wsl, platform) => {
const { dependencies, dockerRun } = gpuModeDependencies();
const seed = authority("openclaw");
const originalInput = compatibilityLifecycleInput(seed, dependencies);
const input = {
...originalInput,
sandboxGpuConfig: { ...originalInput.sandboxGpuConfig, hostGpuPlatform: platform },
};
sandboxCreateMocks.isDockerDesktopWslRuntime.mockReturnValue(wsl);
dockerAdapterMocks.imageInspect.mockReturnValue({ status: 1 });
dockerAdapterMocks.pullWithProgressWatchdog.mockResolvedValue({
status: 124,
signal: "SIGTERM",
output: "pull stopped",
timedOut: true,
timeoutKind: "max",
});
const lifecycle = createDockerManagedBootstrapSurface().createLifecycle(input);

await expect(
lifecycle.runCreate(async () => ({ value: "created", receipt: seed.handle.createReceipt })),
).rejects.toThrow(
"Docker managed sandbox image pull failed before GPU mode selection: exceeded the 30-minute safety limit.",
);
expect(dockerRun).not.toHaveBeenCalled();
expect(adapterMocks.prepare).not.toHaveBeenCalled();
});
await expect(
lifecycle.runCreate(async () => ({ value: "created", receipt: seed.handle.createReceipt })),
).rejects.toThrow(
"Docker managed sandbox image pull failed before GPU mode selection: exceeded the 30-minute safety limit.",
);
expect(dockerRun).not.toHaveBeenCalled();
expect(adapterMocks.prepare).not.toHaveBeenCalled();
},
);

it.each([
{
Expand Down
15 changes: 8 additions & 7 deletions src/lib/onboard/managed-bootstrap/docker-runtime.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -167,7 +167,7 @@ function selectedDockerMode(
device: input.sandboxGpuConfig.sandboxGpuDevice,
backend,
dockerDesktopWsl,
...(dockerDesktopWsl ? { pullPolicy: "never" as const } : {}),
...(dockerDesktopWsl || backend === "jetson" ? { pullPolicy: "never" as const } : {}),
},
withDockerClientEnvDeps(input.dependencies as DockerGpuPatchDeps, prepared),
);
Expand Down Expand Up @@ -248,8 +248,13 @@ function createDockerLifecycle(
}
const dockerDesktopWsl =
input.route === "compatibility" ? isDockerDesktopWslRuntime() : undefined;
const preselectedMode = dockerDesktopWsl ? null : selectedDockerMode(input, dockerDesktopWsl);
const backend = input.sandboxGpuConfig.hostGpuPlatform === "jetson" ? "jetson" : "generic";
const prepareGpuProbeImage =
input.route === "compatibility" &&
input.sandboxGpuConfig.sandboxGpuEnabled &&
(dockerDesktopWsl || backend === "jetson");
// Cold image downloads must not consume the bounded GPU-mode probe timeout.
const preselectedMode = prepareGpuProbeImage ? null : selectedDockerMode(input, dockerDesktopWsl);
const persistStartupCommand =
input.persistStartupCommand && (input.route !== "native" || input.requiredLimits.length > 0);
const commandExecutor = input.dependencies.commandExecutor;
Expand Down Expand Up @@ -350,11 +355,7 @@ function createDockerLifecycle(
readonly bootstrapIdentity: string;
}) => Promise<ManagedBootstrapRuntimeCreateLaunchResult<T>>,
): Promise<T> {
if (
dockerDesktopWsl &&
input.route === "compatibility" &&
input.sandboxGpuConfig.sandboxGpuEnabled
) {
if (prepareGpuProbeImage) {
await prepareDockerManagedBootstrapGpuProbeImage(
managedBootstrapImageReference(input),
input.dockerClientEnv,
Expand Down
Loading