feat(distributed): point-to-point @addressed calls - #6
Closed
celve wants to merge 1 commit into
Closed
Conversation
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 way to express that. @addressed marks a method point-to-point: bound on Handle.worker(k) rather than on Handle, skipping dispatch (args go whole) and collect (one worker's result, unmerged). PendingWorkerCall wraps a single ref so ready() means what it says; wait_any batches a probe of N pendings into one ray.wait. Correctness the slab path gets for free and this one has to do deliberately: localize the args to the target (a Sample produced on another worker is otherwise unresolvable), rebind results onto the worker that ran the call (the wrong actor sends the decref to a store with no such key, and that decref swallows its error, so the real owner leaks silently), strip Broadcast wrappers that dispatch would have consumed, and refuse under a GradContext, where a slab-wide backward answering a single-worker forward yields a zero-gradient or an IndexError depending on worker ordering. Handle.engine_replicas exposes 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 engine shells whose generate returns None; the slab collect filters that, a point-to-point call would not. _bind_methods now rejects a method carrying both markers, and rejects a @distributed name that collides with the Handle API — its setattr silently wins against a plain attribute, which would have made Handle.worker unreachable. Lands with the manager-queue change that consumes it; on its own it has no caller. Deliberately NOT migrating the existing raw Worker.call sites in weight_sync: they are cross-slab and pool-free, so they would get the same stringly-typed call with no localize/rebind benefit. Test plan: compileall passes; an AST sweep confirms no existing @distributed method collides with the reserved names. Behaviour needs Ray (not installed in the analysis environment).
2 tasks
Owner
Author
|
Closing: opened against the fork by mistake. Re-opening against Tencent-Hunyuan/UniRL, with the @addressed primitive and the manager move combined into one PR. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a point-to-point call primitive to
unirl/distributed/group/.Every
Dispatchmode targets the whole slab, andPendingHandleCall.ready()isall-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 way to express that — it either
runs the slab in lockstep (the straggler barrier) or hand-rolls
w.call.remote(role, "name", args, kwargs), which three subsystems already do.@addressedmarks a method point-to-point: bound onHandle.worker(k)ratherthan on
Handle, skipping dispatch (args go whole to one worker) and collect(that worker's result, unmerged).
PendingWorkerCallwraps a single ref, soready()means what it says;wait_anybatches a probe of N pendings into oneray.wait, because probing them individually is N round trips through the objectstore and defeats the point.
The interesting part is what the slab path does for free and this one has to do
deliberately. Each of these is a silent failure if missed:
Worker.callskips this, and aSampleproduced on another worker is then unresolvable(
colocate_store/transport.py:46-50)._resolve_callrebindspositionally with
self.workers[i]. The wrong actor sends the decref to a storewith no such key, and
GPUTensorHandle._releaseis fire-and-forget with aswallowed exception — so the real owner leaks for the process lifetime, silently.
Broadcastwrappers, which only the dispatch fns consume.GradContext. A slab-wide_auto_backwardanswering asingle-worker forward finds no saved
call_idon the other workers; whether thatyields a zero-gradient or an
IndexErrordepends on which worker lands atresults[0]. Order-dependent silence, so it raises instead.Also adds
Handle.engine_replicas— the DP heads (tp_rank == 0 and pp_rank == 0),the ranks a whole unit of work may be addressed to. A non-tp-zero rank is an engine
shell whose
generatereturnsNone; the slab collect filters that out, apoint-to-point call would not.
_bind_methodsnow rejects a method carrying both markers, and rejects a@distributedname that collides with the Handle API — itssetattrsilently winsagainst a plain attribute, which would otherwise have made
Handle.workerunreachable with no diagnostic.
Related Issue
LIN-693 — prerequisite for moving the agentic rollout task queue to the driver.
Test Plan
python -m compileall— passes.unirl/confirming no existing@distributedmethod collideswith the new reserved-name set: no hits.
Not run; reason: behaviour requires Ray, which is not installed in the environment
this was authored in (nor is torch, so the modules do not import there). Nothing
in this PR has been executed. The consuming PR is where it first runs.
Compatibility / Risk
Additive. No existing call path changes:
handle_fn,launch_nowaitandPendingHandleCallare untouched, and the new code is reached only throughHandle.worker(k).Two new failure modes are deliberately loud rather than silent — a method marked
both
@distributedand@addressed, and a@distributedmethod whose namecollides with the Handle API. Both raise at bind time. The AST sweep says neither
fires on the current tree, but a downstream branch adding such a method would now
fail fast where it previously would have shadowed silently.
Reviewer Notes
This has no caller until the manager-queue PR that consumes it, and should not
merge alone. Stacked PR:
LIN-693/manager-queueis branched off this one.I deliberately did not migrate the five existing raw
Worker.callsites inweight_sync/as a proof-of-fit, although an earlier plan called for it. Thosesites run inside a Worker process, cross-slab, with no
DevicePool, and aretensor-free — so they would get the same stringly-typed call with no
localize/rebind benefit. That is churn, not validation.
The
functools.wrapsdependency in both decorators is load-bearing (it copies__dict__, which is what lets the markers stack in either order) and is nowcommented; it was previously silent.
Prepared with AI assistance; every line needs human review before merge, and this
one in particular — the three silent-leak paths above are not covered by any test
that exists. Checked for overlapping open PRs — there are none.
Checklist