Context

API reference for urun.Context — distributed execution context via dependency injection.

Public API at a Glance

urun.Context

class urun.Context

Distributed execution context, injected via type annotation. Add ctx: urun.Context as a parameter to any @app.function and urun injects the context automatically at runtime.

import urun

@app.function(gpus="h100:8")
def world(ctx: urun.Context):
    print(f"Rank {ctx.rank} of {ctx.world_size}")
    model = model.to(ctx.device)  # each rank holds its shard of the pipeline

Use urun.Context as the primary interface for distributed execution info and streaming pipeline construction inside GPU functions.

Properties

All properties are read-only. Values are determined by the runtime environment and GPU allocation.

ctx.session_live()

while ctx.session_live():
    step()

Cooperative, lockstep "should this session keep running?" gate — call once per frame in a session loop. On a multi-rank group it broadcasts the rank-0-owned liveness flag over the SPMD communicator, so every rank ends the session on the same frame (no desync). On world_size == 1 it is a pure local read. Returns False once the session has been asked to stop (lease lapse, authoritative close, consumer abandonment).

ctx.completing_work()

async with ctx.completing_work(max_seconds=30):
    image = await pipeline.generate(prompt)
    await ctx.stream("out").send(image)   # emits → DVR records → S3

Mark a bounded discrete work item (one image render, one generation) so a lost browser leg cannot discard its result mid-flight. While the guard is open, a leg-lost verdict (browser disconnected) is deferred so the item finishes and streams — the DVR records it, and a later reconnect replays the result with zero GPU work. Platform-terminal verdicts (lease lapse, authoritative close) are never deferred, and max_seconds bounds the guard so a hung item ages out. Usable as with or async with.

ctx.run_all()

await ctx.run_all([p1, p2, p3])   # canonical; varargs also accepted

Run several pipelines concurrently and keep them all alive until session end — the multi-pipeline counterpart of await pipeline.run(). Returns when every pipeline's run() returns (session disconnected); propagates the first failure. Never hand-roll asyncio.gather for this.

ctx.channel()

ctx.channel(name: str) -> EventStream

Create or retrieve a session-scoped bidirectional channel — the control/data-oriented alias of ctx.stream(name) (same namespacing, same transport).

ctx.requests()

rpc = ctx.requests()          # bound to ctx.doc(); requests_key="requests", responses_key="responses"

@rpc.on_request("echo")
def echo(req):
    return {"echoed": req.payload}

rpc.start()                   # observes until the session ends (auto-stopped at teardown)

The standard doc-based request/response surface: a client writes a request register under the requests key of the session doc; the responder dispatches by kind and writes the correlated reply under responses of the same doc. No new transport — it rides the ctx.doc() CRDT primitive. Kwargs: key= (which session doc), requests_key=, responses_key=.

ctx.rpc()

async for req in ctx.rpc("generate", kind="image", codec="jpeg"):
    await req.stream.emit(image)     # ndarray/PIL → one binary JPEG turn
    await req.respond(body)          # or req.error("...") on failure

The async-iterator request/response primitive the TS client's session.request / requestStream / complete speaks. Replies ride the request's own response stream — responses never touch the doc. Legacy adapter: apps should never keystone on ctx.rpc; prefer ctx.requests() and doc desired-state, with outputs on streams.

ctx.serve()

@app.function(gpus="l4:1", lease=14400)
async def serve(ctx: urun.Context):
    await ctx.serve("qwen-coder:fp8@l4:1")

Serve a model-catalog model — the whole serving loop inside a normal @app.function body: resolves the spec against the catalog at boot, loads the engine, subscribes to the shared request registry on ctx.doc("llm"), and streams token deltas + the final OpenAI-compatible body to each requester's own consumer-addressed stream (request A's tokens never reach consumer B). spec is a model string "<id>[:<variant>][@<gpu>]" or a resolved ServeSpec.

ctx.input()

driven = ctx.input(field="input", hz=30.0) | ctx.window(...) | ctx.pipe(step) | ...
# or imperatively:
async for sample in ctx.input():
    ...

Continuous ephemeral input (held keys, pointer counters) published by the browser into the control doc's awareness (via @urun-sh/react useInputPresence), sampled at hz as per-sample deltas. Usable imperatively or as a pipeline head; starves when input is quiet. Split rule: continuous ephemeral input rides awareness; discrete semantic events ride the doc or named data streams.

ctx.recordings, ctx.artifacts, ctx.session_state

ctx.stream()

ctx.stream(
    name: str = "default",
    *,
    format: str | None = None,
    codec: str | None = None,
    audio_codec: str | None = None,
    kind: Literal["audio", "video", "data"] | None = None,
    fps: int | None = None,
    gop: int | None = None,
    bitrate: int | None = None,
    host: str = "",
    port: int = 0,
    rank: int | None = None,
    rank_policy: str | None = None,
    broadcast: bool = False,
    record: bool = True,
    retention: str | None = None,
    delivery: Literal["paced", "latest"] = "paced",
    direction: Literal["auto", "in", "out"] = "auto",
) -> EventStream | SinkStage

Create a session-scoped stream. Has two modes depending on parameters:

Data stream mode (no codec/bitrate): Returns an EventStream for bidirectional JSON/event data. Session streams are automatically namespaced to the current execution session.

Pipeline sink mode (codec or bitrate set): Returns a SinkStage that encodes and streams media (e.g., H.264 video via RTP). Chainable with the | operator as the final stage in a pipeline.

Returns: EventStream (data mode) or pipeline sink stage (codec/bitrate mode).

# CRDT document — bidirectional state (see ctx.doc())
control = ctx.doc("control")
listener = control.listen(keys=["prompt"])
while ctx.connected:
    change = listener.wait(timeout=0.1)
    if change:
        handle(change.snapshot.get("prompt"))

# Pipeline sink — video output with encoding
video = ctx.stream("video", codec="h264", bitrate=8_000_000, fps=24, rank=1)

pipeline = (
    ctx.source(denoise, device="cuda", rank=0)
    | ctx.pipe(vae_decode, device="cuda", rank=1)
    | ctx.pace(fps=24, rank=1)
    | video
)

ctx.doc()

ctx.doc(key: str)

Return one keyed replicated document for the current session. Documents are CRDT-backed (Yjs) and automatically synced from the browser via transport.sendDoc(key, patch). The key is app-defined — urun does not reserve any keys.

@app.function(gpus="b200:4")
def render(ctx: urun.Context):
    control = ctx.doc("control")
    prompt = control.get("prompt.text", "a cinematic landscape")
    keyboard = control.get("keyboard.keys", {})
    camera = control.get("camera.pose", {})

For reactive loops:

control = ctx.doc("control")
listener = control.listen(keys=["keyboard", "camera"])

while ctx.connected:
    change = listener.wait(timeout=0.016)
    if change is not None:
        latest = change.snapshot
        update_camera(latest["keyboard"]["keys"], latest["camera"]["pose"])

ctx.docs

ctx.docs

Process-local session document registry. Use for multi-session worker code that observes all active sessions. Normal handlers should use ctx.doc(key).

Streaming pipeline

The pipeline builder methods on ctx create a composable, rank-aware streaming pipeline. Stages chain with the | operator and automatically insert IPC transport at rank boundaries.

ctx.source()

ctx.source(
    fn: Callable,
    name: str | None = None,
    device: str = "cuda",
    priority: int = 0,
    rank: int | tuple[int, ...] | None = None,
    rank_policy: str | None = None,
) -> SourceStage

Create the head of a pipeline. fn is a generator that yields items. It receives stop_event and config as keyword arguments — config is updated live when pipeline.update(prompt=...) is called.

def denoise(*, stop_event, config):
    prompt = config.get("prompt", "default scene")
    while not stop_event.is_set():
        latents = model.generate(prompt)
        yield latents
        prompt = config.get("prompt", prompt)

pipeline = ctx.source(denoise, name="denoise", device="cuda", rank=0)
pipeline.update(prompt="A sunset over the ocean")

ctx.pipe()

ctx.pipe(
    fn: Callable,
    name: str | None = None,
    device: str = "cpu",
    priority: int = 0,
    rank: int | None = None,
    rank_policy: str | None = None,
    max_depth: int = 8,
    backpressure_at: int = 6,
) -> Stage

Create a processing stage. Calls fn(item) on each pipeline item and forwards the result downstream.

When a stage runs on a different rank than its predecessor, urun automatically inserts shared-memory IPC transport with CUDA IPC handles for zero-copy GPU tensor transfer.

pipeline = (
    ctx.source(denoise, device="cuda", rank=0)
    | ctx.pipe(vae_decode, name="vae", device="cuda", rank=1)
    | ctx.pipe(encode_h264, name="encode", device="cuda", rank=1, priority=-1)
)

ctx.tap()

ctx.tap(
    fn: Callable,
    name: str | None = None,
    device: str = "cpu",
    priority: int = 0,
    rank: int | None = None,
    rank_policy: str | None = None,
) -> TapStage

Create a side-effect stage. Calls fn(item) but does not forward a return value downstream — the original item passes through unchanged. Use for logging, metrics, saving snapshots, or any observation that should not alter the pipeline data.

def log_frame(item):
    print(f"Frame generation={item.metadata.prompt_generation}")

pipeline = (
    ctx.source(generate, device="cuda")
    | ctx.tap(log_frame, name="metrics")
    | ctx.pipe(encode, device="cuda")
    | video
)

ctx.pace()

ctx.pace(
    fps: float = 60.0,
    name: str = "pace",
    rank: int | None = None,
    rank_policy: str | None = None,
    max_depth: int = 256,
    backpressure_at: int = 192,
    min_buffer_frames: int = 0,
    target_buffer_frames: int | None = None,
    max_buffer_frames: int | None = None,
) -> PacedStage

