refactor(rollout): engine-level sync/async contracts for async rollout

.github/CODEOWNERS

@@ -28,7 +28,7 @@
 
 # --- rollout ---------------------------------------------------------------
 /unirl/rollout/                         @celve @zzhuoxin1508 @leviking98z-rgb
-/unirl/rollout/async_runtime.py         @CjhHa1 @celve
+/unirl/rollout/engine/asynchronous.py         @CjhHa1 @celve
 /unirl/rollout/engine/sglang/           @celve @leviking98z-rgb
 /unirl/rollout/engine/sglang_diffusion/ @celve @leviking98z-rgb
 /unirl/rollout/engine/vllm_omni/        @celve @zzhuoxin1508

examples/diffusion/bagel/bagel_vllmomni_async.yaml

@@ -55,9 +55,9 @@ weight_sync_interval: 4
 # max_inflight: concurrent generations. MUST be 1: the trajectory-segment
 # cross-slab transfer (NCCL send) runs on the rollout worker; a second in-flight
 # generation co-tenanting that worker blocks the send behind it (~150s/rollout).
-# With max_inflight=1 the shared async runtime (reap_before_launch) reaps and
-# transfers each generation in the idle window before launching the next, then
-# overlaps that next generation with the train step. Real overlap needs
+# With max_inflight=1 the trainer polls (reaps) before topping up launches, so
+# each generation transfers in the idle window before the next launch, then
+# that next generation overlaps the train step. Real overlap needs
 # weight_sync_interval>1 (interval=1 drains every step).
 # buffer_max_staleness: how many weight syncs a buffered group may cross.
 #   0 = never crosses a regular rollout-weight sync (~174s/rollout on BAGEL 4+4).

unirl/distributed/group/handle.py

@@ -317,6 +317,49 @@ class HandleRef:
     ep_size: int = 1
 
 
+class PendingHandleCall:
+    """Future-like result of :meth:`Handle.launch_nowait`: launched, not yet collected.
+
+    ``ready()`` probes without blocking; ``wait()`` blocks without collecting;
+    ``result()`` blocks if needed, then runs the rebind + collect half of
+    ``handle_fn`` and returns the method's collected value. The collect half
+    runs at most once — rebind registers GC finalizers on the result refs — so
+    a successful ``result()`` caches its value and later calls return it; a
+    ``result()`` that raised may be retried.
+    """
+
+    def __init__(self, handle: "Handle", method_name: str, refs: List[Any], worker_local: bool) -> None:
+        self._handle = handle
+        self._method_name = method_name
+        self._refs = refs
+        self._worker_local = worker_local
+        self._consumed = False
+        self._value: Any = None
+
+    def ready(self) -> bool:
+        """True once every worker's ref is resolved (non-blocking probe)."""
+        done, _ = ray.wait(self._refs, num_returns=len(self._refs), timeout=0)
+        return len(done) == len(self._refs)
+
+    def wait(self) -> None:
+        """Block until every worker finishes, without collecting; re-raises worker errors."""
+        ray.get(self._refs)
+
+    def result(self) -> Any:
+        """Block if needed, then rebind + collect: the method's collected return value."""
+        if self._consumed:
+            return self._value
+        handle = self._handle
+        results = ray.get(self._refs)
+        results = [
+            handle._rebind_tree(r, handle.workers[i], worker_local=self._worker_local) for i, r in enumerate(results)
+        ]
+        _, _, collect_fn, _ = handle._method_configs[self._method_name]
+        self._value = collect_fn(handle, results)
+        self._consumed = True
+        return self._value
+
+
 class Handle:
     """Controller-side SPMD handle.
 
@@ -453,6 +496,7 @@ def _rank_init_kwargs(i: int) -> Dict[str, Any]:
         )
 
         # Bind @distributed methods as handle functions
+        self._method_configs: Dict[str, tuple] = {}
         self._bind_methods(role_cls)
 
         # Counter for unique call_id generation within enable_grad contexts.
@@ -542,6 +586,7 @@ def _bind_methods(self, role_cls) -> None:
             else:
                 execute_fn = self._execute_rank_zero
 
+            self._method_configs[name] = (config["dispatch_mode"], dispatch_fn, collect_fn, execute_fn)
             bound = self._make_handle_fn(name, config["dispatch_mode"], dispatch_fn, collect_fn, execute_fn)
             setattr(self, name, bound)
 
@@ -559,6 +604,9 @@ def _make_handle_fn(
         TensorMetas and append an RPCBackwardNode for later auto-backward.
         grad_mode and call_id are passed as dedicated parameters to Worker.call
         (not via kwargs) so dispatch internals remain unaware of grad state.
+
+        The non-blocking twin is :meth:`launch_nowait` +
+        :meth:`PendingHandleCall.result` below — keep the halves in parity.
         """
 
         def handle_fn(*args, **kwargs):
@@ -628,6 +676,40 @@ def handle_fn(*args, **kwargs):
         handle_fn.__doc__ = f"SPMD handle: {method_name} (dispatch={dispatch_fn.__name__})"
         return handle_fn
 
