unionai-sandbox 0.0.1b2

Local and remote code-execution sandbox for Flyte / Union workflows.

unionai-sandbox

A small, embeddable code-execution sandbox for Union / Flyte workflows. One Python package, two transports:

• union.sandbox.local — runs sandboxed children inside the current container via a Rust extension. On Linux it picks the best available isolation (bubblewrap → userns → none); on macOS it prefers sandbox-exec and otherwise falls back to none.
• union.sandbox.remote — Connect-RPC client to a long-running sandbox-server process elsewhere. Used by union.sandbox.SandboxSession, which launches a sandbox-server pod via a Flyte task and connects to it.

Both speak the same session contract: open → many run() / put_bytes / get_bytes calls → close.

Process output helpers:

• await proc.communicate() returns (stdout_bytes, stderr_bytes)
• await proc.communicate_text() returns decoded (stdout, stderr)
• async for line in proc.iter_stdout_lines(): ... streams decoded lines

Quickstart

Local (in-process)

import asyncio
import flyte
from union import sandbox as sb

async def main():
    async with sb.local.session(resources=flyte.Resources(cpu="500m", memory="512Mi")) as sbx:
        proc = await sbx.run("python3 -c 'print(2+2)'", stdout=True, timeout=10)
        out, _ = await proc.communicate_text()
        print(out)

asyncio.run(main())

Remote (sandbox-server pod, launched as a Flyte task)

Deploy the default sandbox task envs once per cluster:

flyte deploy --all python/union/sandbox/task/_server.py

Then in your code:

from datetime import timedelta
from union import sandbox as sb

async with await sb.session(timeout=timedelta(minutes=30)) as sbx:
    proc = await sbx.run("uname -a", stdout=True)
    out, _ = await proc.communicate_text()
    print(sbx.name, sbx.ip, sbx.created_at)

sb.session() submits the deployed union-sandbox task as a Flyte run, waits for the pod to be addressable, opens a Connect-RPC transport, and returns a SandboxSession. Closing the session aborts the run.

Note on defaults: the local transport auto-detects the strongest backend available on the current Linux host, preferring bubblewrap. The remote sb.session() path is separate: if you do not pass environment=, it launches sb.DEFAULT_SANDBOX_ENV, which is the userns-backed remote environment. Use environment=sb.DEFAULT_SANDBOX_ENV_BWRAP (or a custom SandboxEnvironment) if you want a bwrap-backed remote pod.

SandboxEnvironment

SandboxEnvironment is a declarative spec for the sandbox-server pod the session runs in. The default is registered for you; build a custom one if you need extra packages, a tighter resource budget, secrets, or a different isolation mode.

import flyte
from union import sandbox as sb

my_sandbox = sb.SandboxEnvironment(
    name="ml-sandbox",
    image=sb.base_sandbox_image.with_pip_packages("torch", "transformers"),
    resources=flyte.Resources(cpu="8", memory="32Gi", gpu="L4:1"),
    secrets=[flyte.Secret(group="hf", key="HF_TOKEN")],
    env_vars={"HF_HOME": "/tmp/hf"},
    sandbox_mode="userns",     # "userns" | "bwrap"
    runtime="container",       # "container" | "gvisor"
)

# Deploy: `flyte deploy --all path/to/your_module.py` finds
# `my_sandbox.task_env` and registers the task.

async with await sb.session(environment=my_sandbox) as sbx:
    ...

Param	Type	Notes
name	str	Task-env identifier; session resolves {name}.sandbox_server.
image	flyte.Image	Defaults to base_sandbox_image. Extend via base_sandbox_image.with_*().
resources	flyte.Resources	Default per-session resources. Override per launch with sb.session(resources=...) — the platform task spec is rewritten via TaskDetails.override(resources=...), and the in-pod sandbox cgroup ceiling is derived from the same value (minus a small supervisor reserve).
secrets	list[flyte.Secret]	Forwarded to the TaskEnvironment.
env_vars	dict[str, str]	Forwarded.
include	tuple[str, ...]	Globs of source files to include.
interruptible	bool	Allow preemptible nodes.
description	str	Free text shown in the UI.
sandbox_mode	"userns" | "bwrap"	"bwrap" adds SYS_ADMIN + AppArmor unconfined.
runtime	"container" | "gvisor"	"gvisor" sets runtimeClassName: gvisor on the pod.

The pod_template is not exposed — it's derived from sandbox_mode and runtime.

The repo also exports two ready-built defaults:

• sb.DEFAULT_SANDBOX_ENV — userns isolation, container runtime, name union-sandbox (used when you call session() with no environment=).
• sb.DEFAULT_SANDBOX_ENV_BWRAP — bubblewrap isolation, container runtime, name union-sandbox-bwrap.

SandboxSession API

The SandboxSession returned by sb.session() is a dataclass (serialisable across Flyte task boundaries via mashumaro).

Fields (mashumaro-serialised)

