Skip to content

[P2] [feature-request] Support RunPod as an alternative sandbox backend to Modal #49

Description

@lucastononro

Summary

Today every sandboxed code execution and Jupyter kernel goes through Modal: modal.Sandbox, modal.Volume, modal.Image, modal.App are imported directly inside services/sandbox.py, services/kernel_manager.py, and services/volume.py. That couples us to one provider for pricing, GPU availability, and credit pools.

Add RunPod as a swappable second backend, selectable per-deployment (and ideally per-project). Same SDK semantics for the agent (trainable.log(...), code execution, persistent kernel) — just a different compute provider underneath.

Why RunPod

  • Cheaper GPU-hour pricing on most tiers, and a different GPU inventory (e.g. RTX 4090, MI300X) that Modal doesn't expose. Some users have RunPod credits but not Modal credits.
  • Long-running pods are first-class on RunPod, which is a better fit for the persistent-Jupyter-kernel use case (today we keep a Modal Sandbox alive for up to 2h via timeouts + idle reaper).
  • Provider redundancy — if Modal has a regional outage we can fall back.
  • Per-project provider choice — heavy training on RunPod (cost), interactive iteration on Modal (faster cold starts).

Modal vs RunPod — primitives mapping

Concern Modal (today) RunPod (target)
One-shot code exec modal.Sandbox.create.aio() with python -u -c <code> RunPod Serverless endpoint: POST /v2/{endpoint_id}/run with {"input": {"code": ...}}, stream via GET /v2/{endpoint_id}/stream/{job_id}
Persistent Jupyter kernel Modal Sandbox with kernel proxy on stdin/stdout, idle-reaped RunPod Pod (long-running container) created via GraphQL/REST, exposes HTTP/WS for the kernel proxy; torn down by reaper
Image build modal.Image.debian_slim().pip_install(...) (programmatic, lazy) Pre-built Docker image pushed to a registry; serverless workers and pods reference it by tag
Persistent storage modal.Volume.from_name("trainable-data"), mounted at /data RunPod Network Volume attached to workers/pods at /data (similar UX, different reach)
Backend-side file read/write vol.read_file() / vol.batch_upload() directly from API server No direct equivalent — must route through a worker handler or mirror state to S3 (codebase already has services/s3_client.py)
Stdout streaming async for chunk in sb.stdout Job stream endpoint yields JSON chunks; or pod-side WebSocket; same parse pipeline downstream
Stdin (kernel commands) sb.stdin.write() Pod HTTP endpoint posts {"cmd": ...} to the kernel proxy; serverless can't do stdin (use streaming generator with input updates instead)
Cold start ~1s (Modal Sandbox) Serverless: 1–10s without FlashBoot, ~1s with FlashBoot warm; Pods: ~30–60s but stay hot
Auth MODAL_TOKEN_ID + MODAL_TOKEN_SECRET, or ~/.modal.toml RUNPOD_API_KEY (single token); endpoint/volume IDs in env
Pricing model Per-second on every primitive Per-second serverless; per-hour pods (regardless of idle) — cost tracking has to differ
Cost API None (we compute from local pricing table) GET /v2/{endpoint_id}/jobs/{job_id} returns billable seconds — can use ground-truth

The big asymmetries we need to design around:

  1. Backend-driven volume reads — Modal lets us read the volume from the API server. RunPod doesn't, so the abstraction needs a "file gateway" worker, OR we shift the source of truth to S3 and treat the network volume as scratch.
  2. Pods bill by the hour — the idle-reaper has to be more aggressive on RunPod to avoid burning credits while a user goes to lunch with a kernel open.
  3. Serverless workers can't accept stdin — the kernel proxy model needs HTTP/WebSocket on a Pod, not Serverless.

Current state — anchors in the code

