unionai-sandbox 0.0.1b8

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.on_device — runs sandboxed children inside the current process/container via a Rust extension. You pick the isolation backend explicitly — session(backend="bubblewrap") (default, strongest), "userns" (lite), or "none" (dev only). There is no auto-detection and no silent downgrade: if the chosen backend isn't available here, run() fails loudly. Use this to embed a sandbox in any Python process.
• 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

Standalone (embed in any Python process)

The on-device transport runs the sandbox in your current process — no Flyte run, no server pod. Choose the backend explicitly (bubblewrap is the default and needs CAP_SYS_ADMIN; use userns on a vanilla host, or none for unisolated dev):

import asyncio
from union import sandbox as sb

async def main():
    async with sb.on_device.session(backend="userns") as sbx:
        proc = await sbx.run("python3 -c 'print(2+2)'", stdout=True)
        out, _ = await proc.communicate_text()
        print(out)  # -> 4

asyncio.run(main())

Resources, the network policy, and the filesystem allow-list are per session:

import flyte
async with sb.on_device.session(
    backend="bubblewrap",
    resources=flyte.Resources(cpu="2", memory="2Gi"),
    network_mode="allowlist",                 # "blocked" (default) | "open" | "allowlist"
    network_allowlist=["pypi.org", "*.pythonhosted.org"],
    read_only_paths=["/opt/models"],          # add to the FS allow-list (read)
) as sbx:
    ...

With Flyte

1. On-device inside a task — the sandbox runs in the task pod (no extra pod). For backend="bubblewrap" the pod needs CAP_SYS_ADMIN + unconfined AppArmor, granted by flyte.PodTemplate.allow_nested_sandboxing(); otherwise use backend="userns", which works in a vanilla pod.

import flyte
from union import sandbox as sb

env = flyte.TaskEnvironment(name="my-task", image=sb.base_sandbox_image)

@env.task
async def main() -> str:
    async with sb.on_device.session(backend="userns") as sbx:
        proc = await sbx.run("uname -a", stdout=True)
        out, _ = await proc.communicate_text()
        return out

2. Remote sandbox-server pod — sb.session() launches a dedicated sandbox-server pod as a Flyte run and connects over Connect-RPC. Deploy the default task envs once per cluster:

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

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 the transport, and returns a SandboxSession. Closing it aborts the run. Without environment= it launches sb.DEFAULT_SANDBOX_ENV (the userns remote env); pass sandbox_mode="bwrap" or environment=sb.DEFAULT_SANDBOX_ENV_BWRAP for bubblewrap, or build a custom SandboxEnvironment and call env.new_session(...).

Runnable versions of all of the above are in examples/on_device/ and examples/remote/ — hello, network allow-list, put/get bytes, and PyTorch (CPU + GPU).

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 (see below).
sys_cap_admin	bool | None	Explicit override of the CAP_SYS_ADMIN grant. None (default) = granted iff sandbox_mode="bwrap"; True = always; False = never (even for bwrap).
runtime	"container" | "gvisor"	"gvisor" sets runtimeClassName: gvisor on the pod.

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

Pod security: CAP_SYS_ADMIN for the bubblewrap backend

The bubblewrap backend runs as a non-root user via unprivileged user namespaces (no privileged: true, no setuid). But the containerd default seccomp profile only permits the mount / pivot_root / setns / unshare syscalls bwrap needs when the container's capability set includes CAP_SYS_ADMIN, and the default AppArmor profile must be unconfined so those calls aren't blocked. sandbox_mode="bwrap" (or sys_cap_admin=True) makes the deployed sandbox-server pod carry exactly this — and nothing more (allowPrivilegeEscalation: false, no other caps).

The SandboxEnvironment emits it for you; if you build the pod yourself, this is the equivalent Kubernetes spec (the primary container is named primary in Flyte 2):

apiVersion: v1
kind: Pod
metadata:
  annotations:
    # AppArmor must be unconfined or the default profile blocks bwrap's
    # mount/pivot_root/unshare. (K8s 1.30+ alternative: the
    # `securityContext.appArmorProfile` field below.)
    container.apparmor.security.beta.kubernetes.io/primary: unconfined
spec:
  containers:
    - name: primary
      securityContext:
        capabilities:
          add: ["SYS_ADMIN"]          # only cap added — needed for the seccomp profile
        allowPrivilegeEscalation: false
        # K8s >= 1.30 equivalent of the annotation above:
        # appArmorProfile:
        #   type: Unconfined

The userns backend needs none of this — it runs in a vanilla unprivileged pod. Use sys_cap_admin=False to deploy a bubblewrap env on a cluster that already permits unprivileged user namespaces without the cap (keeps the pod's host attack surface minimal), or sys_cap_admin=True to grant it for a userns env on a cluster whose seccomp profile blocks the userns-lite syscalls.

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.on_device.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/
│           ├── on_device/  OnDeviceSession (Rust-backed)
│           ├── remote/   RemoteSession + generated stubs
│           └── task/     Sandbox-server packaged as a Flyte task
├── proto/sandbox/v2/sandbox.proto
├── examples/
│   ├── on_device/{apps,tasks}/  on-device-transport examples (sb.on_device)
│   └── 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 is always explicit — sb.on_device.session(backend="…") (default "bubblewrap"), or the sandbox-server --backend flag. There is no env var and no auto-detection: a requested backend that isn't available here fails loudly rather than downgrading to weaker isolation.

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.on_device.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.