+    # ── Non-blocking launch ──
+
+    def launch_nowait(self, method_name: str, *args, **kwargs) -> PendingHandleCall:
+        """Launch a @distributed method without blocking: the dispatch → localize →
+        execute half of ``handle_fn``, stopping before ``ray.get``.
+
+        Always ``grad_mode=False`` / ``call_id=None`` (a pending call is never
+        valid under a GradContext, so the ``_grad_call_counter`` single-thread
+        assumption is untouched). Kept in line-parity with ``handle_fn`` above —
+        same divisibility gate, same localize. ``result()`` on the returned
+        :class:`PendingHandleCall` runs the collect half.
+        """
+        try:
+            dispatch_mode, dispatch_fn, _, execute_fn = self._method_configs[method_name]
+        except KeyError:
+            raise AttributeError(
+                f"{method_name!r} is not a @distributed method of {_owning_class(self.role_cls).__name__}"
+            ) from None
+
+        batch_size = infer_batch_size(args, kwargs)
+        if (
+            dispatch_mode in (Dispatch.DP_SCATTER, Dispatch.DP_SCATTER_HEAD)
+            and batch_size is not None
+            and batch_size % self.dp_size != 0
+        ):
+            raise ValueError(f"batch_size={batch_size} not divisible by dp_size={self.dp_size}")
+
+        shards = dispatch_fn(self, args, kwargs, batch_size)
+        transport_cls = self.pool.transport_cls
+        worker_local = issubclass(transport_cls, WorkerLocalTransport)
+        shards = transport_cls.localize(shards, self.pool, self.device_ids, self.worker_ids)
+        refs = execute_fn(method_name, shards, grad_mode=False, call_id=None)
+        return PendingHandleCall(self, method_name, refs, worker_local)
+
     # ── Execute strategies ──
 
     def _execute_all(self, method_name: str, shards: List, grad_mode: bool = False, call_id=None) -> List:

unirl/rollout/README.md

@@ -35,7 +35,7 @@ wrong objective.
 
 ## How it works
 
-- **One synchronous generation interface.** `BaseRolloutEngine` (`engine/base.py`)
+- **One synchronous generation interface.** `BaseRolloutEngine` (`engine/synchronous.py`)
   is a `Remote` whose concrete engines implement synchronous `generate(sample)`;
   each keeps its native batching/runtime path. Single-turn engines return one
   `Sample` and dispatch `generate` with `DP_SCATTER`; the agentic coordinator
@@ -63,10 +63,19 @@ wrong objective.
   ratio is 1 on the first update; *separate* — a dedicated engine on its own GPUs
   plus a `sync:` block; *colocate* — a dedicated engine sharing GPUs with train,
   plus offload/onload and `sync:`.
+- **Driver-side async engines** (`engine/asynchronous.py`, the driver-side half next
+  to `engine/synchronous.py`'s worker-side sync contracts). Both engines expose the
+  same consumer verbs the async trainers program against: `poll` / `drain_freshest` /
+  `pop_evicted` / `quiesce` + engine-owned `weight_version`. `AsyncBatchRolloutEngine`
+  (batch granularity; non-blocking `Handle.launch_nowait` generations, stamps
+  versions at launch, used by `AsyncARTrainer`/`AsyncDiffusionTrainer`) and
+  `AsyncAgenticRolloutEngine` (trajectory granularity over the agentic rank-0
+  coordinator; normalizes the `[0]` unwraps, assembles n-sibling GRPO groups,
+  stamps versions at completion, used by the partial/async agentic trainers).
 
 **Extending it:** a new single-turn engine adds `engine/<name>/config.py` (a
 `BaseEngineConfig` whose `make_engine(**deps)` lazily imports and builds it) and
-`engine/<name>/engine.py` (subclass `BaseSingleTurnRolloutEngine`, implement
+`engine/<name>/engine.py` (subclass `SyncRolloutEngine`, implement
 synchronous generation over the whole-`Sample` contract — thread-safe for
 concurrent callers if it should serve as an agentic inner, else serialized
 internally — and dispatch `generate` with `DP_SCATTER`). A dedicated engine also
@@ -83,6 +92,12 @@ implements its weight-receive method and a matching `sync:` handler in
   intentional exception: `BROADCAST + RANK_ZERO` returns its trajectory list.
 - **Direct sampling forbids a `sync:` block; dedicated requires one.** The trainside
   engine also can't live on a `layout: separate` slab — `_build_rollout` raises.
+- **Quiesce before weight sync / eval / checkpoint on the batch async path** —
+  `AsyncBatchRolloutEngine.quiesce()` drains every in-flight generation; a
+  weight + KV update corrupts one mid-flight. The agentic quiesce is a
+  turn-boundary `abort` + final poll, folded into
+  `AsyncAgenticRolloutEngine.quiesce()`. Reap-vs-launch ordering is trainer
+  statement order (diffusion polls before topping up; see its `_next_step`).
 - **Reward/advantage methods are not engine code** — `Part.compute_advantages` and
   `Sample.propagate_rewards` are called by the trainer after scoring. An engine
   fills generation fields such as `segment`, `conditions`, `primitive`, and