Direct Modal coupling we need to abstract:

  • backend/services/sandbox.py:9import modal. Owns one-shot run_code() (L123-260), the SDK preamble (L55-69), the global App (L28-40) and Image (L72-117).
  • backend/services/kernel_manager.py:23import modal. Owns the persistent kernel (_spawn at L297-306, shutdown at L545-561, stdin proxy at L503-540).
  • backend/services/volume.py:10import modal. Owns volume RPC (get_volume, read_volume_file_async, listdir_async, upload_to_volume, write_to_volume).
  • backend/config.py:30-31modal_app_name, modal_volume_name. Provider-named.
  • backend/config.py:44-45sandbox_timeout (already neutral, keep).
  • backend/services/usage.py:51-62_SANDBOX_PRICING table (Modal-rate-keyed).
  • backend/services/usage.py:248provider="modal" hard-coded when recording usage.
  • backend/models.py:217-250UsageEvent already has a provider column. Good — multi-backend cost attribution is half-built.
  • backend/schemas.py:23-30SandboxProfile { gpu, timeout }, SandboxConfig { default, training } — provider-neutral, keep.
  • backend/models.py:38Project.sandbox_config JSON column — extend with backend field (see §UI below).
  • backend/tools/execute_code.py:96-128 — Reads heavy flag, picks profile, calls run_code(...). The integration point.
  • backend/tools/run_notebook_cell.py — Same pattern via kernel_manager. Same integration point.
  • backend/routers/files.py:40-80 — Serves /files/list, /files/read, /files/raw from Modal volume. This is the layer that breaks under RunPod and needs the storage abstraction.
  • backend/tests/conftest.py:204-275MockVolume + mock_volume_patches already mock the Modal volume in tests; the abstraction means we keep one mock instead of mocking two providers.
  • docker-compose.yml:51 + .env.example:12-14MODAL_TOKEN_* env wiring.

Stage handling, observability spans, metrics parsing, and the SDK preamble are all backend-neutral and can be reused 1:1 — they live above the sandbox-creation call.

Proposed design

1. SandboxBackend protocol

New module: backend/services/sandbox/. Move existing Modal logic into sandbox/modal_backend.py. Define an abstract protocol:

# backend/services/sandbox/base.py
class SandboxBackend(Protocol):
    name: Literal["modal", "runpod"]

    async def run_code(
        self,
        code: str,
        session_id: str,
        *,
        stage: str | None = None,
        gpu: str | None = None,
        timeout: int | None = None,
        agent_type: str | None = None,
        agent_id: str | None = None,
    ) -> SandboxResult:
        """One-shot code exec. Streams stdout to broadcaster, returns
        {stdout, stderr, returncode, elapsed_s}."""

    async def spawn_kernel(self, session_id: str) -> KernelHandle:
        """Persistent Jupyter kernel. KernelHandle exposes execute()/interrupt()/
        shutdown() — implementation owns transport (Modal stdin vs RunPod HTTP)."""

    def get_storage(self) -> StorageBackend: ...
    def gpu_pricing(self) -> dict[str, float]: ...

class StorageBackend(Protocol):
    async def read_file(self, path: str) -> bytes: ...
    async def listdir(self, path: str, recursive: bool = False) -> list[FileNode]: ...
    async def upload(self, local_path: str, remote_path: str) -> None: ...
    async def write(self, content: bytes, remote_path: str) -> None: ...
    async def reload(self) -> None: ...

SandboxResult and KernelHandle are both already implicit in today's code — formalize them as dataclasses.

2. Factory + selection

# backend/services/sandbox/__init__.py
def get_backend(name: str | None = None) -> SandboxBackend:
    name = name or settings.sandbox_backend  # default "modal"
    return {
        "modal": ModalBackend,
        "runpod": RunpodBackend,
    }[name]()

config.py additions:

  • sandbox_backend: str = "modal" — default.
  • Rename modal_app_name → keep as modal_app_name (provider-specific is fine); add runpod_endpoint_id, runpod_pod_template_id, runpod_network_volume_id, runpod_api_key.
  • Keep modal_volume_name and runpod_network_volume_id as separate provider-keyed settings.
  • Add per-project override: Project.sandbox_config gains an optional "backend": "modal" | "runpod" field. routers/sessions.py:133-136 already reads sandbox_config — extend to pass backend through to the agent loop and into execute_code's factory.

backend/tools/execute_code.py:126 becomes:

backend = get_backend(_sandbox_config.get("backend"))
result = await backend.run_code(code, session_id, stage=stage, gpu=gpu, timeout=timeout, ...)

Same shape, different implementation underneath.

3. RunPod Serverless implementation (run_code)

