Skip to content

feat(distributed): point-to-point @addressed calls - #6

Closed
celve wants to merge 1 commit into
mainfrom
LIN-693/addressed-calls
Closed

feat(distributed): point-to-point @addressed calls#6
celve wants to merge 1 commit into
mainfrom
LIN-693/addressed-calls

Conversation

@celve

@celve celve commented Aug 4, 2026

Copy link
Copy Markdown
Owner

Summary

Adds a point-to-point call primitive to unirl/distributed/group/.

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 — 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.

@addressed marks a method point-to-point: bound on Handle.worker(k) rather
than on Handle, skipping dispatch (args go whole to one worker) and collect
(that 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, because probing them individually is N round trips through the object
store 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:

  • Localize the args to the target. A raw Worker.call skips this, and a
    Sample produced on another worker is then unresolvable
    (colocate_store/transport.py:46-50).
  • Rebind results onto the worker that ran the call. _resolve_call rebinds
    positionally with self.workers[i]. The wrong actor sends the decref to a store
    with no such key, and GPUTensorHandle._release is fire-and-forget with a
    swallowed exception — so the real owner leaks for the process lifetime, silently.
  • Strip Broadcast wrappers, which only the dispatch fns consume.
  • Refuse under a GradContext. A slab-wide _auto_backward answering a
    single-worker forward finds no saved call_id on the other workers; whether that
    yields a zero-gradient or an IndexError depends on which worker lands at
    results[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 generate returns None; the slab collect filters that out, 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 otherwise have made Handle.worker
unreachable with no diagnostic.

Related Issue

LIN-693 — prerequisite for moving the agentic rollout task queue to the driver.

Test Plan

  • python -m compileall — passes.
  • AST sweep over unirl/ confirming no existing @distributed method collides
    with 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_nowait and
PendingHandleCall are untouched, and the new code is reached only through
Handle.worker(k).

Two new failure modes are deliberately loud rather than silent — a method marked
both @distributed and @addressed, and a @distributed method whose name
collides 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-queue is branched off this one.

I deliberately did not migrate the five existing raw Worker.call sites in
weight_sync/ as a proof-of-fit, although an earlier plan called for it. Those
sites run inside a Worker process, cross-slab, with no DevicePool, and are
tensor-free — so they would get the same stringly-typed call with no
localize/rebind benefit. That is churn, not validation.

The functools.wraps dependency in both decorators is load-bearing (it copies
__dict__, which is what lets the markers stack in either order) and is now
commented; 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

  • I reviewed the changed code and removed unrelated/generated artifacts.
  • I updated tests, docs, and configs where needed, or explained why not.

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).
@celve

celve commented Aug 4, 2026

Copy link
Copy Markdown
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.

@celve celve closed this Aug 4, 2026
@github-actions github-actions Bot removed the wip label Aug 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant