Add: support mixed local and remote L4 CommDomains

[sunkaixuan2018]

Summary

• Add an L4-brokered Global CommDomain flow that collects rank-ordered L2 export descriptors through L3 and returns the complete table for L2 import, without mpirun.
• Support one domain spanning both forked local L3 workers (add_worker) and TCP-connected remote L3 workers (add_remote_worker).
• Report only backend capabilities that are actually implemented: sim platforms accept sim, real A2/A3 accepts a3-fabric-v1, and unsupported platform/profile pairs are rejected before PREPARE and again at L3 COMM_INIT.
• Harden the full lifecycle: device-backed remote buffers remain exportable, partial node/backend imports roll back staged mappings, partial release invalidates committed views before fan-out, configured control timeouts are honored, and Global handles release behind the correct run fence.
• When a rollback leg itself fails, refuse further work on that worker. The domain is never registered on that path, so no run fence and no close() sweep can reach what the failed rollback leaves mapped; both the ALLOC_DOMAIN ABORT fan-out and the L3-node close path route through Worker._record_unreclaimable.
• Add two-machine pod scene tests global_tload_mixed_l3 (L4-brokered peer TLOAD) and compute_then_tload_mixed_l3 (an L2 compute round, then cross-machine communication through the same retained domain), registered in _st-pod.yml.
• Add focused codec, capability, transaction, node-side rollback, remote-only sim, and mixed local/remote sim coverage.

Testing

• Head f1b8bd9f is a single commit on top of upstream/main.
• Python unit tests: 1216 passed, 13 skipped. tests/ut/py/test_global_comm_domain.py is 25 passed.
• test_global_domain_abort_failure_preserves_primary_error_and_poisons_admission is the regression barrier for the rollback-failure path: it asserts the primary error still propagates, every prepared node still receives ABORT, and the worker then refuses further work with the injected failure as __cause__. Confirmed failing against the pre-fix code and passing after.
• Two end-to-end tests run without mocks: test_two_remote_daemons_build_and_copy_global_domain_without_mpirun and test_local_and_remote_l3_build_and_copy_global_domain_without_mpirun spawn real simpler.remote_l3_worker daemons over TCP and fork a real local L3, then build, copy into, read back and release a domain across the mixed topology.
• Other regression coverage: HostBuffer-backed remote export, local-rank import rollback, committed-view invalidation and retry after partial release, run-owned deferred release, and a mixed local/remote control path where remote import failure must roll back the already-imported local L3 node.
• Repository hooks passed (headers, English-only content, large files, EOF/whitespace, clang-format, cpplint, Markdown, Ruff, Pyright).

A3 Fabric coverage

• The real a3-fabric-v1 backend is covered by repository CI. The two-machine st-pod-onboard-a2a3 job runs both new examples with their default --platform a2a3 --comm-profile a3-fabric-v1; pod-run-example injects only the endpoint, device lists and session settings, so the Fabric profile is the one actually exercised.
• Observed on hardware, not inferred from a green check: [global-tload-mixed-l3] local rank=0 max_diff=0.000e+00 / remote rank=1 max_diff=0.000e+00.
• Control plane and data plane are separate here: descriptors are relayed over the TCP control plane, while peer access goes through aclrtMemFabricHandle VMM import, so the imported window is read directly by the AICore kernel.
• The sim backend still covers the transaction, rollback and mixed local/remote paths in the GitHub-hosted jobs.

[coderabbitai]

📝 Walkthrough

Walkthrough

This change adds Global CommDomain support across remote L3 workers, including wire codecs, transactional allocation and cleanup, simulation and A3 backend implementations, Python APIs, remote control handling, tests, and end-to-end smoke workflows.

Changes

Global CommDomain