Worker side (separate Docker image, lives in the repo at backend/runpod_worker/):

# handler.py
import runpod, subprocess, json, sys
SDK_PREAMBLE = "..."  # copied from services/sandbox.py:55-69

def handler(event):
    code = SDK_PREAMBLE + event["input"]["code"]
    proc = subprocess.Popen(
        ["python", "-u", "-c", code],
        stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        cwd=f"/data/sessions/{event['input']['session_id']}",
    )
    for line in proc.stdout:
        yield {"stream": "stdout", "data": line.decode()}
    yield {"returncode": proc.wait()}

runpod.serverless.start({"handler": handler, "return_aggregate_stream": True})

The Dockerfile bakes in the same package set today's get_image() provides (backend/services/sandbox.py:78-106). Network Volume mounts at /data.

Backend side (runpod_backend.py):

async def run_code(self, code, session_id, **kw):
    job = await self._post(f"/v2/{settings.runpod_endpoint_id}/run", {"input": {...}})
    async for chunk in self._stream(f"/v2/{settings.runpod_endpoint_id}/stream/{job['id']}"):
        await broadcaster.publish(session_id, "code_output", {...})
        # same JSON metric parsing as sandbox.py:203-221
    final = await self._poll(f"/v2/{settings.runpod_endpoint_id}/status/{job['id']}")
    return SandboxResult(stdout=..., stderr=..., returncode=final["..."], elapsed_s=final["executionTime"])

GPU mapping: today's gpu="T4" becomes a RunPod GPU type filter on the endpoint config (set at deploy time, not per-request — RunPod serverless endpoints are typed). To handle Modal's runtime gpu param, deploy two endpoints (trainable-cpu, trainable-gpu-h100, etc.) and route based on the gpu arg. Capture the endpoint mapping in config.

4. RunPod Pod implementation (persistent kernel)

spawn_kernel on RunPod creates a Pod from a template image (same Docker image as the worker, plus the kernel proxy entrypoint), attaches the Network Volume, and exposes a single HTTP port for the kernel proxy. Backend talks to the pod via that HTTP endpoint:

  • POST /execute with {cell_id, code} → streams output via SSE.
  • POST /interrupt → cancels current cell.
  • POST /shutdown → graceful exit; we then call mutation podStop to free GPU.

Idle reaper at backend/services/kernel_manager.py:36-38 ports over directly — just call pod.stop() instead of sandbox.terminate.aio().

The kernel proxy script in kernel_manager.py:85-219 becomes a small FastAPI app baked into the pod image, replacing the stdin-JSON protocol with HTTP. Same outputs, different transport.

5. Storage abstraction

This is the trickiest piece because Modal lets the API server read the volume directly; RunPod doesn't.

Two options, decide per-deployment:

(a) "File gateway" worker — a tiny RunPod Serverless endpoint that wraps the network volume and exposes read/list/write handlers. The backend's StorageBackend impl calls those handlers. Pro: keeps storage local to RunPod's network. Con: extra hop, must keep the gateway warm-ish.

(b) S3 mirror — treat the network volume as scratch only, mirror writes to S3 (which services/s3_client.py already supports). The backend reads exclusively from S3. Pro: simpler API. Con: doubles write IO inside the sandbox; eventual consistency.

Recommendation: (a) for v1 because (b) requires changing how the SDK preamble writes files (currently writes to /data directly). The gateway worker is ~50 lines.

Either way, backend/routers/files.py:40-80 becomes:

storage = get_backend().get_storage()
return await storage.listdir(path, recursive=...)

6. Cost tracking

backend/services/usage.py:51-62 — split the pricing table by provider:

_SANDBOX_PRICING: dict[str, dict[str, float]] = {
    "modal": { "cpu": ..., "T4": ..., "H100": ... },
    "runpod": { "cpu": ..., "RTX4090": ..., "H100": ..., "MI300X": ... },
}

compute_sandbox_cost(seconds, gpu, provider) looks up the right table.

record_sandbox_usage at L209-269 already takes a provider — currently hard-coded to "modal" at L248. Change to use backend.name. The DB column already exists (UsageEvent.provider, models.py:236).

