diff --git a/unirl/distributed/group/dispatch.py b/unirl/distributed/group/dispatch.py index 7de4f35b..c6b5cbc5 100644 --- a/unirl/distributed/group/dispatch.py +++ b/unirl/distributed/group/dispatch.py @@ -241,6 +241,7 @@ def resolve_backward_dispatch_mode( DISTRIBUTED_CONFIG_ATTR = "_distributed_config" +ADDRESSED_CONFIG_ATTR = "_addressed_config" def distributed( @@ -283,3 +284,39 @@ def wrapper(*args, **kwargs): if _func is not None: return decorator(_func) return decorator + + +def addressed(_func: Callable = None) -> Callable: + """Declare a Role method point-to-point: ONE target worker, ONE result. + + The counterpart to :func:`distributed`. Every ``Dispatch`` mode targets the whole + slab and ``PendingHandleCall.ready()`` is all-or-nothing across it, so a caller + that must place one unit of work on one worker and reap it the moment it lands + has no dispatch mode to express that. An ``@addressed`` method skips dispatch + (args go whole to one worker) and collect (the result is that worker's, unmerged), + and is bound on ``Handle.worker(k)`` rather than on ``Handle``. + + A method may carry ``@addressed`` or ``@distributed``, never both — ``Handle`` + rejects that at bind time, because the two answer different questions about the + same name. + + Usage: + class AgenticRolloutEngine(Remote): + @addressed + def run_trajectory(self, task): ... + """ + + def decorator(func: Callable) -> Callable: + # ``wraps`` copies ``__dict__``, which is what lets this decorator and + # ``distributed`` stack in either order without the inner one's marker + # being lost. Removing it silently deletes the other marker. + @wraps(func) + def wrapper(*args, **kwargs): + return func(*args, **kwargs) + + setattr(wrapper, ADDRESSED_CONFIG_ATTR, {"addressed": True}) + return wrapper + + if _func is not None: + return decorator(_func) + return decorator diff --git a/unirl/distributed/group/handle.py b/unirl/distributed/group/handle.py index 7bd35c59..3d3ee7a8 100644 --- a/unirl/distributed/group/handle.py +++ b/unirl/distributed/group/handle.py @@ -35,10 +35,12 @@ import ray from unirl.distributed.group.dispatch import ( + ADDRESSED_CONFIG_ATTR, DISPATCH_MODE_REGISTRY, DISTRIBUTED_CONFIG_ATTR, Dispatch, Execute, + _unwrap_broadcast, resolve_backward_dispatch_mode, ) from unirl.distributed.group.remote import RankInfo, Remote @@ -60,6 +62,33 @@ _role_name_counter: Dict[str, int] = {} +# Handle attributes a bound @distributed method must not overwrite. ``_bind_methods`` +# ends in ``setattr(self, name, bound)``: against a read-only property that raises +# opaquely, and against a plain method or attribute it silently wins. Both deserve a +# named error instead. +_HANDLE_RESERVED_NAMES = frozenset( + { + "worker", + "workers", + "worker_ids", + "role_name", + "role_cls", + "rank_infos", + "device_ids", + "world_size", + "pool", + "initialize", + "launch_nowait", + "engine_replicas", + "dp_size", + "sp_size", + "tp_size", + "pp_size", + "ep_size", + "tp_zero_workers", + } +) + def _owning_class(role_cls) -> Type[Remote]: """Return the class for Handle method-binding and role naming. @@ -350,6 +379,128 @@ def result(self) -> Any: return self._value +class PendingWorkerCall: + """Future-like result of :meth:`WorkerHandle.launch_nowait`: ONE worker, ONE ref. + + The point-to-point twin of :class:`PendingHandleCall`, whose ``ready()`` is + all-or-nothing across the slab. Here ``ready()`` means exactly what it says, + which is what lets a caller reap completions as they land rather than at the + slowest worker's pace. + + ``result()`` runs the resolution phase — ``ray.get`` then rebind onto the worker + that actually ran the call. It is cached after the first success, because rebind + registers GC finalizers on the result handles and asserts once-only. + """ + + def __init__(self, owner: "WorkerHandle", ref: Any) -> None: + self._owner = owner + self._ref = ref + self._consumed = False + self._value: Any = None + + @property + def ref(self) -> Any: + """The underlying ObjectRef — :func:`wait_any` batches these into one ``ray.wait``.""" + return self._ref + + def ready(self) -> bool: + """True once this call's ref is resolved (non-blocking probe).""" + done, _ = ray.wait([self._ref], num_returns=1, timeout=0) + return bool(done) + + def wait(self) -> None: + """Block until the worker finishes, without collecting; re-raises worker errors.""" + ray.get(self._ref) + + def result(self) -> Any: + """Block if needed, then resolve: the method's return value, rebound locally.""" + if self._consumed: + return self._value + self._value = self._owner._resolve_one(self._ref) + self._consumed = True + return self._value + + +def wait_any(pendings: Sequence[PendingWorkerCall], *, timeout: float = 0) -> List[PendingWorkerCall]: + """The ready subset of *pendings*, in ONE ``ray.wait``. + + Probing N pendings via :meth:`PendingWorkerCall.ready` is N round trips through + the object store, which is exactly the cost a per-call pending exists to avoid. + Duplicate pendings collapse (``ray.wait`` requires unique refs) and a pending + that is already consumed still reports ready, so callers may re-probe safely. + """ + by_ref: Dict[Any, PendingWorkerCall] = {p.ref: p for p in pendings} + if not by_ref: + return [] + done, _ = ray.wait(list(by_ref), num_returns=len(by_ref), timeout=timeout) + return [by_ref[ref] for ref in done] + + +class WorkerHandle: + """Point-to-point calls at ONE worker of a slab. + + Answers *call worker k*. It must never grow *which worker*, *how many at once*, + or *retry* — placement and admission are the caller's policy, and putting them + here would rebuild a scheduler inside the transport layer. + + Obtained from :meth:`Handle.worker`, which supplies the localize/rebind context + so ``TensorRef`` arguments are transferred to this worker and returned handles + are bound to it. ``@addressed`` methods of the role are available as attributes; + ``launch_nowait`` is the non-blocking form. + """ + + def __init__(self, handle: "Handle", index: int) -> None: + self._handle = handle + self._index = index + self._actor = handle.workers[index] + self._role = handle.role_name + for name in handle._addressed_methods: + setattr(self, name, self._make_fn(name)) + + def __repr__(self) -> str: # pragma: no cover - diagnostics only + return f"WorkerHandle(role={self._role!r}, index={self._index})" + + def _make_fn(self, method_name: str) -> Callable: + def worker_fn(*args, **kwargs): + return self.launch_nowait(method_name, *args, **kwargs).result() + + worker_fn.__name__ = method_name + return worker_fn + + def launch_nowait(self, method_name: str, *args, **kwargs) -> PendingWorkerCall: + """Launch an ``@addressed`` method at this worker without blocking.""" + if method_name not in self._handle._addressed_methods: + raise AttributeError( + f"{method_name!r} is not an @addressed method of " + f"{_owning_class(self._handle.role_cls).__name__}" + ) + # A single-worker forward under a GradContext would be answered by the + # slab-wide _auto_backward, which finds no saved call_id on the other + # workers; whether that yields a silent zero-gradient or an IndexError + # depends on which worker lands first in the collect. Refuse instead. + if current_grad_context() is not None: + raise RuntimeError( + f"@addressed {method_name!r} called under a GradContext; backward is " + "slab-wide and cannot answer a single-worker forward." + ) + # Broadcast is a controller-side annotation consumed by the dispatch fns. + # The addressed path skips dispatch, so it must strip them itself or the + # wrapper object reaches the worker. + args, kwargs = _unwrap_broadcast(args, kwargs) + (args, kwargs) = self._handle._localize_one((args, kwargs), self._index) + ref = self._actor.call.remote(self._role, method_name, args, kwargs) + return PendingWorkerCall(self, ref) + + def _resolve_one(self, ref: Any) -> Any: + """``ray.get`` + rebind onto THIS worker (never ``workers[0]``). + + Rebinding to the wrong actor sends the decref to a store that has no such + key; ``GPUTensorHandle._release`` is fire-and-forget and swallows the error, + so the real owner leaks for the process lifetime with no diagnostic. + """ + return self._handle._rebind_tree(ray.get(ref), self._actor, worker_local=self._handle._worker_local) + + class Handle: """Controller-side SPMD handle. @@ -466,6 +617,8 @@ def _rank_init_kwargs(i: int) -> Dict[str, Any]: ) self._method_configs: Dict[str, tuple] = {} + self._addressed_methods: set = set() + self._worker_handles: Dict[int, "WorkerHandle"] = {} self._bind_methods(role_cls) self._grad_call_counter = count() @@ -508,6 +661,50 @@ def tp_zero_workers(self) -> List[Any]: returns every worker (identical to ``self.workers``).""" return [w for w, ri in zip(self.workers, self.rank_infos) if ri.tp_rank == 0] + @property + def engine_replicas(self) -> List[int]: + """Indices of the workers that host one engine each — the DP heads. + + ``tp_rank == 0 and pp_rank == 0``: the ranks a whole unit of work may be + addressed to. Non-tp-zero ranks are no-op shells whose ``generate`` returns + ``None`` (the slab collect filters that; a point-to-point call would not), + and the pp ranks of one group are one engine, not several. Identical to + ``range(world_size)`` while tp_size == pp_size == 1.""" + return [i for i, ri in enumerate(self.rank_infos) if ri.tp_rank == 0 and ri.pp_rank == 0] + + @property + def _worker_local(self) -> bool: + """Whether returned handles must be rebound to their producing worker. + + Same value ``_launch_call`` computes; a wrong ``False`` on a worker-local + backend means no GC finalizer is registered and every returned tensor pins + its worker's store for the process lifetime, with no error.""" + return issubclass(self.pool.transport_cls, WorkerLocalTransport) + + def worker(self, index: int) -> "WorkerHandle": + """A point-to-point handle for one worker of this slab (``@addressed`` methods). + + ``index`` is a position in ``self.workers``; use :attr:`engine_replicas` to + pick one that actually hosts an engine. Cached, because callers address the + same worker repeatedly.""" + if not 0 <= index < self.world_size: + raise IndexError(f"worker index {index} out of range for world_size={self.world_size}") + cached = self._worker_handles.get(index) + if cached is None: + cached = self._worker_handles[index] = WorkerHandle(self, index) + return cached + + def _localize_one(self, shard: Tuple[tuple, dict], index: int) -> Tuple[tuple, dict]: + """Make every ref in one shard resolvable on worker ``index``. + + The single-target slice of what ``_launch_call`` does for the whole slab; + ``localize`` zips shards against ``(worker_ids, device_ids)`` elementwise, so + one-element lists are a correct slice. Skipping this is what makes a task + produced on one worker unresolvable on another.""" + return self.pool.transport_cls.localize( + [shard], self.pool, [self.device_ids[index]], [self.worker_ids[index]] + )[0] + def initialize(self, *args, **kwargs) -> None: """Call role.initialize(*args, **kwargs) on all workers. @@ -521,12 +718,17 @@ def initialize(self, *args, **kwargs) -> None: self.rank_infos = ray.get([w.get_rank_info.remote(self.role_name) for w in self.workers]) def _bind_methods(self, role_cls) -> None: - """Scan role_cls for @distributed methods and create handle functions. + """Scan role_cls for @distributed / @addressed methods and bind them. For classmethod ``role_cls`` (e.g. ``SD3Bundle.from_config``) we scan the owning class instead — the constructed instance is of that class, so its ``@distributed`` methods are the ones callers will dispatch through this Handle. + + ``@distributed`` methods are bound on this Handle; ``@addressed`` ones are + recorded for :meth:`worker` to bind on a :class:`WorkerHandle`. A method + carrying both is rejected here rather than silently resolving to whichever + marker is checked first. """ role_cls = _owning_class(role_cls) for name in dir(role_cls): @@ -534,8 +736,23 @@ def _bind_methods(self, role_cls) -> None: if method is None: continue config = getattr(method, DISTRIBUTED_CONFIG_ATTR, None) + addressed_config = getattr(method, ADDRESSED_CONFIG_ATTR, None) + if config is not None and addressed_config is not None: + raise TypeError( + f"{role_cls.__name__}.{name} is both @distributed and @addressed; " + "a method is either a slab collective or point-to-point, not both." + ) + if addressed_config is not None: + self._addressed_methods.add(name) + continue if config is None: continue + if name in _HANDLE_RESERVED_NAMES: + # setattr below would shadow the Handle API silently (or raise + # opaquely on the read-only properties). Fail with the reason. + raise TypeError( + f"@distributed {role_cls.__name__}.{name} collides with the Handle API; rename it." + ) fns = DISPATCH_MODE_REGISTRY[config["dispatch_mode"]] dispatch_fn = fns["dispatch_fn"]