Field	Type	Meaning
name	str	Session name. Equals the Flyte run name.
endpoint	str	HTTP(S) URL the transport opens against. Empty until the owner has waited for the pod to surface.
ip	str	Pod IP, when surfaced by the runtime; otherwise empty.
created_at	datetime	UTC timestamp of construction.
network_mode	"blocked" | "open" | "allowlist"	Default network posture applied to every run().
network_allowlist	list[str]	Default CIDR / DNS-pattern allow-list, used only with network_mode="allowlist".

Lifecycle

Bringup is split into two phases so user setup work can overlap with pod startup:

sbx = await sb.session(...)             # submit run; returns instantly.
                                         # A background task watches for
                                         # phase=RUNNING + a pod IP.

# ... do other work here while the pod comes up ...

async with sbx:                          # awaits the endpoint task. After
                                         # this, sbx.endpoint / sbx.ip are
                                         # populated. NO Health-check yet.
    proc = await sbx.run("uname -a")     # first send: Health-checks the
                                         # transport (exponential backoff,
                                         # usually first try), then sends.
                                         # Subsequent sends just send.
# context exit closes the local transport and aborts the run (owner side).

await sbx is equivalent to async with sbx: for phase 1: both wait for the pod to be addressable. The Health-check is intentionally deferred to the first run / put_bytes / get_bytes call, so the time between "session created" and "ready to send" overlaps with whatever the caller does next.

Methods

# Run a command in the sandbox.
proc = await sbx.run(
    "python3 -c 'import os; print(os.uname())'",
    stdout=True,           # PIPE | INHERIT | DEVNULL | False
    stderr=True,
    env={"FOO": "bar"},
    cwd="/tmp",            # under the sandbox work dir
    script_type="shell",   # "shell" | "python"
    network_mode="allowlist",
    network_allowlist=["pypi.org", "*.pythonhosted.org"],
)
out, err = await proc.communicate()
# proc.returncode, proc.runtime_ms, proc.backend, proc.termination_reason
#
# communicate() / wait() raise SandboxExecutionError if the process never
# ran to a real exit — a server-side spawn failure or a stream that died
# before termination (the cases that used to surface as a silent rc=-1).
# A non-zero exit of *your* code (any code, incl. 128+signo for a signal)
# is not an error: it returns normally and you branch on proc.returncode.
#
#   from union.sandbox import SandboxExecutionError
#   try:
#       out, err = await proc.communicate()
#   except SandboxExecutionError as e:
#       print("could not run:", e.reason)   # + e.returncode / e.backend

# Push raw bytes into a file under the sandbox work dir.
n = await sbx.put_bytes("/tmp/sandbox-work/input.json", b'{"x": 1}')

# Pull raw bytes back out.
data = await sbx.get_bytes("/tmp/sandbox-work/result.json", max_bytes=10 * 1024 * 1024)

# Inspect.
sbx.name        # str
sbx.endpoint    # str
sbx.ip          # str  ("" if not yet known)
sbx.created_at  # datetime (UTC)
sbx.is_owner    # bool — True iff this side created the run (can abort it)
sbx.url         # Optional[str] — Flyte console URL when owner

Owner vs reference mode

Only the owner side (the task that called sb.session()) can abort the run. When you pass a SandboxSession as an input to another task, mashumaro carries the dataclass fields (name, endpoint, ip, created_at, network_mode, network_allowlist) across the boundary; the in-memory _run / _client / _endpoint_task are not serialised. The receiver lands in reference mode:

@env.task
async def child(sbx: sb.SandboxSession):
    # No `async with` needed — endpoint + ip already round-tripped via
    # mashumaro, and the first .run() lazily opens the local transport.
    proc = await sbx.run("uname -a", stdout=True)
    out, _ = await proc.communicate()

close() on the receiver only shuts the receiver's transport; the run keeps going until the owner aborts.

Network policy

A sandboxed call gets one of three network postures:

# 1. Blocked (default): fresh netns, only `lo`.
await sbx.run("curl -fsS https://example.com/")    # fails

# 2. Open: full host network.
await sbx.run("curl -fsS https://example.com/", network_mode="open")

# 3. Allow-list: shares host net, but a per-call HTTP CONNECT proxy
#    403s anything outside the list.
await sbx.run(
    "pip install requests",
    network_mode="allowlist",
    network_allowlist=["pypi.org", "*.pythonhosted.org", "10.0.0.0/8"],
)

Read the two knobs as:

• network_mode — the posture selector: "blocked", "open", or "allowlist".
• network_allowlist — the allow-list entries, used only with network_mode="allowlist".

network_allowlist constrains clients that honour HTTPS_PROXY (pip, curl, requests, boto3, huggingface_hub), not adversarial code. See docs/network-allow-list.md for the threat model.

Without a context manager

The SandboxSession returned by sb.session(...) doesn't require async with — you can keep the handle and own its lifetime. await it once to wait for the pod to surface (optional; the transport also opens lazily on the first run), then call await sbx.close() yourself (which closes the transport and aborts the run).

from union import sandbox as sb

