sandbox: dynamic directory mount & snapshot (#239)

Port of modal-labs/modal-client#3791.

Author: ehdr

modal-client

@@ -0,0 +1,68 @@
+// This example shows how to mount Images in the Sandbox filesystem and take snapshots
+// of them.
+//
+// The feature is still experimental in the sense that the API is subject to change.
+//
+// High level, it allows you to:
+// - Mount any Modal Image at a specific directory within the Sandbox filesystem.
+// - Take a snapshot of that directory, which will create a new Modal Image with
+//   the updated contents of the directory.
+//
+// You can only snapshot directories that have previously been mounted using
+// `Sandbox.experimentalMountImage`. If you want to mount an empty directory,
+// you can pass undefined as the image parameter.
+//
+// For exmaple, you can use this to mount user specific dependencies into a running
+// Sandbox, that is started with a base Image with shared system dependencies. This
+// way, you can update system dependencies and user projects independently.
+
+import { ModalClient } from "modal";
+
+const modal = new ModalClient();
+
+const app = await modal.apps.fromName("libmodal-example", {
+  createIfMissing: true,
+});
+// The base Image you use for the Sandbox must have a /usr/bin/mount binary.
+const baseImage = modal.images.fromRegistry("debian:12-slim");
+
+const sb = await modal.sandboxes.create(app, baseImage);
+
+// You must mount an Image at a directory in the Sandbox filesystem before you
+// can snapshot it. You can pass undefined as the image parameter to mount an
+// empty directory.
+//
+// The target directory must exist before you can mount it:
+await (await sb.exec(["mkdir", "-p", "/repo"])).wait();
+await sb.experimentalMountImage("/repo");
+
+const gitClone = await sb.exec([
+  "git",
+  "clone",
+  "https://github.com/modal-labs/libmodal.git",
+  "/repo",
+]);
+await gitClone.wait();
+
+const repoSnapshot = await sb.experimentalSnapshotDirectory("/repo");
+console.log(
+  "Took a snapshot of the /repo directory, Image ID:",
+  repoSnapshot.imageId,
+);
+
+await sb.terminate();
+
+// Start a new Sandbox, and mount the repo directory:
+const sb2 = await modal.sandboxes.create(app, baseImage);
+
+await (await sb2.exec(["mkdir", "-p", "/repo"])).wait();
+await sb2.experimentalMountImage("/repo", repoSnapshot);
+
+const repoLs = await sb2.exec(["ls", "/repo"]);
+console.log(
+  "Contents of /repo directory in new Sandbox sb2:\n",
+  await repoLs.stdout.readText(),
+);
+
+await sb2.terminate();
+await modal.images.delete(repoSnapshot.imageId);

modal-js/examples/sandbox-image-mount.ts

@@ -21,6 +21,8 @@ import {
   TaskExecStartRequest,
   TaskExecStdoutConfig,
   TaskExecStderrConfig,
+  TaskMountDirectoryRequest,
+  TaskSnapshotDirectoryRequest,
 } from "../proto/modal_proto/task_command_router";
 import { TaskCommandRouterClientImpl } from "./task_command_router_client";
 import { v4 as uuidv4 } from "uuid";
@@ -1029,6 +1031,60 @@ export class Sandbox {
     return new Image(this.#client, resp.imageId, "");
   }
 
+  /**
+   * [Alpha] Mount an {@link Image} at a path in the Sandbox filesystem.
+   *
+   * @alpha
+   * @param path - The path where the directory should be mounted
+   * @param image - Optional {@link Image} to mount. If undefined, mounts an empty directory.
+   */
+  async experimentalMountImage(path: string, image?: Image): Promise<void> {
+    const taskId = await this.#getTaskId();
+    const commandRouterClient =
+      await this.#getOrCreateCommandRouterClient(taskId);
+
+    if (image && !image.imageId) {
+      throw new Error(
+        "Image must be built before mounting. Call `image.build(app)` first.",
+      );
+    }
+
+    const pathBytes = new TextEncoder().encode(path);
+    const imageId = image?.imageId ?? "";
+    const request = TaskMountDirectoryRequest.create({
+      taskId,
+      path: pathBytes,
+      imageId,
+    });
+    await commandRouterClient.mountDirectory(request);
+  }
+
+  /**
+   * [Alpha] Snapshot local changes to a previously mounted {@link Image} into a new {@link Image}.
+   *
+   * @alpha
+   * @param path - The path of the directory to snapshot
+   * @returns Promise that resolves to an {@link Image}
+   */
+  async experimentalSnapshotDirectory(path: string): Promise<Image> {
+    const taskId = await this.#getTaskId();
+    const commandRouterClient =
+      await this.#getOrCreateCommandRouterClient(taskId);
+
+    const pathBytes = new TextEncoder().encode(path);
+    const request = TaskSnapshotDirectoryRequest.create({
+      taskId,
+      path: pathBytes,
+    });
+    const response = await commandRouterClient.snapshotDirectory(request);
+
+    if (!response.imageId) {
+      throw new Error("Sandbox snapshot directory response missing `imageId`");
+    }
+
+    return new Image(this.#client, response.imageId, "");
+  }
+
   /**
    * Check if the Sandbox has finished running.
    *

modal-js/src/sandbox.ts

@@ -22,6 +22,9 @@ import {
   TaskExecStdioReadResponse,
   TaskExecWaitRequest,
   TaskExecWaitResponse,
+  TaskMountDirectoryRequest,
+  TaskSnapshotDirectoryRequest,
+  TaskSnapshotDirectoryResponse,
 } from "../proto/modal_proto/task_command_router";
 import {
   TaskGetCommandRouterAccessRequest,
@@ -342,6 +345,20 @@ export class TaskCommandRouterClientImpl {
     }
   }
 
+  async mountDirectory(request: TaskMountDirectoryRequest): Promise<void> {
+    await callWithRetriesOnTransientErrors(() =>
+      this.callWithAuthRetry(() => this.stub.taskMountDirectory(request)),
+    );
+  }
+
+  async snapshotDirectory(
+    request: TaskSnapshotDirectoryRequest,
+  ): Promise<TaskSnapshotDirectoryResponse> {
+    return await callWithRetriesOnTransientErrors(() =>
+      this.callWithAuthRetry(() => this.stub.taskSnapshotDirectory(request)),
+    );
+  }
+
   private async refreshJwt(): Promise<void> {
     this.jwtRefreshLock = this.jwtRefreshLock.then(async () => {
       if (this.closed) {

modal-js/src/task_command_router_client.ts

@@ -980,7 +980,7 @@ test("SandboxExecOutputTimeout", async () => {
 
     const elapsed = Date.now() - t0;
     expect(elapsed).toBeGreaterThan(1000);
-    expect(elapsed).toBeLessThan(3000);
+    expect(elapsed).toBeLessThan(4000);
 
     const exitCode = await p.wait();
     expect(exitCode).toBe(0);
@@ -989,6 +989,6 @@ test("SandboxExecOutputTimeout", async () => {
 
     const elapsed = Date.now() - t0;
     expect(elapsed).toBeGreaterThan(1000);
-    expect(elapsed).toBeLessThan(3000);
+    expect(elapsed).toBeLessThan(4000);
   }
 });

modal-js/test/sandbox.test.ts

@@ -0,0 +1,104 @@
+import { tc } from "../test-support/test-client";
+import { expect, test, onTestFinished } from "vitest";
+
+test.skip("SandboxMountDirectoryEmpty", async () => {
+  const app = await tc.apps.fromName("libmodal-test", {
+    createIfMissing: true,
+  });
+  const image = tc.images.fromRegistry("debian:12-slim");
+
+  const sb = await tc.sandboxes.create(app, image);
+  onTestFinished(async () => await sb.terminate());
+
+  await (await sb.exec(["mkdir", "-p", "/mnt/empty"])).wait();
+  await sb.experimentalMountImage("/mnt/empty");
+
+  const dirCheck = await sb.exec(["test", "-d", "/mnt/empty"]);
+  expect(await dirCheck.wait()).toBe(0);
+});
+
+test.skip("SandboxMountDirectoryWithImage", async () => {
+  const app = await tc.apps.fromName("libmodal-test", {
+    createIfMissing: true,
+  });
+  const baseImage = tc.images.fromRegistry("debian:12-slim");
+
+  const sb1 = await tc.sandboxes.create(app, baseImage);
+  onTestFinished(async () => await sb1.terminate());
+
+  const echoProc = await sb1.exec([
+    "sh",
+    "-c",
+    "echo -n 'mounted content' > /tmp/test.txt",
+  ]);
+  await echoProc.wait();
+
+  const mountImage = await sb1.snapshotFilesystem();
+  expect(mountImage.imageId).toMatch(/^im-/);
+
+  await sb1.terminate();
+
+  const sb2 = await tc.sandboxes.create(app, baseImage);
+  onTestFinished(async () => await sb2.terminate());
+
+  await (await sb2.exec(["mkdir", "-p", "/mnt/data"])).wait();
+  await sb2.experimentalMountImage("/mnt/data", mountImage);
+
+  const catProc = await sb2.exec(["cat", "/mnt/data/tmp/test.txt"]);
+  const output = await catProc.stdout.readText();
+  expect(output).toBe("mounted content");
+});
+
+test.skip("SandboxSnapshotDirectory", async () => {
+  const app = await tc.apps.fromName("libmodal-test", {
+    createIfMissing: true,
+  });
+  const baseImage = tc.images.fromRegistry("debian:12-slim");
+
+  const sb1 = await tc.sandboxes.create(app, baseImage);
+  onTestFinished(async () => await sb1.terminate());
+
+  await (await sb1.exec(["mkdir", "-p", "/mnt/data"])).wait();
+  await sb1.experimentalMountImage("/mnt/data");
+
+  const echoProc = await sb1.exec([
+    "sh",
+    "-c",
+    "echo -n 'snapshot test content' > /mnt/data/snapshot.txt",
+  ]);
+  await echoProc.wait();
+
+  const snapshotImage = await sb1.experimentalSnapshotDirectory("/mnt/data");
+  expect(snapshotImage.imageId).toMatch(/^im-/);
+
+  await sb1.terminate();
+
+  const sb2 = await tc.sandboxes.create(app, baseImage);
+  onTestFinished(async () => await sb2.terminate());
+
+  await (await sb2.exec(["mkdir", "-p", "/mnt/data"])).wait();
+  await sb2.experimentalMountImage("/mnt/data", snapshotImage);
+
+  const catProc = await sb2.exec(["cat", "/mnt/data/snapshot.txt"]);
+  const output = await catProc.stdout.readText();
+  expect(output).toBe("snapshot test content");
+});
+
+test.skip("SandboxMountDirectoryWithUnbuiltImageThrows", async () => {
+  const app = await tc.apps.fromName("libmodal-test", {
+    createIfMissing: true,
+  });
+  const baseImage = tc.images.fromRegistry("debian:12-slim");
+
+  const sb = await tc.sandboxes.create(app, baseImage);
+  onTestFinished(async () => await sb.terminate());
+
+  await (await sb.exec(["mkdir", "-p", "/mnt/data"])).wait();
+
+  const unbuiltImage = tc.images.fromRegistry("alpine:3.21");
+  expect(unbuiltImage.imageId).toBe("");
+
+  await expect(
+    sb.experimentalMountImage("/mnt/data", unbuiltImage),
+  ).rejects.toThrow("Image must be built before mounting");
+});