Add: dispatch MPI worker groups through mailbox - #1624
Conversation
📝 WalkthroughWalkthroughChangesGlobal CommDomain contracts and backends
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Sequence Diagram(s)sequenceDiagram
participant Orchestrator
participant Worker
participant Mailbox
participant MpiDispatcher
participant L3Rank
Orchestrator->>Worker: allocate_global_domain
Worker->>Mailbox: publish domain control
Mailbox->>MpiDispatcher: dispatch ordered request
MpiDispatcher->>L3Rank: prepare or import domain
L3Rank-->>MpiDispatcher: descriptor or context result
MpiDispatcher-->>Mailbox: publish completion
Mailbox-->>Worker: return control response
Worker-->>Orchestrator: commit domain view
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 11
🧹 Nitpick comments (19)
python/simpler/global_comm_smoke.py (1)
22-167: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the shared prologue and tensor builders.
The four callbacks repeat the same steps: validate
scalar_count() == 6, importget_inner_handle, decode the digest, and look up the domain. The group variants then repeat the body of their single-worker counterparts inside a loop. Two small helpers, one for the prologue and one per kernel argument shape, would remove most of this duplication and keep the four entry points to a few lines each.♻️ Sketch
+def _resolve(orch, args: TaskArgs, message: str): + from .remote_l3_session import get_inner_handle # noqa: PLC0415 + + if args.scalar_count() != 6: + raise ValueError(message) + return ( + get_inner_handle(_digest_from_scalars(args, 2).hex()), + orch.get_global_domain(int(args.scalar(0))), + int(args.scalar(1)), + ) + + +def _compute_args(context) -> TaskArgs: + chip_args = TaskArgs() + for buffer_name in ("lhs", "rhs"): + chip_args.add_tensor(_domain_tensor(context, buffer_name), TensorArgType.INPUT) + chip_args.add_tensor(_domain_tensor(context, "input"), TensorArgType.OUTPUT_EXISTING) + return chip_args🤖 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 `@python/simpler/global_comm_smoke.py` around lines 22 - 167, Refactor remote_compute_orch, remote_rank_orch, remote_compute_group_orch, and remote_rank_group_orch to share a helper for importing get_inner_handle, validating the six scalars, decoding the digest, and resolving the domain. Add reusable tensor-argument builders for the compute and TLOAD shapes, then have the group callbacks reuse the corresponding single-worker argument construction inside their loops while preserving worker selection and submission behavior.tools/a3_l4_tcp_smoke/mpirun_compute_then_tload_2x2_smoke.py (1)
51-60: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: use iterable unpacking instead of concatenation.
Static analysis flags
list(include_dirs) + [str(...)]on line 53. Use unpacking for a more idiomatic construction.♻️ Proposed fix
- kernel_include_dirs = list(include_dirs) + [str(compiler.project_root / "src" / "common")] + kernel_include_dirs = [*include_dirs, str(compiler.project_root / "src" / "common")]🤖 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 `@tools/a3_l4_tcp_smoke/mpirun_compute_then_tload_2x2_smoke.py` around lines 51 - 60, Update the kernel_include_dirs construction in _compile_aiv to use iterable unpacking when combining include_dirs with the common source directory, preserving the existing ordering and values.Source: Linters/SAST tools
python/simpler/mpi_l3_session.py (1)
504-513: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueReplace the
asserton the shutdown payload with an explicit check.Line 507 uses
assert payload is not None. Python removes asserts under-O, and_rewrite_frame_identitywould then fail onNone. Raise an explicit error instead.🤖 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 `@python/simpler/mpi_l3_session.py` around lines 504 - 513, Replace the assert in the MailboxOpcode.SHUTDOWN branch of the request handling flow with an explicit payload None check that raises an appropriate error before calling _rewrite_frame_identity; preserve the existing shutdown behavior when payload is present.tools/mpi_group_mailbox_smoke.py (1)
34-37: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winAdd a sleep to the wait helper.
_wait_untilre-evaluates the predicate with no pause. Each mailbox state read maps to a shared-memory load plus an import lookup, so this pins a core for the whole wait. Sleep for a short interval between checks.♻️ Proposed fix
def _wait_until(predicate, *, deadline: float, label: str) -> None: while not predicate(): if time.monotonic() >= deadline: raise TimeoutError(f"timed out waiting for {label}") + time.sleep(0.001)🤖 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 `@tools/mpi_group_mailbox_smoke.py` around lines 34 - 37, Update _wait_until to pause briefly between predicate evaluations, while retaining the existing deadline check and TimeoutError behavior; add the sleep inside the loop after a failed predicate check.tests/ut/py/test_global_comm_domain.py (1)
259-278: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer allocated ports over the fixed
19073 + indexrange.The test never binds these endpoints, so it passes today. The file already provides
_free_tcp_ports. Using it removes the fixed range and keeps the endpoint construction consistent across the file.🤖 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 259 - 278, Update _failure_injection_worker to obtain two ports through the existing _free_tcp_ports helper, then build each RemoteWorkerSpec endpoint from those allocated ports instead of the fixed 19073 + index range. Preserve the current node ordering and endpoint construction format.tests/ut/cpp/hierarchical/test_remote_endpoint.cpp (1)
659-661: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winBound the state wait so a regression fails instead of hanging.
Line 661 spins until
RequestState::REQUEST_READYwith no deadline. Line 697 uses the same pattern. Ifexchange_group_taskstops publishing a request, the test hangs and the CI job times out with no diagnostic. Add a deadline andFAIL()when it expires.♻️ Proposed change for the helper
-void respond_with_payloads(std::vector<uint8_t> &mailbox, const std::vector<std::vector<uint8_t>> &payloads) { +void wait_for_request_ready(const std::vector<uint8_t> &mailbox) { + using namespace mpi_group_mailbox; + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(10); + while (mailbox_state(mailbox, OFF_REQUEST_STATE) != static_cast<int32_t>(RequestState::REQUEST_READY)) { + ASSERT_LT(std::chrono::steady_clock::now(), deadline) << "mailbox request was never published"; + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } +} + +void respond_with_payloads(std::vector<uint8_t> &mailbox, const std::vector<std::vector<uint8_t>> &payloads) { using namespace mpi_group_mailbox; - while (mailbox_state(mailbox, OFF_REQUEST_STATE) != static_cast<int32_t>(RequestState::REQUEST_READY)) {} + wait_for_request_ready(mailbox);🤖 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/cpp/hierarchical/test_remote_endpoint.cpp` around lines 659 - 661, Bound the polling loops in both respond_with_payloads and the matching wait near exchange_group_task with a deadline; when RequestState::REQUEST_READY is not observed before expiration, call FAIL() with a diagnostic, while preserving the existing behavior when the state becomes ready.python/simpler/mpi_group_mailbox.py (3)
320-329: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTruncated error JSON becomes unparsable.
If the encoded rank errors exceed
MAILBOX_ERROR_BYTES, line 324 cuts the JSON mid-document.read_resultthen failsjson.loadsand falls back to a raw replacement decode, so the rank attribution is lost. Truncate each message before encoding instead, so the document stays valid.♻️ Proposed fix
- data = json.dumps([asdict(error) for error in errors], sort_keys=True).encode("utf-8") - if len(data) > MAILBOX_ERROR_BYTES: - data = data[: MAILBOX_ERROR_BYTES - 1] + entries = [asdict(error) for error in errors] + data = json.dumps(entries, sort_keys=True).encode("utf-8") + while len(data) > MAILBOX_ERROR_BYTES and entries: + budget = max(0, len(entries[-1]["message"]) // 2) + if budget == 0: + entries.pop() + else: + entries[-1]["message"] = entries[-1]["message"][:budget] + data = json.dumps(entries, sort_keys=True).encode("utf-8") + data = data[:MAILBOX_ERROR_BYTES]🤖 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 `@python/simpler/mpi_group_mailbox.py` around lines 320 - 329, Update the error serialization flow in the mailbox failure-writing method so oversized data remains valid JSON. Truncate individual rank-error message fields before `json.dumps`, then encode and write the complete serialized document without slicing the encoded JSON at `MAILBOX_ERROR_BYTES`; preserve rank attribution and the existing state updates.
345-357: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winNarrow the exception scope around the error decode.
Line 353 catches
BaseException, so aKeyboardInterruptorSystemExitraised during the decode is converted into a plain message. Catch the decode and lookup errors only. Ruff also reports BLE001 here.♻️ Proposed fix
- except BaseException: + except (ValueError, TypeError, KeyError, UnicodeDecodeError): message = raw.decode("utf-8", errors="replace") or "MPI group request failed"🤖 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 `@python/simpler/mpi_group_mailbox.py` around lines 345 - 357, The TASK_FAILED error decoding in the mailbox request handling must not catch control-flow exceptions. Narrow the try/except around json decoding and entry field access to the specific decode and lookup/type errors that can occur, replacing the broad BaseException handler so Ruff BLE001 is resolved while preserving the raw-message fallback.Source: Linters/SAST tools
373-375: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMove the test-only mutator out of the shipped protocol class.
overwrite_request_payload_for_testwrites arbitrary bytes into the request region with no state or capacity check. It is reachable in production. The single caller istests/ut/py/test_mpi_group_mailbox.pyline 96. Write throughmailbox._bufferfrom the test, or guard the helper with a length check againstMAILBOX_PAYLOAD_BYTES.🤖 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 `@python/simpler/mpi_group_mailbox.py` around lines 373 - 375, The test-only method overwrite_request_payload_for_test should not remain exposed on the shipped mailbox protocol class. Remove it and update the sole caller in test_mpi_group_mailbox.py to write the payload directly through mailbox._buffer, or otherwise enforce MAILBOX_PAYLOAD_BYTES capacity before writing if the helper must remain.tests/ut/py/test_mpi_group_mailbox.py (1)
25-211: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a capacity-guard case to this suite.
The suite covers targets, sequencing, failure, and shutdown. It does not cover the capacity guard in
_encode_payloads(MAILBOX_PAYLOAD_BYTES) or the truncation path infail_request. Both are wire-protocol limits. Add one test that writes an oversized payload vector and expectsValueError.🤖 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_mpi_group_mailbox.py` around lines 25 - 211, Add a test covering the mailbox payload capacity guard by attempting to write a request whose encoded payloads exceed MAILBOX_PAYLOAD_BYTES and asserting ValueError. Exercise the write_request path and ensure the oversized payload vector is rejected before acceptance or state progression.tools/mpi_l3_group_smoke.py (1)
41-43: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRestore
SIMPLER_MPI_SMOKE_DIRafter the run.Line 42 sets the variable and never removes it. The temporary directory is deleted when the
withblock exits, so the variable then points at a missing path. Ifrunis ever called from another module or twice, the stale value leaks. Usetry/finallyoros.environ.popat the end of the block.🤖 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 `@tools/mpi_l3_group_smoke.py` around lines 41 - 43, Restore the SIMPLER_MPI_SMOKE_DIR environment variable after the temporary-directory run in the surrounding run flow: save any prior value before assigning output_dir, then restore it in a finally block (or remove it when absent) so repeated or nested calls never retain the deleted path.python/simpler/worker.py (2)
3515-3523: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unreachable node-identity fallback.
_build_remote_manifestruns only for ids inself._remote_worker_ids, and_remote_like_worker_ids()is the union of_remote_worker_idsand_mpi_worker_ids. The condition at Line 3515 is therefore always true, so lines 3520-3523 never execute. Drop the branch, or state the caller contract that makes it reachable.♻️ Proposed simplification
- if worker_id in self._remote_like_worker_ids(): - runtime = self._resolved_global_nodes()[int(worker_id)] - node_rank = runtime.node_rank - node_count = runtime.node_count - global_device_ranks = runtime.global_device_ranks - else: - node_rank = 0 - node_count = 1 - global_device_ranks = spec.global_device_ranks or tuple(range(len(spec.device_ids))) + runtime = self._resolved_global_nodes()[int(worker_id)] + node_rank = runtime.node_rank + node_count = runtime.node_count + global_device_ranks = runtime.global_device_ranks🤖 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 `@python/simpler/worker.py` around lines 3515 - 3523, Remove the unreachable else fallback in _build_remote_manifest and rely directly on the runtime values from _resolved_global_nodes()[int(worker_id)] for node_rank, node_count, and global_device_ranks. Preserve the existing remote-worker caller contract and eliminate the redundant _remote_like_worker_ids() condition.
7295-7310: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winBind the release to the run that allocated the handle.
_release_global_domain_handlereadsself._building_run_resources, so the pending-release queue it appends to is the run whose graph is being built at release time, not the run that allocated the handle. The local path avoids this:_allocate_domaincaptures the owning_RunResourcesin the_release_fnclosure. For aretain_after_run=Truedomain released inside a later run, the fence that frees it is that later run's fence, and no queue entry exists on the allocating run. Double free is prevented by_release_global_domain_nowmemoization, so this is an ordering concern, not a corruption. Bind the owning resources at allocation for symmetry with_release_domain_handle.♻️ Proposed binding at allocation
- _release_fn=self._release_global_domain_handle, + _release_fn=lambda released, owner=resources: self._release_global_domain_handle(released, owner),Then accept the owning resources explicitly:
- def _release_global_domain_handle(self, handle: GlobalCommDomainHandle) -> None: + def _release_global_domain_handle( + self, handle: GlobalCommDomainHandle, resources: _RunResources | None = None + ) -> None: if self._worker is None: return - resources = self._building_run_resources🤖 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 `@python/simpler/worker.py` around lines 7295 - 7310, Bind each global domain handle to its allocating _RunResources when created, mirroring _allocate_domain’s _release_fn closure. Update _release_global_domain_handle to accept and use the owning resources rather than reading self._building_run_resources, so pending releases and fences are associated with the allocating run; preserve the existing cleanup and _release_global_domain_now memoization behavior.python/simpler/orchestrator.py (1)
432-441: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCorrect the docstring: MPI groups are also supported.
The summary line says "without MPI", and the body lists only
Worker.add_workerandWorker.add_remote_worker.Worker._allocate_global_domainalso routes a complete MPI group through_mpi_group_control, anddocs/comm-domain.mddocumentsadd_mpirun_worker_groupmembers. Update the text so users of MPI groups find this API.📝 Proposed docstring fix
- """Create a CommDomain across local and/or remote L3 nodes without MPI. + """Create a CommDomain across local, remote, and MPI-launched L3 nodes. Each member is ``(l3_worker_id, local_l2_worker_id)``. The L3 worker - may have been registered by ``Worker.add_worker`` or - ``Worker.add_remote_worker``. L4 collects every L2 export descriptor, + may have been registered by ``Worker.add_worker``, + ``Worker.add_remote_worker``, or ``Worker.add_mpirun_worker_group``. + L4 collects every L2 export descriptor,🤖 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 `@python/simpler/orchestrator.py` around lines 432 - 441, Update the docstring for the CommDomain creation method to state that MPI groups are supported, removing the “without MPI” limitation and mentioning MPI group registration alongside Worker.add_worker and Worker.add_remote_worker. Ensure the member description reflects add_mpirun_worker_group usage while preserving the existing lifecycle and commit behavior documentation.src/common/platform_comm/comm_sim.cpp (1)
199-216: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueConsider suppressing implicit moves on
GlobalDomainAllocation.
GlobalDomainAllocationownslocal_baseandshm_nameas raw members and defines a destructor that releases both. The implicitly generated move constructor copieslocal_baseand leaves the source pointer non-null, so a moved-from object wouldmunmapthe same address again. The current code always stores the allocation in astd::unique_ptrand never moves it, so no live path is affected. Deleting the copy/move operations, or reusingGlobalPeerMappingfor the local mapping, removes the hazard for future changes.♻️ Optional hardening
struct GlobalDomainAllocation { + GlobalDomainAllocation() = default; ~GlobalDomainAllocation() { if (local_base != nullptr) { munmap(local_base, mapping_size); } if (!shm_name.empty()) { shm_unlink(shm_name.c_str()); } } + GlobalDomainAllocation(const GlobalDomainAllocation &) = delete; + GlobalDomainAllocation &operator=(const GlobalDomainAllocation &) = delete; + GlobalDomainAllocation(GlobalDomainAllocation &&) = delete; + GlobalDomainAllocation &operator=(GlobalDomainAllocation &&) = delete;🤖 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_sim.cpp` around lines 199 - 216, Make GlobalDomainAllocation non-copyable and non-movable by explicitly deleting its copy and move constructors and assignment operators, preventing duplicated ownership of local_base and shm_name while preserving its current unique_ptr-based usage.src/common/worker/chip_worker.h (1)
161-165: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the new public global-domain APIs.
Every neighbouring
comm_*declaration in this header carries a doc comment that states the ownership and pairing contract. The three new declarations carry none. State whatcomm_global_domain_preparereturns (descriptor bytes, local window base, actual mapping size), that the actual mapping size can exceed the requestedwindow_size, thatcomm_global_domain_importrequires a rank-ordered complete table, and thatcomm_global_domain_releasepairs withprepareand also runs fromfinalize().♻️ Suggested doc comment
+ /// Global CommDomain lifecycle (L4-brokered, independent of the + /// comm_init sessions above). `prepare` creates this rank's local window + /// and returns (descriptor_bytes, local_window_base, mapping_size); the + /// mapping size may exceed `window_size` after backend alignment. + /// `import` takes the complete rank-ordered descriptor table and returns + /// the device CommContext. `release` pairs with `prepare` and is also + /// driven for every tracked domain by `finalize()`. std::tuple<std::vector<uint8_t>, uint64_t, size_t> comm_global_domain_prepare(🤖 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/worker/chip_worker.h` around lines 161 - 165, Document the public APIs comm_global_domain_prepare, comm_global_domain_import, and comm_global_domain_release in chip_worker.h. Specify that prepare returns descriptor bytes, the local window base, and actual mapping size, which may exceed window_size; import requires a complete rank-ordered descriptor table; and release pairs with prepare and is also invoked by finalize().python/simpler/global_comm_domain.py (1)
371-384: 🗄️ Data Integrity & Integration | 🔵 Trivial | 💤 Low valueConsider validating the decoded capability result.
Every other decoder in this module re-validates its fields.
decode_comm_init_resultaccepts anyprofile,max_ranks, anddescriptor_bytes. A peer that reports a different descriptor ABI size passes silently, and the mismatch surfaces later during prepare/import. A cheap check here fails fast.♻️ Optional validation
reader.done("COMM_INIT result") + if profile not in GLOBAL_DOMAIN_PROFILE_IDS: + raise ValueError(f"unsupported global domain profile {profile!r}") + if descriptor_bytes != GLOBAL_DOMAIN_DESCRIPTOR_BYTES or max_ranks == 0 or max_ranks > GLOBAL_DOMAIN_MAX_RANKS: + raise ValueError("global comm init result capability is invalid") result = GlobalCommInitResult(🤖 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 `@python/simpler/global_comm_domain.py` around lines 371 - 384, Update decode_comm_init_result to validate the decoded profile, max_ranks, and descriptor_bytes before constructing GlobalCommInitResult, reusing the module’s existing validation helpers or conventions. Reject unsupported values, including descriptor ABI sizes that do not match the expected value, while preserving the existing successful decode flow.src/common/worker/chip_worker.cpp (1)
868-875: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winHarden the descriptor ABI and the release-pointer use.
Two points in this block:
- Line 869 calls
comm_global_domain_release_fn_without a null check, unlikecomm_global_domain_releaseat Line 900.init()resolves both symbols with the requiredload_symboland clears both together, so the pointer is non-null whenevercomm_global_domain_prepare_fn_is non-null. The guard is still worth adding for consistency with the surrounding code.- Lines 873-875 ship the raw
CommGlobalDomainDescriptorbytes to Python, which decodes them with the fixed little-endian layout"<IIIIQII256s". Only the total size is asserted (sizeof(...) == 288in both backends). A field reorder that keeps the size constant would silently misdecode. Add per-fieldoffsetofassertions next to the existing size assertion.♻️ Proposed hardening
if (local_window_base == 0 || descriptor.mapping_size == 0) { - comm_global_domain_release_fn_(domain_id); + if (comm_global_domain_release_fn_ != nullptr) { + comm_global_domain_release_fn_(domain_id); + } global_domain_ids_.erase(domain_id); throw std::runtime_error("comm_global_domain_prepare returned an invalid window"); }Add next to the existing size assertion (for example in
src/common/platform_comm/comm.h):static_assert(offsetof(CommGlobalDomainDescriptor, mapping_size) == 16, "descriptor layout changed"); static_assert(offsetof(CommGlobalDomainDescriptor, handle_size) == 24, "descriptor layout changed"); static_assert(offsetof(CommGlobalDomainDescriptor, handle) == 32, "descriptor layout changed");🤖 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/worker/chip_worker.cpp` around lines 868 - 875, In the invalid-window cleanup within the descriptor preparation flow, guard comm_global_domain_release_fn_ before invoking it, matching the existing comm_global_domain_release handling. Also strengthen CommGlobalDomainDescriptor ABI validation beside its existing size assertion by adding static_assert checks for the mapping_size, handle_size, and handle field offsets required by the Python little-endian decoder.python/bindings/task_interface.cpp (1)
1588-1617: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winConsider releasing the GIL for the global-domain calls.
The three new bindings hold the GIL for the whole native call.
comm_global_domain_preparereserves and maps a VMM/Fabric window, exports a handle, and zeroes the entire window;comm_global_domain_importimports one peer window per rank and copies aCommContextto the device. Both can run for a long time and block every other Python thread in the process. The adjacent device-side helpers in this file already usenb::call_guard<nb::gil_scoped_release>()for exactly this reason, for example_l3_child_onboard_region_createat Line 1839 and_ChipWorker.initat Line 1387.Note that the lambda for
comm_global_domain_prepareconstructsnb::bytesandnb::make_tuplefrom its result, so a whole-lambda call guard is not correct there. Wrap only the native call, or move the Python object construction after the guard scope ends.♻️ Sketch for the prepare binding
[](ChipWorker &self, uint64_t domain_id, uint32_t domain_rank, uint32_t rank_count, size_t window_size, uint32_t profile) { - auto [descriptor, local_window_base, actual_window_size] = - self.comm_global_domain_prepare(domain_id, domain_rank, rank_count, window_size, profile); + std::vector<uint8_t> descriptor; + uint64_t local_window_base = 0; + size_t actual_window_size = 0; + { + nb::gil_scoped_release release; + std::tie(descriptor, local_window_base, actual_window_size) = + self.comm_global_domain_prepare(domain_id, domain_rank, rank_count, window_size, profile); + } return nb::make_tuple(
comm_global_domain_importandcomm_global_domain_releasereturn plain integers, so anb::call_guard<nb::gil_scoped_release>()on the.def(...)is sufficient for those two.🤖 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 `@python/bindings/task_interface.cpp` around lines 1588 - 1617, Release the GIL while executing the native calls in comm_global_domain_prepare, comm_global_domain_import, and comm_global_domain_release. For comm_global_domain_prepare, scope the GIL release only around self.comm_global_domain_prepare so nb::bytes and nb::make_tuple construction still runs with the GIL held; add a whole-binding call guard for the plain-integer import and release methods.
🤖 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/mpi-l3-mailbox.md`:
- Line 1: Add the mpi-l3-mailbox documentation page to the nav configuration in
mkdocs.yml, using the existing title and navigation structure so strict MkDocs
builds include docs/mpi-l3-mailbox.md.
In `@docs/remote-l3-worker-design/implementation-record.md`:
- Around line 95-96: Update the implementation record text to refer to the
shipped A3 Fabric profile as “A3 Fabric V1” rather than “A3 Fabric V2,” matching
the identifiers GLOBAL_DOMAIN_PROFILE_A3_FABRIC and
COMM_GLOBAL_DOMAIN_PROFILE_A3_FABRIC.
In `@python/simpler/global_comm_domain.py`:
- Around line 395-400: Enforce the 64-buffer limit in encode_domain_command
before serializing buffers, using a shared GLOBAL_DOMAIN_MAX_BUFFERS constant.
Update decode_domain_command to use the same constant instead of the existing
hardcoded or unrelated limit, and raise the established validation error when
command.buffers exceeds the bound.
In `@python/simpler/mpi_l3_session.py`:
- Around line 476-495: Add a short sleep/backoff to the mailbox.request_state
polling loop in the MPI session receive flow, and enforce a deadline when the
group remains terminal so it cannot spin indefinitely; also add the same pause
after each request.test() call in the dispatcher loop at
python/simpler/mpi_l3_session.py lines 71-84. Apply both changes in
python/simpler/mpi_l3_session.py:476-495 and
python/simpler/mpi_l3_session.py:71-84, preserving the existing request dispatch
behavior.
- Around line 500-517: Move the _payload_for_rank(request, rank) call inside the
existing try block in the per-rank dispatch logic, before opcode handling.
Preserve the existing MpiRankError conversion so IndexError and ValueError are
captured and included in dispatch_comm.gather rather than escaping the loop.
In `@python/simpler/task_interface.py`:
- Around line 1105-1109: Update TaskInterface.release to set _released before
invoking _release_fn(self), while preserving the existing early return for
already released handles. Keep the callback invocation unchanged so release
failures still propagate, but ensure member() and buffer_range() observe the
handle as released even when the callback raises.
In `@src/common/hierarchical/remote_endpoint.cpp`:
- Around line 904-945: Update both timeout branches in exchange_group_task,
including the dispatch-timeout path after the leader’s run_exchange, to
increment group_departed_ while holding group_mu_ before marking the rendezvous
complete and throwing. Preserve the existing timeout error, terminal marking,
group termination, notification, and exception behavior so the normal reset
logic can clear group_active_ and group_frames_.
- Around line 766-830: Update the polling loop in the request-waiting method
containing OFF_REQUEST_STATE to use bounded backoff instead of continuously
spinning. Add an escalating yield or short sleep on each iteration while
preserving prompt handling of TASK_DONE, SHUTDOWN_DONE, TASK_FAILED, terminal,
and timeout states; reset or initialize the backoff appropriately for each
request.
In `@src/common/platform_comm/comm_sim.cpp`:
- Line 219: Protect accesses to global_domain_allocations in
comm_global_domain_prepare, comm_global_domain_import, and
comm_global_domain_release with a mutex consistently across the sim and HCCL
backends, covering all reads, writes, and erases; alternatively, explicitly
document and enforce the required single-thread caller contract.
In `@task.md`:
- Around line 7-11: Remove personal and internal identifiers from the task
record: replace the branch name, myserver host, and validation host references
in the affected entries with generic placeholders while preserving the remaining
handoff metadata and structure.
In `@tools/a3_l4_tcp_smoke/mpirun_compute_then_tload_2x2_smoke.py`:
- Around line 310-313: Replace the concrete network defaults in
tools/a3_l4_tcp_smoke/mpirun_compute_then_tload_2x2_smoke.py lines 310-313 for
the --host-37, --host-35, --roce-37, and --roce-35 arguments with
documentation-reserved placeholder addresses, or make them required without
defaults. Update the example command in tools/a3_l4_tcp_smoke/README.md lines
17-23 to use the same placeholders.
---
Nitpick comments:
In `@python/bindings/task_interface.cpp`:
- Around line 1588-1617: Release the GIL while executing the native calls in
comm_global_domain_prepare, comm_global_domain_import, and
comm_global_domain_release. For comm_global_domain_prepare, scope the GIL
release only around self.comm_global_domain_prepare so nb::bytes and
nb::make_tuple construction still runs with the GIL held; add a whole-binding
call guard for the plain-integer import and release methods.
In `@python/simpler/global_comm_domain.py`:
- Around line 371-384: Update decode_comm_init_result to validate the decoded
profile, max_ranks, and descriptor_bytes before constructing
GlobalCommInitResult, reusing the module’s existing validation helpers or
conventions. Reject unsupported values, including descriptor ABI sizes that do
not match the expected value, while preserving the existing successful decode
flow.
In `@python/simpler/global_comm_smoke.py`:
- Around line 22-167: Refactor remote_compute_orch, remote_rank_orch,
remote_compute_group_orch, and remote_rank_group_orch to share a helper for
importing get_inner_handle, validating the six scalars, decoding the digest, and
resolving the domain. Add reusable tensor-argument builders for the compute and
TLOAD shapes, then have the group callbacks reuse the corresponding
single-worker argument construction inside their loops while preserving worker
selection and submission behavior.
In `@python/simpler/mpi_group_mailbox.py`:
- Around line 320-329: Update the error serialization flow in the mailbox
failure-writing method so oversized data remains valid JSON. Truncate individual
rank-error message fields before `json.dumps`, then encode and write the
complete serialized document without slicing the encoded JSON at
`MAILBOX_ERROR_BYTES`; preserve rank attribution and the existing state updates.
- Around line 345-357: The TASK_FAILED error decoding in the mailbox request
handling must not catch control-flow exceptions. Narrow the try/except around
json decoding and entry field access to the specific decode and lookup/type
errors that can occur, replacing the broad BaseException handler so Ruff BLE001
is resolved while preserving the raw-message fallback.
- Around line 373-375: The test-only method overwrite_request_payload_for_test
should not remain exposed on the shipped mailbox protocol class. Remove it and
update the sole caller in test_mpi_group_mailbox.py to write the payload
directly through mailbox._buffer, or otherwise enforce MAILBOX_PAYLOAD_BYTES
capacity before writing if the helper must remain.
In `@python/simpler/mpi_l3_session.py`:
- Around line 504-513: Replace the assert in the MailboxOpcode.SHUTDOWN branch
of the request handling flow with an explicit payload None check that raises an
appropriate error before calling _rewrite_frame_identity; preserve the existing
shutdown behavior when payload is present.
In `@python/simpler/orchestrator.py`:
- Around line 432-441: Update the docstring for the CommDomain creation method
to state that MPI groups are supported, removing the “without MPI” limitation
and mentioning MPI group registration alongside Worker.add_worker and
Worker.add_remote_worker. Ensure the member description reflects
add_mpirun_worker_group usage while preserving the existing lifecycle and commit
behavior documentation.
In `@python/simpler/worker.py`:
- Around line 3515-3523: Remove the unreachable else fallback in
_build_remote_manifest and rely directly on the runtime values from
_resolved_global_nodes()[int(worker_id)] for node_rank, node_count, and
global_device_ranks. Preserve the existing remote-worker caller contract and
eliminate the redundant _remote_like_worker_ids() condition.
- Around line 7295-7310: Bind each global domain handle to its allocating
_RunResources when created, mirroring _allocate_domain’s _release_fn closure.
Update _release_global_domain_handle to accept and use the owning resources
rather than reading self._building_run_resources, so pending releases and fences
are associated with the allocating run; preserve the existing cleanup and
_release_global_domain_now memoization behavior.
In `@src/common/platform_comm/comm_sim.cpp`:
- Around line 199-216: Make GlobalDomainAllocation non-copyable and non-movable
by explicitly deleting its copy and move constructors and assignment operators,
preventing duplicated ownership of local_base and shm_name while preserving its
current unique_ptr-based usage.
In `@src/common/worker/chip_worker.cpp`:
- Around line 868-875: In the invalid-window cleanup within the descriptor
preparation flow, guard comm_global_domain_release_fn_ before invoking it,
matching the existing comm_global_domain_release handling. Also strengthen
CommGlobalDomainDescriptor ABI validation beside its existing size assertion by
adding static_assert checks for the mapping_size, handle_size, and handle field
offsets required by the Python little-endian decoder.
In `@src/common/worker/chip_worker.h`:
- Around line 161-165: Document the public APIs comm_global_domain_prepare,
comm_global_domain_import, and comm_global_domain_release in chip_worker.h.
Specify that prepare returns descriptor bytes, the local window base, and actual
mapping size, which may exceed window_size; import requires a complete
rank-ordered descriptor table; and release pairs with prepare and is also
invoked by finalize().
In `@tests/ut/cpp/hierarchical/test_remote_endpoint.cpp`:
- Around line 659-661: Bound the polling loops in both respond_with_payloads and
the matching wait near exchange_group_task with a deadline; when
RequestState::REQUEST_READY is not observed before expiration, call FAIL() with
a diagnostic, while preserving the existing behavior when the state becomes
ready.
In `@tests/ut/py/test_global_comm_domain.py`:
- Around line 259-278: Update _failure_injection_worker to obtain two ports
through the existing _free_tcp_ports helper, then build each RemoteWorkerSpec
endpoint from those allocated ports instead of the fixed 19073 + index range.
Preserve the current node ordering and endpoint construction format.
In `@tests/ut/py/test_mpi_group_mailbox.py`:
- Around line 25-211: Add a test covering the mailbox payload capacity guard by
attempting to write a request whose encoded payloads exceed
MAILBOX_PAYLOAD_BYTES and asserting ValueError. Exercise the write_request path
and ensure the oversized payload vector is rejected before acceptance or state
progression.
In `@tools/a3_l4_tcp_smoke/mpirun_compute_then_tload_2x2_smoke.py`:
- Around line 51-60: Update the kernel_include_dirs construction in _compile_aiv
to use iterable unpacking when combining include_dirs with the common source
directory, preserving the existing ordering and values.
In `@tools/mpi_group_mailbox_smoke.py`:
- Around line 34-37: Update _wait_until to pause briefly between predicate
evaluations, while retaining the existing deadline check and TimeoutError
behavior; add the sleep inside the loop after a failed predicate check.
In `@tools/mpi_l3_group_smoke.py`:
- Around line 41-43: Restore the SIMPLER_MPI_SMOKE_DIR environment variable
after the temporary-directory run in the surrounding run flow: save any prior
value before assigning output_dir, then restore it in a finally block (or remove
it when absent) so repeated or nested calls never retain the deleted path.
🪄 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: a870b115-1c58-4057-9696-095762c08e11
📒 Files selected for processing (50)
docs/comm-domain.mddocs/mpi-l3-mailbox.mddocs/remote-l3-worker-design.mddocs/remote-l3-worker-design/implementation-record.mddocs/remote-l3-worker-design/protocol.mdpython/bindings/CMakeLists.txtpython/bindings/task_interface.cpppython/bindings/worker_bind.hpython/simpler/global_comm_domain.pypython/simpler/global_comm_smoke.pypython/simpler/mpi_group_mailbox.pypython/simpler/mpi_group_smoke.pypython/simpler/mpi_l3_session.pypython/simpler/orchestrator.pypython/simpler/remote_l3_protocol.pypython/simpler/remote_l3_session.pypython/simpler/remote_l3_worker.pypython/simpler/task_interface.pypython/simpler/worker.pysrc/a2a3/platform/onboard/host/comm_hccl.cppsrc/a5/platform/onboard/host/comm_hccl.cppsrc/common/hierarchical/mpi_group_mailbox.hsrc/common/hierarchical/remote_endpoint.cppsrc/common/hierarchical/remote_endpoint.hsrc/common/hierarchical/remote_wire.cppsrc/common/hierarchical/remote_wire.hsrc/common/hierarchical/worker.cppsrc/common/hierarchical/worker.hsrc/common/hierarchical/worker_manager.cppsrc/common/hierarchical/worker_manager.hsrc/common/platform_comm/comm.hsrc/common/platform_comm/comm_sim.cppsrc/common/worker/chip_worker.cppsrc/common/worker/chip_worker.htask.mdtests/ut/cpp/CMakeLists.txttests/ut/cpp/hierarchical/test_remote_endpoint.cpptests/ut/py/test_callable_identity.pytests/ut/py/test_global_comm_domain.pytests/ut/py/test_mpi_group_mailbox.pytests/ut/py/test_mpi_l3_group.pytests/ut/py/test_worker/test_startup_readiness.pytools/a3_l4_tcp_smoke/README.mdtools/a3_l4_tcp_smoke/kernels/aiv/global_tload_kernel.cpptools/a3_l4_tcp_smoke/kernels/aiv/local_add_kernel.cpptools/a3_l4_tcp_smoke/kernels/orchestration/global_tload_orch.cpptools/a3_l4_tcp_smoke/kernels/orchestration/local_add_orch.cpptools/a3_l4_tcp_smoke/mpirun_compute_then_tload_2x2_smoke.pytools/mpi_group_mailbox_smoke.pytools/mpi_l3_group_smoke.py
8e8a7c4 to
a348313
Compare
6918ac8 to
9811bdf
Compare
9811bdf to
c536dfe
Compare
- Route MPI group task and control traffic through a named rank-0 shared-memory mailbox: requests reach the ranks over one MPI bcast and ranked results return via gather, with no per-rank TCP sessions - Mark group-wide domain controls with an explicit FRAME_FLAG_GROUP_TARGET frame flag; unmarked controls stay rank-targeted, so partial-group Global CommDomains keep their per-node fallback semantics - Batch a full-group submit_next_level_group into one PER_RANK envelope through per-transport progress helper threads, keeping the scheduler's submit/poll non-blocking; rendezvous timeouts fail only the waiting task and never kill mpirun, while the mailbox round trip under lane_mu_ owns terminal and kill semantics - Reject group selections that mix MPI ranks with other workers at the group's world size; reap mpirun children and unlink the mailbox even when the native worker close fails - Preserve the TCP transport for ordinary non-MPI Remote L3 workers - Cover the batched path in the mpirun pod ST example and transport-level unit tests Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
c536dfe to
cd51a88
Compare
|
Reviewed at Stated vs. real goalThe body says: replace the MPI group's TCP command/health path with one L4-owned named shared-memory mailbox attached by local rank 0, distribute ordered task/control envelopes over dedicated MPI communicators, keep ordinary remote L3 on TCP. Reading the code independently, that is what it does — no scope mismatch. Change breakdown❌ Oversized (Core 2375 lines). GitHub's 37 files / +4051 is against a stale base; the real three-dot diff is the above. Still well past the 1000-line mark, and it is a new cross-language wire protocol plus its two endpoints plus a transport swap. A reviewer cannot hold this in one pass — the protocol header ( Mechanism briefThe problem: an MPI group previously needed each rank to bind fixed, pre-agreed command and health TCP ports so L4 could attach to each one as an ordinary remote L3. That means N sockets, N port reservations that survive a killed run, and a control lane that duplicates what MPI already provides between ranks. The new shape: L4 creates one POSIX shm mailbox ( Layout is a 256-byte header plus two 16 MiB payload regions plus a 64 KiB error region. Two state machines: Ownership: L4 owns and unlinks; Failure model: FindingsMust fix1. The atomic accessors silently fall back to non-atomic reads — the fallback is always taken. try:
from .task_interface import _mailbox_load_i32
except (ImportError, AttributeError):
return self._read_i32(offset)
So every state transition on this mailbox uses 2. Rank 0's request-wait loop is an unbounded busy-spin. In
Should fix3. The mailbox layout is duplicated in Python and C++ with nothing checking they agree. 4. The PR body's testing section is stale in a way that undersells the PR. It lists the hardware items as unchecked with "The hardware and MPI integration items remain explicitly unvalidated until the server-37 agent returns logs." But 5. No stated version-evolution rule for a brand-new cross-process ABI. Consider6. 7. Error truncation is silent. pto-isa pin check — advisoryℹ️ Verification I ran
VerdictRequest changes — on findings 1 and 2 only. Both are small and local: one import line, and one blocking primitive in the idle wait. The design itself is good and I'd approve it otherwise. Collapsing N sockets to one mailbox plus MPI's own collectives is the right call, the sequence-id replay guard and sticky TERMINAL state show the failure modes were thought through, and the resource-tracker workaround is the kind of detail that usually gets found in production instead. Finding 1 is worth blocking on specifically because it is invisible: the protocol is designed around acquire/release semantics and currently doesn't have them, with no signal that anything is wrong. Given Core is 2375 lines, I'd also suggest that any future protocol like this land as two PRs — header/codec first with its tests, then the endpoints — so the wire format gets reviewed on its own terms before two implementations depend on it. |
- Import the mailbox acquire/release helpers from _task_interface and fail the import when the extension is missing: the cross-process handshake depends on their publication barriers, and the previous simpler.task_interface import always fell back to plain struct reads because that module does not re-export the helpers - Park rank 0 on the request-state word between requests -- a shared futex on Linux, woken by every L4 publish -- with a bounded chunk that re-checks group state, so an idle group no longer holds a full host core and a lost wake cannot stall the loop - Validate reserved header bytes [80, 256) as zero on both attach paths and document the version-evolution rule, so a version-1 peer cannot silently carry fields it does not understand - Export the C++ wire layout as _mpi_mailbox_layout and assert the Python declaration matches it, so an edit to one side of the hand-mirrored layout fails a unit test instead of corrupting the lane at runtime - Cap per-rank error messages and replace dropped tail entries with a sentinel entry, so the gathered error blob always parses as JSON - Remove overwrite_request_payload_for_test from the production surface; the test writes through the shared-memory buffer directly
|
Re-reviewed at Finding 1 — atomics: fixed, and I confirmed the fix actually resolvesThe silent fallback is gone; it's now a hard module-level import from Worth noting how this played out: on first checkout the import raised Finding 2 — idle spin: fixed with a futex, not a sleep
This is the shape One asymmetry I checked and am not flagging: Finding 3 — layout duplication: fixed, and the guard is real
I verified it is wired in the direction that catches divergence (the dict comes from Finding 5 — reserved bytes: fixed on both sides
Consider items 6 and 7 — both taken
BodyAlso fixed — the stale "hardware and MPI integration items remain explicitly unvalidated" section is replaced with the actual evidence, including the re-validation on this head. That was my finding 4 and it is fully addressed. Verification I ran on
|
2 participants
Summary
mpirunprocess-group cleanupReview-round hardening (
87489e1): the mailbox state words now use the_task_interfaceacquire/release helpers unconditionally (the previous import path silently fell back to plain struct reads), rank 0 parks on a shared futex instead of busy-polling the request lane, reserved header bytes [80, 256) are validated as zero with a documented version-evolution rule, a_mpi_mailbox_layoutbinding plus unit test pins the Python and C++ layout declarations together, gathered rank errors stay valid JSON under truncation, and the test-only payload-overwrite API is removed.Dependency
Stacked on #1623, which is now merged; this branch is rebased onto its merge commit.
Testing
cd51a88, includingst-pod-onboard-a2a3— the two-machinetest_global_tload_mpirun_l3example runs the mailbox path end to end (PASS 15.0s, devices=[12, 13])tests/ut/pyon Linux: 1440 passed, 13 skipped (includes this PR's 46 mailbox/MPI tests)87489e1: 18 checks green, includingst-pod-onboard-a2a3re-runningtest_global_tload_mpirun_l3on the futex-park path (PASS 15.1s, devices=[12, 13])