sbx = await sb.session(timeout=timedelta(minutes=30))
await sbx   # wait for the pod to surface — fail fast on a bad launch
try:
    proc = await sbx.run("uname -a", stdout=True)
    out, _ = await proc.communicate_text()
finally:
    await sbx.close()

To attach to a sandbox-server you started yourself, call sb.remote.session(endpoint=...); for in-process, sb.local.session(...).

Installing

The Python wheel bundles the native extension. From a checkout:

make develop                  # creates .venv-build/, builds the Rust ext into it
.venv-build/bin/pip install -e ".[remote]"   # if you want the remote transport

For a release build of the server binary:

make server                   # produces target/release/sandbox-server

To cross-build the linux/amd64 binary that base_sandbox_image ships into Flyte:

make deploy-binary            # docker-runs rust:1-bookworm

What's in the repo

.
├── crates/
│   ├── sandbox-core/     Isolation backends (bubblewrap, userns, sandbox-exec, none)
│   ├── sandbox-server/   Connect-RPC server
│   └── sandbox-py/       PyO3 bindings → union.sandbox._native
├── python/
│   └── union/
│       └── sandbox/
│           ├── local/    LocalSession (Rust-backed)
│           ├── remote/   RemoteSession + generated stubs
│           └── task/     Sandbox-server packaged as a Flyte task
├── proto/sandbox/v2/sandbox.proto
├── examples/
│   ├── local/{apps,tasks}/   local-transport examples
│   └── remote/{apps,tasks}/  SandboxSession examples + Streamlit app
└── docs/
    ├── changes.md             Original proposal that drove v2
    ├── design.md              Architecture + threat model
    ├── sandbox-as-task.md     `sb.session(...)` lifecycle, dataproxy hop
    ├── network-allow-list.md  Per-host/CIDR allow-list design
    ├── filesystem-allow-list.md  Default FS allow-list + how to extend it
    ├── vs-sandbox-rs.md       Comparison to ErickJ3/sandbox-rs (Rust lib)
    └── vs-openshell.md        Comparison to NVIDIA OpenShell (agent sandbox)

Isolation backends

Backend	How it works	Unprivileged on K8s	Default?
bubblewrap	bwrap(1) with --unshare-all --die-with-parent --cap-drop ALL, plus a Landlock ruleset applied in pre_exec as a kernel-side backstop	Yes	Yes
userns	Direct unshare(2) + prctl(NO_NEW_PRIVS) + capset + setrlimit, plus Landlock and a seccomp BPF deny-list (mount/eBPF/keyring/kexec/raw I/O)	Yes	Fallback
sandbox-exec	macOS-only wrapper around Apple's deprecated sandbox-exec; restricts writes to the sandbox work dir and can deny outbound sockets	n/a	macOS auto
none	setpgid + best-effort setrlimit; logs a warning	n/a	macOS/Windows dev only

Backend selection for the local Rust-backed path: SANDBOX_BACKEND=<name> if set, else auto-detect.

Across the Linux backends, the sandbox layers four kernel-enforced control families. Both bubblewrap and userns use the first three; the current seccomp deny-list is userns-only:

• Namespaces (user / mount / UTS / IPC, plus net unless the caller opted in to host networking).
• Capabilities dropped from every set (bounding, ambient, effective/permitted/inheritable).
• Landlock (5.13+) — inode-level read-only access to /usr, /lib, /etc, /proc, /sys, /dev/null, …; read-write to /tmp, /dev/shm, the per-session work dir (and the NVIDIA device nodes when a GPU is requested). Built once in the parent (landlock_create_ruleset + landlock_add_rule per allowed path) and applied by the child via a single async-signal-safe landlock_restrict_self syscall in pre_exec. Older kernels gracefully degrade. Callers can add paths to this allow-list (never replacing the defaults) via read_only_paths / read_write_paths on sb.local.session(...) — see docs/filesystem-allow-list.md.
• Seccomp BPF deny-list (userns backend only — bubblewrap's own setup needs mount/unshare, so it'd need bwrap's --seccomp FD path; that's a follow-up). Pre-compiled with seccompiler, cached in a OnceLock, installed via one prctl(PR_SET_SECCOMP) in pre_exec. Denies syscalls with no legitimate purpose in a sandbox: mount / namespace control, eBPF, kernel keyring, kexec, raw port I/O, file-handle reopen, perf counters, etc. Returns EPERM, not SIGSYS.

Building and testing

make test              # Rust workspace + Python suites
cargo test             # just the Rust crates
make py-test           # just the Python tests
make examples-local    # run every local-task example end-to-end
make deploy-sandbox    # flyte deploy --all python/union/sandbox/task/_server.py

License

Business Source License 1.1 (BSL). Each released version converts to the Apache License 2.0 on its Change Date (four years after release). See LICENSE and the licenses/ directory for the full BSL and Apache-2.0 texts.

Releasing

See RELEASING.md. Releases are cut from the GitHub UI (Actions → Release → Run workflow); the PyPI upload uses tokenless Trusted Publishing from the release environment.