Layer / File(s)	Summary
Protocol contracts and communication backends python/simpler/global_comm_domain.py, src/common/platform_comm/*, src/a2a3/.../comm_hccl.cpp, src/a5/.../comm_hccl.cpp	Defines validated domain messages and descriptors, then implements prepare/import/release for simulation and A3 Fabric backends; A5 exports unsupported stubs.
Control transport and runtime bindings src/common/hierarchical/*, src/common/worker/*, python/bindings/*	Routes domain controls through workers and remote endpoints, exposes byte-oriented Python bindings, and dynamically loads the new communicator APIs.
Remote session and L2 control handling python/simpler/remote_l3_session.py, python/simpler/remote_l3_worker.py, python/simpler/worker.py	Validates node manifests, handles COMM_INIT and domain phases, performs copy operations, manages host buffers, and sweeps child-side domain state.
Python domain API and transaction orchestration python/simpler/orchestrator.py, python/simpler/task_interface.py, python/simpler/worker.py	Adds domain handles, committed views, allocation/release/copy APIs, remote manifest fields, transactional commit/abort logic, and run/close cleanup.
Validation, smoke workflows, and documentation tests/ut/py/*, tools/a3_l4_tcp_smoke/*, tools/remote_l4_npu/*, docs/*	Adds codec and integration tests, A3 compute/TLOAD kernels and drivers, remote NPU smoke scripts, and protocol/design documentation.

Estimated code review effort: 5 (Critical) | ~120 minutes

Possibly related PRs

• hw-native-sys/simpler#1424: Refactors the Fabric V2 window handle lifecycle used by the new A3 Global CommDomain exchange.

Poem

I’m a rabbit with domains in a row,
Descriptors hop where the remote winds blow.
Prepare, import, commit with delight,
Then TLOAD makes the tensors unite.
Copy and release, cleanup so neat—
A carrot-powered control-plane feat!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 7.66% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the main change: support for mixed local and remote L4 CommDomains.
Description check	✅ Passed	The description directly explains the Global CommDomain flow, mixed local and remote support, testing, rollback, and A3 coverage.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 3

🧹 Nitpick comments (2)

tests/ut/py/test_global_comm_domain.py (1)

176-189: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win

Fixed time.sleep(0.3) daemon warmup is racy. worker.init() connects to the freshly spawned daemons, but there is no guarantee they have bound their ports within 300 ms on a loaded CI host, which can cause intermittent connection failures. Consider polling each (127.0.0.1, port) for accept-readiness (short connect-retry loop with a bounded deadline) before add_remote_worker/init instead of a static sleep.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/ut/py/test_global_comm_domain.py` around lines 176 - 189, Replace the
fixed time.sleep(0.3) in the worker setup test with bounded polling that
attempts short TCP connections to each 127.0.0.1 port until all daemons accept
connections or a deadline is reached. Perform this readiness check before the
add_remote_worker calls and worker.init(), and fail clearly if any daemon
remains unavailable.

src/common/platform_comm/comm.h (1)

48-57: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win

Pin the descriptor ABI used by the wire codec.

ChipWorker::comm_global_domain_prepare serializes sizeof(CommGlobalDomainDescriptor), while the Python wire contract fixes this descriptor at 288 bytes. Add size/offset static_asserts or explicit field serialization so compiler packing changes cannot silently break remote imports.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/common/platform_comm/comm.h` around lines 48 - 57, Pin the
CommGlobalDomainDescriptor wire layout used by
ChipWorker::comm_global_domain_prepare: add compile-time size and field-offset
assertions matching the 288-byte Python contract, including the handle field, so
packing changes fail at build time rather than altering serialized imports.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@docs/remote-l3-worker-design/implementation-record.md`:
- Around line 91-94: Update the L4-brokered Global CommDomain setup description
to state that domains are released after the L4 DAG drain by default, while
explicitly noting that retain_after_run=True preserves them for a later run.

In `@python/simpler/task_interface.py`:
- Around line 910-921: Sort the entries alphabetically in both __slots__ tuples
in the relevant classes, including the declaration containing name and members
and the second declaration around the referenced location, while preserving
every existing slot name.
- Around line 975-979: Update release so _released is set to True only after
_release_fn(self) completes successfully. If the remote cleanup raises,
propagate the error while leaving the handle unreleased so a later release()
retries the cleanup.

---

Nitpick comments:
In `@src/common/platform_comm/comm.h`:
- Around line 48-57: Pin the CommGlobalDomainDescriptor wire layout used by
ChipWorker::comm_global_domain_prepare: add compile-time size and field-offset
assertions matching the 288-byte Python contract, including the handle field, so
packing changes fail at build time rather than altering serialized imports.

In `@tests/ut/py/test_global_comm_domain.py`:
- Around line 176-189: Replace the fixed time.sleep(0.3) in the worker setup
test with bounded polling that attempts short TCP connections to each 127.0.0.1
port until all daemons accept connections or a deadline is reached. Perform this
readiness check before the add_remote_worker calls and worker.init(), and fail
clearly if any daemon remains unavailable.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: f217e7f8-c970-4de9-9e33-aa5539b441f5

📥 Commits

Reviewing files that changed from the base of the PR and between d0bc661 and 73c4500.

📒 Files selected for processing (41)

• docs/comm-domain.md
• docs/remote-l3-worker-design.md
• docs/remote-l3-worker-design/implementation-record.md
• docs/remote-l3-worker-design/protocol.md
• python/bindings/CMakeLists.txt
• python/bindings/task_interface.cpp
• python/bindings/worker_bind.h
• python/simpler/global_comm_domain.py
• python/simpler/global_comm_smoke.py
• python/simpler/orchestrator.py
• python/simpler/remote_l3_protocol.py
• python/simpler/remote_l3_session.py
• python/simpler/remote_l3_worker.py
• python/simpler/task_interface.py
• python/simpler/worker.py
• src/a2a3/platform/onboard/host/comm_hccl.cpp
• src/a5/platform/onboard/host/comm_hccl.cpp
• src/common/hierarchical/remote_endpoint.cpp
• src/common/hierarchical/remote_endpoint.h
• src/common/hierarchical/remote_wire.cpp
• src/common/hierarchical/remote_wire.h
• src/common/hierarchical/worker.h
• src/common/hierarchical/worker_manager.cpp
• src/common/hierarchical/worker_manager.h
• src/common/platform_comm/comm.h
• src/common/platform_comm/comm_sim.cpp
• src/common/worker/chip_worker.cpp
• src/common/worker/chip_worker.h
• tests/ut/py/test_callable_identity.py
• tests/ut/py/test_global_comm_domain.py
• tools/a3_l4_tcp_smoke/README.md
• tools/a3_l4_tcp_smoke/compute_then_tload_smoke.py
• tools/a3_l4_tcp_smoke/global_tload_smoke.py
• tools/a3_l4_tcp_smoke/kernels/aiv/global_tload_kernel.cpp
• tools/a3_l4_tcp_smoke/kernels/aiv/local_add_kernel.cpp
• tools/a3_l4_tcp_smoke/kernels/orchestration/global_tload_orch.cpp
• tools/a3_l4_tcp_smoke/kernels/orchestration/local_add_orch.cpp
• tools/remote_l4_npu/README.md
• tools/remote_l4_npu/remote_l4_npu_smoke.py
• tools/remote_l4_npu/run_parent_smoke.sh
• tools/remote_l4_npu/start_machine_daemon.sh

[sunkaixuan2018]

Also addressed the two CodeRabbit nitpicks listed in the review summary:

• Replaced the fixed daemon warmup sleep with bounded TCP readiness polling.
• Pinned the Global CommDomain descriptor to the 288-byte wire ABI with compile-time size and field-offset checks.

[ChaoWao]

Re-reviewed at d1eff1d2 (previous review was against 34c8b742). Scope tightening and the rebase look good; the four blocking defects from that review are all still present.

What changed since 34c8b742

Normalising both revisions against their own merge-bases, the author-authored delta is:

• tools/ removed entirely (−1258 lines: a3_l4_tcp_smoke/, remote_l4_npu/). Commit message updated accordingly ("local/remote sim integration coverage" replacing the smoke bullet).
• .github/workflows/ci.yml change dropped — the a5 path: source + -p no:cacheprovider hunks are gone. Thanks for pulling that out; it was unrelated to this feature and undisclosed in the body.
• Rebased onto current main (37 → 23 commits behind), which accounts for the rest: comm_hccl.cpp is 3 lines of LOG_INFO_V0 → LOG_INFO, and worker.py is net +14/−1 — _CTRL_GLOBAL_DOMAIN_NODE moved 18 → 24 to avoid the new chip-child control codes, plus a provenance comment refresh.

global_comm_domain.py, comm.h, comm_sim.cpp and remote_l3_session.py are byte-identical between the two revisions.

CI is green, including st-onboard-a2a3 / ut-a2a3 / ut-a5.

Blocking issues — status unchanged

All four are verified against d1eff1d2 itself, not carried over from the earlier read:

M1 — remote_export fails on any remote node that has chip devices. HostBuffer still has only token/data_ptr/nbytes/buffer — no shm_name. _RemoteBufferEntry.shm_name still raises for a non-SharedMemory backing (remote_l3_session.py:142), and export still calls it unconditionally (:824). A node configured with device_ids — the configuration this feature requires — takes the create_host_buffer branch and gets a HostBuffer, so remote_malloc + remote_export raises ValueError. The _NoHostBufferChildrenError change fixed the string-matching branch selection, but not this.

M2 — release does not invalidate the committed view before fanning out. _release_global_domain_node is unchanged: generation check → drop provenance → straight into the member loop, with no phase/view invalidation. On partial failure the state is not popped, phase stays COMMIT, and _get_global_domain keeps handing out a view whose device_ctx points at windows already freed on the members that succeeded.

M3 — comm_global_domain_import leaks on failure and double-maps on retry. Both backends still return -1 without release_domain_windows(...) / munmap, contrary to the convention domain_alloc_via_fabric follows and documents in the same file. The L2 handler re-enters on device_ctx == 0, so a retry maps ranks 0..k-1 a second time.

M4 — the Global handle's release is not bound to the run that allocated it.

# local CommDomain, worker.py:6008
_release_fn=lambda released, owner=resources: self._release_domain_handle(released, owner),
# Global, worker.py:6797
_release_fn=self._release_global_domain_handle,

_release_global_domain_handle reads self._building_run_resources at release time, and that is only non-None inside the graph-building window. So calling domain.release() after submit() returns — while the DAG is still executing — takes the immediate-free path and tears down the backend windows under a running kernel. That is exactly the retain_after_run=True usage the new docs/comm-domain.md section recommends. The local path is immune only because it binds owner=resources at allocation; copying that shape fixes it.

One consequence of removing tools/

Dropping the smokes is right for PR scope, but it leaves the A3 Fabric path with no verification of any kind — it has no CI coverage (sim-only) and now no in-repo manual harness either. Worth saying explicitly in the body where that validation now lives, so a future reader doesn't assume the a3-fabric-v1 profile is exercised somewhere.

Also still open from the previous review and unaddressed here: the local/remote timeout asymmetry (local leg hardcodes _PY_CONTROL_TIMEOUT_S at 5 call sites, bypassing the configured py_control_timeout_s, while the remote leg honours remote_session_timeout_s), silent contextlib.suppress on the ABORT fan-out, and the failure-injection tests patching _global_domain_control at the L4 dispatch point so no node-side code runs on any failure path — which means the mixed local+remote rollback, this revision's headline capability, has no coverage.

[sunkaixuan2018]

@ChaoWao Thanks for the detailed re-review. Addressed the four blockers and the remaining follow-ups in a05b944f (rebased onto upstream/main 2a650f2d):

• M1: remote_export now resolves the backing shm name for device-backed HostBuffer allocations through the owning Worker; the regression test now configures a device child so it exercises this path rather than the plain SharedMemory fallback.
• M2: node release sets the state to ABORT and marks the committed view unusable before the first destructive fan-out. A partial-release test verifies that both _get_global_domain() and a previously held view reject access, then verifies a successful retry.
• M3: sim peer mappings and A3 VMM peer windows are staged in RAII-owned local vectors and moved into the allocation only after the complete import succeeds, so partial failures/retries do not leak or double-map earlier peers.
• M4: the Global handle release callback now captures its allocation run resources, marks the run as cleanup-bearing, and defers backend release behind the owning/current run fence as appropriate for retained cross-run handles.

Also fixed the five local control calls to honor py_control_timeout_s, made ABORT cleanup failures visible while preserving the primary transaction exception, added local-rank rollback tests, and added a mixed local/remote control-path test where the remote import failure must actually release the already-imported local L3 node.

The PR body and docs now state the A3 scope explicitly: repository CI covers sim; the earlier two-machine A3 evidence is external and is not current-head A3 validation after the smoke tools were removed. All executable local hooks passed except Linux-only clang-tidy/pytest entry points blocked on Windows by POSIX fcntl/native extensions. Fresh CI is now running for this head.

[sunkaixuan2018]

CI follow-up for a05b944f: both A5 hardware jobs failed before pytest started. On runners a5ci8p and a5ci8p-2, ${DEVICE_RANGE} was empty, so the workflow's first int(s) conversion raised ValueError: invalid literal for int() with base 10: '' immediately after package installation. An unrelated PR's A5 job has the same runner-side traceback. This is not caused by the CommDomain changes; I did not modify the workflow or rerun the jobs. The code-bearing checks completed so far are green, including pre-commit, Linux/macOS UT, A2/A3 UT/onboard/sim, and Ubuntu A5 sim.

[ChaoWao]

@sunkaixuan2018 Re-reviewed at adc312de and pushed one fix onto this branch (cafbe063, amended into your commit — the PR stays a single commit). Details below.

Your four blockers are all closed — verified against adc312de, not carried over

	Fix
M1 remote_export on a node with chip devices	_RemoteBufferEntry.shm_name falls back to owner._host_buffer_shm_name(...) for a HostBuffer backing
M2 release did not invalidate the committed view	_release_global_domain_node sets phase = ABORT and clears view._committed before the first destructive child call
M3 import leaked on failure	GlobalPeerMapping / VmmWindow are now move-only RAII, so every return -1 in the peer loop unwinds
M4 handle release not bound to its allocating run	_release_fn=lambda released, owner=resources: ..., plus handling for a retained handle released while a later run builds

The RAII approach in M3 is better than the per-branch cleanup I suggested — a failure path can't forget to unmap. The timeout asymmetry is closed too, and the failure-injection tests got real depth: test_two_remote_daemons_... and test_local_and_remote_l3_build_and_copy_global_domain_without_mpirun spawn actual simpler.remote_l3_worker daemons over TCP and fork a real local L3, so node-side code genuinely runs on the mixed topology.

Removing tools/ also no longer leaves A3 Fabric unverified — wiring the two examples into _st-pod.yml means CI now exercises a3-fabric-v1 on real two-machine hardware. I read the job log rather than trusting the green check: [global-tload-mixed-l3] remote rank=1 max_diff=0.000e+00, job 92743892751.

What I pushed: the ABORT path now refuses admission instead of reporting

My earlier comment called this "silent contextlib.suppress on the ABORT fan-out". That wording was on me — it emphasised silence, so making it loud (sys.stderr.write) was a reasonable reading. But visibility was never the requirement here, which is why I fixed it rather than just re-raising it:

On the ABORT path the domain is never registered — self._live_global_domains[name] and resources.live_global_domains[name] are both assigned after the try, and requires_ordered_cleanup is never set. So when an ABORT leg fails, there is no handle for the run's fence to fail on and nothing for close() to sweep, yet the run is still classified as not cleanup-bearing and the next run is admitted at full pipeline depth as if teardown had succeeded.

Both sites now route through Worker._record_unreclaimable(...), whose docstring describes exactly this case — "For device state that no cleanup can reach: nothing tracks it, so there is no handle for a run's fence to fail on and no sweep that could retry it." It is first-wins and returns the error, so the control flow is unchanged:

• Worker._allocate_global_domain — ABORT fan-out except
• Worker._release_all_global_domain_nodes — L3-node close path (this one also didn't re-raise)

Worth noting how mechanical this is: _record_unreclaimable had 10 call sites in worker.py and the whole global-domain section had none, so these two were the exception rather than a new requirement.

test_global_domain_abort_failure_preserves_primary_error_and_reports_cleanup asserted on captured stderr, so it is now ..._poisons_admission and asserts the worker refuses further work with the injected error as __cause__. I confirmed it fails against the pre-fix code (assert None is not None, with the swallowed warning visible in captured stderr) and passes after. Full py UT: 1216 passed / 13 skipped.

Two heads-ups — explicitly not change requests

1. GlobalDomainMember will need a deployment component — please don't add one here. encode_domain_command puts len(members) plus each _put_member on the cross-node control-plane wire, and members are (node_worker_id, local_worker_id, ...) — positions in the orchestration tree. Whether an endpoint can join a comm domain is ultimately a function of where it is deployed and what it is physically connected to, not of its tree position, so that field is expected to grow a deployment component.

The registry that owns this vocabulary is in flight in #1696; adding a second independent enum now would be worse than waiting. Flagging it only because #1623/#1624 stack on this same encoding — adding a field is cheap while the stack is still moving and a protocol migration afterwards.

2. One backing per domain, and the domain-wide comm_profile, are known future breaking changes — CommBufferSpec becoming an independent Buffer rather than an offset slice of one window, and transport/medium selection moving from a node-level string to a per-(member, buffer) decision. Both need the same registry first. Nothing to do here; noting it so the current shape isn't mistaken for the final one.

One question

Are the two pod machines in the same A3 super-pod / fabric? Your peer TLOAD result is the first in-repo evidence that cross-machine device-peer direct-map works, and I'd like to record it accurately — specifically whether "cross-machine" here means within one fabric (control plane over TCP, data plane over fabric) or something broader. I can read the code path but not your pod topology.