Bonus: RunPod's job status response includes executionTime in seconds — use that as the source of truth for billable time on RunPod and fall back to wall-clock only on parse failure. Keeps cost numbers honest when the worker has cold-start overhead.

7. Auth / secrets

Add to .env.example:

SANDBOX_BACKEND=modal
RUNPOD_API_KEY=
RUNPOD_ENDPOINT_ID_CPU=
RUNPOD_ENDPOINT_ID_GPU=
RUNPOD_POD_TEMPLATE_ID=
RUNPOD_NETWORK_VOLUME_ID=

docker-compose.yml:40,51 — pass these through the same way Modal vars flow today.

8. Tests / dev mode

Today's MockVolume (backend/tests/conftest.py:204-228) becomes a MockStorageBackend implementing the new StorageBackend protocol. Add a MockSandboxBackend that runs code in a local subprocess — gives us a true dry-run mode (SANDBOX_BACKEND=mock or SANDBOX_BACKEND=local) for CI and local dev where neither Modal nor RunPod credentials are present. This is independently valuable.

9. UI exposure

frontend/ — extend the project settings UI (where sandbox profiles are edited today, if any) with a backend selector: modal | runpod. Stored in Project.sandbox_config.backend. Default for new projects: whatever settings.sandbox_backend is.

If no UI for sandbox config exists yet, this can be deferred — the backend-level switch via env is enough for v1.

Rollout plan

Phase 1 — refactor only (no behavior change):

  • Extract SandboxBackend / StorageBackend protocols in services/sandbox/base.py.
  • Move existing Modal code into services/sandbox/modal_backend.py and services/sandbox/modal_storage.py behind the new interface.
  • Update call sites in tools/execute_code.py, tools/run_notebook_cell.py, routers/files.py, services/usage.py to go through get_backend().
  • All tests still green; default backend remains Modal.

Phase 2 — RunPod one-shot:

  • Build the worker image (backend/runpod_worker/{Dockerfile,handler.py,requirements.txt}).
  • Implement RunpodBackend.run_code using the Serverless REST + stream API.
  • Document the deploy step (build worker image, push to registry, create endpoints, set env vars).
  • Smoke test: SANDBOX_BACKEND=runpod runs execute_code end-to-end on a single endpoint.

Phase 3 — RunPod storage gateway:

  • Add the file-gateway endpoint + handler.
  • Implement RunpodStorage against the gateway.
  • routers/files.py works on RunPod.

Phase 4 — RunPod persistent kernel:

  • Pod-based kernel manager.
  • Replace stdin/stdout proxy with HTTP/SSE.
  • Idle reaper + cost-aware shutdown (more aggressive than Modal because pods bill hourly).

Phase 5 — Mock backend + UI:

  • MockSandboxBackend for CI / local dev with no Modal/RunPod account.
  • Per-project backend selector in the UI.

Acceptance criteria

  • SANDBOX_BACKEND=modal (default) reproduces today's behavior bit-for-bit. No regressions in execute_code, kernel, file routes, cost tracking.
  • SANDBOX_BACKEND=runpod runs a "hello world" code execution end-to-end — code submitted, stdout streams to the chat, file written to /data is readable via /api/files/raw.
  • SANDBOX_BACKEND=runpod with gpu="A100" actually lands on an A100 worker.
  • Kernel manager works on RunPod with the same run_notebook_cell agent prompt — no agent-side changes needed.
  • UsageEvent.provider is correctly set per-execution ("modal" vs "runpod"); cost figures are within ±5% of the provider's billing dashboard.
  • Idle reaper kills RunPod pods within 5 min of inactivity (vs Modal's 15 min) to avoid runaway hourly billing.
  • SANDBOX_BACKEND=mock lets the test suite run with no provider credentials.
  • Per-project backend override works: setting Project.sandbox_config.backend = "runpod" routes that project's executions to RunPod even when the global default is Modal.

Out of scope

  • Hyperscaler backends (AWS Batch, GCP Cloud Run) — separate epic.
  • Live migration of an in-flight session between backends.
  • Cross-backend cost optimization (auto-pick cheapest backend for a given GPU at request time).
  • Image build pipeline UX — for v1 the worker image is rebuilt and pushed manually; automated build/push is a follow-up.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions