[REA-5060] Give a model a way to run on more than one GPU

[Orion-Zheng]

The abstraction and its tests. End-to-end validation against a real multi-GPU
model happens against a published release rather than in this PR, since a model
pins a runtime version rather than a branch.

Why

Some models do not fit on one GPU, or fit and miss their latency target. The
usual answer is SPMD: run N copies of the program, one per device, and let them
split the forward pass between themselves over a process group. The usual
launcher for that is torchrun, and it assumes your program is a script that
runs main() and exits.

A model here is not a script. It is a long-lived server with one event loop, one
WebRTC connection per client, a session lifecycle, inbound commands that mutate
state mid-stream, and an output stream paced to a frame rate. You cannot start
four copies of that: you would get four servers contending for one port, four
sessions, and four sets of frames racing to the same client.

So there was no way to write a multi-GPU model against this runtime at all. This
adds one, by splitting the process into two roles rather than replicating it: the
model you already write stays single, owns the network and the session, and
touches no CUDA; underneath it sits a plain SPMD world of worker processes, each
pinned to one device, all running identical code.

The design is deliberately thin, and the two acts of restraint are the point.
The framework issues no collectives of its own — the process group is created
for you and then belongs entirely to the model, because collectives must be
entered by every rank in the same order and a framework that interleaved its own
would need to understand your parallelism strategy to stay out of the way. What
synchronises the ranks instead is that collecting one reply from every rank is
the only synchronization mechanism: when a call has heard from all N ranks, all
N have finished, including their writes. And frames never cross the command
channel — pickling a chunk of video through a queue costs tens of milliseconds,
so there is one shared-memory buffer that workers write into in place, and what
travels over the queue is an integer.

What Changed

Two commits. The first moves six modules in byte-identical, so it is the baseline
the rest is reviewed against and "nothing was lost in transit" is checkable
rather than asserted. The second adapts them to this repo. Reviewing them in that
order is much easier than reading the result cold.

A model subclasses DistributedWorker and implements setup / warmup /
start_session / generate_chunk / end_session; the framework runs one
instance per GPU in its own process with rank, world_size, device and a
frames handle already populated, and the process group already initialized. The
model holds a WorkerGroup, built in load(), and drives it from its own loop.
Every method broadcasts one verb to every rank and waits for the full reply set,
so a command succeeds only if it succeeded everywhere, and every blocking wait
polls liveness — a rank that dies silently surfaces as WorkerCrashed in seconds
instead of hanging in a collective.

generate_chunk may return an array of frames, an int end row, or None, and
the chunk's frame count is the maximum end row across ranks. That one choice is
what makes three different write patterns work with no coordination between
ranks: a leader assembling the whole chunk, ranks sharding the frames axis, or
every rank writing its own pixel band. Taking the maximum, all three come out
right, and no rank needs to know what the others did.

The two failure types differ in whether the worker managed to say anything.
WorkerError means a rank reported; it is recoverable only from
start_session, where all ranks rendezvous on the outcome and stay alive for a
retry. WorkerCrashed means a process died without a word, and is never
recoverable. Teardown follows from that: a healthy group gets the exit verb and
runs a clean barrier-and-destroy, while a group already broken by a fail-fast
failure skips the clean path entirely — a survivor entering the exit barrier
would wait forever on its dead peer — and goes straight to termination, with
anything still alive after its grace period terminated and then killed.

generate() blocks the caller for the chunk's compute time. That is deliberate,
not an oversight: it is what lets the emission rate follow measured throughput.
The practical consequence for an author is to keep each turn short — yield while
waiting, and do not park a long await inside one turn.

Adapting to this repo turned up one thing worth naming, because it would have
failed at runtime rather than in review: get_logger here returns a structured
logger whose level methods take fields as keywords and reject positional
arguments, and two calls passed them — both on failure paths, which is the worst
place to learn that a logging call raises. Those now use a static message plus
fields, and a spawned worker installs the runtime's own formatter so its lines
share the runtime's shape and carry rank as a field rather than a string
prefix.

torch stays an optional dependency. It is imported lazily, inside the functions
that need it, so importing the runtime pulls in nothing heavier than numpy and a
model image without torch is unaffected. That rests on discipline rather than
structure, so a test now pins it.

API surface

Five names, re-exported from the package root:

from reactor_runtime import DistributedWorker, WorkerGroup

class MyWorker(DistributedWorker):
    def setup(self, *, weights_path: str) -> None:  # ty: ignore[invalid-method-override]
        self.model = load_sharded(weights_path, device=self.device, shards=self.world_size)

    def start_session(self, params: dict) -> None:
        self.cache = self.model.allocate_cache(params["resolution"])

    def generate_chunk(self, index: int, controls: dict):
        latents = self.model.denoise(self.cache, controls)
        return self.decode(latents) if self.is_leader else None

    def end_session(self) -> None:
        self.cache = None          # keep self.model — that is the point


