diff --git a/js/lib/index.ts b/js/lib/index.ts index 54f85e9..e6514c6 100644 --- a/js/lib/index.ts +++ b/js/lib/index.ts @@ -86,6 +86,7 @@ export { SandboxFileSystemError, SandboxProcessError, } from "./resources/abstraction/sandbox"; +export type { SandboxCreateOptions } from "./resources/abstraction/sandbox"; // Export Image classes and types export { Image } from "./resources/abstraction/image"; diff --git a/js/lib/resources/abstraction/sandbox-files.ts b/js/lib/resources/abstraction/sandbox-files.ts new file mode 100644 index 0000000..984b9ce --- /dev/null +++ b/js/lib/resources/abstraction/sandbox-files.ts @@ -0,0 +1,6 @@ +export function sandboxFileContentUrl( + containerId: string, + sandboxPath: string, +): string { + return `api/v1/gateway/pods/${containerId}/files/download/${encodeURIComponent(sandboxPath)}`; +} diff --git a/js/lib/resources/abstraction/sandbox.ts b/js/lib/resources/abstraction/sandbox.ts index 52bccd8..db401fb 100644 --- a/js/lib/resources/abstraction/sandbox.ts +++ b/js/lib/resources/abstraction/sandbox.ts @@ -1,5 +1,6 @@ import * as fs from "fs"; import { Pod, PodInstance } from "./pod"; +import { sandboxFileContentUrl } from "./sandbox-files"; import { CreateStubConfig } from "./stub"; import { EStubType } from "../../types/stub"; import type { @@ -25,6 +26,11 @@ export class SandboxFileSystemError extends Error {} /** Error thrown for sandbox process operations. */ export class SandboxProcessError extends Error {} +export interface SandboxCreateOptions { + entrypoint?: string[]; + waitForReady?: boolean; +} + function shellQuote(arg: string): string { if (arg === "") return "''"; // Simple POSIX single-quote escaping @@ -109,6 +115,16 @@ export class Sandbox extends Pod { ); } + /** Terminate a sandbox by ID without connecting to it first. */ + public static async terminate(id: string): Promise { + const response = await beamClient.request({ + method: "POST", + url: `api/v1/gateway/containers/${id}/stop`, + data: {}, + }); + return Boolean(response.data?.ok); + } + /** * Create a sandbox instance from a filesystem snapshot. * @@ -175,6 +191,25 @@ export class Sandbox extends Pod { ); } + private async createContainer(preparationCacheKey?: string): Promise<{ + ok: boolean; + containerId: string; + errorMsg?: string; + stubId?: string; + }> { + const response = await beamClient.request({ + method: "POST", + url: "api/v1/gateway/pods", + data: this.stub.stubId ? { stubId: this.stub.stubId } : {}, + headers: preparationCacheKey + ? { + "Grpc-Metadata-Preparation-Cache-Key": preparationCacheKey, + } + : undefined, + }); + return response.data; + } + /** * Create a new sandbox instance. * @@ -185,7 +220,14 @@ export class Sandbox extends Pod { * * Throws: SandboxConnectionError if the sandbox creation fails. */ - public async create(entrypoint?: string[]): Promise { + public async create( + entrypointOrOptions?: string[] | SandboxCreateOptions, + ): Promise { + const options = Array.isArray(entrypointOrOptions) + ? { entrypoint: entrypointOrOptions } + : entrypointOrOptions; + const entrypoint = options?.entrypoint; + this.stub.config.entrypoint = ["tail", "-f", "/dev/null"]; if (entrypoint && entrypoint.length) { this.stub.config.entrypoint = entrypoint; @@ -193,47 +235,47 @@ export class Sandbox extends Pod { const ignorePatterns = this.syncLocalDir ? undefined : ["*"]; - if (!this.runtimePreparation) { - this.runtimePreparation = this.stub.prepareRuntime( - undefined, - EStubType.Sandbox, - true, - ignorePatterns, - ); + const preparationCacheKey = + !this.syncLocalDir && !this.stub.runtimeReady + ? this.stub.preparationCacheKey(EStubType.Sandbox, ignorePatterns) + : undefined; + let body = await this.createContainer(preparationCacheKey); + if (body.ok && body.stubId) { + this.stub.stubCreated = true; + this.stub.stubId = body.stubId; + this.stub.runtimeReady = true; } - const currentPreparation = this.runtimePreparation; - let prepared: boolean; - try { - prepared = await currentPreparation; - } catch (error) { - if (this.runtimePreparation === currentPreparation) { - this.runtimePreparation = undefined; + if (!body.ok && !body.stubId) { + if (!this.runtimePreparation) { + this.runtimePreparation = this.stub.prepareRuntime( + undefined, + EStubType.Sandbox, + true, + ignorePatterns, + ); } - throw error; - } - if (!prepared && this.runtimePreparation === currentPreparation) { - this.runtimePreparation = undefined; - } - if (!prepared) { - const detail = this.stub.lastError?.message ?? "unknown reason"; - throw new SandboxConnectionError(`Failed to prepare runtime: ${detail}`); - } - - // eslint-disable-next-line no-console - console.log("Creating sandbox"); + const currentPreparation = this.runtimePreparation; + let prepared: boolean; + try { + prepared = await currentPreparation; + } catch (error) { + if (this.runtimePreparation === currentPreparation) { + this.runtimePreparation = undefined; + } + throw error; + } - const createResp = await beamClient.request({ - method: "POST", - url: `api/v1/gateway/pods`, - data: { stubId: this.stub.stubId }, - }); - const body = createResp.data as { - ok: boolean; - containerId: string; - errorMsg?: string; - }; + if (!prepared && this.runtimePreparation === currentPreparation) { + this.runtimePreparation = undefined; + } + if (!prepared) { + const detail = this.stub.lastError?.message ?? "unknown reason"; + throw new SandboxConnectionError(`Failed to prepare runtime: ${detail}`); + } + body = await this.createContainer(); + } if (!body.ok) { throw new SandboxConnectionError( @@ -241,35 +283,21 @@ export class Sandbox extends Pod { ); } - // eslint-disable-next-line no-console - console.log(`Sandbox created successfully ===> ${body.containerId}`); - - // Connect to the sandbox to ensure it's ready - const connectResp = await beamClient.request({ - method: "POST", - url: `api/v1/gateway/pods/${body.containerId}/connect`, - data: {}, - }); - const connectData = connectResp.data as { - ok: boolean; - errorMsg?: string; - }; - if (!connectData.ok) { - throw new SandboxConnectionError( - connectData.errorMsg || "Failed to connect to sandbox", - ); - } - - if ((this.stub.config.keepWarmSeconds as number) < 0) { - // eslint-disable-next-line no-console - console.log( - "This sandbox has no timeout, it will run until it is shut down manually.", - ); - } else { - // eslint-disable-next-line no-console - console.log( - `This sandbox will timeout after ${this.stub.config.keepWarmSeconds} seconds.`, - ); + if (options?.waitForReady !== false) { + const connectResp = await beamClient.request({ + method: "POST", + url: `api/v1/gateway/pods/${body.containerId}/connect`, + data: {}, + }); + const connectData = connectResp.data as { + ok: boolean; + errorMsg?: string; + }; + if (!connectData.ok) { + throw new SandboxConnectionError( + connectData.errorMsg || "Failed to connect to sandbox", + ); + } } return new SandboxInstance( @@ -449,7 +477,7 @@ export class SandboxInstance extends PodInstance { * Terminate the sandbox instance. */ public async terminate(): Promise { - const result = await super.terminate(); + const result = await Sandbox.terminate(this.containerId); if (result) { this.terminated = true; } @@ -1012,7 +1040,7 @@ export class SandboxFileSearchResult { /** * File system interface for managing files within a sandbox. * - * Upload, download, stat, list, and manage files and directories. + * Upload, stat, list, and manage files and directories. */ export class SandboxFileSystem { private sandbox_instance: SandboxInstance; @@ -1065,32 +1093,18 @@ export class SandboxFileSystem { return this.writeBytes(sandboxPath, Buffer.from(content, "utf8"), mode); } - /** Download a file from the sandbox and return its bytes. */ - public async download(sandboxPath: string): Promise { - return this.readBytes(sandboxPath); - } - - /** Download a file from the sandbox to a local path. */ - public async downloadFile( - sandboxPath: string, - localPath: string, - ): Promise { - fs.writeFileSync(localPath, await this.readBytes(sandboxPath)); - } - /** Read a file from the sandbox as bytes. */ public async readBytes(sandboxPath: string): Promise { const resp = await beamClient.request({ method: "GET", - url: `api/v1/gateway/pods/${ - this.sandbox_instance.containerId - }/files/download/${encodeURIComponent(sandboxPath)}`, + url: sandboxFileContentUrl( + this.sandbox_instance.containerId, + sandboxPath, + ), }); const data = resp.data as { ok: boolean; errorMsg?: string; data?: string }; if (!data.ok || !data.data) - throw new SandboxFileSystemError( - data.errorMsg || "Failed to download file", - ); + throw new SandboxFileSystemError(data.errorMsg || "Failed to read file"); return Buffer.from(data.data, "base64"); } diff --git a/js/lib/resources/abstraction/stub.ts b/js/lib/resources/abstraction/stub.ts index 1f25bfe..f8051ee 100644 --- a/js/lib/resources/abstraction/stub.ts +++ b/js/lib/resources/abstraction/stub.ts @@ -1,4 +1,5 @@ import * as path from "path"; +import { createHash } from "crypto"; import beamClient, { GpuType, GpuTypeAlias } from "../.."; import { Image } from "./image"; import { Volume } from "../volume"; @@ -56,9 +57,10 @@ export interface StubConfig { allowList?: string[]; } -export interface CreateStubConfig extends Partial { +export type CreateStubConfig = Omit, "image"> & { name: string; -} + image?: Image | string; +}; // Global stub creation state management let _stubCreatedForWorkspace = false; @@ -120,7 +122,8 @@ export class StubBuilder { this.config.name = name; this.config.app = app || name; this.config.authorized = authorized; - this.config.image = image || new Image({}); + this.config.image = + typeof image === "string" ? Image.fromRegistry(image) : image || new Image({}); this.config.callbackUrl = callbackUrl; this.config.cpu = cpu; this.config.memory = memory; @@ -239,6 +242,11 @@ export class StubBuilder { return true; } + const preparationCacheKey = + ignorePatterns?.length === 1 && ignorePatterns[0] === "*" + ? this.preparationCacheKey(stubType, ignorePatterns) + : undefined; + // Build image if not available if (!this.imageAvailable) { try { @@ -263,20 +271,25 @@ export class StubBuilder { // Sync files if not already synced if (!this.filesSynced) { - try { - const syncResult = await this.syncer.sync(ignorePatterns); - if (syncResult.success) { - this.filesSynced = true; - this.objectId = syncResult.objectId; - } else { - this.lastError = new Error("File sync failed"); - console.error("File sync failed"); + if (ignorePatterns?.length === 1 && ignorePatterns[0] === "*") { + this.filesSynced = true; + this.objectId = ""; + } else { + try { + const syncResult = await this.syncer.sync(ignorePatterns); + if (syncResult.success) { + this.filesSynced = true; + this.objectId = syncResult.objectId; + } else { + this.lastError = new Error("File sync failed"); + console.error("File sync failed"); + return false; + } + } catch (error) { + this.lastError = error instanceof Error ? error : new Error(String(error)); + console.error("File sync failed:", error); return false; } - } catch (error) { - this.lastError = error instanceof Error ? error : new Error(String(error)); - console.error("File sync failed:", error); - return false; } } @@ -372,6 +385,9 @@ export class StubBuilder { method: "POST", url: "/api/v1/gateway/stubs", data: camelCaseToSnakeCaseKeys(stubRequest), + headers: preparationCacheKey + ? { "Grpc-Metadata-Preparation-Cache-Key": preparationCacheKey } + : undefined, }); stubResponse = response.data; } else { @@ -392,6 +408,9 @@ export class StubBuilder { method: "POST", url: "/api/v1/gateway/stubs", data: camelCaseToSnakeCaseKeys(stubRequest), + headers: preparationCacheKey + ? { "Grpc-Metadata-Preparation-Cache-Key": preparationCacheKey } + : undefined, }); stubResponse = response.data; setStubCreatedForWorkspace(true); @@ -423,6 +442,27 @@ export class StubBuilder { return true; } + public preparationCacheKey( + stubType: string, + ignorePatterns?: string[] + ): string { + return createHash("sha256") + .update( + JSON.stringify({ + version: 1, + stubType, + config: { + ...this.config, + image: this.config.image.config, + volumes: this.config.volumes.map((volume) => volume.export()), + }, + extra: this.extra, + ignorePatterns, + }), + ) + .digest("hex"); + } + public async deployStub( request: DeployStubRequest ): Promise { diff --git a/js/package-lock.json b/js/package-lock.json index 77adff1..4c6b0c1 100644 --- a/js/package-lock.json +++ b/js/package-lock.json @@ -1,12 +1,12 @@ { "name": "@beamcloud/beam-js", - "version": "1.0.13", + "version": "1.0.17", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "@beamcloud/beam-js", - "version": "1.0.13", + "version": "1.0.17", "license": "MIT", "dependencies": { "axios": "^1.16.0", diff --git a/js/package.json b/js/package.json index 81a73e1..7c18035 100644 --- a/js/package.json +++ b/js/package.json @@ -1,6 +1,6 @@ { "name": "@beamcloud/beam-js", - "version": "1.0.13", + "version": "1.0.17", "description": "TypeScript and JavaScript SDK for Beam", "main": "dist/index.js", "module": "dist/index.mjs", diff --git a/js/tests/sandbox-network.test.ts b/js/tests/sandbox-network.test.ts index b724d30..771fac8 100644 --- a/js/tests/sandbox-network.test.ts +++ b/js/tests/sandbox-network.test.ts @@ -1,4 +1,5 @@ import beamClient from "../lib"; +import { Image } from "../lib/resources/abstraction/image"; import { Sandbox, SandboxConnectionError, SandboxInstance } from "../lib/resources/abstraction/sandbox"; import { EStubType } from "../lib/types/stub"; @@ -59,6 +60,151 @@ describe("Sandbox network parity", () => { ); }); + test("creates from a prepared stub without rebuilding or syncing", async () => { + const sandbox = new Sandbox({ name: "cached-sandbox" }); + const buildMock = jest.spyOn(sandbox.stub.config.image, "build"); + const syncMock = jest.spyOn(sandbox.stub.syncer, "sync"); + const requestMock = jest + .spyOn(beamClient, "request") + .mockImplementation(async (config) => { + if (config.url?.endsWith("/connect")) { + return { data: { ok: true } }; + } + return { + data: { + ok: true, + containerId: "sandbox-1", + stubId: "stub-cached", + }, + }; + }); + + await expect(sandbox.create()).resolves.toMatchObject({ + containerId: "sandbox-1", + stubId: "stub-cached", + }); + + expect(sandbox.stub.stubId).toBe("stub-cached"); + expect(buildMock).not.toHaveBeenCalled(); + expect(syncMock).not.toHaveBeenCalled(); + expect(requestMock).toHaveBeenCalledTimes(2); + expect(requestMock).toHaveBeenCalledWith( + expect.objectContaining({ + method: "POST", + url: "api/v1/gateway/pods", + data: {}, + headers: { + "Grpc-Metadata-Preparation-Cache-Key": expect.stringMatching( + /^[0-9a-f]{64}$/ + ), + }, + }) + ); + }); + + test("uses distinct preparation keys for distinct images", () => { + const node = new Sandbox({ name: "image-key", image: "node:20" }); + const python = new Sandbox({ name: "image-key", image: "python:3.12" }); + + expect(node.stub.config.image).toBeInstanceOf(Image); + expect( + node.stub.preparationCacheKey(EStubType.Sandbox, ["*"]) + ).not.toBe(python.stub.preparationCacheKey(EStubType.Sandbox, ["*"])); + }); + + test("does not use a prepared stub when syncing local files", async () => { + const requestMock = jest.spyOn(beamClient, "request").mockResolvedValue({ + data: { ok: true, containerId: "sandbox-1", stubId: "stub-1" }, + }); + + await new Sandbox({ name: "synced-sandbox" }, true).create({ + waitForReady: false, + }); + + expect(requestMock).toHaveBeenCalledWith( + expect.objectContaining({ + url: "api/v1/gateway/pods", + headers: undefined, + }) + ); + }); + + test("can return before readiness when the next operation waits for it", async () => { + const requestMock = jest.spyOn(beamClient, "request").mockResolvedValue({ + data: { + ok: true, + containerId: "sandbox-1", + stubId: "stub-cached", + }, + }); + + await expect( + new Sandbox({ name: "cached-sandbox" }).create({ waitForReady: false }) + ).resolves.toMatchObject({ containerId: "sandbox-1" }); + + expect(requestMock).toHaveBeenCalledTimes(1); + expect(requestMock).toHaveBeenCalledWith( + expect.objectContaining({ url: "api/v1/gateway/pods" }) + ); + }); + + test("terminates by ID without connecting first", async () => { + const requestMock = jest.spyOn(beamClient, "request").mockResolvedValue({ + data: { ok: true }, + }); + + await expect(Sandbox.terminate("sandbox-1")).resolves.toBe(true); + expect(requestMock).toHaveBeenCalledTimes(1); + expect(requestMock).toHaveBeenCalledWith({ + method: "POST", + url: "api/v1/gateway/containers/sandbox-1/stop", + data: {}, + }); + }); + + test("does not prepare again when a prepared sandbox cannot be scheduled", async () => { + const sandbox = new Sandbox({ name: "cached-sandbox" }); + const prepareMock = jest.spyOn(sandbox.stub, "prepareRuntime"); + jest.spyOn(beamClient, "request").mockResolvedValue({ + data: { + ok: false, + errorMsg: "cpu quota exceeded", + stubId: "stub-cached", + }, + }); + + await expect(sandbox.create()).rejects.toThrow("cpu quota exceeded"); + expect(prepareMock).not.toHaveBeenCalled(); + }); + + test("skips file sync for an ignored workspace and caches the stub", async () => { + const sandbox = new Sandbox({ name: "empty-workspace" }); + sandbox.stub.imageAvailable = true; + sandbox.stub.config.image.id = "image-123"; + const syncMock = jest.spyOn(sandbox.stub.syncer, "sync"); + const requestMock = jest + .spyOn(beamClient, "request") + .mockResolvedValue({ data: { ok: true, stubId: "stub-new" } }); + + await expect( + sandbox.stub.prepareRuntime(undefined, EStubType.Sandbox, true, ["*"]) + ).resolves.toBe(true); + + expect(syncMock).not.toHaveBeenCalled(); + expect(requestMock).toHaveBeenCalledWith( + expect.objectContaining({ + method: "POST", + url: "/api/v1/gateway/stubs", + data: expect.objectContaining({ object_id: "" }), + headers: { + "Grpc-Metadata-Preparation-Cache-Key": expect.stringMatching( + /^[0-9a-f]{64}$/ + ), + }, + }) + ); + }); + test("updates network permissions with the sandbox update endpoint", async () => { const requestMock = jest.spyOn(beamClient, "request").mockResolvedValue({ data: { @@ -170,7 +316,11 @@ describe("Sandbox network parity", () => { const sandbox = new Sandbox({ name: "concurrent-sandbox" }); let releasePreparation!: (prepared: boolean) => void; const preparation = new Promise((resolve) => { - releasePreparation = resolve; + releasePreparation = (prepared) => { + sandbox.stub.stubId = "stub-1"; + sandbox.stub.runtimeReady = prepared; + resolve(prepared); + }; }); const prepareRuntimeMock = jest .spyOn(sandbox.stub, "prepareRuntime") @@ -180,6 +330,9 @@ describe("Sandbox network parity", () => { .spyOn(beamClient, "request") .mockImplementation(async (config) => { if (config.url === "api/v1/gateway/pods") { + if (!config.data?.stubId) { + return { data: { ok: false } }; + } nextContainer += 1; return { data: { @@ -197,6 +350,7 @@ describe("Sandbox network parity", () => { const firstCreate = sandbox.create(); const secondCreate = sandbox.create(); + await new Promise((resolve) => setImmediate(resolve)); expect(prepareRuntimeMock).toHaveBeenCalledTimes(1); releasePreparation(true); @@ -205,7 +359,7 @@ describe("Sandbox network parity", () => { "sandbox-1", "sandbox-2", ]); - expect(requestMock).toHaveBeenCalledTimes(4); + expect(requestMock).toHaveBeenCalledTimes(6); }); test("returns inline exec results without follow-up requests", async () => { @@ -290,8 +444,11 @@ describe("prepareRuntime surfaces real errors via lastError", () => { }); test("file sync exception is surfaced in SandboxConnectionError", async () => { - const sandbox = new Sandbox({ name: "test-sandbox" }); + const sandbox = new Sandbox({ name: "test-sandbox" }, true); sandbox.stub.imageAvailable = true; + jest.spyOn(beamClient, "request").mockResolvedValue({ + data: { ok: false }, + }); jest .spyOn(sandbox.stub.syncer, "sync")