Emit items at a fixed frame rate with adaptive buffering. Smooths bursty upstream output into steady playback.

The pacer adapts dynamically: it grows the buffer when frame underruns are detected and shrinks it when the upstream is keeping up. This prevents both stalls (too few buffered frames) and latency creep (too many).

pipeline = (
    ctx.source(denoise, device="cuda", rank=0)
    | ctx.pipe(vae_decode, device="cuda", rank=1)
    | ctx.unbatch(rank=1)
    | ctx.pace(fps=24, rank=1, min_buffer_frames=4, target_buffer_frames=8)
    | video
)

ctx.unbatch()

ctx.unbatch(
    name: str = "unbatch",
    rank: int | None = None,
    rank_policy: str | None = None,
) -> UnbatchStage

Receive a list/batch from upstream, emit each element individually downstream. Preserves pipeline metadata (prompt generation, flush markers) across the unbatch boundary.

Use when an upstream stage produces batches (e.g., a VAE decoding a window of latent frames into multiple video frames).

ctx.window()

ctx.window(
    window: int,
    stride: int | None = None,
    idle: Any = ...,            # pad value when the head has nothing buffered
    name: str = "window",
    rank: int | None = None,
    rank_policy: str | None = None,
    max_depth: int = 256,
    backpressure_at: int = 192,
) -> WindowStage

Clips/buffers the inbound stream into windows — the streaming dual of ctx.pace. Where the pacer spaces items out in time, the window accumulates window inbound items and emits them as one clip every stride (default non-overlapping). Insert it after an inbound ctx.stream(...) for a model that wants windows rather than the raw stream:

pipeline = (
    ctx.stream("mic")
    | ctx.window(window=24)
    | ctx.source(make_model(model), device="cuda")
    | ctx.pace(fps=24)
    | ctx.stream("captions")
)

idle= supplies a pad value used when the inbound head has nothing buffered (idle-park). A pattern that seems to need hand-rolled buffering is a missing windowing stage — use ctx.window, don't roll your own loop.

ctx.self_test()

with ctx.self_test():
    ...

A context manager — the one end-to-end gate. It is self-driving: the app's block drives one representative pass through the composed pipeline (with pipeline.drive(frames=N)) and asserts its own success criteria. That single block simultaneously compiles every component (precompile), warms caches/allocator (warmup), and validates (self-test).

with ctx.self_test():
    out = pipeline.drive(frames=1)      # the app drives its own representative pass
    _assert_frames_shaped(out)          # the app asserts its own criteria — shape/dtype, no value read

Mandatory-gate semantics are the block's natural exception flow:

• Clean exit = validated. A passed self-test is recorded (the readiness gate reads it) and, in prod, the weight-excluded snapshot is triggered at this validated point, then serving proceeds.
• A raise inside the block = self-test failed. A failed report is recorded, readiness is forced not-open (nothing serves, nothing is snapshotted), the last-validated snapshot keeps serving, and the exception propagates so the deploy/boot path sees the real failure.

The same block is your CPU test's pass/fail check: under MockContext with the mocked store, a clean exit passes and a raise fails. See Strict mode & CPU testing.

Reactive control

A live control doc steers a running pipeline in exactly two ways — cheap config is read inline; structural change reconstructs a stage. There is no field-mapping table.

Cheap / live config (inline read)

doc.desired.<prop> is a live reactive accessor over the CRDT: its .value reads the current snapshot each time. A stage reads it inline at a chunk boundary — no machinery, no bind:

def make_denoise(models, doc):
    def denoise(*, stop_event, config):
        while not stop_event.is_set():
            yield models.denoise(doc.desired.prompt.value)   # live read each chunk
    return denoise

doc.bind()

doc.bind(
    construct: Callable[[Doc], Stage],
    deps: list[LiveProp] | None = None,
) -> Stage

useMemo for stages. construct is (doc) -> stage (it builds a composable stage, reading what it needs from the doc); deps is a list of live doc properties ([doc.desired.vae.model]). bind evaluates construct(doc) once and, when any dep property value changes (shallow compare, like useMemo), re-runs construct → reconstructs the stage → the pipeline drains and swaps it at a safe boundary. It returns a composable stage, so it slots straight into the | chain:

doc = ctx.doc("control")
pipeline = (
    ctx.stream("mic") | ctx.window(window=24)
    | ctx.source(make_denoise(models, doc))                                   # cheap reads inline
    | doc.bind(lambda d: ctx.pipe(make_vae(models, d)), [doc.desired.vae.model])  # useMemo
    | ctx.pace(fps=24) | ctx.stream("video", codec="h264")
)
await pipeline.run()

Use doc.bind ONLY for structural change (model / resolution / topology) that needs a rebuild; cheap config stays an inline live doc.desired.* read.

Auto tensor conversion: when a pipeline stage outputs a float tensor (NHWC, values in [0,1]), the ctx.stream sink auto-converts it to encoded frames before sending. The app emits a tensor and never touches PIL/numpy itself.