class MyModel(ReactorModel):
    def load(self, config_path: Path | None) -> None:
        self.workers = WorkerGroup(
            MyWorker,
            frame_shape=(12, 720, 1280, 3),        # worst-case chunk
            setup_kwargs={"weights_path": get_weights_path()},
        )
        self.workers.start()

    async def run(self) -> None:
        while True:
            await self.connected.wait()
            self.workers.start_session({"resolution": "720p"}, seed=self.seed)
            index = 0
            while self.connected.is_set():
                frames = self.workers.generate(index, {"prompt": self.prompt})
                await self.emit(MyOutput(main_video=frames))
                index += 1
            self.workers.end_session()

SharedFrameBuffer, WorkerError and WorkerCrashed come along for the authors
who write slices themselves or catch the failures.

Nothing needs a process group to be useful in tests: init_process_group=False
gives the same protocol with no torch involved.

A model is free to carve its own sub-groups — tensor-, sequence-, or
context-parallel — by calling new_group or init_device_mesh inside setup,
where every rank reaches the call in the same order. The framework never touches
the group after creating it, so there is nothing for a model's own collectives to
interleave with.

Verification

The ten protocol tests that came with the abstraction are ported unchanged, and
they are the reason "no behavioural change" is a claim rather than a hope. They
drive the real machinery — real spawned processes, real queues, real POSIX shared
memory — at world_size=2 with no torch installed, and cover every failure path:
retryable session init, fail-fast chunk failure, exactly one error posted per
failure, SIGKILL surfacing as WorkerCrashed, the shutdown escalation ladder,
and all three frame-write patterns. Green on 3.12 and 3.13 alongside the rest of
the suite.

Those tests all pass init_process_group=False, so they never form a group. That
gap is closed separately, since it cannot live in this suite: a probe forms a
real group and runs over gloo (torch present, CUDA absent) and then NCCL
on 4×B200, at world_size 2 and 4. At world_size=4 split into two halves, a
world all_reduce of ones returns 4 while the sub-group's returns 2 — a pair of
values only two genuinely distinct groups produce, which is what makes
"a model can build its own mesh on this" a demonstrated claim. The same run
confirms a thread created inside generate_chunk allocates on its own rank's
device, that host-memory pinning actually succeeds rather than silently degrading
to pageable copies, and that the clean-exit barrier and destroy complete without
hanging.

One thing surfaced there and is left alone on purpose: torch warns that
init_process_group is called without device_id, so the clean-exit barrier
infers its device from the ambient context. It is a warning, and passing
device_id would change what this code asks of torch, which is a decision rather
than something to slip into a move.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[tempusfrangit]

The engine is a good approach. Failure taxonomy, the max-end-row trick, the shared-memory path, liveness, the teardown ladder, the ten protocol tests are solid.

The changes I'd like to see are the author surface. We need to be careful as once the names are importable they are minimally an implicit contract and likely to be taken as an explicit contract.

Today multi-GPU is a second way to write a model (DistributedWorker plus a hand-driven WorkerGroup) alongside ReactorModel/ReactorPipeline. I'd rather it be a capability of the existing lineage. Here's the bar. Taking a single-GPU model to two GPUs should be a small in-place change: read self.rank/self.world_size, guard the emit with self.is_leader, shard a tensor.

Distributed by default, single-GPU is world_size=1:

• One class in the existing lineage. The runtime injects rank/world_size/device/is_leader, defaulting to a world of one. A model that never shards never reads them.
• world_size=1 runs in-process. No spawn, no shared-memory round-trip, today's latency. The spawn and shared-memory path is for world_size > 1 only.
• world_size comes from config/manifest, not hardcoded in load(). Device count is a deployment fact like fps. RuntimeConfig already reaches the runner whole, and config_path is the analogue.
• WorkerGroup, SharedFrameBuffer, and the protocol stay as internals the base class drives, off the package root. Exposing SharedFrameBuffer pins one transport. Exposing WorkerGroup commits the hand-driven shape.
• Mark the surface experimental so we can refine it without a break.

The base class is the paved road that covers essentially every model. The raw primitives stay available as experimental internals for the rare model that genuinely needs to hand-drive. That demotes them from "the API" to "the escape hatch" instead of removing the capability.

Think about DX ergonomics when deciding to either keep generate_chunk(index, controls), or converge on inference() so single-GPU is literally today's pipeline at world_size=1. I lean inference(), though there's a real trade with the shared-memory writes.

The README documents DistributedWorker/WorkerGroup as public API. Repoint it at the base-class and injected-context shape, and flag it experimental.

This is moving the driving into a base class and thinning the surface.

If my recommendations are wildly out of left field or otherwise would pose too many ergonomic challenges, please let me know.