From 34c8b74258f40112b4b3f0c7465ee4d3627b60ee Mon Sep 17 00:00:00 2001 From: sunkaixuan2018 Date: Fri, 24 Jul 2026 14:41:36 +0800 Subject: [PATCH] Add: support mixed local and remote L4 CommDomains - Relay rank-ordered Fabric descriptors through L4 to local and remote L3 nodes - Route local Global CommDomain control through NEXT_LEVEL mailboxes - Reject unsupported platform/profile capabilities before allocation and at COMM_INIT - Make same-name release tracking identity-safe and release failures terminal - Replace message matching with a dedicated childless-buffer exception - Add transaction failure injection plus remote, mixed-node, compute, and communication smokes - Document no-mpirun setup and retained Global CommDomain lifetimes Co-authored-by: Leaf-Salix <2503954024@qq.com> --- .github/workflows/ci.yml | 28 +- docs/comm-domain.md | 62 + docs/remote-l3-worker-design.md | 16 +- .../implementation-record.md | 32 +- docs/remote-l3-worker-design/protocol.md | 24 +- python/bindings/CMakeLists.txt | 1 + python/bindings/task_interface.cpp | 30 + python/bindings/worker_bind.h | 38 + python/simpler/global_comm_domain.py | 574 +++++++ python/simpler/global_comm_smoke.py | 88 ++ python/simpler/orchestrator.py | 85 + python/simpler/remote_l3_protocol.py | 2 + python/simpler/remote_l3_session.py | 191 ++- python/simpler/remote_l3_worker.py | 17 + python/simpler/task_interface.py | 130 ++ python/simpler/worker.py | 1378 ++++++++++++++++- src/a2a3/platform/onboard/host/comm_hccl.cpp | 191 +++ src/a5/platform/onboard/host/comm_hccl.cpp | 16 + src/common/hierarchical/remote_endpoint.cpp | 6 + src/common/hierarchical/remote_endpoint.h | 2 + src/common/hierarchical/remote_wire.cpp | 2 + src/common/hierarchical/remote_wire.h | 2 + src/common/hierarchical/worker.h | 10 + src/common/hierarchical/worker_manager.cpp | 46 + src/common/hierarchical/worker_manager.h | 10 + src/common/platform_comm/comm.h | 79 + src/common/platform_comm/comm_sim.cpp | 178 +++ src/common/worker/chip_worker.cpp | 82 + src/common/worker/chip_worker.h | 16 + tests/ut/cpp/CMakeLists.txt | 1 + tests/ut/py/test_callable_identity.py | 78 + tests/ut/py/test_global_comm_domain.py | 709 +++++++++ tools/a3_l4_tcp_smoke/README.md | 80 + .../compute_then_tload_smoke.py | 282 ++++ tools/a3_l4_tcp_smoke/global_tload_smoke.py | 178 +++ .../kernels/aiv/global_tload_kernel.cpp | 86 + .../kernels/aiv/local_add_kernel.cpp | 61 + .../orchestration/global_tload_orch.cpp | 35 + .../kernels/orchestration/local_add_orch.cpp | 34 + .../mixed_global_tload_smoke.py | 212 +++ tools/remote_l4_npu/README.md | 60 + tools/remote_l4_npu/remote_l4_npu_smoke.py | 282 ++++ tools/remote_l4_npu/run_parent_smoke.sh | 34 + tools/remote_l4_npu/start_machine_daemon.sh | 26 + 44 files changed, 5458 insertions(+), 36 deletions(-) create mode 100644 python/simpler/global_comm_domain.py create mode 100644 python/simpler/global_comm_smoke.py create mode 100644 tests/ut/py/test_global_comm_domain.py create mode 100644 tools/a3_l4_tcp_smoke/README.md create mode 100644 tools/a3_l4_tcp_smoke/compute_then_tload_smoke.py create mode 100644 tools/a3_l4_tcp_smoke/global_tload_smoke.py create mode 100644 tools/a3_l4_tcp_smoke/kernels/aiv/global_tload_kernel.cpp create mode 100644 tools/a3_l4_tcp_smoke/kernels/aiv/local_add_kernel.cpp create mode 100644 tools/a3_l4_tcp_smoke/kernels/orchestration/global_tload_orch.cpp create mode 100644 tools/a3_l4_tcp_smoke/kernels/orchestration/local_add_orch.cpp create mode 100644 tools/a3_l4_tcp_smoke/mixed_global_tload_smoke.py create mode 100644 tools/remote_l4_npu/README.md create mode 100644 tools/remote_l4_npu/remote_l4_npu_smoke.py create mode 100755 tools/remote_l4_npu/run_parent_smoke.sh create mode 100755 tools/remote_l4_npu/start_machine_daemon.sh diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3e2f0307fe..b7bd747bf7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -789,6 +789,9 @@ jobs: if: false && needs.detect-changes.outputs.docs_only != 'true' runs-on: [self-hosted, a5] timeout-minutes: 30 + defaults: + run: + working-directory: source env: SIMPLER_SCHEDULER_TIMEOUT_MS: "2000" SIMPLER_OP_EXECUTE_TIMEOUT_US: "3000000" @@ -797,6 +800,11 @@ jobs: steps: - name: Checkout repository uses: actions/checkout@v5 + with: + # The legacy workspace root contains an undeletable privileged + # .pytest_cache. Use a clean child checkout instead. + path: source + - name: Set up environment run: | python3 -m venv --system-site-packages .venv @@ -810,7 +818,7 @@ jobs: run: | source .venv/bin/activate DEVICE_LIST=$(python -c "p='${DEVICE_RANGE}'.split('-'); s,e=p[0],p[-1]; print(','.join(str(i) for i in range(int(s),int(e)+1)))") - task-submit --timeout 1800 --max-time 1800 --device "$DEVICE_LIST" --run "python -m pytest tests -m requires_hardware --platform a5 --device ${DEVICE_RANGE} -v" + task-submit --timeout 1800 --max-time 1800 --device "$DEVICE_LIST" --run "python -m pytest tests -m requires_hardware -p no:cacheprovider --platform a5 --device ${DEVICE_RANGE} -v" - name: Build and run C++ hardware unit tests (a5) run: | @@ -860,6 +868,9 @@ jobs: if: false && needs.detect-changes.outputs.a5_changed == 'true' runs-on: [self-hosted, a5] timeout-minutes: 60 + defaults: + run: + working-directory: source env: SIMPLER_SCHEDULER_TIMEOUT_MS: "2000" SIMPLER_OP_EXECUTE_TIMEOUT_US: "3000000" @@ -868,6 +879,11 @@ jobs: steps: - name: Checkout repository uses: actions/checkout@v5 + with: + # The legacy workspace root contains an undeletable privileged + # .pytest_cache. Use a clean child checkout instead. + path: source + - name: Set up environment run: | python3 -m venv --system-site-packages .venv @@ -881,7 +897,7 @@ jobs: run: | source .venv/bin/activate DEVICE_LIST=$(python -c "p='${DEVICE_RANGE}'.split('-'); s,e=p[0],p[-1]; print(','.join(str(i) for i in range(int(s),int(e)+1)))") - PYTEST="python -m pytest examples tests/st --platform a5 --device ${DEVICE_RANGE} -v --require-pto-isa" + PYTEST="python -m pytest examples tests/st -p no:cacheprovider --platform a5 --device ${DEVICE_RANGE} -v --require-pto-isa" task-submit --timeout 1800 --max-time 1800 --device "$DEVICE_LIST" --run "$PYTEST --pto-session-timeout 1200" # DFX per-feature smokes — hardware mirror of the st-sim-a5 set. The @@ -893,7 +909,7 @@ jobs: source .venv/bin/activate DEVICE_LIST=$(python -c "p='${DEVICE_RANGE}'.split('-'); s,e=p[0],p[-1]; print(','.join(str(i) for i in range(int(s),int(e)+1)))") PYTEST="python -m pytest tests/st/a5/tensormap_and_ringbuffer/dfx/dep_gen/test_dep_gen.py \ - --platform a5 --device ${DEVICE_RANGE} -p no:xdist --pto-session-timeout 600 \ + --platform a5 --device ${DEVICE_RANGE} -p no:xdist -p no:cacheprovider --pto-session-timeout 600 \ --require-pto-isa --enable-dep-gen" task-submit --timeout 1800 --max-time 1800 --device "$DEVICE_LIST" --run "$PYTEST" @@ -902,7 +918,7 @@ jobs: source .venv/bin/activate DEVICE_LIST=$(python -c "p='${DEVICE_RANGE}'.split('-'); s,e=p[0],p[-1]; print(','.join(str(i) for i in range(int(s),int(e)+1)))") PYTEST="python -m pytest tests/st/a5/tensormap_and_ringbuffer/dfx/l2_swimlane/ \ - --platform a5 --device ${DEVICE_RANGE} -p no:xdist --pto-session-timeout 600 \ + --platform a5 --device ${DEVICE_RANGE} -p no:xdist -p no:cacheprovider --pto-session-timeout 600 \ --require-pto-isa --enable-l2-swimlane --enable-dep-gen" task-submit --timeout 1800 --max-time 1800 --device "$DEVICE_LIST" --run "$PYTEST" @@ -911,7 +927,7 @@ jobs: source .venv/bin/activate DEVICE_LIST=$(python -c "p='${DEVICE_RANGE}'.split('-'); s,e=p[0],p[-1]; print(','.join(str(i) for i in range(int(s),int(e)+1)))") PYTEST="python -m pytest tests/st/a5/tensormap_and_ringbuffer/dfx/pmu/test_pmu.py \ - --platform a5 --device ${DEVICE_RANGE} -p no:xdist --pto-session-timeout 600 \ + --platform a5 --device ${DEVICE_RANGE} -p no:xdist -p no:cacheprovider --pto-session-timeout 600 \ --require-pto-isa --enable-pmu 2" task-submit --timeout 1800 --max-time 1800 --device "$DEVICE_LIST" --run "$PYTEST" @@ -920,6 +936,6 @@ jobs: source .venv/bin/activate DEVICE_LIST=$(python -c "p='${DEVICE_RANGE}'.split('-'); s,e=p[0],p[-1]; print(','.join(str(i) for i in range(int(s),int(e)+1)))") PYTEST="python -m pytest tests/st/a5/tensormap_and_ringbuffer/dfx/args_dump/test_args_dump.py \ - --platform a5 --device ${DEVICE_RANGE} -p no:xdist --pto-session-timeout 600 \ + --platform a5 --device ${DEVICE_RANGE} -p no:xdist -p no:cacheprovider --pto-session-timeout 600 \ --require-pto-isa --dump-args" task-submit --timeout 1800 --max-time 1800 --device "$DEVICE_LIST" --run "$PYTEST" diff --git a/docs/comm-domain.md b/docs/comm-domain.md index e0b3f8bfc4..cfef07f4f0 100644 --- a/docs/comm-domain.md +++ b/docs/comm-domain.md @@ -48,6 +48,68 @@ allocation: if `sum(b.nbytes) > window_size`, `allocate_domain` raises Kernels read peer windows through `device_ctx` (which holds every rank's window base, local + imported peer); `buffer_ptrs[name]` is the local slice. +### Global CommDomain across local and remote L3 nodes + +An L4 worker can build the same `CommContext` shape across any combination of +forked local L3 workers (`add_worker`) and TCP-connected L3 workers +(`add_remote_worker`) without `mpirun`: + +```python +with orch.allocate_global_domain( + name="tp", + members=[(node0_worker_id, 0), (node1_worker_id, 0)], + window_size=4096, + buffers=[CommBufferSpec("payload", "uint8", 4096, 4096)], +) as domain: + ... +``` + +Each member is `(l3_worker_id, local_l2_worker_id)`. The order defines dense +domain ranks. A remote node reads `comm_profile` and `global_device_ranks` +from `RemoteWorkerSpec`; a local L3 reads the same fields from its `Worker` +configuration. All participating nodes must use the same profile. + +Global CommDomain capability follows the backend that the node actually +loads: a platform ending in `sim` supports the `sim` profile, and a real +`a2a3` platform supports `a3-fabric-v1`. Real A5 and any other +platform/profile combination currently reject allocation before `PREPARE`. +Each local or remote L3 repeats the same check during `COMM_INIT`, so an +unsupported backend never advertises a usable descriptor capability. + +The control flow is: + +1. L4 sends `COMM_INIT` with cluster, node, global-device, and domain-rank + identities. +2. Each L3 asks its participating L2 children to create a local window and + export a transport descriptor. +3. L4 validates and assembles one complete rank-ordered descriptor table. +4. L4 returns that table to every L3, which forwards it to each L2 for import. +5. L4 commits only after all imports succeed. Any earlier failure sends + `ABORT` and releases every prepared local window. + +The descriptor reports the backend's actual mapped size. A3 Fabric may align +the requested size to its VMM granularity; buffer carving and bounds checks +therefore use the returned mapping size. Handles and device pointers never +cross the public Python API. Remote orchestration code calls +`orch.get_global_domain(domain_id)` to obtain only its committed L3-local +contexts. + +`copy_to_global_domain` and `copy_from_global_domain` provide bounded +control-plane staging and smoke checks. Normal communication still runs in +L2 kernels through the imported `CommContext`. + +By default a live Global CommDomain is swept after the current `Worker.run` +drains. Set `retain_after_run=True` when a communication kernel writes results +into the window and a second L4 run must inspect them. The later run should +call `domain.release()` after copying the results; `Worker.close()` is the +final safety net. + +The remote-only end-to-end A3 smoke is +`tools/a3_l4_tcp_smoke/global_tload_smoke.py`; the local-plus-remote variant is +`tools/a3_l4_tcp_smoke/mixed_global_tload_smoke.py`. They do not use `mpirun` +and pass only when every L2 kernel successfully performs peer `TLOAD` through +the L4-brokered Fabric descriptor table. + --- ## 2. Lifetime model diff --git a/docs/remote-l3-worker-design.md b/docs/remote-l3-worker-design.md index 607a6f9bc4..1324911b62 100644 --- a/docs/remote-l3-worker-design.md +++ b/docs/remote-l3-worker-design.md @@ -67,11 +67,18 @@ Implemented: imported-handle scheduling eligibility, and deferred owner free. - Registry-scope-aware remote callable manifest/control install for dispatcher `PYTHON_IMPORT`, inner `PYTHON_IMPORT`, and inner inline `CHIP_CALLABLE`. + Pre-init `ChipCallable` registrations on an L4 worker are serialized into + each remote session manifest and installed on that L3's L2 children. +- A no-`mpirun` A3 TCP smoke that keeps a Global CommDomain across two L4 + runs, executes peer `TLOAD` from each remote L2, and verifies the reduced + values before release. +- Two-server hardware validation covers L4-brokered peer `TLOAD`, one L2 + compute followed by cross-machine communication, and two-NPU-per-node + remote L3 group compute. Still pending: - A2 RoCE, A3 HCCS, and A5 UB HCOMM profiles. -- Remote `CommDomain` allocation/import and hardware-gated validation. - Negotiated `PYTHON_SERIALIZED` remote callable payloads and staged `CHIP_CALLABLE` blob adapters. @@ -449,10 +456,9 @@ Session execution rules: the current one-`WorkerThread`-per-child local scheduling model and keeps ordering, buffer lifetime, and callable visibility simple. - State-changing CONTROL frames such as register, unregister, buffer free, - copy, export/import, and import release serialize with TASK execution on the - ordered command lane. They are not applied concurrently with a running TASK - on the same endpoint. Future Remote CommDomain controls follow the same - ordering rule when they enter scope. + copy, export/import, import release, and Global CommDomain transactions + serialize with TASK execution on the ordered command lane. They are not + applied concurrently with a running TASK on the same endpoint. - Bulk data movement may use a separate data plane, but the state change that makes staged bytes, callable payloads, or imported handles visible is ordered by the command lane. diff --git a/docs/remote-l3-worker-design/implementation-record.md b/docs/remote-l3-worker-design/implementation-record.md index 9958df3816..0b59682ffa 100644 --- a/docs/remote-l3-worker-design/implementation-record.md +++ b/docs/remote-l3-worker-design/implementation-record.md @@ -16,7 +16,7 @@ It is updated as each documented feature is completed and verified. | 5 | Versioned remote frame codec | In progress | TASK/COMPLETION/CONTROL_REPLY/HELLO/CONTROL/HEALTH exist; core fuzz/bounds coverage is present, with more exhaustive corpus testing still possible. | | 6 | Remote callable registry | In progress | Dispatcher `PYTHON_IMPORT`, inner manifest/control `PYTHON_IMPORT`, and inner manifest/control inline `CHIP_CALLABLE` are implemented; serialized payloads and staged chip blobs remain negotiated extensions. | | 7 | Fork-safe simulation session runner | In progress | Daemon/session bootstrap and HELLO READY barrier are implemented for sim transport. | -| 8 | Remote control-plane parity | In progress | Registry, alloc/free/copy, export/import/release-import controls are implemented for sim; Remote CommDomain controls are reserved/unsupported. | +| 8 | Remote control-plane parity | In progress | Registry, remote buffers, and Global CommDomain prepare/import/commit/release/copy controls are implemented. | | 9 | Remote buffer registry | In progress | Sim owner/imported buffers, TASK materialization, public memory API, opaque handles, slot/import-ref capture, and deferred free/release-import are implemented. | | 10 | A2 RoCE HCOMM profile | Pending | Hardware-gated profile. | | 11 | A3 HCCS HCOMM profile | Pending | Hardware-gated profile. | @@ -88,6 +88,26 @@ It is updated as each documented feature is completed and verified. Imports use shared-memory backed mappings in the session runner, imported handles remain opaque on the parent, and owner frees wait for live imports and slot refs to drain. +- Added L4-brokered Global CommDomain setup without MPI. L2 export + descriptors are collected by L3, assembled by L4, returned to every L3/L2 + for import, and released after the L4 DAG drain by default. Domains created + with `retain_after_run=True` remain live for a later run until explicitly + released or the Worker closes. Sim shm and A3 Fabric V2 use the same + descriptor ABI. +- Added startup-manifest delivery for pre-registered inner `CHIP_CALLABLE` + payloads and `tools/a3_l4_tcp_smoke/global_tload_smoke.py`. The smoke uses + TCP daemons for launch/control, executes peer Fabric `TLOAD` in L2 kernels, + and verifies results without `mpirun`. +- Added `tools/remote_l4_npu/remote_l4_npu_smoke.py` as the real-device + control/data-path baseline. L4 submits one task to each remote L3, each L3 + runs a two-NPU local group, and L4 copies the results back for a golden + check. Remote buffers now use L3-owned child-visible host buffers whenever + the L3 has forked chip children, while childless sim sessions keep the + shared-memory fallback. +- Added `tools/a3_l4_tcp_smoke/compute_then_tload_smoke.py`. One persistent L4 + session first dispatches a local L2 vector-add through every remote L3, then + dispatches a cross-machine peer `TLOAD` over the computed values. The smoke + checks both the per-rank compute output and the final communication output. - Documented the v1 remote registry target/kind matrix, inner `INNER_L3_WORKER` visibility rules, remote `CHIP_CALLABLE` staged/inline payload contract, partial-register cleanup outcomes, and health-expiry @@ -95,6 +115,16 @@ It is updated as each documented feature is completed and verified. ## Verification +- Global CommDomain codec/validation tests and the Linux two-daemon sim + transaction test live in `tests/ut/py/test_global_comm_domain.py`. +- The A3 two-server hardware commands in + `tools/a3_l4_tcp_smoke/README.md` passed for both peer `TLOAD` and + compute-then-`TLOAD`; every rank reported zero maximum difference. +- The two-machine real-NPU compute baseline is documented in + `tools/remote_l4_npu/README.md`. It intentionally verifies remote L4 + dispatch and per-machine L3/L2 compute separately from the cross-machine + Global CommDomain smoke. The two-NPU-per-node run passed with zero maximum + difference for every output. - Python focused sidecar/callable tests: `tests/ut/py/test_task_interface.py tests/ut/py/test_callable_identity.py` passed with `145 passed`. diff --git a/docs/remote-l3-worker-design/protocol.md b/docs/remote-l3-worker-design/protocol.md index d4e3c24da9..a9cf9f3f4b 100644 --- a/docs/remote-l3-worker-design/protocol.md +++ b/docs/remote-l3-worker-design/protocol.md @@ -296,15 +296,29 @@ Required remote controls: - `IMPORT_BUFFER` - `RELEASE_IMPORT` -Reserved future controls for Remote CommDomain: +Required Global CommDomain controls: - `COMM_INIT` - `ALLOC_DOMAIN` - `RELEASE_DOMAIN` - -The first Remote L3 task-dispatch cut rejects the reserved domain controls -with an unsupported-control reply. They become required only when Remote -CommDomain enters scope. +- `COPY_TO_DOMAIN` +- `COPY_FROM_DOMAIN` + +`COMM_INIT` validates the cluster id, node identity, communication profile, +global device ranks, and dense domain-rank table. `ALLOC_DOMAIN` is a +transaction with `PREPARE_EXPORT`, `IMPORT`, `COMMIT`, and `ABORT` phases. +Each L2 exports its local transport descriptor during prepare. L4 assembles +the complete rank-ordered table and sends it to every L3; each L3 forwards it +to its L2 children for import. No domain becomes visible to a remote task +before every node acknowledges `COMMIT`. + +`RELEASE_DOMAIN` is idempotent and normally runs after the L4 DAG drain. +An allocation marked `retain_after_run` may remain live for a later L4 run +that reads kernel results; explicit release or session shutdown then performs +the same teardown. +`COPY_TO_DOMAIN` and `COPY_FROM_DOMAIN` are bounded smoke/control data +operations for a committed local window. They do not replace kernel data +movement through `CommContext`. The register-family controls are registry-scope-aware. `PREPARE_REGISTER_CALLABLE` carries: diff --git a/python/bindings/CMakeLists.txt b/python/bindings/CMakeLists.txt index 75d24dff82..f1df0c248b 100644 --- a/python/bindings/CMakeLists.txt +++ b/python/bindings/CMakeLists.txt @@ -61,6 +61,7 @@ target_include_directories(_task_interface PRIVATE ${CMAKE_SOURCE_DIR}/src/common/task_interface ${CMAKE_SOURCE_DIR}/src/common/worker ${CMAKE_SOURCE_DIR}/src/common/hierarchical + ${CMAKE_SOURCE_DIR}/src/common/platform/include ${CMAKE_SOURCE_DIR}/src/common/platform/include/common ${CMAKE_SOURCE_DIR}/src/common/platform/include/host ${CMAKE_CURRENT_SOURCE_DIR} diff --git a/python/bindings/task_interface.cpp b/python/bindings/task_interface.cpp index 1aac74a091..f462060c15 100644 --- a/python/bindings/task_interface.cpp +++ b/python/bindings/task_interface.cpp @@ -1530,6 +1530,36 @@ NB_MODULE(_task_interface, m) { nb::arg("allocation_id"), nb::arg("rank_count"), nb::arg("domain_rank"), "Pair to comm_alloc_domain_windows: collectively release the per-rank pool." ) + .def( + "comm_global_domain_prepare", + [](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); + return nb::make_tuple( + nb::bytes(reinterpret_cast(descriptor.data()), descriptor.size()), local_window_base, + actual_window_size + ); + }, + nb::arg("domain_id"), nb::arg("domain_rank"), nb::arg("rank_count"), nb::arg("window_size"), + nb::arg("profile"), "Create a Global CommDomain local window and return its transport descriptor." + ) + .def( + "comm_global_domain_import", + [](ChipWorker &self, uint64_t domain_id, nb::bytes descriptors) { + std::vector descriptor_bytes( + reinterpret_cast(descriptors.c_str()), + reinterpret_cast(descriptors.c_str()) + descriptors.size() + ); + return self.comm_global_domain_import(domain_id, descriptor_bytes); + }, + nb::arg("domain_id"), nb::arg("descriptors"), + "Import a rank-ordered Global CommDomain descriptor table and return the device context." + ) + .def( + "comm_global_domain_release", &ChipWorker::comm_global_domain_release, nb::arg("domain_id"), + "Release a prepared or imported Global CommDomain." + ) .def("comm_barrier", &ChipWorker::comm_barrier, nb::arg("comm_handle"), "Synchronize all ranks.") .def( "comm_destroy", &ChipWorker::comm_destroy, nb::arg("comm_handle"), diff --git a/python/bindings/worker_bind.h b/python/bindings/worker_bind.h index f110c9a164..426366664e 100644 --- a/python/bindings/worker_bind.h +++ b/python/bindings/worker_bind.h @@ -692,6 +692,25 @@ inline void bind_worker(nb::module_ &m) { nb::arg("importer_worker_id"), nb::arg("owner_worker_id"), nb::arg("buffer_id"), nb::arg("generation"), nb::arg("import_id"), "Release an imported remote buffer mapping." ) + .def( + "remote_domain_control", + [](Worker &self, int worker_id, uint32_t control_name, nb::bytes command) { + std::vector command_bytes( + reinterpret_cast(command.c_str()), + reinterpret_cast(command.c_str()) + command.size() + ); + std::vector result; + { + nb::gil_scoped_release release; + result = self.remote_domain_control( + worker_id, static_cast(control_name), command_bytes + ); + } + return nb::bytes(reinterpret_cast(result.data()), result.size()); + }, + nb::arg("worker_id"), nb::arg("control_name"), nb::arg("command"), + "Send one Global CommDomain control to a remote L3 endpoint." + ) .def( "broadcast_unregister_all", [](Worker &self, nb::object digest) { @@ -733,6 +752,25 @@ inline void bind_worker(nb::module_ &m) { "If payload is a Python buffer, C++ stages it in POSIX shm and writes the shm name " "into the mailbox. Returns per-child ControlResult entries." ) + .def( + "control_payload", + [](Worker &self, WorkerType worker_type, int worker_id, uint64_t sub_cmd, nb::object payload, + nb::object timeout_s) { + std::string payload_bytes = buffer_to_string(payload, "payload"); + double timeout_val = timeout_s.is_none() ? -1.0 : nb::cast(timeout_s); + std::vector result; + { + nb::gil_scoped_release release; + result = self.control_payload( + worker_type, worker_id, sub_cmd, payload_bytes.data(), payload_bytes.size(), timeout_val + ); + } + return nb::bytes(reinterpret_cast(result.data()), result.size()); + }, + nb::arg("worker_type"), nb::arg("worker_id"), nb::arg("sub_cmd"), nb::arg("payload"), + nb::arg("timeout_s") = nb::none(), + "Drive one local worker control with a mutable staged payload and return its final bytes." + ) .def( "control_alloc_domain", &Worker::control_alloc_domain, nb::arg("worker_id"), nb::arg("request_shm_name"), nb::arg("reply_shm_name"), nb::call_guard(), diff --git a/python/simpler/global_comm_domain.py b/python/simpler/global_comm_domain.py new file mode 100644 index 0000000000..ec1ca55539 --- /dev/null +++ b/python/simpler/global_comm_domain.py @@ -0,0 +1,574 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Versioned codecs and data models for L4-brokered Global CommDomains.""" + +from __future__ import annotations + +import enum +import struct +from dataclasses import dataclass + +GLOBAL_DOMAIN_VERSION = 1 +GLOBAL_DOMAIN_MAX_RANKS = 64 +GLOBAL_DOMAIN_HANDLE_BYTES = 256 +GLOBAL_DOMAIN_DESCRIPTOR = struct.Struct(" bytes: + _validate_descriptor(self) + handle = bytes(self.handle) + return GLOBAL_DOMAIN_DESCRIPTOR.pack( + int(self.version), + int(self.profile_id), + int(self.domain_rank), + int(self.rank_count), + int(self.mapping_size), + len(handle), + 0, + handle + b"\x00" * (GLOBAL_DOMAIN_HANDLE_BYTES - len(handle)), + ) + + @classmethod + def decode(cls, data: bytes) -> GlobalDomainDescriptor: + if len(data) != GLOBAL_DOMAIN_DESCRIPTOR_BYTES: + raise ValueError("global domain descriptor size mismatch") + version, profile_id, domain_rank, rank_count, mapping_size, handle_size, reserved, handle = ( + GLOBAL_DOMAIN_DESCRIPTOR.unpack(data) + ) + if reserved != 0: + raise ValueError("global domain descriptor reserved field must be zero") + if handle_size > GLOBAL_DOMAIN_HANDLE_BYTES: + raise ValueError("global domain descriptor handle size exceeds maximum") + descriptor = cls( + version=int(version), + profile_id=int(profile_id), + domain_rank=int(domain_rank), + rank_count=int(rank_count), + mapping_size=int(mapping_size), + handle=bytes(handle[:handle_size]), + ) + _validate_descriptor(descriptor) + return descriptor + + +@dataclass(frozen=True) +class GlobalCommInitCommand: + cluster_id: str + topology_hash: str + profile: str + node_rank: int + node_count: int + members: tuple[GlobalDomainMember, ...] + + +@dataclass(frozen=True) +class GlobalCommInitResult: + profile: str + max_ranks: int + descriptor_bytes: int + local_device_count: int + + +def resolve_global_comm_capability(*, platform: str, profile: str, local_device_count: int) -> GlobalCommInitResult: + """Return the capability implemented by the selected platform backend.""" + platform = str(platform) + profile = str(profile) + supported = (platform.endswith("sim") and profile == GLOBAL_DOMAIN_PROFILE_SIM) or ( + platform.startswith("a2a3") and not platform.endswith("sim") and profile == GLOBAL_DOMAIN_PROFILE_A3_FABRIC + ) + if not supported: + raise ValueError(f"Global CommDomain is not supported by platform {platform!r} with comm_profile {profile!r}") + if local_device_count <= 0: + raise ValueError("Global CommDomain capability requires at least one local device") + return GlobalCommInitResult( + profile=profile, + max_ranks=GLOBAL_DOMAIN_MAX_RANKS, + descriptor_bytes=GLOBAL_DOMAIN_DESCRIPTOR_BYTES, + local_device_count=int(local_device_count), + ) + + +@dataclass(frozen=True) +class GlobalDomainCommand: + phase: GlobalDomainPhase + domain_id: int + generation: int + name: str + profile: str + window_size: int + members: tuple[GlobalDomainMember, ...] + buffers: tuple[GlobalDomainBuffer, ...] + descriptors: tuple[GlobalDomainDescriptor, ...] = () + + +@dataclass(frozen=True) +class GlobalDomainReleaseCommand: + domain_id: int + generation: int + + +@dataclass(frozen=True) +class GlobalDomainCopyCommand: + domain_id: int + generation: int + domain_rank: int + offset: int + nbytes: int + data: bytes = b"" + + +class _Reader: + def __init__(self, data: bytes) -> None: + self._data = data + self._offset = 0 + + def _take(self, size: int, field: str) -> bytes: + if size < 0 or self._offset > len(self._data) or size > len(self._data) - self._offset: + raise ValueError(f"global domain wire truncated {field}") + result = self._data[self._offset : self._offset + size] + self._offset += size + return result + + def u32(self) -> int: + return int(struct.unpack(" int: + return int(struct.unpack(" int: + return int(struct.unpack(" str: + size = self.u32() + if size > GLOBAL_DOMAIN_MAX_STRING_BYTES: + raise ValueError(f"global domain wire {field} exceeds maximum") + return self._take(size, field).decode("utf-8") + + def blob(self, maximum: int, field: str) -> bytes: + size = self.u32() + if size > maximum: + raise ValueError(f"global domain wire {field} exceeds maximum") + return self._take(size, field) + + def fixed(self, size: int, field: str) -> bytes: + return self._take(size, field) + + def done(self, field: str) -> None: + if self._offset != len(self._data): + raise ValueError(f"global domain wire trailing bytes after {field}") + + +def _put_string(out: bytearray, value: str, field: str) -> None: + encoded = str(value).encode("utf-8") + if len(encoded) > GLOBAL_DOMAIN_MAX_STRING_BYTES: + raise ValueError(f"global domain wire {field} exceeds maximum") + out.extend(struct.pack(" None: + encoded = bytes(value) + if len(encoded) > maximum: + raise ValueError(f"global domain wire {field} exceeds maximum") + out.extend(struct.pack(" None: + if descriptor.version != GLOBAL_DOMAIN_VERSION: + raise ValueError("global domain descriptor version mismatch") + if descriptor.profile_id not in GLOBAL_DOMAIN_PROFILE_IDS.values(): + raise ValueError("global domain descriptor profile is unknown") + if descriptor.rank_count <= 0 or descriptor.rank_count > GLOBAL_DOMAIN_MAX_RANKS: + raise ValueError("global domain descriptor rank_count is invalid") + if descriptor.domain_rank < 0 or descriptor.domain_rank >= descriptor.rank_count: + raise ValueError("global domain descriptor domain_rank is invalid") + if descriptor.mapping_size <= 0: + raise ValueError("global domain descriptor mapping_size must be positive") + if not descriptor.handle or len(descriptor.handle) > GLOBAL_DOMAIN_HANDLE_BYTES: + raise ValueError("global domain descriptor handle size is invalid") + + +def validate_member_table(members: tuple[GlobalDomainMember, ...]) -> None: + if not members or len(members) > GLOBAL_DOMAIN_MAX_RANKS: + raise ValueError("global domain members must contain between 1 and 64 devices") + ranks = [member.domain_rank for member in members] + if ranks != list(range(len(members))): + raise ValueError("global domain members must be in dense domain-rank order") + devices = [(member.node_worker_id, member.local_worker_id) for member in members] + if len(set(devices)) != len(devices): + raise ValueError("global domain members contain duplicate node/local devices") + if any(node < 0 or local < 0 for node, local in devices): + raise ValueError("global domain member node/local ids must be non-negative") + global_ranks = [member.global_device_rank for member in members] + if len(set(global_ranks)) != len(global_ranks) or any(rank < 0 for rank in global_ranks): + raise ValueError("global domain members require unique non-negative global device ranks") + + +def validate_descriptor_table( + descriptors: tuple[GlobalDomainDescriptor, ...], *, rank_count: int, profile: str +) -> None: + if profile not in GLOBAL_DOMAIN_PROFILE_IDS: + raise ValueError(f"unsupported global domain profile {profile!r}") + if len(descriptors) != rank_count: + raise ValueError("global domain descriptor table is incomplete") + expected_profile = GLOBAL_DOMAIN_PROFILE_IDS[profile] + ranks: set[int] = set() + mapping_size: int | None = None + for descriptor in descriptors: + _validate_descriptor(descriptor) + if descriptor.profile_id != expected_profile or descriptor.rank_count != rank_count: + raise ValueError("global domain descriptor profile or rank_count mismatch") + if descriptor.domain_rank in ranks: + raise ValueError("global domain descriptor table contains a duplicate rank") + ranks.add(descriptor.domain_rank) + if mapping_size is None: + mapping_size = descriptor.mapping_size + elif mapping_size != descriptor.mapping_size: + raise ValueError("global domain descriptor mapping sizes differ") + if ranks != set(range(rank_count)): + raise ValueError("global domain descriptor table has missing ranks") + + +def _put_member(out: bytearray, member: GlobalDomainMember) -> None: + out.extend( + struct.pack( + " GlobalDomainMember: + return GlobalDomainMember( + node_worker_id=reader.i32(), + local_worker_id=reader.u32(), + global_device_rank=reader.u32(), + domain_rank=reader.u32(), + ) + + +def encode_comm_init(command: GlobalCommInitCommand) -> bytes: + validate_member_table(command.members) + if not command.cluster_id or not command.topology_hash: + raise ValueError("global comm init cluster_id and topology_hash must be non-empty") + if command.profile not in GLOBAL_DOMAIN_PROFILE_IDS: + raise ValueError(f"unsupported global domain profile {command.profile!r}") + if command.node_rank < 0 or command.node_count <= 0 or command.node_rank >= command.node_count: + raise ValueError("global comm init node identity is invalid") + out = bytearray(struct.pack(" GlobalCommInitCommand: + reader = _Reader(data) + version = reader.u32() + if version != GLOBAL_DOMAIN_VERSION: + raise ValueError("global comm init version mismatch") + node_rank = reader.u32() + node_count = reader.u32() + cluster_id = reader.string("cluster_id") + topology_hash = reader.string("topology_hash") + profile = reader.string("profile") + member_count = reader.u32() + if member_count > GLOBAL_DOMAIN_MAX_RANKS: + raise ValueError("global comm init member count exceeds maximum") + members = tuple(_read_member(reader) for _ in range(member_count)) + reader.done("COMM_INIT") + command = GlobalCommInitCommand(cluster_id, topology_hash, profile, node_rank, node_count, members) + validate_member_table(command.members) + if not command.cluster_id or not command.topology_hash: + raise ValueError("global comm init cluster_id and topology_hash must be non-empty") + if command.profile not in GLOBAL_DOMAIN_PROFILE_IDS: + raise ValueError(f"unsupported global domain profile {command.profile!r}") + if node_count <= 0 or node_rank >= node_count: + raise ValueError("global comm init node identity is invalid") + return command + + +def encode_comm_init_result(result: GlobalCommInitResult) -> bytes: + out = bytearray( + struct.pack( + " GlobalCommInitResult: + reader = _Reader(data) + result = GlobalCommInitResult( + profile="", + max_ranks=reader.u32(), + descriptor_bytes=reader.u32(), + local_device_count=reader.u32(), + ) + result = GlobalCommInitResult( + profile=reader.string("profile"), + max_ranks=result.max_ranks, + descriptor_bytes=result.descriptor_bytes, + local_device_count=result.local_device_count, + ) + reader.done("COMM_INIT result") + return result + + +def encode_domain_command(command: GlobalDomainCommand) -> bytes: + validate_member_table(command.members) + if command.domain_id == 0 or command.generation == 0 or command.window_size <= 0: + raise ValueError("global domain command identity and window_size must be positive") + if not command.name: + raise ValueError("global domain command name must be non-empty") + if command.profile not in GLOBAL_DOMAIN_PROFILE_IDS: + raise ValueError(f"unsupported global domain profile {command.profile!r}") + if len({buffer.name for buffer in command.buffers}) != len(command.buffers): + raise ValueError("global domain command contains duplicate buffer names") + if any(not buffer.name or buffer.nbytes <= 0 for buffer in command.buffers): + raise ValueError("global domain buffers require a name and positive size") + if sum(buffer.nbytes for buffer in command.buffers) > command.window_size: + raise ValueError("global domain buffers exceed the requested window") + if command.descriptors: + validate_descriptor_table(command.descriptors, rank_count=len(command.members), profile=command.profile) + if command.phase in (GlobalDomainPhase.IMPORT, GlobalDomainPhase.COMMIT): + if len(command.descriptors) != len(command.members): + raise ValueError("global domain IMPORT/COMMIT requires a complete descriptor table") + elif command.descriptors: + raise ValueError("global domain PREPARE/ABORT must not carry descriptors") + + out = bytearray( + struct.pack( + " GlobalDomainCommand: + reader = _Reader(data) + version = reader.u32() + if version != GLOBAL_DOMAIN_VERSION: + raise ValueError("global domain command version mismatch") + try: + phase = GlobalDomainPhase(reader.u32()) + except ValueError as exc: + raise ValueError("global domain command phase is unknown") from exc + domain_id = reader.u64() + generation = reader.u64() + window_size = reader.u64() + name = reader.string("name") + profile = reader.string("profile") + member_count = reader.u32() + if member_count > GLOBAL_DOMAIN_MAX_RANKS: + raise ValueError("global domain command member count exceeds maximum") + members = tuple(_read_member(reader) for _ in range(member_count)) + buffer_count = reader.u32() + if buffer_count > GLOBAL_DOMAIN_MAX_RANKS: + raise ValueError("global domain command buffer count exceeds maximum") + buffers = tuple(GlobalDomainBuffer(reader.string("buffer.name"), reader.u64()) for _ in range(buffer_count)) + descriptor_count = reader.u32() + if descriptor_count > GLOBAL_DOMAIN_MAX_RANKS: + raise ValueError("global domain command descriptor count exceeds maximum") + descriptors = tuple( + GlobalDomainDescriptor.decode(reader.fixed(GLOBAL_DOMAIN_DESCRIPTOR_BYTES, "descriptor")) + for _ in range(descriptor_count) + ) + reader.done("ALLOC_DOMAIN") + command = GlobalDomainCommand( + phase=phase, + domain_id=domain_id, + generation=generation, + name=name, + profile=profile, + window_size=window_size, + members=members, + buffers=buffers, + descriptors=descriptors, + ) + encode_domain_command(command) + return command + + +def encode_descriptor_table(descriptors: tuple[GlobalDomainDescriptor, ...]) -> bytes: + if len(descriptors) > GLOBAL_DOMAIN_MAX_RANKS: + raise ValueError("global domain descriptor table exceeds maximum") + return struct.pack(" tuple[GlobalDomainDescriptor, ...]: + reader = _Reader(data) + count = reader.u32() + if count > GLOBAL_DOMAIN_MAX_RANKS: + raise ValueError("global domain descriptor table exceeds maximum") + descriptors = tuple( + GlobalDomainDescriptor.decode(reader.fixed(GLOBAL_DOMAIN_DESCRIPTOR_BYTES, "descriptor")) for _ in range(count) + ) + reader.done("descriptor table") + return descriptors + + +def encode_release_command(command: GlobalDomainReleaseCommand) -> bytes: + if command.domain_id == 0 or command.generation == 0: + raise ValueError("global domain release identity must be positive") + return struct.pack(" GlobalDomainReleaseCommand: + if len(data) != struct.calcsize(" bytes: + if ( + command.domain_id == 0 + or command.generation == 0 + or command.domain_rank < 0 + or command.offset < 0 + or command.nbytes <= 0 + or command.nbytes > GLOBAL_DOMAIN_MAX_COPY_BYTES + ): + raise ValueError("global domain copy fields are invalid") + if include_data and len(command.data) != command.nbytes: + raise ValueError("global domain copy payload size mismatch") + if not include_data and command.data: + raise ValueError("global domain copy-from request must not contain data") + out = bytearray( + struct.pack( + " GlobalDomainCopyCommand: + header_size = struct.calcsize(" bytes: + out = bytearray() + _put_blob(out, data, GLOBAL_DOMAIN_MAX_COPY_BYTES, "copy result") + return bytes(out) + + +def decode_copy_result(data: bytes) -> bytes: + reader = _Reader(data) + result = reader.blob(GLOBAL_DOMAIN_MAX_COPY_BYTES, "copy result") + reader.done("copy result") + return result diff --git a/python/simpler/global_comm_smoke.py b/python/simpler/global_comm_smoke.py new file mode 100644 index 0000000000..12a17193c8 --- /dev/null +++ b/python/simpler/global_comm_smoke.py @@ -0,0 +1,88 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Remote L3 callback used by the no-mpirun Global CommDomain TLOAD smoke.""" + +from __future__ import annotations + +from .task_interface import CallConfig, DataType, TaskArgs, Tensor, TensorArgType + +_SMOKE_COUNT = 256 + + +def _digest_from_scalars(args: TaskArgs, start: int) -> bytes: + return b"".join(int(args.scalar(start + i)).to_bytes(8, "little") for i in range(4)) + + +def remote_compute_orch(orch, args: TaskArgs, cfg: CallConfig) -> None: + """Submit one local vector-add task from a remote L3 to its L2.""" + from .remote_l3_session import get_inner_handle # noqa: PLC0415 + + if args.scalar_count() != 6: + raise ValueError("remote compute task expects domain_id, local_worker_id, and four digest scalars") + domain_id = int(args.scalar(0)) + local_worker_id = int(args.scalar(1)) + chip_handle = get_inner_handle(_digest_from_scalars(args, 2).hex()) + context = orch.get_global_domain(domain_id)[local_worker_id] + + chip_args = TaskArgs() + for buffer_name in ("lhs", "rhs"): + chip_args.add_tensor( + Tensor.make( + data=context.buffer_ptrs[buffer_name], + shapes=(_SMOKE_COUNT,), + dtype=DataType.FLOAT32, + child_memory=True, + ), + TensorArgType.INPUT, + ) + chip_args.add_tensor( + Tensor.make( + data=context.buffer_ptrs["input"], + shapes=(_SMOKE_COUNT,), + dtype=DataType.FLOAT32, + child_memory=True, + ), + TensorArgType.OUTPUT_EXISTING, + ) + orch.submit_next_level(chip_handle, chip_args, cfg, worker=local_worker_id) + + +def remote_rank_orch(orch, args: TaskArgs, cfg: CallConfig) -> None: + """Submit the registered TLOAD kernel from one remote L3 to its local L2.""" + from .remote_l3_session import get_inner_handle # noqa: PLC0415 + + if args.scalar_count() != 6: + raise ValueError("global TLOAD remote task expects domain_id, local_worker_id, and four digest scalars") + domain_id = int(args.scalar(0)) + local_worker_id = int(args.scalar(1)) + chip_handle = get_inner_handle(_digest_from_scalars(args, 2).hex()) + context = orch.get_global_domain(domain_id)[local_worker_id] + + chip_args = TaskArgs() + chip_args.add_tensor( + Tensor.make( + data=context.buffer_ptrs["input"], + shapes=(_SMOKE_COUNT,), + dtype=DataType.FLOAT32, + child_memory=True, + ), + TensorArgType.INPUT, + ) + chip_args.add_tensor( + Tensor.make( + data=context.buffer_ptrs["result"], + shapes=(_SMOKE_COUNT,), + dtype=DataType.FLOAT32, + child_memory=True, + ), + TensorArgType.OUTPUT_EXISTING, + ) + chip_args.add_scalar(context.domain_size) + chip_args.add_scalar(context.device_ctx) + orch.submit_next_level(chip_handle, chip_args, cfg, worker=local_worker_id) diff --git a/python/simpler/orchestrator.py b/python/simpler/orchestrator.py index a79c619264..f1ef896334 100644 --- a/python/simpler/orchestrator.py +++ b/python/simpler/orchestrator.py @@ -48,6 +48,8 @@ def my_orch(orch, args, cfg): CommBufferSpec, CommDomainHandle, DataType, + GlobalCommDomainHandle, + GlobalCommDomainView, RemoteAddressSpace, TaskArgs, Tensor, @@ -414,6 +416,89 @@ def release_domain(self, handle: CommDomainHandle) -> None: """Collective release. Equivalent to ``handle.release()``.""" handle.release() + def allocate_global_domain( + self, + *, + name: str, + members: Sequence[tuple[int, int]], + window_size: int, + buffers: Sequence[CommBufferSpec] = (), + retain_after_run: bool = False, + ) -> GlobalCommDomainHandle: + """Create a CommDomain across local and/or remote L3 nodes without MPI. + + 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, + sends the complete rank-ordered table back to every L3, and commits + only after all L2 imports succeed. + ``retain_after_run=True`` keeps the domain live after the current DAG + drains so a later run can inspect communication results; explicit + release or ``Worker.close()`` still tears it down. + """ + if self._worker is None: + raise RuntimeError("allocate_global_domain requires an Orchestrator bound to a Worker") + return self._worker._allocate_global_domain( + name=str(name), + members=tuple((int(node), int(local)) for node, local in members), + window_size=int(window_size), + buffers=list(buffers), + retain_after_run=bool(retain_after_run), + ) + + def release_global_domain(self, handle: GlobalCommDomainHandle) -> None: + handle.release() + + def get_global_domain(self, domain_id: int) -> GlobalCommDomainView: + """Return the committed L3-local view for a domain created by L4.""" + if self._worker is None: + raise RuntimeError("get_global_domain requires an Orchestrator bound to a Worker") + return self._worker._get_global_domain(int(domain_id)) + + @staticmethod + def _global_copy_range(handle: GlobalCommDomainHandle, *, buffer: str | None, offset: int, nbytes: int) -> int: + absolute = int(offset) + if absolute < 0 or nbytes <= 0: + raise ValueError("Global CommDomain copy offset must be non-negative and size must be positive") + limit = handle.mapping_size + if buffer is not None: + buffer_offset, buffer_nbytes = handle.buffer_range(str(buffer)) + if absolute > buffer_nbytes or nbytes > buffer_nbytes - absolute: + raise ValueError(f"Global CommDomain copy exceeds buffer {buffer!r}") + absolute += buffer_offset + elif absolute > limit or nbytes > limit - absolute: + raise ValueError("Global CommDomain copy exceeds the mapped window") + return absolute + + def copy_to_global_domain( + self, + handle: GlobalCommDomainHandle, + domain_rank: int, + data: bytes, + *, + buffer: str | None = None, + offset: int = 0, + ) -> None: + payload = bytes(data) + absolute = self._global_copy_range(handle, buffer=buffer, offset=int(offset), nbytes=len(payload)) + if self._worker is None: + raise RuntimeError("copy_to_global_domain requires an Orchestrator bound to a Worker") + self._worker._copy_to_global_domain(handle, int(domain_rank), payload, absolute) + + def copy_from_global_domain( + self, + handle: GlobalCommDomainHandle, + domain_rank: int, + nbytes: int, + *, + buffer: str | None = None, + offset: int = 0, + ) -> bytes: + absolute = self._global_copy_range(handle, buffer=buffer, offset=int(offset), nbytes=int(nbytes)) + if self._worker is None: + raise RuntimeError("copy_from_global_domain requires an Orchestrator bound to a Worker") + return self._worker._copy_from_global_domain(handle, int(domain_rank), int(nbytes), absolute) + def create_l3_l2_region(self, *, worker_id: int, payload_bytes: int, counter_bytes: int): """Create an L3-L2 communication region on one NEXT_LEVEL chip worker.""" if self._worker is None: diff --git a/python/simpler/remote_l3_protocol.py b/python/simpler/remote_l3_protocol.py index 24ca8ab402..f1263b80f3 100644 --- a/python/simpler/remote_l3_protocol.py +++ b/python/simpler/remote_l3_protocol.py @@ -64,6 +64,8 @@ class ControlName(enum.IntEnum): COMM_INIT = 13 ALLOC_DOMAIN = 14 RELEASE_DOMAIN = 15 + COPY_TO_DOMAIN = 16 + COPY_FROM_DOMAIN = 17 class RemoteRegistryTarget(enum.IntEnum): diff --git a/python/simpler/remote_l3_session.py b/python/simpler/remote_l3_session.py index c62a613779..4049d35f0f 100644 --- a/python/simpler/remote_l3_session.py +++ b/python/simpler/remote_l3_session.py @@ -42,6 +42,18 @@ parse_python_import_target, validate_hashid, ) +from .global_comm_domain import ( + GlobalDomainPhase, + GlobalDomainReleaseCommand, + decode_comm_init, + decode_copy_command, + decode_domain_command, + decode_release_command, + encode_comm_init_result, + encode_copy_result, + encode_descriptor_table, + resolve_global_comm_capability, +) from .remote_l3_protocol import ( PROTOCOL_VERSION, CallableKind, @@ -75,7 +87,7 @@ send_frame, ) from .task_interface import ChipCallable, TaskArgs, Tensor -from .worker import Worker +from .worker import Worker, _NoHostBufferChildrenError sys.modules.setdefault("simpler.remote_l3_session", sys.modules[__name__]) @@ -111,6 +123,7 @@ class _RemoteBufferEntry: nbytes: int generation: int address_space: RemoteAddressSpace + owner: Worker | None = None offset: int = 0 released: bool = False @@ -120,6 +133,8 @@ def addr(self) -> int: buf = self.data.buf assert buf is not None return ctypes.addressof(ctypes.c_char.from_buffer(buf)) + if hasattr(self.data, "data_ptr"): + return int(self.data.data_ptr) return ctypes.addressof(self.data) @property @@ -130,6 +145,13 @@ def shm_name(self) -> str: def close(self, *, unlink: bool = False) -> None: if not isinstance(self.data, shared_memory.SharedMemory): + if self.owner is not None: + owner, self.owner = self.owner, None + # HostBuffer itself owns a memoryview. Release it before asking + # Worker to close the backing shm so teardown is prompt and does + # not report a false live-view warning. + self.data.buffer.release() + owner.free_host_buffer(self.data) return self.data.close() if unlink: @@ -461,6 +483,20 @@ def _control_reply( send_frame(conn, FrameHeader(FrameType.CONTROL_REPLY, session_id, worker_id, sequence), payload) +def _control_result_reply( + conn: socket.socket, + manifest: dict[str, Any], + sequence: int, + control_name: ControlName, + version: int, + result: bytes, +) -> None: + session_id = int(manifest["session_id"]) + worker_id = int(manifest["worker_id"]) + payload = encode_control_reply(sequence, control_name, version, 0, "", bytes(result)) + send_frame(conn, FrameHeader(FrameType.CONTROL_REPLY, session_id, worker_id, sequence), payload) + + def _copy_command_header(data: bytes) -> tuple[int, int, int, int, int, bytes]: if len(data) < 36: raise ValueError("remote buffer copy command is truncated") @@ -544,12 +580,13 @@ def _run_command_loop( # noqa: PLR0912, PLR0915 next_export_id = 1 next_import_id = 1 buffers: dict[tuple[int, ...], _RemoteBufferEntry] = {} + global_comm_inits: dict[str, Any] = {} hello = HelloPayload( session_id=session_id, worker_id=worker_id, protocol_version=PROTOCOL_VERSION, - comm_profile=str(manifest["transport"]), + comm_profile=str(manifest.get("comm_profile", manifest["transport"])), feature_flags=0, ready_state=ReadyState.READY, ) @@ -656,9 +693,24 @@ def _run_command_loop( # noqa: PLR0912, PLR0915 buffer_id = next_buffer_id next_buffer_id += 1 generation = 1 - buf = shared_memory.SharedMemory(create=True, size=int(nbytes)) + try: + buf = inner_worker.create_host_buffer(int(nbytes)) + entry = _RemoteBufferEntry( + buf, + int(nbytes), + generation, + RemoteAddressSpace.REMOTE_DEVICE, + owner=inner_worker, + ) + except _NoHostBufferChildrenError: + buf = shared_memory.SharedMemory(create=True, size=int(nbytes)) + entry = _RemoteBufferEntry( + buf, + int(nbytes), + generation, + RemoteAddressSpace.REMOTE_DEVICE, + ) key = _buffer_key(buffer_id, generation) - entry = _RemoteBufferEntry(buf, int(nbytes), generation, RemoteAddressSpace.REMOTE_DEVICE) buffers[key] = entry remote_addr = entry.addr result = struct.pack( @@ -846,19 +898,136 @@ def _run_command_loop( # noqa: PLR0912, PLR0915 _control_reply( conn, manifest, header.sequence, control.control_name, control.control_version, 0, "" ) - elif control.control_name in ( - ControlName.COMM_INIT, - ControlName.ALLOC_DOMAIN, - ControlName.RELEASE_DOMAIN, - ): + elif control.control_name == ControlName.COMM_INIT: + command = decode_comm_init(control.command_bytes) + manifest_profile = str(manifest.get("comm_profile", manifest["transport"])) + if command.cluster_id != str(manifest.get("cluster_id", "")): + raise ValueError("COMM_INIT cluster_id does not match the session manifest") + if command.profile != manifest_profile: + raise ValueError("COMM_INIT profile does not match the session manifest") + if command.node_rank != int(manifest.get("node_rank", 0)) or command.node_count != int( + manifest.get("node_count", 1) + ): + raise ValueError("COMM_INIT node identity does not match the session manifest") + global_ranks = tuple(int(rank) for rank in manifest.get("global_device_ranks", ())) + local_members = tuple( + member for member in command.members if member.node_worker_id == worker_id + ) + if not local_members: + raise ValueError("COMM_INIT topology has no local members") + for member in local_members: + if member.local_worker_id < 0 or member.local_worker_id >= len(global_ranks): + raise ValueError("COMM_INIT local worker id exceeds the session device list") + if member.global_device_rank != global_ranks[member.local_worker_id]: + raise ValueError("COMM_INIT global device rank does not match the manifest") + prior = global_comm_inits.get(command.topology_hash) + if prior is not None and prior != command: + raise ValueError("COMM_INIT topology hash conflicts with an earlier command") + result = resolve_global_comm_capability( + platform=str(manifest["platform"]), + profile=manifest_profile, + local_device_count=len(global_ranks), + ) + global_comm_inits[command.topology_hash] = command + _control_result_reply( + conn, + manifest, + header.sequence, + control.control_name, + control.control_version, + encode_comm_init_result(result), + ) + elif control.control_name == ControlName.ALLOC_DOMAIN: + command = decode_domain_command(control.command_bytes) + if not any( + init.profile == command.profile and init.members == command.members + for init in global_comm_inits.values() + ): + raise RuntimeError("ALLOC_DOMAIN requires a matching COMM_INIT topology") + if command.phase is GlobalDomainPhase.PREPARE_EXPORT: + if command.descriptors: + raise ValueError("PREPARE_EXPORT must not carry descriptors") + descriptors = inner_worker._prepare_global_domain_node(command, worker_id) # noqa: SLF001 + _control_result_reply( + conn, + manifest, + header.sequence, + control.control_name, + control.control_version, + encode_descriptor_table(descriptors), + ) + elif command.phase is GlobalDomainPhase.IMPORT: + inner_worker._import_global_domain_node(command, worker_id) # noqa: SLF001 + _control_reply( + conn, + manifest, + header.sequence, + control.control_name, + control.control_version, + 0, + "", + ) + elif command.phase is GlobalDomainPhase.COMMIT: + inner_worker._commit_global_domain_node(command) # noqa: SLF001 + _control_reply( + conn, + manifest, + header.sequence, + control.control_name, + control.control_version, + 0, + "", + ) + elif command.phase is GlobalDomainPhase.ABORT: + inner_worker._release_global_domain_node( # noqa: SLF001 + GlobalDomainReleaseCommand(command.domain_id, command.generation), + suppress_errors=True, + ) + _control_reply( + conn, + manifest, + header.sequence, + control.control_name, + control.control_version, + 0, + "", + ) + else: + raise ValueError("ALLOC_DOMAIN phase is not supported") + elif control.control_name == ControlName.RELEASE_DOMAIN: + command = decode_release_command(control.command_bytes) + inner_worker._release_global_domain_node(command) # noqa: SLF001 _control_reply( conn, manifest, header.sequence, control.control_name, control.control_version, - 1, - f"unsupported reserved remote domain control {control.control_name.name}", + 0, + "", + ) + elif control.control_name == ControlName.COPY_TO_DOMAIN: + command = decode_copy_command(control.command_bytes, include_data=True) + inner_worker._copy_global_domain_node(command, copy_to_device=True) # noqa: SLF001 + _control_reply( + conn, + manifest, + header.sequence, + control.control_name, + control.control_version, + 0, + "", + ) + elif control.control_name == ControlName.COPY_FROM_DOMAIN: + command = decode_copy_command(control.command_bytes, include_data=False) + result = inner_worker._copy_global_domain_node(command, copy_to_device=False) # noqa: SLF001 + _control_result_reply( + conn, + manifest, + header.sequence, + control.control_name, + control.control_version, + encode_copy_result(result), ) else: _control_reply( diff --git a/python/simpler/remote_l3_worker.py b/python/simpler/remote_l3_worker.py index 3561ee95c4..4937b424c3 100644 --- a/python/simpler/remote_l3_worker.py +++ b/python/simpler/remote_l3_worker.py @@ -64,6 +64,23 @@ def _validate_manifest(manifest: dict[str, Any]) -> None: raise ValueError("manifest platform must be non-empty") if str(manifest["transport"]) != "sim": raise ValueError("only sim transport is accepted by simpler-remote-worker") + comm_profile = str(manifest.get("comm_profile", manifest["transport"])) + if comm_profile not in ("sim", "a3-fabric-v1"): + raise ValueError("manifest comm_profile is not supported") + if comm_profile == "a3-fabric-v1" and not str(manifest["platform"]).startswith("a2a3"): + raise ValueError("manifest a3-fabric-v1 comm_profile requires an a2a3 platform") + if comm_profile == "a3-fabric-v1" and str(manifest["platform"]).endswith("sim"): + raise ValueError("manifest a3-fabric-v1 comm_profile requires real A3 devices") + node_rank = int(manifest.get("node_rank", 0)) + node_count = int(manifest.get("node_count", 1)) + if node_count <= 0 or node_rank < 0 or node_rank >= node_count: + raise ValueError("manifest node identity is invalid") + device_ids = [int(device_id) for device_id in manifest.get("device_ids", [])] + global_device_ranks = [int(rank) for rank in manifest.get("global_device_ranks", range(len(device_ids)))] + if len(global_device_ranks) != len(device_ids): + raise ValueError("manifest global_device_ranks must match device_ids length") + if any(rank < 0 for rank in global_device_ranks) or len(set(global_device_ranks)) != len(global_device_ranks): + raise ValueError("manifest global_device_ranks must be unique and non-negative") def _session_timeout_s(manifest: dict[str, Any]) -> float: diff --git a/python/simpler/task_interface.py b/python/simpler/task_interface.py index e29bec4c58..9bb3c64ade 100644 --- a/python/simpler/task_interface.py +++ b/python/simpler/task_interface.py @@ -130,6 +130,8 @@ def _assert_bindings_match_source_tree() -> None: _assert_bindings_match_source_tree() +from .global_comm_domain import GlobalDomainBuffer, GlobalDomainMember # noqa: E402 + __all__ = [ "DataType", "get_element_size", @@ -163,6 +165,8 @@ def _assert_bindings_match_source_tree() -> None: "CommBufferSpec", "ChipDomainContext", "CommDomainHandle", + "GlobalCommDomainHandle", + "GlobalCommDomainView", ] COMM_MAX_RANK_NUM = 64 @@ -1017,6 +1021,132 @@ def __repr__(self) -> str: return f"CommDomainHandle(name={self.name!r}, workers={self.workers}, {state})" +class GlobalCommDomainHandle: + """L4-owned handle for one CommDomain spanning local and/or remote L3 nodes. + + The handle contains stable topology and buffer offsets only. Device + addresses remain in the L3/L2 process that imported the transport handles. + """ + + __slots__ = ( + "_freed", + "_release_fn", + "_released", + "buffers", + "domain_id", + "generation", + "mapping_size", + "members", + "name", + "retain_after_run", + ) + + def __init__( + self, + *, + name: str, + members: tuple[GlobalDomainMember, ...], + buffers: tuple[GlobalDomainBuffer, ...], + domain_id: int, + generation: int, + mapping_size: int, + retain_after_run: bool, + _release_fn, + ) -> None: + self.name = str(name) + self.members = tuple(members) + self.buffers = tuple(buffers) + self.domain_id = int(domain_id) + self.generation = int(generation) + self.mapping_size = int(mapping_size) + self.retain_after_run = bool(retain_after_run) + self._release_fn = _release_fn + self._released = False + self._freed = False + + def member(self, domain_rank: int) -> GlobalDomainMember: + if self._released: + raise RuntimeError(f"GlobalCommDomainHandle({self.name!r}) is already released") + rank = int(domain_rank) + if rank < 0 or rank >= len(self.members): + raise IndexError(f"global domain rank {rank} is out of range") + member = self.members[rank] + if member.domain_rank != rank: + raise RuntimeError("global domain member table is not rank ordered") + return member + + def buffer_range(self, name: str) -> tuple[int, int]: + if self._released: + raise RuntimeError(f"GlobalCommDomainHandle({self.name!r}) is already released") + offset = 0 + for buffer in self.buffers: + if buffer.name == name: + return offset, buffer.nbytes + offset += buffer.nbytes + raise KeyError(f"global domain {self.name!r} has no buffer {name!r}") + + @property + def released(self) -> bool: + return self._released + + @property + def freed(self) -> bool: + return self._freed + + def release(self) -> None: + if self._released: + return + self._release_fn(self) + self._released = True + + def __enter__(self) -> GlobalCommDomainHandle: + return self + + def __exit__(self, *_): + self.release() + + +class GlobalCommDomainView: + """L3-local imported view exposed to remote orchestration callables.""" + + __slots__ = ( + "_committed", + "contexts", + "domain_id", + "generation", + "mapping_size", + "members", + "name", + ) + + def __init__( + self, + *, + name: str, + members: tuple[GlobalDomainMember, ...], + contexts: dict[int, ChipDomainContext], + domain_id: int, + generation: int, + mapping_size: int, + ) -> None: + self.name = str(name) + self.members = tuple(members) + self.contexts = dict(contexts) + self.domain_id = int(domain_id) + self.generation = int(generation) + self.mapping_size = int(mapping_size) + self._committed = False + + def __getitem__(self, local_worker_id: int) -> ChipDomainContext: + if not self._committed: + raise RuntimeError(f"GlobalCommDomainView({self.name!r}) is not committed") + return self.contexts[int(local_worker_id)] + + @property + def committed(self) -> bool: + return self._committed + + # Process-wide RTLD_GLOBAL preload registry. host_runtime.so resolves its # undefined HostLogger / unified_log_* (and, on sim, sim_context_*) symbols # against these globals, so they must be loaded — exactly once — before any diff --git a/python/simpler/worker.py b/python/simpler/worker.py index ed9a99e42b..83566c4c28 100644 --- a/python/simpler/worker.py +++ b/python/simpler/worker.py @@ -63,6 +63,7 @@ def my_l4_orch(orch, args, config): import contextlib import ctypes import enum +import hashlib import importlib import json import math @@ -107,6 +108,51 @@ def my_l4_orch(orch, args, config): parse_python_callable_payload, parse_python_import_target, ) +from .global_comm_domain import ( + CTRL_GLOBAL_DOMAIN_COPY_FROM, + CTRL_GLOBAL_DOMAIN_COPY_TO, + CTRL_GLOBAL_DOMAIN_IMPORT, + CTRL_GLOBAL_DOMAIN_PREPARE, + CTRL_GLOBAL_DOMAIN_RELEASE, + GLOBAL_DOMAIN_DESCRIPTOR_BYTES, + GLOBAL_DOMAIN_MAX_COPY_BYTES, + GLOBAL_DOMAIN_MAX_RANKS, + GLOBAL_DOMAIN_MAX_STRING_BYTES, + GLOBAL_DOMAIN_PROFILE_IDS, + GLOBAL_DOMAIN_VERSION, + LOCAL_COPY_REPLY, + LOCAL_COPY_REQUEST, + LOCAL_DOMAIN_MAGIC, + LOCAL_IMPORT_REPLY, + LOCAL_IMPORT_REQUEST, + LOCAL_PREPARE_REPLY, + LOCAL_PREPARE_REQUEST, + LOCAL_RELEASE_REQUEST, + GlobalCommInitCommand, + GlobalDomainBuffer, + GlobalDomainCommand, + GlobalDomainCopyCommand, + GlobalDomainDescriptor, + GlobalDomainMember, + GlobalDomainPhase, + GlobalDomainReleaseCommand, + decode_comm_init, + decode_comm_init_result, + decode_copy_command, + decode_copy_result, + decode_descriptor_table, + decode_domain_command, + decode_release_command, + encode_comm_init, + encode_comm_init_result, + encode_copy_command, + encode_copy_result, + encode_descriptor_table, + encode_domain_command, + encode_release_command, + resolve_global_comm_capability, + validate_descriptor_table, +) from .l3_l2_orch_comm import ( _CTRL_SHM_TOKEN_BYTES, _REGION_CREATE_REPLY, @@ -136,6 +182,8 @@ def my_l4_orch(orch, args, config): ChipWorker, CommBufferSpec, CommDomainHandle, + GlobalCommDomainHandle, + GlobalCommDomainView, RemoteAddressSpace, RemoteBufferExport, RemoteBufferHandle, @@ -312,6 +360,12 @@ def my_l4_orch(orch, args, config): _BLOB_HEADER_BYTES = 8 _CTRL_L3_L2_REGION_CREATE = 16 _CTRL_L3_L2_REGION_RELEASE = 17 +# L4-to-local-L3 envelope for the Global CommDomain control protocol. The +# enclosed command uses remote_l3_protocol.ControlName; values 13-17 already +# belong to other local child controls. +_CTRL_GLOBAL_DOMAIN_NODE = 18 +_LOCAL_GLOBAL_CONTROL_HEADER = struct.Struct(" None: object.__setattr__(self, "platform", str(self.platform)) object.__setattr__(self, "runtime", str(self.runtime)) object.__setattr__(self, "transport", str(self.transport)) + object.__setattr__(self, "comm_profile", str(self.comm_profile)) object.__setattr__( self, "session_listen_host", @@ -430,9 +487,24 @@ def __post_init__(self) -> None: ) object.__setattr__(self, "allow_wildcard_session_bind", bool(self.allow_wildcard_session_bind)) object.__setattr__(self, "device_ids", tuple(int(x) for x in self.device_ids)) + object.__setattr__(self, "global_device_ranks", tuple(int(x) for x in self.global_device_ranks)) object.__setattr__(self, "num_sub_workers", int(self.num_sub_workers)) if self.num_sub_workers < 0: raise ValueError("RemoteWorkerSpec.num_sub_workers must be non-negative") + if self.transport != "sim": + raise ValueError("RemoteWorkerSpec.transport must be 'sim' for the TCP daemon control plane") + if self.comm_profile not in GLOBAL_DOMAIN_PROFILE_IDS: + raise ValueError(f"RemoteWorkerSpec.comm_profile {self.comm_profile!r} is not supported") + if self.comm_profile == "a3-fabric-v1" and not self.platform.startswith("a2a3"): + raise ValueError("RemoteWorkerSpec.comm_profile 'a3-fabric-v1' requires an a2a3 platform") + if self.comm_profile == "a3-fabric-v1" and self.platform.endswith("sim"): + raise ValueError("RemoteWorkerSpec.comm_profile 'a3-fabric-v1' requires real A3 devices") + if self.global_device_ranks and len(self.global_device_ranks) != len(self.device_ids): + raise ValueError("RemoteWorkerSpec.global_device_ranks must match device_ids length") + if any(rank < 0 for rank in self.global_device_ranks) or len(set(self.global_device_ranks)) != len( + self.global_device_ranks + ): + raise ValueError("RemoteWorkerSpec.global_device_ranks must be unique and non-negative") @dataclass(frozen=True) @@ -446,6 +518,31 @@ class _RemoteSession: pid: int +@dataclass +class _GlobalNodeDomainState: + command: GlobalDomainCommand + prepared_domain_ranks: set[int] = field(default_factory=set) + descriptors: dict[int, GlobalDomainDescriptor] = field(default_factory=dict) + local_window_bases: dict[int, int] = field(default_factory=dict) + mapping_sizes: dict[int, int] = field(default_factory=dict) + contexts: dict[int, ChipDomainContext] = field(default_factory=dict) + view: GlobalCommDomainView | None = None + phase: GlobalDomainPhase = GlobalDomainPhase.PREPARE_EXPORT + + +@dataclass(frozen=True) +class _GlobalNodeRuntime: + worker_id: int + device_ids: tuple[int, ...] + platform: str + comm_profile: str + global_device_ranks: tuple[int, ...] + node_rank: int + node_count: int + cluster_id: str + is_remote: bool + + _IdentitySnapshotEntry = tuple[bytes, Any, int, str, str] @@ -516,6 +613,10 @@ class HostBuffer: buffer: memoryview +class _NoHostBufferChildrenError(RuntimeError): + """The Worker has no process child that can attach a host buffer.""" + + def _rewrite_blob_host_addrs(buf: memoryview, blob_off: int, ranges: list[tuple[int, int, int]]) -> None: """Redirect registered host pointers in a task-args blob to child mappings. @@ -713,14 +814,15 @@ def _pack_py_callable_payload(target) -> bytes: def _chip_descriptor_context(worker: Worker) -> tuple[str, str]: platform = str(worker._config.get("platform", "")) runtime = str(worker._config.get("runtime", "")) - if platform or runtime: - return platform, runtime - contexts: list[tuple[str, str]] = [] + if platform or runtime: + contexts.append((platform, runtime)) for child in getattr(worker, "_next_level_workers", []): child_context = _chip_descriptor_context(child) if child_context != ("", ""): contexts.append(child_context) + for spec in getattr(worker, "_remote_worker_specs", []): + contexts.append((str(spec.platform), str(spec.runtime))) if not contexts: return "", "" first = contexts[0] @@ -1279,6 +1381,25 @@ class _L2HostL3L2RegionStore: next_region_id: int = 1 +@dataclass +class _L2GlobalDomain: + domain_id: int + generation: int + domain_rank: int + rank_count: int + descriptor: GlobalDomainDescriptor + local_window_base: int + mapping_size: int + requested_window_size: int + device_ctx: int = 0 + descriptor_table: bytes = b"" + + +@dataclass +class _L2GlobalDomainStore: + domains: dict[int, _L2GlobalDomain] = field(default_factory=dict) + + @dataclass(frozen=True) class _L2HostL3L2RegionReplyMeta: payload_base: int @@ -1440,6 +1561,223 @@ def _sweep_l2_host_l3_l2_regions(store: _L2HostL3L2RegionStore) -> None: pass +def _open_global_domain_payload(buf: memoryview) -> tuple[SharedMemory, memoryview, int]: + payload_size = int(struct.unpack_from("Q", buf, _CTRL_OFF_ARG0)[0]) + if payload_size <= 0: + raise RuntimeError("Global CommDomain control payload must be non-empty") + staged = SharedMemory(name=_read_ctrl_staged_shm_name(buf)) + staged_buf = cast(memoryview, staged.buf) + if payload_size > staged.size: + staged_buf.release() + staged.close() + raise RuntimeError("Global CommDomain control payload exceeds staged shm") + return staged, staged_buf, payload_size + + +def _validate_local_global_header( + magic: bytes, version: int, domain_id: int, generation: int, *, operation: str +) -> None: + if magic != LOCAL_DOMAIN_MAGIC or version != GLOBAL_DOMAIN_VERSION: + raise RuntimeError(f"{operation}: local protocol magic or version mismatch") + if domain_id == 0 or generation == 0: + raise RuntimeError(f"{operation}: domain identity must be positive") + + +def _handle_ctrl_global_domain_prepare(cw: ChipWorker, buf: memoryview, store: _L2GlobalDomainStore) -> None: + staged, payload, payload_size = _open_global_domain_payload(buf) + try: + if payload_size < max(LOCAL_PREPARE_REQUEST.size, LOCAL_PREPARE_REPLY.size + GLOBAL_DOMAIN_DESCRIPTOR_BYTES): + raise RuntimeError("Global CommDomain prepare payload is too small") + fields = LOCAL_PREPARE_REQUEST.unpack_from(payload, 0) + magic, version, domain_id, generation, domain_rank, rank_count, profile_id, window_size = fields + _validate_local_global_header(magic, version, domain_id, generation, operation="prepare") + if rank_count <= 0 or rank_count > GLOBAL_DOMAIN_MAX_RANKS or domain_rank >= rank_count: + raise RuntimeError("Global CommDomain prepare rank identity is invalid") + if profile_id not in GLOBAL_DOMAIN_PROFILE_IDS.values() or window_size <= 0: + raise RuntimeError("Global CommDomain prepare profile or window size is invalid") + prior = store.domains.get(int(domain_id)) + if prior is not None: + if ( + prior.generation != generation + or prior.domain_rank != domain_rank + or prior.rank_count != rank_count + or prior.descriptor.profile_id != profile_id + or prior.requested_window_size != window_size + ): + raise RuntimeError("Global CommDomain prepare conflicts with a live domain") + descriptor = prior.descriptor + local_base = prior.local_window_base + mapping_size = prior.mapping_size + else: + descriptor_bytes, local_base, mapping_size = cw._impl.comm_global_domain_prepare( + int(domain_id), + int(domain_rank), + int(rank_count), + int(window_size), + int(profile_id), + ) + descriptor = GlobalDomainDescriptor.decode(bytes(descriptor_bytes)) + if ( + descriptor.domain_rank != domain_rank + or descriptor.rank_count != rank_count + or descriptor.profile_id != profile_id + or descriptor.mapping_size != mapping_size + or descriptor.mapping_size < window_size + ): + cw._impl.comm_global_domain_release(int(domain_id)) + raise RuntimeError("Global CommDomain backend returned an inconsistent descriptor") + store.domains[int(domain_id)] = _L2GlobalDomain( + domain_id=int(domain_id), + generation=int(generation), + domain_rank=int(domain_rank), + rank_count=int(rank_count), + descriptor=descriptor, + local_window_base=int(local_base), + mapping_size=int(mapping_size), + requested_window_size=int(window_size), + ) + LOCAL_PREPARE_REPLY.pack_into( + payload, + 0, + LOCAL_DOMAIN_MAGIC, + GLOBAL_DOMAIN_VERSION, + int(domain_id), + int(generation), + int(local_base), + int(mapping_size), + ) + start = LOCAL_PREPARE_REPLY.size + payload[start : start + GLOBAL_DOMAIN_DESCRIPTOR_BYTES] = descriptor.encode() + finally: + payload.release() + staged.close() + + +def _handle_ctrl_global_domain_import(cw: ChipWorker, buf: memoryview, store: _L2GlobalDomainStore) -> None: + staged, payload, payload_size = _open_global_domain_payload(buf) + try: + if payload_size < LOCAL_IMPORT_REQUEST.size: + raise RuntimeError("Global CommDomain import payload is truncated") + magic, version, domain_id, generation, descriptor_count = LOCAL_IMPORT_REQUEST.unpack_from(payload, 0) + _validate_local_global_header(magic, version, domain_id, generation, operation="import") + expected_size = LOCAL_IMPORT_REQUEST.size + int(descriptor_count) * GLOBAL_DOMAIN_DESCRIPTOR_BYTES + if descriptor_count <= 0 or descriptor_count > GLOBAL_DOMAIN_MAX_RANKS or expected_size > payload_size: + raise RuntimeError("Global CommDomain import descriptor table size is invalid") + entry = store.domains.get(int(domain_id)) + if entry is None or entry.generation != generation: + raise RuntimeError("Global CommDomain import requires a matching prepared domain") + descriptor_bytes = bytes(payload[LOCAL_IMPORT_REQUEST.size : expected_size]) + descriptors = tuple( + GlobalDomainDescriptor.decode(descriptor_bytes[offset : offset + GLOBAL_DOMAIN_DESCRIPTOR_BYTES]) + for offset in range(0, len(descriptor_bytes), GLOBAL_DOMAIN_DESCRIPTOR_BYTES) + ) + profile = next( + name for name, profile_id in GLOBAL_DOMAIN_PROFILE_IDS.items() if profile_id == entry.descriptor.profile_id + ) + validate_descriptor_table(descriptors, rank_count=entry.rank_count, profile=profile) + if descriptors[entry.domain_rank] != entry.descriptor: + raise RuntimeError("Global CommDomain import table does not contain the local exported descriptor") + if entry.descriptor_table and entry.descriptor_table != descriptor_bytes: + raise RuntimeError("Global CommDomain repeated import carries a different descriptor table") + if entry.device_ctx == 0: + entry.device_ctx = int(cw._impl.comm_global_domain_import(int(domain_id), descriptor_bytes)) + if entry.device_ctx == 0: + raise RuntimeError("Global CommDomain backend returned a zero device context") + entry.descriptor_table = descriptor_bytes + if payload_size < LOCAL_IMPORT_REPLY.size: + raise RuntimeError("Global CommDomain import reply capacity is too small") + LOCAL_IMPORT_REPLY.pack_into( + payload, + 0, + LOCAL_DOMAIN_MAGIC, + GLOBAL_DOMAIN_VERSION, + int(domain_id), + int(generation), + entry.device_ctx, + entry.local_window_base, + entry.mapping_size, + ) + finally: + payload.release() + staged.close() + + +def _handle_ctrl_global_domain_release(cw: ChipWorker, buf: memoryview, store: _L2GlobalDomainStore) -> None: + staged, payload, payload_size = _open_global_domain_payload(buf) + try: + if payload_size < LOCAL_RELEASE_REQUEST.size: + raise RuntimeError("Global CommDomain release payload is truncated") + magic, version, domain_id, generation = LOCAL_RELEASE_REQUEST.unpack_from(payload, 0) + _validate_local_global_header(magic, version, domain_id, generation, operation="release") + entry = store.domains.get(int(domain_id)) + if entry is not None and entry.generation != generation: + raise RuntimeError("Global CommDomain release generation mismatch") + if entry is not None: + cw._impl.comm_global_domain_release(int(domain_id)) + store.domains.pop(int(domain_id), None) + finally: + payload.release() + staged.close() + + +def _handle_ctrl_global_domain_copy( + cw: ChipWorker, buf: memoryview, store: _L2GlobalDomainStore, *, copy_to_device: bool +) -> None: + staged, payload, payload_size = _open_global_domain_payload(buf) + try: + if payload_size < LOCAL_COPY_REQUEST.size: + raise RuntimeError("Global CommDomain copy payload is truncated") + magic, version, domain_id, generation, offset, nbytes = LOCAL_COPY_REQUEST.unpack_from(payload, 0) + operation = "copy-to" if copy_to_device else "copy-from" + _validate_local_global_header(magic, version, domain_id, generation, operation=operation) + entry = store.domains.get(int(domain_id)) + if entry is None or entry.generation != generation or entry.device_ctx == 0: + raise RuntimeError(f"Global CommDomain {operation} requires an imported live domain") + if nbytes <= 0 or nbytes > GLOBAL_DOMAIN_MAX_COPY_BYTES: + raise RuntimeError(f"Global CommDomain {operation} size is invalid") + if offset > entry.mapping_size or nbytes > entry.mapping_size - offset: + raise RuntimeError(f"Global CommDomain {operation} range exceeds the local window") + if copy_to_device: + data_offset = LOCAL_COPY_REQUEST.size + if data_offset + nbytes > payload_size: + raise RuntimeError("Global CommDomain copy-to data is truncated") + exported = ctypes.c_char.from_buffer(payload, data_offset) + try: + cw.copy_to(entry.local_window_base + int(offset), ctypes.addressof(exported), int(nbytes)) + finally: + del exported + else: + data_offset = LOCAL_COPY_REPLY.size + if data_offset + nbytes > payload_size: + raise RuntimeError("Global CommDomain copy-from reply capacity is too small") + exported = ctypes.c_char.from_buffer(payload, data_offset) + try: + cw.copy_from(ctypes.addressof(exported), entry.local_window_base + int(offset), int(nbytes)) + finally: + del exported + LOCAL_COPY_REPLY.pack_into( + payload, + 0, + LOCAL_DOMAIN_MAGIC, + GLOBAL_DOMAIN_VERSION, + int(domain_id), + int(generation), + int(nbytes), + ) + finally: + payload.release() + staged.close() + + +def _sweep_l2_global_domains(cw: ChipWorker, store: _L2GlobalDomainStore) -> None: + for domain_id in list(store.domains): + store.domains.pop(domain_id, None) + try: + cw._impl.comm_global_domain_release(int(domain_id)) + except Exception: # noqa: BLE001 + pass + + def _handle_ctrl_release_domain(cw: ChipWorker, buf: memoryview) -> None: """CTRL_RELEASE_DOMAIN handler — collective free for one allocation.""" request_shm_name = _read_shm_name(buf, _OFF_ARGS) @@ -1511,6 +1849,7 @@ def _run_chip_main_loop( # noqa: PLR0913, PLR0915 -- fork-child entry: every de """ prepared = prepared if prepared is not None else set() l3_l2_region_store = _L2HostL3L2RegionStore() + global_domain_store = _L2GlobalDomainStore() # Post-fork host buffers mapped into this child. `host_buf_table` # owns the mmap per token (for unmap + teardown); `host_buf_ranges` is the # parent-VA → child-VA translation table the per-task blob rewrite consults, @@ -1664,6 +2003,26 @@ def handle_control(sub_cmd: int) -> tuple[int, str]: # noqa: PLR0912 -- one bra _handle_ctrl_l3_l2_region_create(cw, buf, chip_platform, l3_l2_region_store) elif sub_cmd == _CTRL_L3_L2_REGION_RELEASE: _handle_ctrl_l3_l2_region_release(buf, l3_l2_region_store) + elif sub_cmd == CTRL_GLOBAL_DOMAIN_PREPARE: + _handle_ctrl_global_domain_prepare(cw, buf, global_domain_store) + elif sub_cmd == CTRL_GLOBAL_DOMAIN_IMPORT: + _handle_ctrl_global_domain_import(cw, buf, global_domain_store) + elif sub_cmd == CTRL_GLOBAL_DOMAIN_RELEASE: + _handle_ctrl_global_domain_release(cw, buf, global_domain_store) + elif sub_cmd == CTRL_GLOBAL_DOMAIN_COPY_TO: + _handle_ctrl_global_domain_copy( + cw, + buf, + global_domain_store, + copy_to_device=True, + ) + elif sub_cmd == CTRL_GLOBAL_DOMAIN_COPY_FROM: + _handle_ctrl_global_domain_copy( + cw, + buf, + global_domain_store, + copy_to_device=False, + ) else: raise RuntimeError(f"unknown control sub-command {int(sub_cmd)}") except Exception as e: # noqa: BLE001 @@ -1678,6 +2037,7 @@ def handle_control(sub_cmd: int) -> tuple[int, str]: # noqa: PLR0912 -- one bra try: _run_mailbox_loop(buf, state_addr, handle_task=handle_task, handle_control=handle_control) finally: + _sweep_l2_global_domains(cw, global_domain_store) _sweep_l2_host_l3_l2_regions(l3_l2_region_store) for host_shm, _lo, _hi, _base in host_buf_table.values(): try: @@ -1806,12 +2166,91 @@ def _read_config_from_mailbox(buf: memoryview) -> CallConfig: return cfg +def _run_local_global_domain_control( # noqa: PLR0912 -- one ordered dispatcher for the Global CommDomain protocol + inner_worker: Worker, + runtime: _GlobalNodeRuntime, + comm_inits: dict[str, GlobalCommInitCommand], + control_name: int, + request: bytes, +) -> bytes: + """Execute one Global CommDomain command inside an add_worker L3 child.""" + from .remote_l3_protocol import ControlName # noqa: PLC0415 + + control = ControlName(control_name) + if control is ControlName.COMM_INIT: + command = decode_comm_init(request) + if command.cluster_id != runtime.cluster_id: + raise ValueError("COMM_INIT cluster_id does not match the local L3 topology") + if command.profile != runtime.comm_profile: + raise ValueError("COMM_INIT profile does not match the local L3 topology") + if command.node_rank != runtime.node_rank or command.node_count != runtime.node_count: + raise ValueError("COMM_INIT node identity does not match the local L3 topology") + local_members = tuple(member for member in command.members if member.node_worker_id == runtime.worker_id) + if not local_members: + raise ValueError("COMM_INIT topology has no local members") + for member in local_members: + if member.local_worker_id < 0 or member.local_worker_id >= len(runtime.global_device_ranks): + raise ValueError("COMM_INIT local worker id exceeds the local L3 device list") + if member.global_device_rank != runtime.global_device_ranks[member.local_worker_id]: + raise ValueError("COMM_INIT global device rank does not match the local L3 topology") + prior = comm_inits.get(command.topology_hash) + if prior is not None and prior != command: + raise ValueError("COMM_INIT topology hash conflicts with an earlier command") + capability = resolve_global_comm_capability( + platform=runtime.platform, + profile=runtime.comm_profile, + local_device_count=len(runtime.global_device_ranks), + ) + comm_inits[command.topology_hash] = command + return encode_comm_init_result(capability) + + if control is ControlName.ALLOC_DOMAIN: + command = decode_domain_command(request) + if not any(init.profile == command.profile and init.members == command.members for init in comm_inits.values()): + raise RuntimeError("ALLOC_DOMAIN requires a matching COMM_INIT topology") + if command.phase is GlobalDomainPhase.PREPARE_EXPORT: + if command.descriptors: + raise ValueError("PREPARE_EXPORT must not carry descriptors") + return encode_descriptor_table(inner_worker._prepare_global_domain_node(command, runtime.worker_id)) + if command.phase is GlobalDomainPhase.IMPORT: + inner_worker._import_global_domain_node(command, runtime.worker_id) + return b"" + if command.phase is GlobalDomainPhase.COMMIT: + inner_worker._commit_global_domain_node(command) + return b"" + if command.phase is GlobalDomainPhase.ABORT: + inner_worker._release_global_domain_node( + GlobalDomainReleaseCommand(command.domain_id, command.generation), + suppress_errors=True, + ) + return b"" + raise ValueError("ALLOC_DOMAIN phase is not supported") + + if control is ControlName.RELEASE_DOMAIN: + inner_worker._release_global_domain_node(decode_release_command(request)) + return b"" + if control is ControlName.COPY_TO_DOMAIN: + inner_worker._copy_global_domain_node( + decode_copy_command(request, include_data=True), + copy_to_device=True, + ) + return b"" + if control is ControlName.COPY_FROM_DOMAIN: + result = inner_worker._copy_global_domain_node( + decode_copy_command(request, include_data=False), + copy_to_device=False, + ) + return encode_copy_result(result) + raise ValueError(f"unsupported local Global CommDomain control {int(control)}") + + def _child_worker_loop( buf: memoryview, registry: dict[int, Any], identity_table: dict[bytes, int], identity_refs: dict[bytes, int], inner_worker: Worker, + global_node: _GlobalNodeRuntime | None = None, ) -> None: """Runs in forked child process. Any-level Worker as child of its parent. @@ -1823,6 +2262,7 @@ def _child_worker_loop( into the inner Worker (see docs section 7). """ state_addr = _buffer_field_addr(buf, _OFF_STATE) + global_comm_inits: dict[str, GlobalCommInitCommand] = {} def handle_task() -> tuple[int, str]: digest = _read_task_digest(buf) @@ -1873,6 +2313,41 @@ def handle_control(sub_cmd: int) -> tuple[int, str]: sub_cmd, context=f"child_worker level={inner_worker.level}", ) + elif sub_cmd == _CTRL_GLOBAL_DOMAIN_NODE: + if global_node is None: + raise RuntimeError("Global CommDomain control requires a local L3 child") + staged, payload, payload_size = _open_global_domain_payload(buf) + try: + if payload_size < _LOCAL_GLOBAL_CONTROL_HEADER.size: + raise RuntimeError("local Global CommDomain control payload is truncated") + control_name, request_size, response_size = _LOCAL_GLOBAL_CONTROL_HEADER.unpack_from(payload, 0) + capacity = payload_size - _LOCAL_GLOBAL_CONTROL_HEADER.size + if request_size > capacity: + raise RuntimeError("local Global CommDomain request exceeds staged payload") + if response_size != 0: + raise RuntimeError("local Global CommDomain request contains a response") + start = _LOCAL_GLOBAL_CONTROL_HEADER.size + request = bytes(payload[start : start + request_size]) + response = _run_local_global_domain_control( + inner_worker, + global_node, + global_comm_inits, + int(control_name), + request, + ) + if len(response) > capacity: + raise RuntimeError("local Global CommDomain response exceeds staged payload") + payload[start : start + len(response)] = response + _LOCAL_GLOBAL_CONTROL_HEADER.pack_into( + payload, + 0, + int(control_name), + int(request_size), + len(response), + ) + finally: + payload.release() + staged.close() else: raise RuntimeError(f"unknown control sub-command {sub_cmd}") except Exception as e: # noqa: BLE001 @@ -1954,6 +2429,8 @@ class _RunResources: remote_slot_refs: list[RemoteBufferHandle] = field(default_factory=list) live_domains: dict[str, CommDomainHandle] = field(default_factory=dict) pending_release_domains: list[CommDomainHandle] = field(default_factory=list) + live_global_domains: dict[str, GlobalCommDomainHandle] = field(default_factory=dict) + pending_release_global_domains: list[GlobalCommDomainHandle] = field(default_factory=list) l3_l2_regions: list[Any] = field(default_factory=list) l3_l2_orch_comm_host_buffers: dict[int, int] = field(default_factory=dict) # True once the owning run's fence has claimed the domains above. A release @@ -2376,6 +2853,10 @@ def __init__( # among live handles). ``orch.allocate_domain`` adds entries here; # ``release()`` removes them and queues a deferred backend free. self._live_domains: dict[str, CommDomainHandle] = {} + self._live_global_domains: dict[str, GlobalCommDomainHandle] = {} + self._failed_global_domain_releases: dict[int, GlobalCommDomainHandle] = {} + self._global_node_domains: dict[int, _GlobalNodeDomainState] = {} + self._global_cluster_id = uuid.uuid4().hex # Monotonic per-Worker counter; mixed into IPC barrier filenames so # two concurrent allocations don't share a marker file. Wraps after # 2^64 allocations — far beyond any realistic Worker lifetime. @@ -2389,6 +2870,8 @@ def __init__( # down under an in-flight release. self._domain_free_mu = threading.Lock() self._domain_free_results: dict[int, BaseException | None] = {} + self._global_domain_free_mu = threading.Lock() + self._global_domain_free_results: dict[int, BaseException | None] = {} self._alloc_id_lock = threading.Lock() # Base HCCL/sim communicator is built lazily on the first # ``orch.allocate_domain`` call (see ``_ensure_comm_base``). We @@ -2630,6 +3113,153 @@ def _remote_dispatcher_entries_for_worker(self, worker_id: int) -> list[dict[str ) return entries + def _inner_registry_entries_for_spec(self, spec: RemoteWorkerSpec) -> list[dict[str, Any]]: + from .remote_l3_protocol import ( # noqa: PLC0415 + ChipCallableBlobLocation, + RemoteChipCallablePayload, + encode_remote_chip_callable_payload, + ) + + entries: list[dict[str, Any]] = [] + with self._registry_lock: + states = list(self._identity_registry.values()) + for state in states: + if state.target_namespace != "LOCAL_CHIP": + continue + if not isinstance(state.target, ChipCallable): + raise RuntimeError(f"inner chip hashid {state.hashid} does not carry a ChipCallable target") + descriptor = build_chip_callable_descriptor( + target=state.target, + platform=spec.platform, + runtime=spec.runtime, + ) + if descriptor != state.descriptor: + raise RuntimeError(f"inner chip hashid {state.hashid} was registered for a different platform/runtime") + blob = ctypes.string_at(int(state.target.buffer_ptr()), int(state.target.buffer_size())) + payload = encode_remote_chip_callable_payload( + RemoteChipCallablePayload( + descriptor_bytes=descriptor, + blob_location=ChipCallableBlobLocation.INLINE_BLOB, + blob_size=len(blob), + blob_sha256=hashlib.sha256(blob).digest(), + inline_blob=blob, + staged_blob_token=b"", + ) + ) + entries.append( + { + "hashid": state.digest.hex(), + "kind": "CHIP_CALLABLE", + "target_registry": "INNER_L3_WORKER", + "payload_version": 1, + "payload_hex": payload.hex(), + } + ) + return entries + + @staticmethod + def _validate_global_node_config( + *, + label: str, + platform: str, + device_ids: tuple[int, ...], + comm_profile: str, + global_device_ranks: tuple[int, ...], + ) -> None: + if comm_profile not in GLOBAL_DOMAIN_PROFILE_IDS: + raise ValueError(f"{label} comm_profile {comm_profile!r} is not supported") + if comm_profile == "a3-fabric-v1" and (not platform.startswith("a2a3") or platform.endswith("sim")): + raise ValueError(f"{label} comm_profile 'a3-fabric-v1' requires real A3 devices") + if global_device_ranks and len(global_device_ranks) != len(device_ids): + raise ValueError(f"{label} global_device_ranks must match device_ids length") + if any(rank < 0 for rank in global_device_ranks) or len(set(global_device_ranks)) != len(global_device_ranks): + raise ValueError(f"{label} global_device_ranks must be unique and non-negative") + + def _resolved_global_nodes(self) -> dict[int, _GlobalNodeRuntime]: + configs: list[tuple[int, tuple[int, ...], str, str, tuple[int, ...], bool]] = [] + for worker_id, spec in zip(self._remote_worker_ids, self._remote_worker_specs, strict=True): + configs.append( + ( + int(worker_id), + tuple(spec.device_ids), + spec.platform, + spec.comm_profile, + tuple(spec.global_device_ranks), + True, + ) + ) + for worker_id, child in zip(self._next_level_worker_ids, self._next_level_workers, strict=True): + if child.level != 3: + continue + device_ids = tuple(int(device_id) for device_id in child._config.get("device_ids", ())) + platform = str(child._config.get("platform", "")) + comm_profile = str(child._config.get("comm_profile", "sim")) + global_device_ranks = tuple(int(rank) for rank in child._config.get("global_device_ranks", ())) + self._validate_global_node_config( + label=f"local L3 worker {worker_id}", + platform=platform, + device_ids=device_ids, + comm_profile=comm_profile, + global_device_ranks=global_device_ranks, + ) + configs.append( + ( + int(worker_id), + device_ids, + platform, + comm_profile, + global_device_ranks, + False, + ) + ) + configs.sort(key=lambda item: item[0]) + + explicit_ranks: set[int] = set() + for worker_id, _device_ids, _platform, _profile, ranks, _is_remote in configs: + overlap = explicit_ranks.intersection(ranks) + if overlap: + raise ValueError( + f"Global CommDomain worker {worker_id} duplicates global_device_ranks {sorted(overlap)}" + ) + explicit_ranks.update(ranks) + + used = set(explicit_ranks) + next_rank = 0 + resolved: dict[int, _GlobalNodeRuntime] = {} + node_count = len(configs) + for node_rank, (worker_id, device_ids, platform, profile, ranks, is_remote) in enumerate(configs): + self._validate_global_node_config( + label=f"{'remote' if is_remote else 'local'} L3 worker {worker_id}", + platform=platform, + device_ids=device_ids, + comm_profile=profile, + global_device_ranks=ranks, + ) + if not ranks: + assigned: list[int] = [] + for _device_id in device_ids: + while next_rank in used: + next_rank += 1 + assigned.append(next_rank) + used.add(next_rank) + next_rank += 1 + ranks = tuple(assigned) + resolved[worker_id] = _GlobalNodeRuntime( + worker_id=worker_id, + device_ids=device_ids, + platform=platform, + comm_profile=profile, + global_device_ranks=ranks, + node_rank=node_rank, + node_count=node_count, + cluster_id=self._global_cluster_id, + is_remote=is_remote, + ) + return resolved + + def _resolved_global_device_ranks(self) -> dict[int, tuple[int, ...]]: + return {worker_id: runtime.global_device_ranks for worker_id, runtime in self._resolved_global_nodes().items()} + def _build_remote_manifest( self, *, spec: RemoteWorkerSpec, worker_id: int, session_id: int, startup_remaining_s: float ) -> dict[str, Any]: @@ -2637,6 +3267,15 @@ def _build_remote_manifest( listen_host = spec.session_listen_host or ("127.0.0.1" if daemon_host == "localhost" else daemon_host) if self._is_wildcard_session_host(listen_host) and not spec.allow_wildcard_session_bind: raise ValueError("RemoteWorkerSpec wildcard session bind requires allow_wildcard_session_bind=True") + if worker_id in self._remote_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))) return { "session_id": int(session_id), "parent_worker_level": int(self.level), @@ -2648,6 +3287,11 @@ def _build_remote_manifest( "num_sub_workers": int(spec.num_sub_workers), "heap_ring_size": self._config.get("remote_heap_ring_size", None), "transport": spec.transport, + "comm_profile": spec.comm_profile, + "cluster_id": self._global_cluster_id, + "node_rank": node_rank, + "node_count": node_count, + "global_device_ranks": list(global_device_ranks), # session_timeout_s bounds the runtime command socket; startup_remaining_s # bounds this session's slice of the single root startup budget. They are # distinct: the remote must not spend runtime-command time as startup time. @@ -2656,7 +3300,7 @@ def _build_remote_manifest( "listen_host": listen_host, "connect_host": daemon_host, "remote_task_dispatcher": self._remote_dispatcher_entries_for_worker(worker_id), - "inner_l3_worker": [], + "inner_l3_worker": self._inner_registry_entries_for_spec(spec), "feature_flags": [], } @@ -4178,7 +4822,15 @@ def _eligible_target_need(self, namespace: str | None, eligible_worker_ids) -> s has_python_child = self._config.get("num_sub_workers", 0) > 0 or bool(self._next_level_workers) return None if has_python_child else "a SUB or next-level child" if namespace == "LOCAL_CHIP": - return None if bool(self._config.get("device_ids")) else "a chip device (device_ids)" + + def has_chip_target(worker: Worker) -> bool: + if worker._config.get("device_ids"): + return True + if any(spec.device_ids for spec in worker._remote_worker_specs): + return True + return any(has_chip_target(child) for child in worker._next_level_workers) + + return None if has_chip_target(self) else "a local or remote chip device" if namespace == "REMOTE_TASK_DISPATCHER": has_remote_workers = set(self._remote_worker_ids) ok = bool(has_remote_workers) and set(eligible_worker_ids) <= has_remote_workers @@ -4452,7 +5104,7 @@ def _activate_remote_sessions(self, deadline: float) -> None: self._worker.add_remote_l3_socket( session.worker_id, session.session_id, - spec.transport, + spec.comm_profile, session.command_host, session.command_port, session.health_host, @@ -4476,6 +5128,7 @@ def _start_hierarchical(self) -> None: # noqa: PLR0912 -- three parallel fork l device_ids = self._config.get("device_ids", []) n_sub = self._config.get("num_sub_workers", 0) deadline = self._startup_deadline + global_nodes = self._resolved_global_nodes() if self.level >= 4 else {} # Freeze the startup registry snapshot. init() already holds the epoch in # the INITIALIZING state, so a concurrent register/unregister is blocked @@ -4590,6 +5243,8 @@ def _setup(): # L3 child → L3's chip/sub grandchildren) and INIT_READY propagates up # only after the whole subtree is ready. for idx, inner_worker in enumerate(self._next_level_workers): + worker_id = self._next_level_worker_ids[idx] + global_node = global_nodes.get(worker_id) pid = os.fork() if pid == 0: buf = self._next_level_shms[idx].buf @@ -4621,7 +5276,12 @@ def _setup(inner=inner_worker): buf, f"next_level worker {idx}", _setup, - lambda tables, b=buf, inner=inner_worker: _child_worker_loop(b, *tables, inner), + lambda tables, b=buf, inner=inner_worker, node=global_node: _child_worker_loop( + b, + *tables, + inner, + node, + ), make_group_leader=self._is_startup_root, ) else: @@ -5387,8 +6047,13 @@ def _retire_run_domains(self, resources: _RunResources) -> None: resources.retired = True stragglers = list(resources.pending_release_domains) resources.pending_release_domains.clear() + global_stragglers = list(resources.pending_release_global_domains) + resources.pending_release_global_domains.clear() + resources.live_global_domains.clear() for handle in stragglers: self._free_domain_after_fence(handle) + for handle in global_stragglers: + self._free_global_domain_after_fence(handle) def _free_domain_after_fence(self, handle: CommDomainHandle) -> None: """Back-end free for a handle whose owning run has retired. @@ -5560,6 +6225,683 @@ def dispatch(chip_idx: int) -> None: f"{len(errors)}/{len(workers)} chips; first error chip={first[0]}: {first[1]}" ) + @staticmethod + def _global_domain_command_identity(command: GlobalDomainCommand) -> tuple[Any, ...]: + return ( + command.domain_id, + command.generation, + command.name, + command.profile, + command.window_size, + command.members, + command.buffers, + ) + + @staticmethod + def _global_domain_provenance_id(domain_id: int) -> int: + # Local CommDomain allocation ids are positive. Keep remote Global + # CommDomains in a disjoint namespace while reusing the same exact-pointer + # provenance table that protects child-memory task submissions. + return -int(domain_id) + + def _global_local_members( + self, command: GlobalDomainCommand, node_worker_id: int + ) -> tuple[GlobalDomainMember, ...]: + members = tuple(member for member in command.members if member.node_worker_id == node_worker_id) + if not members: + raise ValueError(f"Global CommDomain has no members on node worker {node_worker_id}") + local_count = len(self._config.get("device_ids", [])) + for member in members: + if member.local_worker_id < 0 or member.local_worker_id >= local_count: + raise ValueError( + f"Global CommDomain local worker {member.local_worker_id} is outside [0, {local_count})" + ) + return members + + def _prepare_global_domain_node( + self, command: GlobalDomainCommand, node_worker_id: int + ) -> tuple[GlobalDomainDescriptor, ...]: + if self.level != 3 or self._worker is None: + raise RuntimeError("Global CommDomain node prepare requires a ready L3 Worker") + prior = self._global_node_domains.get(command.domain_id) + if prior is not None: + if self._global_domain_command_identity(prior.command) != self._global_domain_command_identity(command): + raise RuntimeError("Global CommDomain prepare conflicts with a live domain") + return tuple(prior.descriptors[rank] for rank in sorted(prior.descriptors)) + + state = _GlobalNodeDomainState(command=command) + self._global_node_domains[command.domain_id] = state + local_members = self._global_local_members(command, node_worker_id) + capacity = max(LOCAL_PREPARE_REQUEST.size, LOCAL_PREPARE_REPLY.size + GLOBAL_DOMAIN_DESCRIPTOR_BYTES) + try: + for member in local_members: + payload = bytearray(capacity) + LOCAL_PREPARE_REQUEST.pack_into( + payload, + 0, + LOCAL_DOMAIN_MAGIC, + GLOBAL_DOMAIN_VERSION, + command.domain_id, + command.generation, + member.domain_rank, + len(command.members), + GLOBAL_DOMAIN_PROFILE_IDS[command.profile], + command.window_size, + ) + state.prepared_domain_ranks.add(member.domain_rank) + reply = bytes( + self._worker.control_payload( + WorkerType.NEXT_LEVEL, + member.local_worker_id, + CTRL_GLOBAL_DOMAIN_PREPARE, + payload, + _PY_CONTROL_TIMEOUT_S, + ) + ) + fields = LOCAL_PREPARE_REPLY.unpack_from(reply, 0) + magic, version, domain_id, generation, local_base, mapping_size = fields + _validate_local_global_header(magic, version, domain_id, generation, operation="prepare reply") + if domain_id != command.domain_id or generation != command.generation: + raise RuntimeError("Global CommDomain prepare reply identity mismatch") + start = LOCAL_PREPARE_REPLY.size + descriptor = GlobalDomainDescriptor.decode(reply[start : start + GLOBAL_DOMAIN_DESCRIPTOR_BYTES]) + if descriptor.domain_rank != member.domain_rank: + raise RuntimeError("Global CommDomain prepare reply rank mismatch") + state.descriptors[member.domain_rank] = descriptor + state.local_window_bases[member.local_worker_id] = int(local_base) + state.mapping_sizes[member.local_worker_id] = int(mapping_size) + return tuple(state.descriptors[rank] for rank in sorted(state.descriptors)) + except BaseException: + self._release_global_domain_node( + GlobalDomainReleaseCommand(command.domain_id, command.generation), + suppress_errors=True, + ) + raise + + def _import_global_domain_node(self, command: GlobalDomainCommand, node_worker_id: int) -> None: + if self.level != 3 or self._worker is None: + raise RuntimeError("Global CommDomain node import requires a ready L3 Worker") + state = self._global_node_domains.get(command.domain_id) + if state is None or state.command.generation != command.generation: + raise RuntimeError("Global CommDomain import requires a matching prepared domain") + if self._global_domain_command_identity(state.command) != self._global_domain_command_identity(command): + raise RuntimeError("Global CommDomain import command conflicts with prepare") + validate_descriptor_table( + command.descriptors, + rank_count=len(command.members), + profile=command.profile, + ) + local_members = self._global_local_members(command, node_worker_id) + descriptor_bytes = b"".join(descriptor.encode() for descriptor in command.descriptors) + request_size = LOCAL_IMPORT_REQUEST.size + len(descriptor_bytes) + capacity = max(request_size, LOCAL_IMPORT_REPLY.size) + for member in local_members: + payload = bytearray(capacity) + LOCAL_IMPORT_REQUEST.pack_into( + payload, + 0, + LOCAL_DOMAIN_MAGIC, + GLOBAL_DOMAIN_VERSION, + command.domain_id, + command.generation, + len(command.descriptors), + ) + payload[LOCAL_IMPORT_REQUEST.size : request_size] = descriptor_bytes + reply = bytes( + self._worker.control_payload( + WorkerType.NEXT_LEVEL, + member.local_worker_id, + CTRL_GLOBAL_DOMAIN_IMPORT, + payload, + _PY_CONTROL_TIMEOUT_S, + ) + ) + fields = LOCAL_IMPORT_REPLY.unpack_from(reply, 0) + magic, version, domain_id, generation, device_ctx, local_base, mapping_size = fields + _validate_local_global_header(magic, version, domain_id, generation, operation="import reply") + if domain_id != command.domain_id or generation != command.generation: + raise RuntimeError("Global CommDomain import reply identity mismatch") + if mapping_size != command.descriptors[member.domain_rank].mapping_size: + raise RuntimeError("Global CommDomain import reply mapping size mismatch") + offset = 0 + buffer_ptrs: dict[str, int] = {} + for buffer in command.buffers: + buffer_ptrs[buffer.name] = int(local_base) + offset + offset += buffer.nbytes + state.contexts[member.local_worker_id] = ChipDomainContext( + name=command.name, + domain_rank=member.domain_rank, + domain_size=len(command.members), + device_ctx=int(device_ctx), + local_window_base=int(local_base), + actual_window_size=int(mapping_size), + buffer_ptrs=buffer_ptrs, + ) + provenance_id = self._global_domain_provenance_id(command.domain_id) + with self._child_prov_lock: + for pointer in {int(local_base), *buffer_ptrs.values()}: + self._child_prov_record_domain(member.local_worker_id, pointer, provenance_id) + state.command = command + state.phase = GlobalDomainPhase.IMPORT + state.view = GlobalCommDomainView( + name=command.name, + members=command.members, + contexts=state.contexts, + domain_id=command.domain_id, + generation=command.generation, + mapping_size=command.descriptors[0].mapping_size, + ) + + def _commit_global_domain_node(self, command: GlobalDomainCommand) -> None: + state = self._global_node_domains.get(command.domain_id) + if state is None or state.command.generation != command.generation: + raise RuntimeError("Global CommDomain commit requires a matching imported domain") + if state.phase is not GlobalDomainPhase.IMPORT or state.view is None: + raise RuntimeError("Global CommDomain commit requires IMPORT completion") + if ( + self._global_domain_command_identity(state.command) != self._global_domain_command_identity(command) + or state.command.descriptors != command.descriptors + ): + raise RuntimeError("Global CommDomain commit command conflicts with IMPORT") + state.phase = GlobalDomainPhase.COMMIT + state.view._committed = True # noqa: SLF001 -- session owns the transaction + + def _release_global_domain_node( + self, command: GlobalDomainReleaseCommand, *, suppress_errors: bool = False + ) -> None: + state = self._global_node_domains.get(command.domain_id) + if state is None: + return + if state.command.generation != command.generation: + raise RuntimeError("Global CommDomain release generation mismatch") + with self._child_prov_lock: + self._child_prov_drop_domain(self._global_domain_provenance_id(command.domain_id)) + if self._worker is None: + return + errors: list[BaseException] = [] + local_members = tuple( + member for member in state.command.members if member.domain_rank in state.prepared_domain_ranks + ) + for member in local_members: + payload = bytearray(LOCAL_RELEASE_REQUEST.size) + LOCAL_RELEASE_REQUEST.pack_into( + payload, + 0, + LOCAL_DOMAIN_MAGIC, + GLOBAL_DOMAIN_VERSION, + command.domain_id, + command.generation, + ) + try: + self._worker.control_payload( + WorkerType.NEXT_LEVEL, + member.local_worker_id, + CTRL_GLOBAL_DOMAIN_RELEASE, + payload, + _PY_CONTROL_TIMEOUT_S, + ) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + if not errors: + self._global_node_domains.pop(command.domain_id, None) + if errors and not suppress_errors: + raise RuntimeError(f"Global CommDomain node release failed: {errors[0]}") from errors[0] + + def _release_all_global_domain_nodes(self) -> None: + for state in list(self._global_node_domains.values())[::-1]: + try: + self._release_global_domain_node( + GlobalDomainReleaseCommand(state.command.domain_id, state.command.generation) + ) + except Exception as exc: # noqa: BLE001 + sys.stderr.write( + f"Worker._release_all_global_domain_nodes: domain_id={state.command.domain_id} " + f"release failed: {type(exc).__name__}: {exc}\n" + ) + sys.stderr.flush() + + def _get_global_domain(self, domain_id: int) -> GlobalCommDomainView: + state = self._global_node_domains.get(int(domain_id)) + if state is None or state.phase is not GlobalDomainPhase.COMMIT or state.view is None: + raise KeyError(f"Global CommDomain {domain_id} is not committed on this L3 node") + return state.view + + def _copy_global_domain_node(self, command: GlobalDomainCopyCommand, *, copy_to_device: bool) -> bytes: + state = self._global_node_domains.get(command.domain_id) + if ( + state is None + or state.command.generation != command.generation + or state.phase is not GlobalDomainPhase.COMMIT + ): + raise RuntimeError("Global CommDomain copy requires a committed live domain") + if command.domain_rank >= len(state.command.members): + raise ValueError("Global CommDomain copy rank is out of range") + member = state.command.members[command.domain_rank] + if member.local_worker_id not in state.contexts: + raise RuntimeError("Global CommDomain copy rank is not local to this L3 node") + request_size = LOCAL_COPY_REQUEST.size + (command.nbytes if copy_to_device else 0) + reply_size = LOCAL_COPY_REPLY.size + (command.nbytes if not copy_to_device else 0) + payload = bytearray(max(request_size, reply_size)) + LOCAL_COPY_REQUEST.pack_into( + payload, + 0, + LOCAL_DOMAIN_MAGIC, + GLOBAL_DOMAIN_VERSION, + command.domain_id, + command.generation, + command.offset, + command.nbytes, + ) + if copy_to_device: + payload[LOCAL_COPY_REQUEST.size : request_size] = command.data + assert self._worker is not None + reply = bytes( + self._worker.control_payload( + WorkerType.NEXT_LEVEL, + member.local_worker_id, + CTRL_GLOBAL_DOMAIN_COPY_TO if copy_to_device else CTRL_GLOBAL_DOMAIN_COPY_FROM, + payload, + _PY_CONTROL_TIMEOUT_S, + ) + ) + magic, version, domain_id, generation, nbytes = LOCAL_COPY_REPLY.unpack_from(reply, 0) + _validate_local_global_header(magic, version, domain_id, generation, operation="copy reply") + if domain_id != command.domain_id or generation != command.generation or nbytes != command.nbytes: + raise RuntimeError("Global CommDomain copy reply mismatch") + if copy_to_device: + return b"" + return reply[LOCAL_COPY_REPLY.size : LOCAL_COPY_REPLY.size + command.nbytes] + + @staticmethod + def _local_global_domain_response_capacity(control_name: int, payload: bytes) -> int: + from .remote_l3_protocol import ControlName # noqa: PLC0415 + + control = ControlName(control_name) + if control is ControlName.COMM_INIT: + return struct.calcsize(" bytes: + if self._worker is None: + raise RuntimeError("Global CommDomain control requires a ready hierarchical Worker") + if worker_id not in self._next_level_worker_ids: + raise ValueError(f"Global CommDomain worker {worker_id} is not a local L3 worker") + response_capacity = self._local_global_domain_response_capacity(control_name, payload) + capacity = max(len(payload), response_capacity) + staged = bytearray(_LOCAL_GLOBAL_CONTROL_HEADER.size + capacity) + _LOCAL_GLOBAL_CONTROL_HEADER.pack_into(staged, 0, int(control_name), len(payload), 0) + start = _LOCAL_GLOBAL_CONTROL_HEADER.size + staged[start : start + len(payload)] = payload + reply = bytes( + self._worker.control_payload( + WorkerType.NEXT_LEVEL, + int(worker_id), + _CTRL_GLOBAL_DOMAIN_NODE, + staged, + _PY_CONTROL_TIMEOUT_S, + ) + ) + reply_control, reply_request_size, response_size = _LOCAL_GLOBAL_CONTROL_HEADER.unpack_from(reply, 0) + if reply_control != int(control_name) or reply_request_size != len(payload) or response_size > capacity: + raise RuntimeError("local Global CommDomain control reply is invalid") + return reply[start : start + response_size] + + def _global_domain_control(self, worker_id: int, control_name: int, payload: bytes) -> bytes: + if self._worker is None: + raise RuntimeError("Global CommDomain control requires a ready hierarchical Worker") + if worker_id in self._remote_worker_ids: + return bytes(self._worker.remote_domain_control(int(worker_id), int(control_name), bytes(payload))) + if worker_id in self._next_level_worker_ids: + return self._local_global_domain_control(worker_id, control_name, payload) + raise ValueError(f"Global CommDomain worker {worker_id} is not a registered L3 worker") + + def _allocate_global_domain( # noqa: PLR0912 -- transaction validation and prepare/import/commit rollback stay ordered + self, + *, + name: str, + members: tuple[tuple[int, int], ...], + window_size: int, + buffers: list[CommBufferSpec], + retain_after_run: bool, + ) -> GlobalCommDomainHandle: + from .remote_l3_protocol import ControlName # noqa: PLC0415 + + if self.level < 4 or self._worker is None: + raise RuntimeError("allocate_global_domain requires a ready L4+ Worker") + resources = self._building_run_resources + if resources is None: + raise RuntimeError("allocate_global_domain is only valid while a run's graph is being built") + if not name: + raise ValueError("allocate_global_domain: name must be non-empty") + if name in self._live_global_domains: + raise ValueError(f"allocate_global_domain: domain {name!r} is already live") + if not members or len(members) > GLOBAL_DOMAIN_MAX_RANKS: + raise ValueError("allocate_global_domain: members must contain between 1 and 64 devices") + if len(set(members)) != len(members): + raise ValueError("allocate_global_domain: members contain duplicate node/local devices") + if window_size <= 0: + raise ValueError("allocate_global_domain: window_size must be positive") + if len({buffer.name for buffer in buffers}) != len(buffers): + raise ValueError("allocate_global_domain: buffer names must be unique") + if any(not buffer.name or int(buffer.nbytes) <= 0 for buffer in buffers): + raise ValueError("allocate_global_domain: buffers require a name and positive nbytes") + if sum(int(buffer.nbytes) for buffer in buffers) > window_size: + raise ValueError("allocate_global_domain: buffers exceed window_size") + + nodes = self._resolved_global_nodes() + profiles: set[str] = set() + domain_members: list[GlobalDomainMember] = [] + for domain_rank, (node_worker_id, local_worker_id) in enumerate(members): + node = nodes.get(int(node_worker_id)) + if node is None: + raise ValueError(f"allocate_global_domain: worker {node_worker_id} is not a registered L3") + if local_worker_id < 0 or local_worker_id >= len(node.device_ids): + raise ValueError( + f"allocate_global_domain: local worker {local_worker_id} is outside " + f"worker {node_worker_id}'s device list" + ) + profiles.add(node.comm_profile) + domain_members.append( + GlobalDomainMember( + node_worker_id=int(node_worker_id), + local_worker_id=int(local_worker_id), + global_device_rank=node.global_device_ranks[int(local_worker_id)], + domain_rank=domain_rank, + ) + ) + if len(profiles) != 1: + raise ValueError("allocate_global_domain: all participating nodes must use the same comm_profile") + profile = next(iter(profiles)) + global_buffers = tuple(GlobalDomainBuffer(buffer.name, int(buffer.nbytes)) for buffer in buffers) + domain_members_tuple = tuple(domain_members) + involved_nodes = tuple(dict.fromkeys(member.node_worker_id for member in domain_members_tuple)) + for node_worker_id in involved_nodes: + node = nodes[node_worker_id] + resolve_global_comm_capability( + platform=node.platform, + profile=node.comm_profile, + local_device_count=len(node.device_ids), + ) + topology_bytes = repr( + ( + self._global_cluster_id, + profile, + tuple( + ( + member.node_worker_id, + member.local_worker_id, + member.global_device_rank, + member.domain_rank, + ) + for member in domain_members_tuple + ), + ) + ).encode() + topology_hash = hashlib.sha256(topology_bytes).hexdigest() + with self._alloc_id_lock: + self._next_alloc_id += 1 + domain_id = self._next_alloc_id + generation = 1 + base_command = GlobalDomainCommand( + phase=GlobalDomainPhase.PREPARE_EXPORT, + domain_id=domain_id, + generation=generation, + name=name, + profile=profile, + window_size=int(window_size), + members=domain_members_tuple, + buffers=global_buffers, + ) + + prepared_nodes: list[int] = [] + try: + for node_worker_id in involved_nodes: + node = nodes[node_worker_id] + init = GlobalCommInitCommand( + cluster_id=self._global_cluster_id, + topology_hash=topology_hash, + profile=profile, + node_rank=node.node_rank, + node_count=node.node_count, + members=domain_members_tuple, + ) + result = decode_comm_init_result( + self._global_domain_control(node_worker_id, ControlName.COMM_INIT, encode_comm_init(init)) + ) + if ( + result.profile != profile + or result.max_ranks < len(domain_members_tuple) + or result.descriptor_bytes != GLOBAL_DOMAIN_DESCRIPTOR_BYTES + or result.local_device_count != len(node.device_ids) + ): + raise RuntimeError(f"Global CommDomain COMM_INIT capability mismatch on node {node_worker_id}") + + descriptor_by_rank: dict[int, GlobalDomainDescriptor] = {} + for node_worker_id in involved_nodes: + prepared_nodes.append(node_worker_id) + reply = self._global_domain_control( + node_worker_id, + ControlName.ALLOC_DOMAIN, + encode_domain_command(base_command), + ) + for descriptor in decode_descriptor_table(reply): + if descriptor.domain_rank in descriptor_by_rank: + raise RuntimeError("Global CommDomain prepare returned a duplicate rank") + descriptor_by_rank[descriptor.domain_rank] = descriptor + descriptors = tuple(descriptor_by_rank[rank] for rank in range(len(domain_members_tuple))) + validate_descriptor_table(descriptors, rank_count=len(domain_members_tuple), profile=profile) + if descriptors[0].mapping_size < window_size: + raise RuntimeError("Global CommDomain backend mapped less than the requested window size") + + import_command = GlobalDomainCommand( + phase=GlobalDomainPhase.IMPORT, + domain_id=domain_id, + generation=generation, + name=name, + profile=profile, + window_size=int(window_size), + members=domain_members_tuple, + buffers=global_buffers, + descriptors=descriptors, + ) + for node_worker_id in involved_nodes: + self._global_domain_control( + node_worker_id, + ControlName.ALLOC_DOMAIN, + encode_domain_command(import_command), + ) + commit_command = GlobalDomainCommand( + phase=GlobalDomainPhase.COMMIT, + domain_id=domain_id, + generation=generation, + name=name, + profile=profile, + window_size=int(window_size), + members=domain_members_tuple, + buffers=global_buffers, + descriptors=descriptors, + ) + for node_worker_id in involved_nodes: + self._global_domain_control( + node_worker_id, + ControlName.ALLOC_DOMAIN, + encode_domain_command(commit_command), + ) + except BaseException: + abort_command = GlobalDomainCommand( + phase=GlobalDomainPhase.ABORT, + domain_id=domain_id, + generation=generation, + name=name, + profile=profile, + window_size=int(window_size), + members=domain_members_tuple, + buffers=global_buffers, + ) + for node_worker_id in prepared_nodes: + with contextlib.suppress(BaseException): + self._global_domain_control( + node_worker_id, + ControlName.ALLOC_DOMAIN, + encode_domain_command(abort_command), + ) + raise + + handle = GlobalCommDomainHandle( + name=name, + members=domain_members_tuple, + buffers=global_buffers, + domain_id=domain_id, + generation=generation, + mapping_size=descriptors[0].mapping_size, + retain_after_run=retain_after_run, + _release_fn=self._release_global_domain_handle, + ) + self._live_global_domains[name] = handle + resources.live_global_domains[name] = handle + return handle + + def _release_global_domain_handle(self, handle: GlobalCommDomainHandle) -> None: + if self._worker is None: + return + resources = self._building_run_resources + if resources is not None: + with resources.domain_lock: + if resources.live_global_domains.get(handle.name) is handle: + resources.live_global_domains.pop(handle.name) + if self._live_global_domains.get(handle.name) is handle: + self._live_global_domains.pop(handle.name) + if not resources.retired: + resources.pending_release_global_domains.append(handle) + return + elif self._live_global_domains.get(handle.name) is handle: + self._live_global_domains.pop(handle.name) + self._free_global_domain_after_fence(handle) + + def _free_global_domain_after_fence(self, handle: GlobalCommDomainHandle) -> None: + if handle.freed: + return + try: + self._release_global_domain_now(handle) + handle._freed = True # noqa: SLF001 -- runtime owns this transition + self._failed_global_domain_releases.pop(handle.domain_id, None) + except Exception: + self._failed_global_domain_releases[handle.domain_id] = handle + raise + + def _release_global_domain_now(self, handle: GlobalCommDomainHandle) -> None: + with self._global_domain_free_mu: + if handle.domain_id in self._global_domain_free_results: + failure = self._global_domain_free_results[handle.domain_id] + if failure is not None: + raise failure + return + try: + self._release_global_domain_claimed(handle) + except BaseException as exc: + self._global_domain_free_results[handle.domain_id] = exc + raise + self._global_domain_free_results[handle.domain_id] = None + + def _release_global_domain_claimed(self, handle: GlobalCommDomainHandle) -> None: + from .remote_l3_protocol import ControlName # noqa: PLC0415 + + command = encode_release_command(GlobalDomainReleaseCommand(handle.domain_id, handle.generation)) + errors: list[BaseException] = [] + for node_worker_id in dict.fromkeys(member.node_worker_id for member in handle.members): + try: + self._global_domain_control(node_worker_id, ControlName.RELEASE_DOMAIN, command) + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + if errors: + raise RuntimeError(f"Global CommDomain release failed: {errors[0]}") from errors[0] + if self._live_global_domains.get(handle.name) is handle: + self._live_global_domains.pop(handle.name) + + def _execute_pending_global_domain_releases(self, resources: _RunResources) -> None: + pending = list(resources.pending_release_global_domains) + resources.pending_release_global_domains.clear() + for handle in pending: + try: + self._free_global_domain_after_fence(handle) + except Exception as exc: # noqa: BLE001 + sys.stderr.write( + f"Worker._execute_pending_global_domain_releases: {handle.name!r} " + f"release failed: {type(exc).__name__}: {exc}\n" + ) + sys.stderr.flush() + + def _release_all_live_global_domains( + self, + resources: _RunResources | None = None, + *, + include_retained: bool = True, + ) -> None: + live_domains = self._live_global_domains if resources is None else resources.live_global_domains + for handle in list(live_domains.values())[::-1]: + if handle.retain_after_run and not include_retained: + continue + try: + handle._released = True # noqa: SLF001 -- runtime owns this transition + self._free_global_domain_after_fence(handle) + if live_domains.get(handle.name) is handle: + live_domains.pop(handle.name) + except Exception as exc: # noqa: BLE001 + sys.stderr.write( + f"Worker._release_all_live_global_domains: {handle.name!r} " + f"release failed: {type(exc).__name__}: {exc}\n" + ) + sys.stderr.flush() + + def _copy_to_global_domain( + self, handle: GlobalCommDomainHandle, domain_rank: int, data: bytes, offset: int + ) -> None: + from .remote_l3_protocol import ControlName # noqa: PLC0415 + + payload = bytes(data) + member = handle.member(domain_rank) + command = GlobalDomainCopyCommand( + domain_id=handle.domain_id, + generation=handle.generation, + domain_rank=int(domain_rank), + offset=int(offset), + nbytes=len(payload), + data=payload, + ) + self._global_domain_control( + member.node_worker_id, + ControlName.COPY_TO_DOMAIN, + encode_copy_command(command, include_data=True), + ) + + def _copy_from_global_domain( + self, handle: GlobalCommDomainHandle, domain_rank: int, nbytes: int, offset: int + ) -> bytes: + from .remote_l3_protocol import ControlName # noqa: PLC0415 + + member = handle.member(domain_rank) + command = GlobalDomainCopyCommand( + domain_id=handle.domain_id, + generation=handle.generation, + domain_rank=int(domain_rank), + offset=int(offset), + nbytes=int(nbytes), + ) + reply = self._global_domain_control( + member.node_worker_id, + ControlName.COPY_FROM_DOMAIN, + encode_copy_command(command, include_data=False), + ) + return decode_copy_result(reply) + def _release_all_live_domains(self, resources: _RunResources | None = None) -> None: """Best-effort release of every still-live domain handle (LIFO). @@ -5826,7 +7168,7 @@ def _create_host_buffer_locked(self, nbytes: int) -> HostBuffer: # and sub alike, via _broadcast_host_control). Only a truly childless L3 # has nowhere to attach it. if not self._chip_shms and not self._sub_shms: - raise RuntimeError( + raise _NoHostBufferChildrenError( "create_host_buffer requires at least one forked chip or sub child (this Worker has none)" ) assert self._worker is not None @@ -6223,6 +7565,14 @@ def _step(fn) -> None: _step(lambda: self._cleanup_l3_l2_regions(resources)) finally: resources.l3_l2_orch_comm_host_buffers.clear() + _step(lambda: self._execute_pending_global_domain_releases(resources)) + if resources.live_global_domains: + _step( + lambda: self._release_all_live_global_domains( + resources, + include_retained=False, + ) + ) _step(lambda: self._execute_pending_domain_releases(resources)) if resources.live_domains: _step(lambda: self._release_all_live_domains(resources)) @@ -6306,6 +7656,8 @@ def _has_live_resources(self) -> bool: or bool(self._sub_shms or self._chip_shms or self._next_level_shms) or bool(self._live_l3_l2_regions) or bool(self._live_domains) + or bool(self._live_global_domains or self._failed_global_domain_releases) + or bool(self._global_node_domains) or bool(self._host_buf_registry) or bool(self._pending_remote_buffer_frees or self._pending_remote_import_releases) ) @@ -6326,6 +7678,12 @@ def _describe_live_resources(self) -> str: parts.append(f"{len(self._live_l3_l2_regions)} L3-L2 region(s)") if self._live_domains: parts.append(f"{len(self._live_domains)} comm domain(s)") + if self._live_global_domains or self._failed_global_domain_releases: + live_global_ids = {handle.domain_id for handle in self._live_global_domains.values()} + live_global_ids.update(self._failed_global_domain_releases) + parts.append(f"{len(live_global_ids)} global comm domain(s)") + if self._global_node_domains: + parts.append(f"{len(self._global_node_domains)} imported global comm domain(s)") if self._host_buf_registry: parts.append(f"{len(self._host_buf_registry)} host buffer(s)") n_remote = len(self._pending_remote_buffer_frees) + len(self._pending_remote_import_releases) @@ -6666,6 +8024,10 @@ def _step(fn) -> None: # C++ scheduler: once `dw.close()` runs the chip mailboxes are unusable # and we can no longer drive CTRL_RELEASE_DOMAIN. _step(self._cleanup_l3_l2_regions) + if self._live_global_domains: + _step(self._release_all_live_global_domains) + if self._global_node_domains: + _step(self._release_all_global_domain_nodes) if self._live_domains: _step(self._release_all_live_domains) _step(self._clear_child_prov) diff --git a/src/a2a3/platform/onboard/host/comm_hccl.cpp b/src/a2a3/platform/onboard/host/comm_hccl.cpp index f6e2e7d815..9f2b190b35 100644 --- a/src/a2a3/platform/onboard/host/comm_hccl.cpp +++ b/src/a2a3/platform/onboard/host/comm_hccl.cpp @@ -84,11 +84,18 @@ struct DomainAllocation { int rank = 0; int nranks = 0; + bool release_failed = false; VmmWindow local_window; std::vector peer_windows; CommContext *device_ctx = nullptr; // aclrtMalloc'd CommContext mirror }; +static_assert(sizeof(CommGlobalDomainDescriptor) == 288, "global domain descriptor ABI changed"); +static_assert( + sizeof(aclrtMemFabricHandle) <= COMM_GLOBAL_DOMAIN_HANDLE_BYTES, "Fabric handle exceeds global descriptor" +); +static std::unordered_map> global_domain_allocations; + struct CommHandle_ { int rank; int nranks; @@ -1533,6 +1540,190 @@ comm_release_domain_windows(CommHandle h, uint64_t allocation_id, size_t rank_co return -1; } +extern "C" int comm_global_domain_prepare( + uint64_t domain_id, uint32_t domain_rank, uint32_t rank_count, size_t window_size, uint32_t profile, + CommGlobalDomainDescriptor *descriptor_out, uint64_t *local_window_base_out +) try { + if (domain_id == 0 || rank_count == 0 || rank_count > COMM_MAX_RANK_NUM || domain_rank >= rank_count || + window_size == 0 || profile != COMM_GLOBAL_DOMAIN_PROFILE_A3_FABRIC || descriptor_out == nullptr || + local_window_base_out == nullptr || global_domain_allocations.count(domain_id) != 0) { + return -1; + } + + int32_t device_id = -1; + if (aclrtGetDevice(&device_id) != ACL_SUCCESS) { + return -1; + } + auto allocation = std::make_unique(); + aclError status = create_local_fabric_window(device_id, window_size, &allocation->local_window); + if (status != ACL_SUCCESS) { + LOG_ERROR("[global domain rank %u] create local Fabric window -> %d", domain_rank, static_cast(status)); + return -1; + } + + aclrtMemFabricHandle fabric_handle{}; + status = export_fabric_window(allocation->local_window, &fabric_handle); + if (status != ACL_SUCCESS) { + LOG_ERROR("[global domain rank %u] export Fabric handle -> %d", domain_rank, static_cast(status)); + release_domain_windows(allocation.get()); + return -1; + } + status = + aclrtMemset(allocation->local_window.base, allocation->local_window.size, 0, allocation->local_window.size); + if (status != ACL_SUCCESS) { + LOG_ERROR("[global domain rank %u] zero local Fabric window -> %d", domain_rank, static_cast(status)); + release_domain_windows(allocation.get()); + return -1; + } + + allocation->rank = static_cast(domain_rank); + allocation->nranks = static_cast(rank_count); + CommGlobalDomainDescriptor descriptor{}; + descriptor.version = COMM_GLOBAL_DOMAIN_VERSION; + descriptor.profile = COMM_GLOBAL_DOMAIN_PROFILE_A3_FABRIC; + descriptor.domain_rank = domain_rank; + descriptor.rank_count = rank_count; + descriptor.mapping_size = allocation->local_window.size; + descriptor.handle_size = sizeof(fabric_handle); + std::memcpy(descriptor.handle, &fabric_handle, sizeof(fabric_handle)); + + *descriptor_out = descriptor; + *local_window_base_out = reinterpret_cast(allocation->local_window.base); + global_domain_allocations.emplace(domain_id, std::move(allocation)); + return 0; +} catch (const std::exception &e) { + LOG_ERROR("[global domain] prepare exception: %s", e.what()); + return -1; +} catch (...) { + LOG_ERROR("[global domain] prepare unknown exception"); + return -1; +} + +extern "C" int comm_global_domain_import( + uint64_t domain_id, const CommGlobalDomainDescriptor *descriptors, size_t descriptor_count, uint64_t *device_ctx_out +) try { + auto it = global_domain_allocations.find(domain_id); + if (it == global_domain_allocations.end() || descriptors == nullptr || device_ctx_out == nullptr) { + return -1; + } + auto &allocation = it->second; + if (descriptor_count != static_cast(allocation->nranks) || allocation->device_ctx != nullptr) { + return -1; + } + + std::vector rank_order(descriptor_count, nullptr); + for (size_t i = 0; i < descriptor_count; ++i) { + const auto &descriptor = descriptors[i]; + if (descriptor.version != COMM_GLOBAL_DOMAIN_VERSION || + descriptor.profile != COMM_GLOBAL_DOMAIN_PROFILE_A3_FABRIC || + descriptor.rank_count != static_cast(allocation->nranks) || + descriptor.domain_rank >= static_cast(allocation->nranks) || + descriptor.mapping_size != allocation->local_window.size || + descriptor.handle_size != sizeof(aclrtMemFabricHandle) || rank_order[descriptor.domain_rank] != nullptr) { + return -1; + } + rank_order[descriptor.domain_rank] = &descriptor; + } + + int32_t device_id = -1; + if (aclrtGetDevice(&device_id) != ACL_SUCCESS) { + return -1; + } + CommContext ctx{}; + ctx.rankId = static_cast(allocation->rank); + ctx.rankNum = static_cast(allocation->nranks); + ctx.winSize = allocation->local_window.size; + allocation->peer_windows.reserve(descriptor_count - 1); + for (uint32_t rank = 0; rank < static_cast(allocation->nranks); ++rank) { + const auto *descriptor = rank_order[rank]; + if (descriptor == nullptr) { + return -1; + } + uint64_t window_addr = reinterpret_cast(allocation->local_window.base); + if (rank != static_cast(allocation->rank)) { + aclrtMemFabricHandle fabric_handle{}; + std::memcpy(&fabric_handle, descriptor->handle, sizeof(fabric_handle)); + VmmWindow peer_window; + aclError status = import_fabric_window(device_id, fabric_handle, descriptor->mapping_size, &peer_window); + if (status != ACL_SUCCESS) { + LOG_ERROR( + "[global domain rank %d] import peer rank %u -> %d", allocation->rank, rank, + static_cast(status) + ); + return -1; + } + window_addr = reinterpret_cast(peer_window.base); + allocation->peer_windows.push_back(std::move(peer_window)); + } + ctx.windowsIn[rank] = window_addr; + ctx.windowsOut[rank] = window_addr; + } + + void *device_ctx = nullptr; + aclError status = aclrtMalloc(&device_ctx, sizeof(CommContext), ACL_MEM_MALLOC_HUGE_FIRST); + if (status != ACL_SUCCESS) { + return -1; + } + status = aclrtMemcpy(device_ctx, sizeof(CommContext), &ctx, sizeof(CommContext), ACL_MEMCPY_HOST_TO_DEVICE); + if (status != ACL_SUCCESS) { + aclrtFree(device_ctx); + return -1; + } + allocation->device_ctx = static_cast(device_ctx); + *device_ctx_out = reinterpret_cast(device_ctx); + return 0; +} catch (const std::exception &e) { + LOG_ERROR("[global domain] import exception: %s", e.what()); + return -1; +} catch (...) { + LOG_ERROR("[global domain] import unknown exception"); + return -1; +} + +extern "C" int comm_global_domain_release(uint64_t domain_id) try { + auto it = global_domain_allocations.find(domain_id); + if (it == global_domain_allocations.end()) { + return 0; + } + auto &allocation = it->second; + if (allocation->release_failed) { + LOG_ERROR( + "[global domain] domain %llu previously failed a partial release", + static_cast(domain_id) + ); + return -1; + } + int rc = 0; + if (allocation->device_ctx != nullptr) { + if (aclrtFree(allocation->device_ctx) != ACL_SUCCESS) { + rc = -1; + } + allocation->device_ctx = nullptr; + } + if (release_domain_windows(allocation.get()) != ACL_SUCCESS) { + rc = -1; + } + if (rc != 0) { + // The release helpers are destructive: some fields may already be + // cleared before a later ACL operation fails. Keep a terminal failure + // record so a repeated call cannot falsely report success. + allocation->release_failed = true; + LOG_ERROR( + "[global domain] domain %llu entered terminal partial-release state", + static_cast(domain_id) + ); + return rc; + } + global_domain_allocations.erase(it); + return 0; +} catch (const std::exception &e) { + LOG_ERROR("[global domain] release exception: %s", e.what()); + return -1; +} catch (...) { + LOG_ERROR("[global domain] release unknown exception"); + return -1; +} + extern "C" int comm_destroy(CommHandle h) try { if (!h) return -1; diff --git a/src/a5/platform/onboard/host/comm_hccl.cpp b/src/a5/platform/onboard/host/comm_hccl.cpp index 4d8be9a529..e1bb66d146 100644 --- a/src/a5/platform/onboard/host/comm_hccl.cpp +++ b/src/a5/platform/onboard/host/comm_hccl.cpp @@ -1396,6 +1396,22 @@ comm_release_domain_windows(CommHandle h, uint64_t allocation_id, size_t rank_co return -1; } +extern "C" int +comm_global_domain_prepare(uint64_t, uint32_t, uint32_t, size_t, uint32_t, CommGlobalDomainDescriptor *, uint64_t *) { + LOG_ERROR("[comm] Global CommDomain is not supported by the a5 backend"); + return -1; +} + +extern "C" int comm_global_domain_import(uint64_t, const CommGlobalDomainDescriptor *, size_t, uint64_t *) { + LOG_ERROR("[comm] Global CommDomain is not supported by the a5 backend"); + return -1; +} + +extern "C" int comm_global_domain_release(uint64_t) { + LOG_ERROR("[comm] Global CommDomain is not supported by the a5 backend"); + return -1; +} + extern "C" int comm_destroy(CommHandle h) try { if (!h) return -1; diff --git a/src/common/hierarchical/remote_endpoint.cpp b/src/common/hierarchical/remote_endpoint.cpp index fb96e3e6f4..f3d73bb270 100644 --- a/src/common/hierarchical/remote_endpoint.cpp +++ b/src/common/hierarchical/remote_endpoint.cpp @@ -934,6 +934,12 @@ void RemoteL3Endpoint::control_remote_release_import(const RemoteBufferHandle &h run_control(remote_l3::ControlName::RELEASE_IMPORT, remote_l3::encode_release_import_request(request)); } +std::vector RemoteL3Endpoint::control_remote_domain( + remote_l3::ControlName control_name, const std::vector &command_bytes +) { + return run_control(control_name, command_bytes).result_bytes; +} + void RemoteL3Endpoint::shutdown_child() { if (!transport_) return; try { diff --git a/src/common/hierarchical/remote_endpoint.h b/src/common/hierarchical/remote_endpoint.h index 39fc86c2c8..aec470d24c 100644 --- a/src/common/hierarchical/remote_endpoint.h +++ b/src/common/hierarchical/remote_endpoint.h @@ -110,6 +110,8 @@ class RemoteL3Endpoint : public WorkerEndpoint { int32_t importer_worker_id, const RemoteBufferExport &export_desc, uint32_t requested_access_flags ) override; void control_remote_release_import(const RemoteBufferHandle &handle) override; + std::vector + control_remote_domain(remote_l3::ControlName control_name, const std::vector &command_bytes) override; private: WorkerEndpointCaps caps_; diff --git a/src/common/hierarchical/remote_wire.cpp b/src/common/hierarchical/remote_wire.cpp index 8e3dbd820e..29d444bb6f 100644 --- a/src/common/hierarchical/remote_wire.cpp +++ b/src/common/hierarchical/remote_wire.cpp @@ -147,6 +147,8 @@ bool valid_control_name(uint32_t v) { case ControlName::COMM_INIT: case ControlName::ALLOC_DOMAIN: case ControlName::RELEASE_DOMAIN: + case ControlName::COPY_TO_DOMAIN: + case ControlName::COPY_FROM_DOMAIN: return true; } return false; diff --git a/src/common/hierarchical/remote_wire.h b/src/common/hierarchical/remote_wire.h index 720ccedba9..e64795adf3 100644 --- a/src/common/hierarchical/remote_wire.h +++ b/src/common/hierarchical/remote_wire.h @@ -63,6 +63,8 @@ enum class ControlName : uint32_t { COMM_INIT = 13, ALLOC_DOMAIN = 14, RELEASE_DOMAIN = 15, + COPY_TO_DOMAIN = 16, + COPY_FROM_DOMAIN = 17, }; enum class ReadyState : uint32_t { diff --git a/src/common/hierarchical/worker.h b/src/common/hierarchical/worker.h index f8e2e82b09..af53c9b234 100644 --- a/src/common/hierarchical/worker.h +++ b/src/common/hierarchical/worker.h @@ -128,6 +128,11 @@ class Worker { control_digest_only(WorkerType type, int worker_id, uint64_t sub_cmd, const uint8_t *digest, double timeout_s) { return manager_.control_digest_only(type, worker_id, sub_cmd, digest, timeout_s); } + std::vector control_payload( + WorkerType type, int worker_id, uint64_t sub_cmd, const void *payload, size_t payload_size, double timeout_s + ) { + return manager_.control_payload(type, worker_id, sub_cmd, payload, payload_size, timeout_s); + } ControlResult remote_prepare_register( int worker_id, remote_l3::RemoteRegistryTarget target_registry, CallableKind callable_kind, const void *payload, size_t payload_size, const uint8_t *digest @@ -175,6 +180,11 @@ class Worker { return manager_.control_remote_import(importer_worker_id, export_desc, requested_access_flags); } void remote_release_import(const RemoteBufferHandle &handle) { manager_.control_remote_release_import(handle); } + std::vector remote_domain_control( + int worker_id, remote_l3::ControlName control_name, const std::vector &command_bytes + ) { + return manager_.control_remote_domain(worker_id, control_name, command_bytes); + } // Broadcast CTRL_REGISTER / CTRL_UNREGISTER for a ChipCallable digest to // every NEXT_LEVEL child in parallel. `blob_ptr`/`blob_size` describe diff --git a/src/common/hierarchical/worker_manager.cpp b/src/common/hierarchical/worker_manager.cpp index ad2466985b..846d1c3dcb 100644 --- a/src/common/hierarchical/worker_manager.cpp +++ b/src/common/hierarchical/worker_manager.cpp @@ -124,6 +124,9 @@ RemoteBufferHandle WorkerEndpoint::control_remote_import(int32_t, const RemoteBu void WorkerEndpoint::control_remote_release_import(const RemoteBufferHandle &) { throw_unsupported_control("control_remote_release_import"); } +std::vector WorkerEndpoint::control_remote_domain(remote_l3::ControlName, const std::vector &) { + throw_unsupported_control("control_remote_domain"); +} void WorkerEndpoint::control_generic(uint64_t, const char *, size_t, double, const uint8_t *) { throw_unsupported_control("control_generic"); } @@ -945,6 +948,12 @@ void WorkerThread::control_remote_release_import(const RemoteBufferHandle &handl endpoint_->control_remote_release_import(handle); } +std::vector +WorkerThread::control_remote_domain(remote_l3::ControlName control_name, const std::vector &command_bytes) { + if (!endpoint_) throw std::runtime_error("control_remote_domain: null endpoint"); + return endpoint_->control_remote_domain(control_name, command_bytes); +} + void WorkerThread::control_generic( uint64_t sub_cmd, const char *shm_name, size_t payload_size, double timeout_s, const uint8_t *digest ) { @@ -1158,6 +1167,24 @@ ControlResult WorkerManager::control_digest_only( return result; } +std::vector WorkerManager::control_payload( + WorkerType type, int worker_id, uint64_t sub_cmd, const void *payload, size_t payload_size, double timeout_s +) { + if (payload == nullptr || payload_size == 0) { + throw std::runtime_error("control_payload: payload must be non-empty"); + } + WorkerThread *wt = get_worker_by_id(type, worker_id); + if (wt == nullptr) { + throw std::runtime_error("control_payload: invalid worker_id " + std::to_string(worker_id)); + } + std::string shm_name = make_shm_name(); + PosixShmHolder shm(shm_name, payload_size); + std::memcpy(shm.addr(), payload, payload_size); + wt->control_generic(sub_cmd, shm_name.c_str(), payload_size, timeout_s, nullptr); + auto *begin = static_cast(shm.addr()); + return {begin, begin + payload_size}; +} + ControlResult WorkerManager::control_remote_prepare_register( int worker_id, remote_l3::RemoteRegistryTarget target_registry, CallableKind callable_kind, const void *payload, size_t payload_size, const uint8_t *digest @@ -1300,6 +1327,25 @@ void WorkerManager::control_remote_release_import(const RemoteBufferHandle &hand wt->control_remote_release_import(handle); } +std::vector WorkerManager::control_remote_domain( + int worker_id, remote_l3::ControlName control_name, const std::vector &command_bytes +) { + WorkerThread *wt = get_worker_by_id(WorkerType::NEXT_LEVEL, worker_id); + if (wt == nullptr) { + throw std::runtime_error("control_remote_domain: invalid worker_id " + std::to_string(worker_id)); + } + switch (control_name) { + case remote_l3::ControlName::COMM_INIT: + case remote_l3::ControlName::ALLOC_DOMAIN: + case remote_l3::ControlName::RELEASE_DOMAIN: + case remote_l3::ControlName::COPY_TO_DOMAIN: + case remote_l3::ControlName::COPY_FROM_DOMAIN: + return wt->control_remote_domain(control_name, command_bytes); + default: + throw std::runtime_error("control_remote_domain: control name is not a domain operation"); + } +} + std::vector WorkerManager::broadcast_register_all(const void *blob_ptr, size_t blob_size, const uint8_t *digest) { std::vector results; diff --git a/src/common/hierarchical/worker_manager.h b/src/common/hierarchical/worker_manager.h index 6d00a95519..e9eb2cdf44 100644 --- a/src/common/hierarchical/worker_manager.h +++ b/src/common/hierarchical/worker_manager.h @@ -248,6 +248,8 @@ class WorkerEndpoint { int32_t importer_worker_id, const RemoteBufferExport &export_desc, uint32_t requested_access_flags ); virtual void control_remote_release_import(const RemoteBufferHandle &handle); + virtual std::vector + control_remote_domain(remote_l3::ControlName control_name, const std::vector &command_bytes); virtual void control_generic( uint64_t sub_cmd, const char *shm_name, size_t payload_size, double timeout_s, const uint8_t *digest ); @@ -446,6 +448,8 @@ class WorkerThread { int32_t importer_worker_id, const RemoteBufferExport &export_desc, uint32_t requested_access_flags ); void control_remote_release_import(const RemoteBufferHandle &handle); + std::vector + control_remote_domain(remote_l3::ControlName control_name, const std::vector &command_bytes); void control_generic( uint64_t sub_cmd, const char *shm_name, size_t payload_size, double timeout_s, const uint8_t *digest ); @@ -532,6 +536,9 @@ class WorkerManager { void control_l3_l2_region_release(int worker_id, uint64_t region_id); ControlResult control_digest_only(WorkerType type, int worker_id, uint64_t sub_cmd, const uint8_t *digest, double timeout_s); + std::vector control_payload( + WorkerType type, int worker_id, uint64_t sub_cmd, const void *payload, size_t payload_size, double timeout_s + ); ControlResult control_remote_prepare_register( int worker_id, remote_l3::RemoteRegistryTarget target_registry, CallableKind callable_kind, const void *payload, size_t payload_size, const uint8_t *digest @@ -560,6 +567,9 @@ class WorkerManager { int32_t importer_worker_id, const RemoteBufferExport &export_desc, uint32_t requested_access_flags ); void control_remote_release_import(const RemoteBufferHandle &handle); + std::vector control_remote_domain( + int worker_id, remote_l3::ControlName control_name, const std::vector &command_bytes + ); // Broadcast CTRL_REGISTER for `digest` to every NEXT_LEVEL worker in // parallel. Stages `blob_size` bytes from `blob_ptr` into a per-call diff --git a/src/common/platform_comm/comm.h b/src/common/platform_comm/comm.h index b43a60675d..11da9b9114 100644 --- a/src/common/platform_comm/comm.h +++ b/src/common/platform_comm/comm.h @@ -37,6 +37,49 @@ extern "C" { typedef struct CommHandle_ *CommHandle; +#define COMM_GLOBAL_DOMAIN_VERSION 1U +#define COMM_GLOBAL_DOMAIN_HANDLE_BYTES 256U +#define COMM_GLOBAL_DOMAIN_DESCRIPTOR_BYTES 288U + +typedef enum CommGlobalDomainProfile { + COMM_GLOBAL_DOMAIN_PROFILE_SIM_SHM = 1, + COMM_GLOBAL_DOMAIN_PROFILE_A3_FABRIC = 2, +} CommGlobalDomainProfile; + +typedef struct CommGlobalDomainDescriptor { + uint32_t version; + uint32_t profile; + uint32_t domain_rank; + uint32_t rank_count; + uint64_t mapping_size; + uint32_t handle_size; + uint32_t reserved; + uint8_t handle[COMM_GLOBAL_DOMAIN_HANDLE_BYTES]; +} CommGlobalDomainDescriptor; + +#ifdef __cplusplus +static_assert( + sizeof(CommGlobalDomainDescriptor) == COMM_GLOBAL_DOMAIN_DESCRIPTOR_BYTES, + "CommGlobalDomainDescriptor wire size drift" +); +static_assert(offsetof(CommGlobalDomainDescriptor, version) == 0, "CommGlobalDomainDescriptor.version offset drift"); +static_assert(offsetof(CommGlobalDomainDescriptor, profile) == 4, "CommGlobalDomainDescriptor.profile offset drift"); +static_assert( + offsetof(CommGlobalDomainDescriptor, domain_rank) == 8, "CommGlobalDomainDescriptor.domain_rank offset drift" +); +static_assert( + offsetof(CommGlobalDomainDescriptor, rank_count) == 12, "CommGlobalDomainDescriptor.rank_count offset drift" +); +static_assert( + offsetof(CommGlobalDomainDescriptor, mapping_size) == 16, "CommGlobalDomainDescriptor.mapping_size offset drift" +); +static_assert( + offsetof(CommGlobalDomainDescriptor, handle_size) == 24, "CommGlobalDomainDescriptor.handle_size offset drift" +); +static_assert(offsetof(CommGlobalDomainDescriptor, reserved) == 28, "CommGlobalDomainDescriptor.reserved offset drift"); +static_assert(offsetof(CommGlobalDomainDescriptor, handle) == 32, "CommGlobalDomainDescriptor.handle offset drift"); +#endif + /** Bit mask of DmaWorkspaceKind values this platform can provision. */ uint32_t dma_workspace_supported_mask(void); @@ -223,6 +266,42 @@ int comm_alloc_domain_windows( */ int comm_release_domain_windows(CommHandle h, uint64_t allocation_id, size_t rank_count, uint32_t domain_rank); +/** + * Create one local window for a Global CommDomain and export its transport + * descriptor. This operation has no + * HCCL or filesystem bootstrap dependency. + * + * The caller relays the descriptor through its control plane. It later + * passes + * the complete rank-ordered table to comm_global_domain_import(). A live + * domain_id is unique within one + * ChipWorker process. + */ +int comm_global_domain_prepare( + uint64_t domain_id, uint32_t domain_rank, uint32_t rank_count, size_t window_size, uint32_t profile, + CommGlobalDomainDescriptor *descriptor_out, uint64_t *local_window_base_out +); + +/** + * Import every peer descriptor and publish a device CommContext. + * + * descriptors must contain exactly one entry + * for every dense domain rank. + * The returned context and all imported mappings remain live until + * + * comm_global_domain_release(). + */ +int comm_global_domain_import( + uint64_t domain_id, const CommGlobalDomainDescriptor *descriptors, size_t descriptor_count, uint64_t *device_ctx_out +); + +/** + * Release a prepared or imported Global CommDomain. This is local teardown. + * The L4 owner is responsible for + * draining all rank tasks before fanout. + */ +int comm_global_domain_release(uint64_t domain_id); + /** * Synchronize all ranks. * diff --git a/src/common/platform_comm/comm_sim.cpp b/src/common/platform_comm/comm_sim.cpp index c57ebe2229..1456df7369 100644 --- a/src/common/platform_comm/comm_sim.cpp +++ b/src/common/platform_comm/comm_sim.cpp @@ -165,6 +165,38 @@ struct DomainAllocation { std::unique_ptr host_ctx; // device_ctx points here on sim }; +struct GlobalPeerMapping { + void *base = nullptr; + size_t size = 0; +}; + +struct GlobalDomainAllocation { + ~GlobalDomainAllocation() { + for (auto &mapping : peer_mappings) { + if (mapping.base != nullptr) { + munmap(mapping.base, mapping.size); + } + } + if (local_base != nullptr) { + munmap(local_base, mapping_size); + } + if (!shm_name.empty()) { + shm_unlink(shm_name.c_str()); + } + } + + uint32_t rank = 0; + uint32_t nranks = 0; + std::string shm_name; + void *local_base = nullptr; + size_t mapping_size = 0; + std::vector peer_mappings; + std::unique_ptr host_ctx; +}; + +static_assert(sizeof(CommGlobalDomainDescriptor) == 288, "global domain descriptor ABI changed"); +static std::unordered_map> global_domain_allocations; + struct CommHandle_ { int rank; int nranks; @@ -685,6 +717,152 @@ comm_release_domain_windows(CommHandle h, uint64_t allocation_id, size_t rank_co return -1; } +extern "C" int comm_global_domain_prepare( + uint64_t domain_id, uint32_t domain_rank, uint32_t rank_count, size_t window_size, uint32_t profile, + CommGlobalDomainDescriptor *descriptor_out, uint64_t *local_window_base_out +) try { + if (domain_id == 0 || rank_count == 0 || rank_count > COMM_MAX_RANK_NUM || domain_rank >= rank_count || + window_size == 0 || profile != COMM_GLOBAL_DOMAIN_PROFILE_SIM_SHM || descriptor_out == nullptr || + local_window_base_out == nullptr) { + return -1; + } + if (global_domain_allocations.count(domain_id) != 0) { + return -1; + } + + std::string identity = + std::to_string(static_cast(domain_id)) + ":" + std::to_string(domain_rank); + std::string shm_name = make_shm_name(static_cast(getpid()), hash_id(identity.c_str())); + if (shm_name.empty() || shm_name.size() > COMM_GLOBAL_DOMAIN_HANDLE_BYTES) { + return -1; + } + + int fd = shm_open(shm_name.c_str(), O_CREAT | O_EXCL | O_RDWR, 0600); + if (fd < 0) { + return -1; + } + if (ftruncate(fd, static_cast(window_size)) != 0) { + close(fd); + shm_unlink(shm_name.c_str()); + return -1; + } + void *base = mmap(nullptr, window_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0); + close(fd); + if (base == MAP_FAILED) { + shm_unlink(shm_name.c_str()); + return -1; + } + std::memset(base, 0, window_size); + + auto allocation = std::make_unique(); + allocation->rank = domain_rank; + allocation->nranks = rank_count; + allocation->shm_name = shm_name; + allocation->local_base = base; + allocation->mapping_size = window_size; + + CommGlobalDomainDescriptor descriptor{}; + descriptor.version = COMM_GLOBAL_DOMAIN_VERSION; + descriptor.profile = COMM_GLOBAL_DOMAIN_PROFILE_SIM_SHM; + descriptor.domain_rank = domain_rank; + descriptor.rank_count = rank_count; + descriptor.mapping_size = window_size; + descriptor.handle_size = static_cast(shm_name.size()); + std::memcpy(descriptor.handle, shm_name.data(), shm_name.size()); + + *descriptor_out = descriptor; + *local_window_base_out = reinterpret_cast(base); + global_domain_allocations.emplace(domain_id, std::move(allocation)); + return 0; +} catch (const std::exception &e) { + std::fprintf(stderr, "[comm_sim] global_domain_prepare: exception: %s\n", e.what()); + return -1; +} catch (...) { + std::fprintf(stderr, "[comm_sim] global_domain_prepare: unknown exception\n"); + return -1; +} + +extern "C" int comm_global_domain_import( + uint64_t domain_id, const CommGlobalDomainDescriptor *descriptors, size_t descriptor_count, uint64_t *device_ctx_out +) try { + auto it = global_domain_allocations.find(domain_id); + if (it == global_domain_allocations.end() || descriptors == nullptr || device_ctx_out == nullptr) { + return -1; + } + auto &allocation = it->second; + if (descriptor_count != allocation->nranks || allocation->host_ctx != nullptr) { + return -1; + } + + std::vector rank_order(descriptor_count, nullptr); + for (size_t i = 0; i < descriptor_count; ++i) { + const auto &descriptor = descriptors[i]; + if (descriptor.version != COMM_GLOBAL_DOMAIN_VERSION || + descriptor.profile != COMM_GLOBAL_DOMAIN_PROFILE_SIM_SHM || descriptor.rank_count != allocation->nranks || + descriptor.domain_rank >= allocation->nranks || descriptor.mapping_size != allocation->mapping_size || + descriptor.handle_size == 0 || descriptor.handle_size > COMM_GLOBAL_DOMAIN_HANDLE_BYTES || + rank_order[descriptor.domain_rank] != nullptr) { + return -1; + } + rank_order[descriptor.domain_rank] = &descriptor; + } + + auto ctx = std::make_unique(); + ctx->rankId = allocation->rank; + ctx->rankNum = allocation->nranks; + ctx->winSize = allocation->mapping_size; + allocation->peer_mappings.reserve(allocation->nranks - 1); + for (uint32_t rank = 0; rank < allocation->nranks; ++rank) { + const auto *descriptor = rank_order[rank]; + if (descriptor == nullptr) { + return -1; + } + void *base = allocation->local_base; + if (rank != allocation->rank) { + std::string peer_name( + reinterpret_cast(descriptor->handle), static_cast(descriptor->handle_size) + ); + int fd = shm_open(peer_name.c_str(), O_RDWR, 0600); + if (fd < 0) { + return -1; + } + base = mmap(nullptr, allocation->mapping_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0); + close(fd); + if (base == MAP_FAILED) { + return -1; + } + allocation->peer_mappings.push_back(GlobalPeerMapping{base, allocation->mapping_size}); + } + ctx->windowsIn[rank] = reinterpret_cast(base); + ctx->windowsOut[rank] = reinterpret_cast(base); + } + + *device_ctx_out = reinterpret_cast(ctx.get()); + allocation->host_ctx = std::move(ctx); + return 0; +} catch (const std::exception &e) { + std::fprintf(stderr, "[comm_sim] global_domain_import: exception: %s\n", e.what()); + return -1; +} catch (...) { + std::fprintf(stderr, "[comm_sim] global_domain_import: unknown exception\n"); + return -1; +} + +extern "C" int comm_global_domain_release(uint64_t domain_id) try { + auto it = global_domain_allocations.find(domain_id); + if (it == global_domain_allocations.end()) { + return 0; + } + global_domain_allocations.erase(it); + return 0; +} catch (const std::exception &e) { + std::fprintf(stderr, "[comm_sim] global_domain_release: exception: %s\n", e.what()); + return -1; +} catch (...) { + std::fprintf(stderr, "[comm_sim] global_domain_release: unknown exception\n"); + return -1; +} + extern "C" int comm_destroy(CommHandle h) try { if (h == nullptr) return -1; diff --git a/src/common/worker/chip_worker.cpp b/src/common/worker/chip_worker.cpp index 68e48e2ce6..a4948ce9a9 100644 --- a/src/common/worker/chip_worker.cpp +++ b/src/common/worker/chip_worker.cpp @@ -144,6 +144,9 @@ void ChipWorker::init( comm_alloc_domain_windows_fn_ = load_symbol(handle, "comm_alloc_domain_windows"); comm_release_domain_windows_fn_ = load_symbol(handle, "comm_release_domain_windows"); + comm_global_domain_prepare_fn_ = load_symbol(handle, "comm_global_domain_prepare"); + comm_global_domain_import_fn_ = load_symbol(handle, "comm_global_domain_import"); + comm_global_domain_release_fn_ = load_symbol(handle, "comm_global_domain_release"); comm_barrier_fn_ = load_symbol(handle, "comm_barrier"); comm_destroy_fn_ = load_symbol(handle, "comm_destroy"); } catch (...) { @@ -234,6 +237,9 @@ void ChipWorker::init( comm_get_window_size_fn_ = nullptr; comm_alloc_domain_windows_fn_ = nullptr; comm_release_domain_windows_fn_ = nullptr; + comm_global_domain_prepare_fn_ = nullptr; + comm_global_domain_import_fn_ = nullptr; + comm_global_domain_release_fn_ = nullptr; comm_barrier_fn_ = nullptr; comm_destroy_fn_ = nullptr; runtime_buf_.clear(); @@ -276,6 +282,9 @@ void ChipWorker::init( comm_derive_context_fn_ = nullptr; comm_alloc_domain_windows_fn_ = nullptr; comm_release_domain_windows_fn_ = nullptr; + comm_global_domain_prepare_fn_ = nullptr; + comm_global_domain_import_fn_ = nullptr; + comm_global_domain_release_fn_ = nullptr; comm_barrier_fn_ = nullptr; comm_destroy_fn_ = nullptr; runtime_buf_.clear(); @@ -306,6 +315,15 @@ void ChipWorker::init( } void ChipWorker::finalize() { + // Global domains are independent of the legacy communicator sessions. + // Release them while the host runtime and device context are still alive. + if (comm_global_domain_release_fn_ != nullptr) { + for (uint64_t domain_id : global_domain_ids_) { + comm_global_domain_release_fn_(domain_id); + } + } + global_domain_ids_.clear(); + // Defensive: if the user never called comm_destroy, reclaim all owned // communicator handles and streams before tearing down the device context. clear_comm_sessions(); @@ -347,6 +365,9 @@ void ChipWorker::finalize() { comm_derive_context_fn_ = nullptr; comm_alloc_domain_windows_fn_ = nullptr; comm_release_domain_windows_fn_ = nullptr; + comm_global_domain_prepare_fn_ = nullptr; + comm_global_domain_import_fn_ = nullptr; + comm_global_domain_release_fn_ = nullptr; comm_barrier_fn_ = nullptr; comm_destroy_fn_ = nullptr; runtime_buf_.clear(); @@ -710,6 +731,67 @@ void ChipWorker::comm_release_domain_windows( } } +std::tuple, uint64_t, size_t> ChipWorker::comm_global_domain_prepare( + uint64_t domain_id, uint32_t domain_rank, uint32_t rank_count, size_t window_size, uint32_t profile +) { + if (comm_global_domain_prepare_fn_ == nullptr) { + throw std::runtime_error("comm_global_domain_prepare is not supported by this runtime"); + } + auto [tracked, inserted] = global_domain_ids_.insert(domain_id); + if (!inserted) { + throw std::runtime_error("comm_global_domain_prepare received a duplicate domain_id"); + } + CommGlobalDomainDescriptor descriptor{}; + uint64_t local_window_base = 0; + int rc = comm_global_domain_prepare_fn_( + domain_id, domain_rank, rank_count, window_size, profile, &descriptor, &local_window_base + ); + if (rc != 0) { + global_domain_ids_.erase(tracked); + throw std::runtime_error("comm_global_domain_prepare failed with code " + std::to_string(rc)); + } + if (local_window_base == 0 || descriptor.mapping_size == 0) { + 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"); + } + const auto *begin = reinterpret_cast(&descriptor); + std::vector descriptor_bytes(begin, begin + sizeof(descriptor)); + return {std::move(descriptor_bytes), local_window_base, static_cast(descriptor.mapping_size)}; +} + +uint64_t ChipWorker::comm_global_domain_import(uint64_t domain_id, const std::vector &descriptors) { + if (comm_global_domain_import_fn_ == nullptr) { + throw std::runtime_error("comm_global_domain_import is not supported by this runtime"); + } + if (descriptors.empty() || descriptors.size() % sizeof(CommGlobalDomainDescriptor) != 0) { + throw std::runtime_error("comm_global_domain_import descriptor table size is invalid"); + } + uint64_t device_ctx = 0; + int rc = comm_global_domain_import_fn_( + domain_id, reinterpret_cast(descriptors.data()), + descriptors.size() / sizeof(CommGlobalDomainDescriptor), &device_ctx + ); + if (rc != 0) { + throw std::runtime_error("comm_global_domain_import failed with code " + std::to_string(rc)); + } + if (device_ctx == 0) { + throw std::runtime_error("comm_global_domain_import returned a null device context"); + } + return device_ctx; +} + +void ChipWorker::comm_global_domain_release(uint64_t domain_id) { + if (comm_global_domain_release_fn_ == nullptr) { + throw std::runtime_error("comm_global_domain_release is not supported by this runtime"); + } + int rc = comm_global_domain_release_fn_(domain_id); + if (rc != 0) { + throw std::runtime_error("comm_global_domain_release failed with code " + std::to_string(rc)); + } + global_domain_ids_.erase(domain_id); +} + void ChipWorker::comm_barrier(uint64_t comm_handle) { int rc = comm_barrier_fn_(reinterpret_cast(comm_handle)); if (rc != 0) { diff --git a/src/common/worker/chip_worker.h b/src/common/worker/chip_worker.h index 7e0dbb84fc..65835f7271 100644 --- a/src/common/worker/chip_worker.h +++ b/src/common/worker/chip_worker.h @@ -14,9 +14,12 @@ #include #include +#include #include +#include #include +#include "../platform_comm/comm.h" #include "../task_interface/call_config.h" #include "../task_interface/task_args.h" #include "pto_runtime_c_api.h" @@ -145,6 +148,11 @@ class ChipWorker { /// inside the backend's per-allocation record). void comm_release_domain_windows(uint64_t comm_handle, uint64_t allocation_id, size_t rank_count, uint32_t domain_rank); + std::tuple, uint64_t, size_t> comm_global_domain_prepare( + uint64_t domain_id, uint32_t domain_rank, uint32_t rank_count, size_t window_size, uint32_t profile + ); + uint64_t comm_global_domain_import(uint64_t domain_id, const std::vector &descriptors); + void comm_global_domain_release(uint64_t domain_id); void comm_barrier(uint64_t comm_handle); void comm_destroy(uint64_t comm_handle); void comm_destroy_all(); @@ -186,6 +194,10 @@ class ChipWorker { using CommAllocDomainWindowsFn = int (*)(void *, uint64_t, const uint32_t *, size_t, uint32_t, size_t, uint64_t *, uint64_t *); using CommReleaseDomainWindowsFn = int (*)(void *, uint64_t, size_t, uint32_t); + using CommGlobalDomainPrepareFn = + int (*)(uint64_t, uint32_t, uint32_t, size_t, uint32_t, CommGlobalDomainDescriptor *, uint64_t *); + using CommGlobalDomainImportFn = int (*)(uint64_t, const CommGlobalDomainDescriptor *, size_t, uint64_t *); + using CommGlobalDomainReleaseFn = int (*)(uint64_t); using CommBarrierFn = int (*)(void *); using CommDestroyFn = int (*)(void *); @@ -234,11 +246,15 @@ class ChipWorker { CommDeriveContextFn comm_derive_context_fn_ = nullptr; CommAllocDomainWindowsFn comm_alloc_domain_windows_fn_ = nullptr; CommReleaseDomainWindowsFn comm_release_domain_windows_fn_ = nullptr; + CommGlobalDomainPrepareFn comm_global_domain_prepare_fn_ = nullptr; + CommGlobalDomainImportFn comm_global_domain_import_fn_ = nullptr; + CommGlobalDomainReleaseFn comm_global_domain_release_fn_ = nullptr; CommBarrierFn comm_barrier_fn_ = nullptr; CommDestroyFn comm_destroy_fn_ = nullptr; void *device_ctx_ = nullptr; std::vector comm_sessions_; std::unordered_map comm_session_index_; + std::unordered_set global_domain_ids_; uint64_t base_comm_handle_ = 0; std::vector runtime_buf_; diff --git a/tests/ut/cpp/CMakeLists.txt b/tests/ut/cpp/CMakeLists.txt index 2e8cf1fcf7..63cb058955 100644 --- a/tests/ut/cpp/CMakeLists.txt +++ b/tests/ut/cpp/CMakeLists.txt @@ -197,6 +197,7 @@ add_library(hierarchical_objs OBJECT ) target_include_directories(hierarchical_objs PUBLIC ${HIERARCHICAL_SRC_DIR} + ${CMAKE_SOURCE_DIR}/../../../src/common/platform/include ${CMAKE_SOURCE_DIR}/../../../src/common/task_interface ${WORKER_SRC_DIR} ) diff --git a/tests/ut/py/test_callable_identity.py b/tests/ut/py/test_callable_identity.py index e8ce434496..1e72e0f682 100644 --- a/tests/ut/py/test_callable_identity.py +++ b/tests/ut/py/test_callable_identity.py @@ -42,10 +42,14 @@ encode_register_callable_command, encode_remote_chip_callable_payload, ) +from simpler.remote_l3_protocol import ( + RemoteAddressSpace as WireRemoteAddressSpace, +) from simpler.remote_l3_session import ( _install_manifest_dispatcher_registry, _install_manifest_inner_registry, _prepare_register_callable, + _RemoteBufferEntry, _unpublish_inner_handle, get_inner_handle, ) @@ -115,6 +119,42 @@ def _remote_sum_u8_orch(orch, args, cfg): dst_data[0] = sum(int(src_data[i]) for i in range(src.nbytes())) & 0xFF +def test_remote_buffer_entry_releases_child_visible_host_buffer(): + raw = bytearray(8) + view = memoryview(raw) + + class FakeHostBuffer: + data_ptr = ctypes.addressof(ctypes.c_char.from_buffer(raw)) + buffer = view + + class FakeOwner: + def __init__(self): + self.freed = [] + + def free_host_buffer(self, handle): + self.freed.append(handle) + + data = FakeHostBuffer() + owner = FakeOwner() + entry = _RemoteBufferEntry( + data=data, + nbytes=len(raw), + generation=1, + address_space=WireRemoteAddressSpace.REMOTE_DEVICE, + owner=cast(Worker, owner), + ) + + assert entry.addr == data.data_ptr + entry.close() + assert owner.freed == [data] + assert entry.owner is None + with pytest.raises(ValueError, match="released memoryview"): + _ = view[0] + + # Session finalization may defensively revisit an already closed entry. + entry.close() + + class _FakeRemoteControlResult: def __init__(self, worker_id: int, ok: bool = True, error_message: str = ""): self.worker_type = "NEXT_LEVEL" @@ -628,6 +668,44 @@ def test_remote_session_manifest_uses_endpoint_host_as_default_bind(): worker.close() +def test_remote_manifest_carries_pre_registered_inner_chip_callable(): + worker = Worker(level=4, num_sub_workers=0) + chip = ChipCallable.build(signature=[], func_name="x", binary=b"\x01", children=[]) + try: + worker_id = worker.add_remote_worker( + RemoteWorkerSpec( + endpoint="127.0.0.1:19073", + platform="a2a3sim", + device_ids=(0,), + ) + ) + handle = worker.register(chip) + manifest = worker._build_remote_manifest( + spec=worker._remote_worker_specs[0], + worker_id=worker_id, + session_id=1, + startup_remaining_s=30.0, + ) + + assert len(manifest["inner_l3_worker"]) == 1 + entry = manifest["inner_l3_worker"][0] + assert entry["hashid"] == handle.digest.hex() + command = encode_register_callable_command( + RemoteRegistryTarget.INNER_L3_WORKER, + CallableKind.CHIP_CALLABLE, + handle.digest, + 1, + bytes.fromhex(entry["payload_hex"]), + ) + digest, kind, registry, target = _prepare_register_callable(command, manifest) + assert digest == handle.digest + assert kind is CallableKind.CHIP_CALLABLE + assert registry is RemoteRegistryTarget.INNER_L3_WORKER + assert isinstance(target, ChipCallable) + finally: + worker.close() + + def test_remote_session_manifest_requires_wildcard_bind_opt_in(): worker = Worker(level=4, num_sub_workers=0) try: diff --git a/tests/ut/py/test_global_comm_domain.py b/tests/ut/py/test_global_comm_domain.py new file mode 100644 index 0000000000..d89a73a467 --- /dev/null +++ b/tests/ut/py/test_global_comm_domain.py @@ -0,0 +1,709 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +from __future__ import annotations + +import os +import socket +import subprocess +import sys +import time + +import pytest +from simpler.global_comm_domain import ( + GLOBAL_DOMAIN_DESCRIPTOR_BYTES, + GLOBAL_DOMAIN_PROFILE_IDS, + GLOBAL_DOMAIN_VERSION, + GlobalCommInitCommand, + GlobalDomainBuffer, + GlobalDomainCommand, + GlobalDomainDescriptor, + GlobalDomainMember, + GlobalDomainPhase, + decode_comm_init, + decode_descriptor_table, + decode_domain_command, + encode_comm_init, + encode_comm_init_result, + encode_descriptor_table, + encode_domain_command, + resolve_global_comm_capability, + validate_descriptor_table, +) + + +def _members() -> tuple[GlobalDomainMember, ...]: + return ( + GlobalDomainMember(0, 0, 3, 0), + GlobalDomainMember(1, 0, 7, 1), + ) + + +def _descriptors() -> tuple[GlobalDomainDescriptor, ...]: + return tuple( + GlobalDomainDescriptor( + version=GLOBAL_DOMAIN_VERSION, + profile_id=GLOBAL_DOMAIN_PROFILE_IDS["sim"], + domain_rank=rank, + rank_count=2, + mapping_size=4096, + handle=f"/simpler-test-{rank}".encode(), + ) + for rank in range(2) + ) + + +@pytest.mark.parametrize( + ("platform", "profile"), + ( + ("a2a3sim", "sim"), + ("a5sim", "sim"), + ("a2a3", "a3-fabric-v1"), + ), +) +def test_global_comm_capability_reports_only_implemented_backends(platform, profile): + result = resolve_global_comm_capability(platform=platform, profile=profile, local_device_count=2) + + assert result.profile == profile + assert result.max_ranks == 64 + assert result.descriptor_bytes == GLOBAL_DOMAIN_DESCRIPTOR_BYTES + assert result.local_device_count == 2 + + +@pytest.mark.parametrize( + ("platform", "profile"), + ( + ("a2a3", "sim"), + ("a2a3sim", "a3-fabric-v1"), + ("a5", "sim"), + ("a5", "a3-fabric-v1"), + ), +) +def test_global_comm_capability_rejects_unimplemented_backends(platform, profile): + with pytest.raises(ValueError, match="Global CommDomain is not supported"): + resolve_global_comm_capability(platform=platform, profile=profile, local_device_count=2) + + +def test_local_l3_comm_init_rejects_unsupported_capability_without_caching_topology(): + from simpler.remote_l3_protocol import ControlName # noqa: PLC0415 + from simpler.worker import Worker, _GlobalNodeRuntime, _run_local_global_domain_control # noqa: PLC0415 + + inner_worker = Worker(level=3, num_sub_workers=0) + runtime = _GlobalNodeRuntime( + worker_id=0, + device_ids=(0,), + platform="a5", + comm_profile="sim", + global_device_ranks=(0,), + node_rank=0, + node_count=1, + cluster_id="cluster", + is_remote=False, + ) + comm_inits = {} + command = GlobalCommInitCommand( + cluster_id="cluster", + topology_hash="topology", + profile="sim", + node_rank=0, + node_count=1, + members=(GlobalDomainMember(0, 0, 0, 0),), + ) + + try: + with pytest.raises(ValueError, match="Global CommDomain is not supported"): + _run_local_global_domain_control( + inner_worker, + runtime, + comm_inits, + ControlName.COMM_INIT, + encode_comm_init(command), + ) + + assert comm_inits == {} + finally: + inner_worker.close() + + +def test_global_domain_wire_round_trips_topology_and_descriptor_table(): + init = GlobalCommInitCommand("cluster", "topology", "sim", 0, 2, _members()) + command = GlobalDomainCommand( + phase=GlobalDomainPhase.IMPORT, + domain_id=11, + generation=1, + name="tp", + profile="sim", + window_size=2048, + members=_members(), + buffers=(GlobalDomainBuffer("payload", 128),), + descriptors=_descriptors(), + ) + + assert decode_comm_init(encode_comm_init(init)) == init + assert decode_domain_command(encode_domain_command(command)) == command + assert decode_descriptor_table(encode_descriptor_table(_descriptors())) == _descriptors() + assert GLOBAL_DOMAIN_DESCRIPTOR_BYTES == 288 + + +def _failure_injection_worker(*, platform: str = "a2a3sim", profile: str = "sim"): + from simpler.worker import RemoteWorkerSpec, Worker, _RunResources # noqa: PLC0415 + + worker = Worker(level=4, num_sub_workers=0) + node_ids = tuple( + worker.add_remote_worker( + RemoteWorkerSpec( + endpoint=f"127.0.0.1:{19073 + index}", + platform=platform, + device_ids=(0,), + comm_profile=profile, + global_device_ranks=(index,), + ) + ) + for index in range(2) + ) + resources = _RunResources() + worker._worker = object() + worker._building_run_resources = resources + return worker, resources, node_ids + + +def _install_global_domain_failure_injector(monkeypatch, worker, *, fail_phase, fail_node): + from simpler.remote_l3_protocol import ControlName # noqa: PLC0415 + + calls = [] + + def control(worker_id, control_name, payload): + control_name = ControlName(control_name) + if control_name is ControlName.COMM_INIT: + init = decode_comm_init(payload) + calls.append(("COMM_INIT", worker_id)) + return encode_comm_init_result( + resolve_global_comm_capability( + platform="a2a3sim", + profile=init.profile, + local_device_count=1, + ) + ) + + assert control_name is ControlName.ALLOC_DOMAIN + command = decode_domain_command(payload) + calls.append((command.phase, worker_id)) + if command.phase is fail_phase and worker_id == fail_node: + raise RuntimeError(f"injected {command.phase.name} failure") + if command.phase is not GlobalDomainPhase.PREPARE_EXPORT: + return b"" + descriptors = tuple( + GlobalDomainDescriptor( + version=GLOBAL_DOMAIN_VERSION, + profile_id=GLOBAL_DOMAIN_PROFILE_IDS[command.profile], + domain_rank=member.domain_rank, + rank_count=len(command.members), + mapping_size=4096, + handle=f"/injected-{member.domain_rank}".encode(), + ) + for member in command.members + if member.node_worker_id == worker_id + ) + return encode_descriptor_table(descriptors) + + monkeypatch.setattr(worker, "_global_domain_control", control) + return calls + + +def _close_failure_injection_worker(worker, resources): + worker._building_run_resources = None + worker._live_global_domains.clear() + resources.live_global_domains.clear() + worker._worker = None + worker.close() + + +@pytest.mark.parametrize( + "fail_phase", + ( + GlobalDomainPhase.PREPARE_EXPORT, + GlobalDomainPhase.IMPORT, + GlobalDomainPhase.COMMIT, + ), +) +def test_global_domain_transaction_aborts_all_prepared_nodes_after_phase_failure(monkeypatch, fail_phase): + from simpler.task_interface import CommBufferSpec # noqa: PLC0415 + + worker, resources, node_ids = _failure_injection_worker() + calls = _install_global_domain_failure_injector( + monkeypatch, + worker, + fail_phase=fail_phase, + fail_node=node_ids[1], + ) + try: + with pytest.raises(RuntimeError, match=f"injected {fail_phase.name} failure"): + worker._allocate_global_domain( + name="failure-injection", + members=((node_ids[0], 0), (node_ids[1], 0)), + window_size=4096, + buffers=[CommBufferSpec("payload", "uint8", 4096, 4096)], + retain_after_run=False, + ) + + abort_nodes = [node_id for phase, node_id in calls if phase is GlobalDomainPhase.ABORT] + assert abort_nodes == list(node_ids) + assert worker._live_global_domains == {} + assert resources.live_global_domains == {} + finally: + _close_failure_injection_worker(worker, resources) + + +def test_global_domain_abort_failure_preserves_primary_error_and_continues_cleanup(monkeypatch): + from simpler.remote_l3_protocol import ControlName # noqa: PLC0415 + from simpler.task_interface import CommBufferSpec # noqa: PLC0415 + + worker, resources, node_ids = _failure_injection_worker() + calls = _install_global_domain_failure_injector( + monkeypatch, + worker, + fail_phase=GlobalDomainPhase.IMPORT, + fail_node=node_ids[1], + ) + original_control = worker._global_domain_control + + def fail_first_abort(worker_id, control_name, payload): + if ControlName(control_name) is ControlName.ALLOC_DOMAIN: + command = decode_domain_command(payload) + if command.phase is GlobalDomainPhase.ABORT and worker_id == node_ids[0]: + calls.append((command.phase, worker_id)) + raise RuntimeError("injected ABORT failure") + return original_control(worker_id, control_name, payload) + + monkeypatch.setattr(worker, "_global_domain_control", fail_first_abort) + try: + with pytest.raises(RuntimeError, match="injected IMPORT failure"): + worker._allocate_global_domain( + name="abort-failure-injection", + members=((node_ids[0], 0), (node_ids[1], 0)), + window_size=4096, + buffers=[CommBufferSpec("payload", "uint8", 4096, 4096)], + retain_after_run=False, + ) + + abort_nodes = [node_id for phase, node_id in calls if phase is GlobalDomainPhase.ABORT] + assert abort_nodes == list(node_ids) + assert worker._live_global_domains == {} + assert resources.live_global_domains == {} + finally: + _close_failure_injection_worker(worker, resources) + + +def test_allocate_global_domain_rejects_unsupported_capability_before_control(monkeypatch): + from simpler.task_interface import CommBufferSpec # noqa: PLC0415 + + worker, resources, node_ids = _failure_injection_worker(platform="a5", profile="sim") + calls = [] + monkeypatch.setattr(worker, "_global_domain_control", lambda *args: calls.append(args)) + try: + with pytest.raises(ValueError, match="Global CommDomain is not supported"): + worker._allocate_global_domain( + name="unsupported", + members=((node_ids[0], 0), (node_ids[1], 0)), + window_size=4096, + buffers=[CommBufferSpec("payload", "uint8", 4096, 4096)], + retain_after_run=False, + ) + + assert calls == [] + assert worker._live_global_domains == {} + assert resources.live_global_domains == {} + finally: + _close_failure_injection_worker(worker, resources) + + +def test_global_domain_descriptor_table_rejects_different_mapping_sizes(): + descriptors = list(_descriptors()) + descriptors[1] = GlobalDomainDescriptor( + version=GLOBAL_DOMAIN_VERSION, + profile_id=GLOBAL_DOMAIN_PROFILE_IDS["sim"], + domain_rank=1, + rank_count=2, + mapping_size=8192, + handle=b"/simpler-test-1", + ) + + with pytest.raises(ValueError, match="mapping sizes differ"): + validate_descriptor_table(tuple(descriptors), rank_count=2, profile="sim") + + +def test_global_domain_release_retries_after_callback_failure(): + from simpler.task_interface import GlobalCommDomainHandle # noqa: PLC0415 + + attempts = 0 + + def release_fn(_handle): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise RuntimeError("transient release failure") + + handle = GlobalCommDomainHandle( + name="retry", + members=(), + buffers=(), + domain_id=17, + generation=1, + mapping_size=4096, + retain_after_run=False, + _release_fn=release_fn, + ) + + with pytest.raises(RuntimeError, match="transient release failure"): + handle.release() + + assert not handle.released + handle.release() + assert handle.released + assert attempts == 2 + + +def test_old_global_domain_release_does_not_remove_same_name_replacement(): + from simpler.task_interface import GlobalCommDomainHandle # noqa: PLC0415 + from simpler.worker import Worker, _RunResources # noqa: PLC0415 + + worker = Worker(level=4, num_sub_workers=0) + resources = _RunResources() + + def make_handle(domain_id: int) -> GlobalCommDomainHandle: + return GlobalCommDomainHandle( + name="reuse", + members=(), + buffers=(), + domain_id=domain_id, + generation=1, + mapping_size=4096, + retain_after_run=False, + _release_fn=worker._release_global_domain_handle, + ) + + first = make_handle(17) + second = make_handle(18) + worker._worker = object() + worker._building_run_resources = resources + worker._live_global_domains[first.name] = first + resources.live_global_domains[first.name] = first + try: + first.release() + worker._live_global_domains[second.name] = second + resources.live_global_domains[second.name] = second + + worker._execute_pending_global_domain_releases(resources) + + assert first.freed + assert worker._live_global_domains[second.name] is second + assert resources.live_global_domains[second.name] is second + finally: + worker._building_run_resources = None + worker._live_global_domains.clear() + resources.live_global_domains.clear() + worker._worker = None + worker.close() + + +def test_global_domain_backend_release_failure_is_terminal(monkeypatch): + from simpler.task_interface import GlobalCommDomainHandle # noqa: PLC0415 + from simpler.worker import Worker # noqa: PLC0415 + + attempts = 0 + worker = Worker(level=4, num_sub_workers=0) + worker._worker = object() + + def fail_control(_worker_id, _control_name, _payload): + nonlocal attempts + attempts += 1 + raise RuntimeError("partial backend release") + + monkeypatch.setattr(worker, "_global_domain_control", fail_control) + handle = GlobalCommDomainHandle( + name="terminal", + members=(_members()[0],), + buffers=(), + domain_id=19, + generation=1, + mapping_size=4096, + retain_after_run=False, + _release_fn=worker._release_global_domain_handle, + ) + try: + with pytest.raises(RuntimeError, match="partial backend release"): + worker._free_global_domain_after_fence(handle) + with pytest.raises(RuntimeError, match="partial backend release"): + worker._free_global_domain_after_fence(handle) + + assert attempts == 1 + assert not handle.freed + assert worker._failed_global_domain_releases[handle.domain_id] is handle + finally: + worker._failed_global_domain_releases.clear() + worker._worker = None + worker.close() + + +def test_childless_host_buffer_uses_dedicated_exception(): + from simpler.worker import Worker, _NoHostBufferChildrenError # noqa: PLC0415 + + worker = Worker(level=3, num_sub_workers=0) + try: + with pytest.raises(_NoHostBufferChildrenError, match="at least one forked chip or sub child"): + worker._create_host_buffer_locked(64) + finally: + worker.close() + + +def test_remote_compute_orch_submits_local_l2_add(monkeypatch): + from simpler.global_comm_smoke import remote_compute_orch # noqa: PLC0415 + from simpler.remote_l3_session import get_inner_handle # noqa: PLC0415 + from simpler.task_interface import CallConfig, TaskArgs, TensorArgType # noqa: PLC0415 + + digest = bytes(range(32)) + chip_handle = object() + monkeypatch.setattr( + "simpler.remote_l3_session.get_inner_handle", + lambda hashid: chip_handle if hashid == digest.hex() else get_inner_handle(hashid), + ) + + class FakeContext: + buffer_ptrs = {"lhs": 0x1000, "rhs": 0x2000, "input": 0x3000} + + class FakeOrchestrator: + def __init__(self): + self.submitted = None + + def get_global_domain(self, domain_id): + assert domain_id == 17 + return {0: FakeContext()} + + def submit_next_level(self, handle, task_args, cfg, *, worker): + self.submitted = (handle, task_args, cfg, worker) + + args = TaskArgs() + args.add_scalar(17) + args.add_scalar(0) + for offset in range(0, 32, 8): + args.add_scalar(int.from_bytes(digest[offset : offset + 8], "little")) + config = CallConfig() + orch = FakeOrchestrator() + + remote_compute_orch(orch, args, config) # type: ignore[arg-type] + + assert orch.submitted is not None + handle, chip_args, submitted_config, worker_id = orch.submitted + assert handle is chip_handle + assert submitted_config is config + assert worker_id == 0 + assert chip_args.tensor_count() == 3 + assert [chip_args.tensor(index).data for index in range(3)] == [0x1000, 0x2000, 0x3000] + assert [chip_args.tensor(index).child_memory for index in range(3)] == [True, True, True] + assert [chip_args.tag(index) for index in range(3)] == [ + TensorArgType.INPUT, + TensorArgType.INPUT, + TensorArgType.OUTPUT_EXISTING, + ] + + +def _free_tcp_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _wait_for_tcp_ports(ports: tuple[int, ...], timeout_s: float = 5.0) -> None: + pending = set(ports) + deadline = time.monotonic() + timeout_s + while pending: + for port in tuple(pending): + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"remote L3 daemons did not become ready on ports {sorted(pending)}") + try: + with socket.create_connection(("127.0.0.1", port), timeout=min(0.1, remaining)): + pending.remove(port) + except OSError: + pass + if pending: + time.sleep(0.01) + + +@pytest.mark.skipif(os.name == "nt", reason="hierarchical workers require fork") +def test_two_remote_daemons_build_and_copy_global_domain_without_mpirun(): + from simpler.task_interface import CommBufferSpec # noqa: PLC0415 + from simpler.worker import RemoteWorkerSpec, Worker # noqa: PLC0415 + + ports = (_free_tcp_port(), _free_tcp_port()) + daemons = [ + subprocess.Popen( + [ + sys.executable, + "-m", + "simpler.remote_l3_worker", + "--host", + "127.0.0.1", + "--port", + str(port), + ], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + for port in ports + ] + worker = Worker(level=4, num_sub_workers=0, remote_session_timeout_s=20) + captured: dict[str, object] = {} + try: + _wait_for_tcp_ports(ports) + node_ids = tuple( + worker.add_remote_worker( + RemoteWorkerSpec( + endpoint=f"127.0.0.1:{port}", + platform="a2a3sim", + device_ids=(0,), + comm_profile="sim", + ) + ) + for port in ports + ) + worker.init() + + def parent_orch(orch, _args, _cfg): + domain = orch.allocate_global_domain( + name="tcp-global", + members=((node_ids[0], 0), (node_ids[1], 0)), + window_size=4096, + buffers=(CommBufferSpec("payload", "uint8", 64, 64),), + retain_after_run=True, + ) + orch.copy_to_global_domain(domain, 0, b"node-zero", buffer="payload") + orch.copy_to_global_domain(domain, 1, b"node-one", buffer="payload") + captured["ranks"] = tuple(member.global_device_rank for member in domain.members) + captured["handle"] = domain + + worker.run(parent_orch) + assert not captured["handle"].freed + + def read_orch(orch, _args, _cfg): + domain = captured["handle"] + try: + captured["rank0"] = orch.copy_from_global_domain(domain, 0, len(b"node-zero"), buffer="payload") + captured["rank1"] = orch.copy_from_global_domain(domain, 1, len(b"node-one"), buffer="payload") + finally: + domain.release() + + worker.run(read_orch) + assert captured["rank0"] == b"node-zero" + assert captured["rank1"] == b"node-one" + assert captured["ranks"] == (0, 1) + assert captured["handle"].freed + finally: + worker.close() + for daemon in daemons: + daemon.terminate() + try: + daemon.wait(timeout=5) + except subprocess.TimeoutExpired: + daemon.kill() + daemon.wait(timeout=5) + + +@pytest.mark.skipif(os.name == "nt", reason="hierarchical workers require fork") +def test_local_and_remote_l3_build_and_copy_global_domain_without_mpirun(): + from simpler.task_interface import CommBufferSpec # noqa: PLC0415 + from simpler.worker import RemoteWorkerSpec, Worker # noqa: PLC0415 + + port = _free_tcp_port() + daemon = subprocess.Popen( + [ + sys.executable, + "-m", + "simpler.remote_l3_worker", + "--host", + "127.0.0.1", + "--port", + str(port), + ], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + worker = Worker(level=4, num_sub_workers=0, remote_session_timeout_s=20) + captured: dict[str, object] = {} + try: + _wait_for_tcp_ports((port,)) + local_node_id = worker.add_worker( + Worker( + level=3, + device_ids=[0], + num_sub_workers=0, + platform="a2a3sim", + runtime="tensormap_and_ringbuffer", + comm_profile="sim", + global_device_ranks=(0,), + ) + ) + remote_node_id = worker.add_remote_worker( + RemoteWorkerSpec( + endpoint=f"127.0.0.1:{port}", + platform="a2a3sim", + device_ids=(0,), + comm_profile="sim", + global_device_ranks=(1,), + ) + ) + worker.init() + + def build_orch(orch, _args, _cfg): + domain = orch.allocate_global_domain( + name="mixed-global", + members=((local_node_id, 0), (remote_node_id, 0)), + window_size=4096, + buffers=(CommBufferSpec("payload", "uint8", 64, 64),), + retain_after_run=True, + ) + orch.copy_to_global_domain(domain, 0, b"local-l3", buffer="payload") + orch.copy_to_global_domain(domain, 1, b"remote-l3", buffer="payload") + captured["ranks"] = tuple(member.global_device_rank for member in domain.members) + captured["domain"] = domain + + worker.run(build_orch) + domain = captured["domain"] + assert not domain.freed + + def read_orch(orch, _args, _cfg): + try: + captured["local"] = orch.copy_from_global_domain( + domain, + 0, + len(b"local-l3"), + buffer="payload", + ) + captured["remote"] = orch.copy_from_global_domain( + domain, + 1, + len(b"remote-l3"), + buffer="payload", + ) + finally: + domain.release() + + worker.run(read_orch) + assert captured["local"] == b"local-l3" + assert captured["remote"] == b"remote-l3" + assert captured["ranks"] == (0, 1) + assert domain.freed + finally: + worker.close() + daemon.terminate() + try: + daemon.wait(timeout=5) + except subprocess.TimeoutExpired: + daemon.kill() + daemon.wait(timeout=5) diff --git a/tools/a3_l4_tcp_smoke/README.md b/tools/a3_l4_tcp_smoke/README.md new file mode 100644 index 0000000000..6d3d3d9601 --- /dev/null +++ b/tools/a3_l4_tcp_smoke/README.md @@ -0,0 +1,80 @@ +# A3 L4 TCP Global CommDomain smoke + +This smoke validates the complete no-`mpirun` path: + +1. L4 connects to the pre-started TCP daemon on each server. +2. Each daemon starts its local L3, and L3 starts the selected L2 worker. +3. L2 exports one Fabric V2 key; L3 returns it to L4. +4. L4 builds the dense rank-ordered key table and sends it to every L3. +5. Every L3 asks its L2 to import the complete table and create `CommContext`. +6. The AIV kernel on every rank executes peer `TLOAD`, sums all rank inputs, + and writes the result to its local Global CommDomain window. + +Start one daemon on each server from the same source revision: + +```bash +python -m simpler.remote_l3_worker --host 10.0.0.1 --port 19073 +python -m simpler.remote_l3_worker --host 10.0.0.2 --port 19073 +``` + +Run the L4 driver on the control host: + +```bash +python tools/a3_l4_tcp_smoke/global_tload_smoke.py \ + --endpoint 10.0.0.1:19073 --device-id 0 \ + --endpoint 10.0.0.2:19073 --device-id 0 +``` + +The test passes only when every result equals the sum of all rank inputs. A +successful descriptor import without working peer `TLOAD` is therefore not +reported as success. + +## Local L3 plus remote L3 case + +`mixed_global_tload_smoke.py` verifies that one L4 process can fork a local L3 +with `add_worker()` and connect to another L3 with `add_remote_worker()`, then +place both L3 nodes in the same Global CommDomain. Only the remote server needs +a pre-started daemon. + +Start the daemon on the remote server: + +```bash +python -m simpler.remote_l3_worker --host 10.0.0.2 --port 19073 +``` + +Run L4 on the local server. Device 0 is owned by its forked local L3; the +remote L3 daemon starts device 0 on `10.0.0.2`: + +```bash +python tools/a3_l4_tcp_smoke/mixed_global_tload_smoke.py \ + --local-device-id 0 \ + --endpoint 10.0.0.2:19073 --remote-device-id 0 +``` + +The case runs the same peer `TLOAD` sum on the local and remote ranks and +checks both results. It does not use `mpirun`. + +## Compute then communicate case + +`compute_then_tload_smoke.py` keeps the same L4/L3/L2 processes alive for two +ordered task rounds: + +1. Every L3 submits one local vector-add task to its L2. +2. After all compute tasks finish, every L3 submits one peer `TLOAD` sum task + to its L2. + +Run it against the same two daemons: + +```bash +python tools/a3_l4_tcp_smoke/compute_then_tload_smoke.py \ + --endpoint 10.0.0.1:19073 --device-id 0 \ + --endpoint 10.0.0.2:19073 --device-id 0 +``` + +The case checks each rank's local addition result before communication, then +checks that every rank receives the sum of all computed rank results. + +For a simpler real-device baseline that first verifies L4 remote dispatch, +per-machine two-NPU L3 group execution, remote-buffer upload/download, and +golden checking without peer-machine memory access, run +[`remote_l4_npu`](../remote_l4_npu/README.md). diff --git a/tools/a3_l4_tcp_smoke/compute_then_tload_smoke.py b/tools/a3_l4_tcp_smoke/compute_then_tload_smoke.py new file mode 100644 index 0000000000..d9d6cf395b --- /dev/null +++ b/tools/a3_l4_tcp_smoke/compute_then_tload_smoke.py @@ -0,0 +1,282 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Run one L2 compute task and one cross-machine communication task from L4.""" + +from __future__ import annotations + +import argparse +import os +import struct +import sys + +from simpler.task_interface import ( + ArgDirection, + CallConfig, + ChipCallable, + CommBufferSpec, + CoreCallable, + TaskArgs, +) +from simpler.worker import RemoteCallable, RemoteWorkerSpec, Worker + +from simpler_setup.elf_parser import extract_text_section +from simpler_setup.kernel_compiler import KernelCompiler +from simpler_setup.pto_isa import ensure_pto_isa_root + +HERE = os.path.dirname(os.path.abspath(__file__)) +LOCAL_ADD_AIV = os.path.join(HERE, "kernels", "aiv", "local_add_kernel.cpp") +LOCAL_ADD_ORCH = os.path.join(HERE, "kernels", "orchestration", "local_add_orch.cpp") +GLOBAL_TLOAD_AIV = os.path.join(HERE, "kernels", "aiv", "global_tload_kernel.cpp") +GLOBAL_TLOAD_ORCH = os.path.join(HERE, "kernels", "orchestration", "global_tload_orch.cpp") +COUNT = 256 +FLOAT_BYTES = 4 +MAX_RANKS = 16 +WINDOW_SIZE = 4096 + + +def _compile_aiv(compiler: KernelCompiler, platform: str, runtime: str, source: str) -> bytes: + include_dirs = compiler.get_orchestration_include_dirs(runtime) + kernel_include_dirs = list(include_dirs) + [str(compiler.project_root / "src" / "common")] + kernel_bytes = compiler.compile_incore( + source_path=source, + core_type="aiv", + pto_isa_root=ensure_pto_isa_root(), + extra_include_dirs=kernel_include_dirs, + ) + return kernel_bytes if platform.endswith("sim") else extract_text_section(kernel_bytes) + + +def _build_chip_callable( + *, + platform: str, + runtime: str, + kernel_source: str, + orchestration_source: str, + signature: list[ArgDirection], + func_name: str, + config_name: str, +) -> ChipCallable: + compiler = KernelCompiler(platform=platform) + kernel_bytes = _compile_aiv(compiler, platform, runtime, kernel_source) + orchestration_bytes = compiler.compile_orchestration( + runtime_name=runtime, + source_path=orchestration_source, + ) + core = CoreCallable.build(signature=signature, binary=kernel_bytes) + return ChipCallable.build( + signature=signature, + func_name=func_name, + config_name=config_name, + binary=orchestration_bytes, + children=[(0, core)], + ) + + +def build_compute_callable(platform: str, runtime: str) -> ChipCallable: + return _build_chip_callable( + platform=platform, + runtime=runtime, + kernel_source=LOCAL_ADD_AIV, + orchestration_source=LOCAL_ADD_ORCH, + signature=[ArgDirection.IN, ArgDirection.IN, ArgDirection.OUT], + func_name="local_add_orchestration", + config_name="local_add_orchestration_config", + ) + + +def build_communication_callable(platform: str, runtime: str) -> ChipCallable: + return _build_chip_callable( + platform=platform, + runtime=runtime, + kernel_source=GLOBAL_TLOAD_AIV, + orchestration_source=GLOBAL_TLOAD_ORCH, + signature=[ArgDirection.IN, ArgDirection.OUT], + func_name="global_tload_orchestration", + config_name="global_tload_orchestration_config", + ) + + +def _digest_scalars(digest: bytes) -> tuple[int, ...]: + if len(digest) != 32: + raise ValueError("callable digest must be 32 bytes") + return tuple(int.from_bytes(digest[offset : offset + 8], "little") for offset in range(0, 32, 8)) + + +def _lhs_values(rank: int) -> tuple[float, ...]: + return tuple(float(rank * 100 + index) for index in range(COUNT)) + + +def _rhs_values(rank: int) -> tuple[float, ...]: + return tuple(float(rank * 10 + 2 * index) for index in range(COUNT)) + + +def _expected_compute(rank: int) -> tuple[float, ...]: + return tuple(lhs + rhs for lhs, rhs in zip(_lhs_values(rank), _rhs_values(rank), strict=True)) + + +def _expected_communication(rank_count: int) -> tuple[float, ...]: + rank_results = tuple(_expected_compute(rank) for rank in range(rank_count)) + return tuple(sum(rank_result[index] for rank_result in rank_results) for index in range(COUNT)) + + +def _remote_args(domain_id: int, chip_digest: bytes) -> TaskArgs: + args = TaskArgs() + args.add_scalar(domain_id) + args.add_scalar(0) + for value in _digest_scalars(chip_digest): + args.add_scalar(value) + return args + + +def _unpack_floats(raw: bytes) -> tuple[float, ...]: + return tuple(float(value) for value in struct.unpack(f"<{COUNT}f", raw)) + + +def _max_diff(actual: tuple[float, ...], expected: tuple[float, ...]) -> float: + return max(abs(observed - wanted) for observed, wanted in zip(actual, expected, strict=True)) + + +def run(endpoints: list[str], device_ids: list[int], platform: str, runtime: str) -> int: + if len(endpoints) != len(device_ids): + raise ValueError("--endpoint and --device-id counts must match") + if not 2 <= len(endpoints) <= MAX_RANKS: + raise ValueError(f"the smoke requires between 2 and {MAX_RANKS} nodes") + + worker = Worker(level=4, num_sub_workers=0, remote_session_timeout_s=120) + node_ids = tuple( + worker.add_remote_worker( + RemoteWorkerSpec( + endpoint=endpoint, + platform=platform, + runtime=runtime, + device_ids=(device_id,), + comm_profile="a3-fabric-v1", + ) + ) + for endpoint, device_id in zip(endpoints, device_ids, strict=True) + ) + print(f"[l4-compute-comm] compiling for {platform}/{runtime}") + compute_handle = worker.register(build_compute_callable(platform, runtime)) + communication_handle = worker.register(build_communication_callable(platform, runtime)) + remote_compute_handle = worker.register( + RemoteCallable("simpler.global_comm_smoke:remote_compute_orch"), + workers=list(node_ids), + ) + remote_communication_handle = worker.register( + RemoteCallable("simpler.global_comm_smoke:remote_rank_orch"), + workers=list(node_ids), + ) + captured: dict[str, object] = {} + observed_compute: list[tuple[float, ...]] = [] + observed_communication: list[tuple[float, ...]] = [] + try: + worker.init() + + def compute_phase(orch, _args, cfg): + domain = orch.allocate_global_domain( + name="a3-l4-compute-then-tload", + members=tuple((node_id, 0) for node_id in node_ids), + window_size=WINDOW_SIZE, + buffers=( + CommBufferSpec("lhs", "float32", COUNT, COUNT * FLOAT_BYTES), + CommBufferSpec("rhs", "float32", COUNT, COUNT * FLOAT_BYTES), + CommBufferSpec("input", "float32", COUNT, COUNT * FLOAT_BYTES), + CommBufferSpec("result", "float32", COUNT, COUNT * FLOAT_BYTES), + ), + retain_after_run=True, + ) + for rank, node_id in enumerate(node_ids): + orch.copy_to_global_domain( + domain, + rank, + struct.pack(f"<{COUNT}f", *_lhs_values(rank)), + buffer="lhs", + ) + orch.copy_to_global_domain( + domain, + rank, + struct.pack(f"<{COUNT}f", *_rhs_values(rank)), + buffer="rhs", + ) + orch.submit_next_level( + remote_compute_handle, + _remote_args(domain.domain_id, compute_handle.digest), + cfg, + worker=node_id, + ) + captured["domain"] = domain + + worker.run(compute_phase, args=None, config=CallConfig()) + domain = captured["domain"] + + def communication_phase(orch, _args, cfg): + for rank, node_id in enumerate(node_ids): + raw = orch.copy_from_global_domain( + domain, + rank, + COUNT * FLOAT_BYTES, + buffer="input", + ) + observed_compute.append(_unpack_floats(raw)) + orch.submit_next_level( + remote_communication_handle, + _remote_args(domain.domain_id, communication_handle.digest), + cfg, + worker=node_id, + ) + + worker.run(communication_phase, args=None, config=CallConfig()) + + def verify_phase(orch, _args, _cfg): + try: + for rank in range(len(node_ids)): + raw = orch.copy_from_global_domain( + domain, + rank, + COUNT * FLOAT_BYTES, + buffer="result", + ) + observed_communication.append(_unpack_floats(raw)) + finally: + domain.release() + + worker.run(verify_phase, args=None, config=CallConfig()) + + for rank, result in enumerate(observed_compute): + max_diff = _max_diff(result, _expected_compute(rank)) + print(f"[l4-compute-comm] compute rank={rank} max_diff={max_diff:.3e}") + if max_diff > 1e-5: + raise AssertionError(f"rank {rank} compute golden mismatch: max_diff={max_diff}") + + expected_communication = _expected_communication(len(node_ids)) + for rank, result in enumerate(observed_communication): + max_diff = _max_diff(result, expected_communication) + print(f"[l4-compute-comm] communication rank={rank} max_diff={max_diff:.3e}") + if max_diff > 1e-3: + raise AssertionError(f"rank {rank} communication golden mismatch: max_diff={max_diff}") + + print("[l4-compute-comm] PASS: L4 -> L3 -> L2 compute and peer TLOAD succeeded") + return 0 + finally: + worker.close() + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--endpoint", action="append", required=True, help="Remote L3 daemon endpoint, HOST:PORT") + parser.add_argument("--device-id", action="append", required=True, type=int, help="One device id per endpoint") + parser.add_argument("--platform", default="a2a3") + parser.add_argument("--runtime", default="tensormap_and_ringbuffer") + args = parser.parse_args() + return run(args.endpoint, args.device_id, args.platform, args.runtime) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/a3_l4_tcp_smoke/global_tload_smoke.py b/tools/a3_l4_tcp_smoke/global_tload_smoke.py new file mode 100644 index 0000000000..f99cc8c3d1 --- /dev/null +++ b/tools/a3_l4_tcp_smoke/global_tload_smoke.py @@ -0,0 +1,178 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""No-mpirun A3 L4 -> remote L3 -> L2 Fabric TLOAD smoke.""" + +from __future__ import annotations + +import argparse +import os +import struct +import sys + +from simpler.task_interface import ( + ArgDirection, + CallConfig, + ChipCallable, + CommBufferSpec, + CoreCallable, + TaskArgs, +) +from simpler.worker import RemoteCallable, RemoteWorkerSpec, Worker + +from simpler_setup.elf_parser import extract_text_section +from simpler_setup.kernel_compiler import KernelCompiler +from simpler_setup.pto_isa import ensure_pto_isa_root + +HERE = os.path.dirname(os.path.abspath(__file__)) +KERNEL_AIV = os.path.join(HERE, "kernels", "aiv", "global_tload_kernel.cpp") +KERNEL_ORCH = os.path.join(HERE, "kernels", "orchestration", "global_tload_orch.cpp") +COUNT = 256 +FLOAT_BYTES = 4 +MAX_RANKS = 16 +WINDOW_SIZE = 4096 + + +def build_chip_callable(platform: str, runtime: str) -> ChipCallable: + compiler = KernelCompiler(platform=platform) + pto_isa_root = ensure_pto_isa_root() + include_dirs = compiler.get_orchestration_include_dirs(runtime) + kernel_include_dirs = list(include_dirs) + [str(compiler.project_root / "src" / "common")] + kernel_bytes = compiler.compile_incore( + source_path=KERNEL_AIV, + core_type="aiv", + pto_isa_root=pto_isa_root, + extra_include_dirs=kernel_include_dirs, + ) + if not platform.endswith("sim"): + kernel_bytes = extract_text_section(kernel_bytes) + orch_bytes = compiler.compile_orchestration(runtime_name=runtime, source_path=KERNEL_ORCH) + core = CoreCallable.build(signature=[ArgDirection.IN, ArgDirection.OUT], binary=kernel_bytes) + return ChipCallable.build( + signature=[ArgDirection.IN, ArgDirection.OUT], + func_name="global_tload_orchestration", + config_name="global_tload_orchestration_config", + binary=orch_bytes, + children=[(0, core)], + ) + + +def _digest_scalars(digest: bytes) -> tuple[int, ...]: + if len(digest) != 32: + raise ValueError("callable digest must be 32 bytes") + return tuple(int.from_bytes(digest[offset : offset + 8], "little") for offset in range(0, 32, 8)) + + +def _input_values(rank: int) -> tuple[float, ...]: + return tuple(float(rank * 100 + index) for index in range(COUNT)) + + +def _expected_values(rank_count: int) -> tuple[float, ...]: + rank_bias = 100 * rank_count * (rank_count - 1) // 2 + return tuple(float(rank_count * index + rank_bias) for index in range(COUNT)) + + +def run(endpoints: list[str], device_ids: list[int], platform: str, runtime: str) -> int: + if len(endpoints) != len(device_ids): + raise ValueError("--endpoint and --device-id counts must match") + if not 2 <= len(endpoints) <= MAX_RANKS: + raise ValueError(f"the smoke requires between 2 and {MAX_RANKS} nodes") + + worker = Worker(level=4, num_sub_workers=0, remote_session_timeout_s=120) + node_ids = tuple( + worker.add_remote_worker( + RemoteWorkerSpec( + endpoint=endpoint, + platform=platform, + runtime=runtime, + device_ids=(device_id,), + comm_profile="a3-fabric-v1", + ) + ) + for endpoint, device_id in zip(endpoints, device_ids, strict=True) + ) + print(f"[l4-global-tload] compiling for {platform}/{runtime}") + chip_handle = worker.register(build_chip_callable(platform, runtime)) + remote_handle = worker.register( + RemoteCallable("simpler.global_comm_smoke:remote_rank_orch"), + workers=list(node_ids), + ) + captured: dict[str, object] = {} + try: + worker.init() + + def build_and_run(orch, _args, cfg): + domain = orch.allocate_global_domain( + name="a3-l4-tload", + members=tuple((node_id, 0) for node_id in node_ids), + window_size=WINDOW_SIZE, + buffers=( + CommBufferSpec("input", "float32", COUNT, COUNT * FLOAT_BYTES), + CommBufferSpec("result", "float32", COUNT, COUNT * FLOAT_BYTES), + ), + retain_after_run=True, + ) + for rank in range(len(node_ids)): + orch.copy_to_global_domain( + domain, + rank, + struct.pack(f"<{COUNT}f", *_input_values(rank)), + buffer="input", + ) + digest_scalars = _digest_scalars(chip_handle.digest) + for node_id in node_ids: + rank_args = TaskArgs() + rank_args.add_scalar(domain.domain_id) + rank_args.add_scalar(0) + for value in digest_scalars: + rank_args.add_scalar(value) + orch.submit_next_level(remote_handle, rank_args, cfg, worker=node_id) + captured["domain"] = domain + + worker.run(build_and_run, args=None, config=CallConfig()) + domain = captured["domain"] + expected = _expected_values(len(node_ids)) + observed: list[tuple[float, ...]] = [] + + def read_and_release(orch, _args, _cfg): + for rank in range(len(node_ids)): + raw = orch.copy_from_global_domain( + domain, + rank, + COUNT * FLOAT_BYTES, + buffer="result", + ) + observed.append(tuple(float(value) for value in struct.unpack(f"<{COUNT}f", raw))) + domain.release() + + worker.run(read_and_release, args=None, config=CallConfig()) + for rank, result in enumerate(observed): + max_diff = max(abs(actual - wanted) for actual, wanted in zip(result, expected, strict=True)) + print(f"[l4-global-tload] rank={rank} max_diff={max_diff:.3e}") + if max_diff > 1e-3: + print("[l4-global-tload] FAILED") + return 1 + print("[l4-global-tload] PASS: L4-brokered Fabric descriptors and peer TLOAD both succeeded") + return 0 + finally: + worker.close() + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--endpoint", action="append", required=True, help="Remote L3 daemon endpoint, HOST:PORT") + parser.add_argument("--device-id", action="append", required=True, type=int, help="One device id per endpoint") + parser.add_argument("--platform", default="a2a3") + parser.add_argument("--runtime", default="tensormap_and_ringbuffer") + args = parser.parse_args() + return run(args.endpoint, args.device_id, args.platform, args.runtime) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/a3_l4_tcp_smoke/kernels/aiv/global_tload_kernel.cpp b/tools/a3_l4_tcp_smoke/kernels/aiv/global_tload_kernel.cpp new file mode 100644 index 0000000000..231a75a6a1 --- /dev/null +++ b/tools/a3_l4_tcp_smoke/kernels/aiv/global_tload_kernel.cpp @@ -0,0 +1,86 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include +#include + +#include "platform_comm/comm_context.h" +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +static constexpr size_t kCount = 256; +static constexpr int kMaxRanks = 16; + +template +AICORE inline __gm__ T *CommRemotePtr(__gm__ CommContext *ctx, __gm__ T *local_ptr, int peer_rank) { + uint64_t local_base = ctx->windowsIn[ctx->rankId]; + uint64_t offset = reinterpret_cast(local_ptr) - local_base; + return reinterpret_cast<__gm__ T *>(ctx->windowsIn[peer_rank] + offset); +} + +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *input_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *result_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + int rank_count = static_cast(args[2]); + __gm__ CommContext *comm_ctx = reinterpret_cast<__gm__ CommContext *>(args[3]); + + if (rank_count <= 0 || rank_count > kMaxRanks) { + pipe_barrier(PIPE_ALL); + return; + } + __gm__ float *input = reinterpret_cast<__gm__ float *>(input_tensor->buffer.addr) + input_tensor->start_offset; + __gm__ float *result = reinterpret_cast<__gm__ float *>(result_tensor->buffer.addr) + result_tensor->start_offset; + + using Shape = pto::Shape; + using Stride = pto::Stride; + using Global = pto::GlobalTensor; + using Tile = pto::Tile; + + Shape shape(1, 1, 1, 1, kCount); + Stride stride(kCount, kCount, kCount, kCount, 1); + Tile accumulator(1, kCount); + Tile peer_tile(1, kCount); + TASSIGN(accumulator, 0x0); + TASSIGN(peer_tile, 0x10000); + + Global local_input(input, shape, stride); + TLOAD(accumulator, local_input); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + int my_rank = static_cast(comm_ctx->rankId); + for (int peer = 0; peer < rank_count; ++peer) { + if (peer == my_rank) continue; + __gm__ float *peer_input = CommRemotePtr(comm_ctx, input, peer); + Global peer_global(peer_input, shape, stride); + TLOAD(peer_tile, peer_global); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + TADD(accumulator, accumulator, peer_tile); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + } + + Global result_global(result, shape, stride); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(result_global, accumulator); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + pipe_barrier(PIPE_ALL); +} diff --git a/tools/a3_l4_tcp_smoke/kernels/aiv/local_add_kernel.cpp b/tools/a3_l4_tcp_smoke/kernels/aiv/local_add_kernel.cpp new file mode 100644 index 0000000000..9d50f4f4a1 --- /dev/null +++ b/tools/a3_l4_tcp_smoke/kernels/aiv/local_add_kernel.cpp @@ -0,0 +1,61 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +static constexpr size_t kCount = 256; + +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *lhs_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *rhs_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *result_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *lhs = reinterpret_cast<__gm__ float *>(lhs_tensor->buffer.addr) + lhs_tensor->start_offset; + __gm__ float *rhs = reinterpret_cast<__gm__ float *>(rhs_tensor->buffer.addr) + rhs_tensor->start_offset; + __gm__ float *result = reinterpret_cast<__gm__ float *>(result_tensor->buffer.addr) + result_tensor->start_offset; + + using Shape = pto::Shape; + using Stride = pto::Stride; + using Global = pto::GlobalTensor; + using Tile = pto::Tile; + + Shape shape(1, 1, 1, 1, kCount); + Stride stride(kCount, kCount, kCount, kCount, 1); + Tile lhs_tile(1, kCount); + Tile rhs_tile(1, kCount); + Tile result_tile(1, kCount); + TASSIGN(lhs_tile, 0x0); + TASSIGN(rhs_tile, 0x10000); + TASSIGN(result_tile, 0x20000); + + Global lhs_global(lhs, shape, stride); + Global rhs_global(rhs, shape, stride); + Global result_global(result, shape, stride); + TLOAD(lhs_tile, lhs_global); + TLOAD(rhs_tile, rhs_global); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TADD(result_tile, lhs_tile, rhs_tile); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(result_global, result_tile); + pipe_barrier(PIPE_ALL); +} diff --git a/tools/a3_l4_tcp_smoke/kernels/orchestration/global_tload_orch.cpp b/tools/a3_l4_tcp_smoke/kernels/orchestration/global_tload_orch.cpp new file mode 100644 index 0000000000..fe9c2da937 --- /dev/null +++ b/tools/a3_l4_tcp_smoke/kernels/orchestration/global_tload_orch.cpp @@ -0,0 +1,35 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include + +#include "pto_orchestration_api.h" + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig +global_tload_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 4, + }; +} + +__attribute__((visibility("default"))) void global_tload_orchestration(const L2TaskArgs &orch_args) { + L0TaskArgs params; + params.add_input(orch_args.tensor(0).ref()); + params.add_output(orch_args.tensor(1).ref()); + params.add_scalar(orch_args.scalar(0)); + params.add_scalar(orch_args.scalar(1)); + rt_submit_aiv_task(0, params); +} + +} // extern "C" diff --git a/tools/a3_l4_tcp_smoke/kernels/orchestration/local_add_orch.cpp b/tools/a3_l4_tcp_smoke/kernels/orchestration/local_add_orch.cpp new file mode 100644 index 0000000000..1c0a2d81e9 --- /dev/null +++ b/tools/a3_l4_tcp_smoke/kernels/orchestration/local_add_orch.cpp @@ -0,0 +1,34 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include + +#include "pto_orchestration_api.h" + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig +local_add_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 3, + }; +} + +__attribute__((visibility("default"))) void local_add_orchestration(const L2TaskArgs &orch_args) { + L0TaskArgs params; + params.add_input(orch_args.tensor(0).ref()); + params.add_input(orch_args.tensor(1).ref()); + params.add_output(orch_args.tensor(2).ref()); + rt_submit_aiv_task(0, params); +} + +} // extern "C" diff --git a/tools/a3_l4_tcp_smoke/mixed_global_tload_smoke.py b/tools/a3_l4_tcp_smoke/mixed_global_tload_smoke.py new file mode 100644 index 0000000000..0e43d0e677 --- /dev/null +++ b/tools/a3_l4_tcp_smoke/mixed_global_tload_smoke.py @@ -0,0 +1,212 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""No-mpirun A3 TLOAD smoke spanning one local L3 and one or more remote L3 nodes.""" + +from __future__ import annotations + +import argparse +import struct +import sys + +from global_tload_smoke import ( + COUNT, + FLOAT_BYTES, + MAX_RANKS, + WINDOW_SIZE, + _digest_scalars, + _expected_values, + _input_values, + build_chip_callable, +) +from simpler.task_interface import CallConfig, CommBufferSpec, DataType, TaskArgs, Tensor, TensorArgType +from simpler.worker import CallableHandle, RemoteCallable, RemoteWorkerSpec, Worker + + +def _make_local_rank_orch(chip_handle: CallableHandle): + def local_rank_orch(orch, args: TaskArgs, cfg: CallConfig) -> None: + if args.scalar_count() != 2: + raise ValueError("local TLOAD task expects domain_id and local_worker_id") + domain_id = int(args.scalar(0)) + local_worker_id = int(args.scalar(1)) + context = orch.get_global_domain(domain_id)[local_worker_id] + + chip_args = TaskArgs() + chip_args.add_tensor( + Tensor.make( + data=context.buffer_ptrs["input"], + shapes=(COUNT,), + dtype=DataType.FLOAT32, + child_memory=True, + ), + TensorArgType.INPUT, + ) + chip_args.add_tensor( + Tensor.make( + data=context.buffer_ptrs["result"], + shapes=(COUNT,), + dtype=DataType.FLOAT32, + child_memory=True, + ), + TensorArgType.OUTPUT_EXISTING, + ) + chip_args.add_scalar(context.domain_size) + chip_args.add_scalar(context.device_ctx) + orch.submit_next_level(chip_handle, chip_args, cfg, worker=local_worker_id) + + return local_rank_orch + + +def _domain_args(domain_id: int, local_worker_id: int) -> TaskArgs: + args = TaskArgs() + args.add_scalar(domain_id) + args.add_scalar(local_worker_id) + return args + + +def run( + endpoints: list[str], + remote_device_ids: list[int], + local_device_id: int, + platform: str, + runtime: str, +) -> int: + if len(endpoints) != len(remote_device_ids): + raise ValueError("--endpoint and --remote-device-id counts must match") + node_count = 1 + len(endpoints) + if not 2 <= node_count <= MAX_RANKS: + raise ValueError(f"the smoke requires between 2 and {MAX_RANKS} nodes including the local L3") + + chip_callable = build_chip_callable(platform, runtime) + local_l3 = Worker( + level=3, + device_ids=[local_device_id], + num_sub_workers=0, + platform=platform, + runtime=runtime, + comm_profile="a3-fabric-v1", + global_device_ranks=(0,), + ) + local_chip_handle = local_l3.register(chip_callable) + + worker = Worker(level=4, num_sub_workers=0, remote_session_timeout_s=120) + local_node_id = worker.add_worker(local_l3) + remote_node_ids = tuple( + worker.add_remote_worker( + RemoteWorkerSpec( + endpoint=endpoints[index], + platform=platform, + runtime=runtime, + device_ids=(remote_device_ids[index],), + comm_profile="a3-fabric-v1", + global_device_ranks=(index + 1,), + ) + ) + for index in range(len(endpoints)) + ) + node_ids = (local_node_id, *remote_node_ids) + + print(f"[l4-mixed-global-tload] compiling for {platform}/{runtime}") + remote_chip_handle = worker.register(chip_callable) + local_orch_handle = worker.register(_make_local_rank_orch(local_chip_handle)) + remote_orch_handle = worker.register( + RemoteCallable("simpler.global_comm_smoke:remote_rank_orch"), + workers=list(remote_node_ids), + ) + captured: dict[str, object] = {} + try: + worker.init() + + def build_and_run(orch, _args, cfg): + domain = orch.allocate_global_domain( + name="a3-l4-mixed-tload", + members=tuple((node_id, 0) for node_id in node_ids), + window_size=WINDOW_SIZE, + buffers=( + CommBufferSpec("input", "float32", COUNT, COUNT * FLOAT_BYTES), + CommBufferSpec("result", "float32", COUNT, COUNT * FLOAT_BYTES), + ), + retain_after_run=True, + ) + for rank in range(node_count): + orch.copy_to_global_domain( + domain, + rank, + struct.pack(f"<{COUNT}f", *_input_values(rank)), + buffer="input", + ) + + orch.submit_next_level( + local_orch_handle, + _domain_args(domain.domain_id, 0), + cfg, + worker=local_node_id, + ) + digest_scalars = _digest_scalars(remote_chip_handle.digest) + for node_id in remote_node_ids: + rank_args = _domain_args(domain.domain_id, 0) + for value in digest_scalars: + rank_args.add_scalar(value) + orch.submit_next_level(remote_orch_handle, rank_args, cfg, worker=node_id) + captured["domain"] = domain + + worker.run(build_and_run, args=None, config=CallConfig()) + domain = captured["domain"] + expected = _expected_values(node_count) + observed: list[tuple[float, ...]] = [] + + def read_and_release(orch, _args, _cfg): + for rank in range(node_count): + raw = orch.copy_from_global_domain( + domain, + rank, + COUNT * FLOAT_BYTES, + buffer="result", + ) + observed.append(tuple(float(value) for value in struct.unpack(f"<{COUNT}f", raw))) + domain.release() + + worker.run(read_and_release, args=None, config=CallConfig()) + for rank, result in enumerate(observed): + max_diff = max(abs(actual - wanted) for actual, wanted in zip(result, expected, strict=True)) + print(f"[l4-mixed-global-tload] rank={rank} max_diff={max_diff:.3e}") + if max_diff > 1e-3: + print("[l4-mixed-global-tload] FAILED") + return 1 + print("[l4-mixed-global-tload] PASS: local and remote L3 ranks completed peer TLOAD") + return 0 + finally: + worker.close() + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--local-device-id", required=True, type=int, help="Device owned by the local L3") + parser.add_argument("--endpoint", action="append", required=True, help="Remote L3 daemon endpoint, HOST:PORT") + parser.add_argument( + "--remote-device-id", + action="append", + required=True, + type=int, + help="One device id per remote endpoint", + ) + parser.add_argument("--platform", default="a2a3") + parser.add_argument("--runtime", default="tensormap_and_ringbuffer") + args = parser.parse_args() + return run( + args.endpoint, + args.remote_device_id, + args.local_device_id, + args.platform, + args.runtime, + ) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/remote_l4_npu/README.md b/tools/remote_l4_npu/README.md new file mode 100644 index 0000000000..04e3c93152 --- /dev/null +++ b/tools/remote_l4_npu/README.md @@ -0,0 +1,60 @@ +# Remote L4 two-machine NPU compute smoke + +This smoke is the real-device baseline imported from +`codex/remote-l4-npu-machineA`. It validates: + +```text +L4 parent -> machine A remote L3 -> local NPU 0 + NPU 1 vector group + -> machine B remote L3 -> local NPU 0 + NPU 1 vector group +``` + +The parent allocates six remote buffers on each machine, uploads different +inputs, dispatches both remote L3 tasks, downloads both outputs, and checks the +vector-example golden result. No `mpirun` is used. + +This baseline proves remote startup, remote buffers, L4 task dispatch, local L3 +group scheduling, NPU execution, and result copy-back. It does not itself read +peer-machine memory. Run the sibling +[`a3_l4_tcp_smoke`](../a3_l4_tcp_smoke/README.md) for the Global CommDomain +Fabric-handle exchange and peer `TLOAD` test. + +## Prepare both NPU machines + +Use the same source revision on both machines: + +```bash +python3 -m venv --system-site-packages .venv +source .venv/bin/activate +pip install --no-build-isolation -e . + +export ASCEND_HOME_PATH=/usr/local/Ascend/ascend-toolkit/latest +export PATH="$ASCEND_HOME_PATH/bin:$PATH" + +SIMPLER_REMOTE_L4_NPU_ROLE=machineA \ + bash tools/remote_l4_npu/start_machine_daemon.sh +``` + +Use `machineB` for the role on the second server. The default daemon address is +`0.0.0.0:19073`; override `SIMPLER_REMOTE_L4_NPU_HOST` or +`SIMPLER_REMOTE_L4_NPU_PORT` when needed. + +The daemon starts a session runner on random TCP command and health ports. +Firewalls between the parent and both NPU machines must allow those returned +ports in addition to the daemon port. + +## Run on the L4 parent + +The parent needs the A3 compilation toolchain because it builds the vector +orchestration and AIV kernels before opening the remote sessions. + +```bash +export SIMPLER_REMOTE_L4_NPU_MACHINE_A=10.0.0.11:19073 +export SIMPLER_REMOTE_L4_NPU_MACHINE_B=10.0.0.12:19073 +export SIMPLER_REMOTE_L4_NPU_MACHINE_A_DEVICES=0,1 +export SIMPLER_REMOTE_L4_NPU_MACHINE_B_DEVICES=0,1 + +bash tools/remote_l4_npu/run_parent_smoke.sh +``` + +Each machine must provide exactly two free A3 device ids. Success requires all +four downloaded outputs to match the vector golden result. diff --git a/tools/remote_l4_npu/remote_l4_npu_smoke.py b/tools/remote_l4_npu/remote_l4_npu_smoke.py new file mode 100644 index 0000000000..791eaf1ec5 --- /dev/null +++ b/tools/remote_l4_npu/remote_l4_npu_smoke.py @@ -0,0 +1,282 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Run L4 -> two remote L3 workers, each executing a two-NPU vector group.""" + +from __future__ import annotations + +import argparse +import ctypes +from pathlib import Path +from typing import Any + +from simpler.task_interface import ( + ArgDirection, + CallConfig, + ChipCallable, + CoreCallable, + DataType, + RemoteBufferHandle, + RemoteTensorRef, + TaskArgs, + TensorArgType, +) +from simpler.worker import RemoteCallable, RemoteWorkerSpec, Worker + +from simpler_setup.elf_parser import extract_text_section +from simpler_setup.kernel_compiler import KernelCompiler +from simpler_setup.pto_isa import ensure_pto_isa_root + +REMOTE_ORCH_TARGET = "tools.remote_l4_npu.remote_l4_npu_smoke:remote_l3_group_orch" +ELEMENTS = 128 * 128 +FLOAT_NBYTES = ctypes.sizeof(ctypes.c_float) +TENSOR_COUNT = 6 +FloatArray = ctypes.c_float * ELEMENTS +_REMOTE_GROUP_KEEPALIVE: list[TaskArgs] = [] + + +def _digest_from_scalars(args: TaskArgs) -> bytes: + return b"".join(int(args.scalar(index)).to_bytes(8, "little") for index in range(4)) + + +def remote_l3_group_orch(orch, args: TaskArgs, cfg: CallConfig) -> None: + """Rebuild the two local chip tasks and submit them as one L3 group.""" + from simpler.remote_l3_session import get_inner_handle # noqa: PLC0415 + + if args.tensor_count() != TENSOR_COUNT or args.scalar_count() != 4: + raise ValueError("remote L3 group task expects six tensors and four digest scalars") + chip_handle = get_inner_handle(_digest_from_scalars(args).hex()) + + chip_args0 = TaskArgs() + chip_args0.add_tensor(args.tensor(0), TensorArgType.INPUT) + chip_args0.add_tensor(args.tensor(1), TensorArgType.INPUT) + chip_args0.add_tensor(args.tensor(2), TensorArgType.OUTPUT_EXISTING) + + chip_args1 = TaskArgs() + chip_args1.add_tensor(args.tensor(3), TensorArgType.INPUT) + chip_args1.add_tensor(args.tensor(4), TensorArgType.INPUT) + chip_args1.add_tensor(args.tensor(5), TensorArgType.OUTPUT_EXISTING) + + # The inner Worker drains after this callback returns, so retain the rebuilt + # TaskArgs until that drain has completed. + _REMOTE_GROUP_KEEPALIVE[:] = [chip_args0, chip_args1] + orch.submit_next_level_group(chip_handle, [chip_args0, chip_args1], cfg, workers=[0, 1]) + + +def _build_vector_chip_callable(platform: str, runtime: str) -> ChipCallable: + root = Path(__file__).resolve().parents[2] + kernels = root / "examples" / "a2a3" / "tensormap_and_ringbuffer" / "vector_example" / "kernels" + orch_source = kernels / "orchestration" / "example_orchestration.cpp" + aiv_sources = ( + kernels / "aiv" / "kernel_add.cpp", + kernels / "aiv" / "kernel_add_scalar.cpp", + kernels / "aiv" / "kernel_mul.cpp", + ) + + compiler = KernelCompiler(platform=platform) + pto_isa_root = ensure_pto_isa_root() + include_dirs = compiler.get_orchestration_include_dirs(runtime) + include_dirs = list(include_dirs) + [str(compiler.project_root / "src" / "common")] + + def compile_aiv(source: Path) -> bytes: + binary = compiler.compile_incore( + source_path=str(source), + core_type="aiv", + pto_isa_root=pto_isa_root, + extra_include_dirs=include_dirs, + ) + return binary if platform.endswith("sim") else extract_text_section(binary) + + children = ( + ( + 0, + CoreCallable.build( + signature=[ArgDirection.IN, ArgDirection.IN, ArgDirection.OUT], + binary=compile_aiv(aiv_sources[0]), + ), + ), + ( + 1, + CoreCallable.build( + signature=[ArgDirection.IN, ArgDirection.OUT], + binary=compile_aiv(aiv_sources[1]), + ), + ), + ( + 2, + CoreCallable.build( + signature=[ArgDirection.IN, ArgDirection.IN, ArgDirection.OUT], + binary=compile_aiv(aiv_sources[2]), + ), + ), + ) + orch_binary = compiler.compile_orchestration(runtime_name=runtime, source_path=str(orch_source)) + return ChipCallable.build( + signature=[ArgDirection.IN, ArgDirection.IN, ArgDirection.OUT], + func_name="aicpu_orchestration_entry", + config_name="aicpu_orchestration_config", + binary=orch_binary, + children=list(children), + ) + + +def _add_digest_scalars(task_args: TaskArgs, digest: bytes) -> None: + if len(digest) != 32: + raise ValueError("inner chip callable digest must be 32 bytes") + for offset in range(0, 32, 8): + task_args.add_scalar(int.from_bytes(digest[offset : offset + 8], "little")) + + +def _parse_device_ids(value: str) -> tuple[int, int]: + device_ids = tuple(int(part.strip()) for part in value.split(",") if part.strip()) + if len(device_ids) != 2: + raise ValueError("each remote L3 group requires exactly two device ids") + if any(device_id < 0 for device_id in device_ids) or len(set(device_ids)) != 2: + raise ValueError("device ids must be distinct and non-negative") + return device_ids + + +def _make_array(value: float) -> Any: + array = FloatArray() + for index in range(ELEMENTS): + array[index] = value + return array + + +def _expected(lhs: float, rhs: float) -> float: + summed = lhs + rhs + return (summed + 1.0) * (summed + 2.0) + summed + + +def _make_remote_group_args(handles: list[RemoteBufferHandle], digest: bytes) -> TaskArgs: + if len(handles) != TENSOR_COUNT: + raise ValueError("remote L3 group requires six remote buffers") + args = TaskArgs() + for index, handle in enumerate(handles): + tag = TensorArgType.OUTPUT_EXISTING if index in (2, 5) else TensorArgType.INPUT + args.add_tensor(RemoteTensorRef(handle, shape=(ELEMENTS,), dtype=DataType.FLOAT32), tag) + _add_digest_scalars(args, digest) + return args + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--machine-a", required=True, help="machine A daemon endpoint, HOST:PORT") + parser.add_argument("--machine-b", required=True, help="machine B daemon endpoint, HOST:PORT") + parser.add_argument("--machine-a-devices", default="0,1") + parser.add_argument("--machine-b-devices", default="0,1") + parser.add_argument("--platform", default="a2a3") + parser.add_argument("--runtime", default="tensormap_and_ringbuffer") + parser.add_argument("--session-timeout", type=float, default=120.0) + parser.add_argument("--session-listen-host", default="0.0.0.0") + return parser.parse_args() + + +def main() -> int: + args = _parse_args() + machine_a_devices = _parse_device_ids(args.machine_a_devices) + machine_b_devices = _parse_device_ids(args.machine_b_devices) + + worker = Worker(level=4, num_sub_workers=0, remote_session_timeout_s=args.session_timeout) + remote_buffers: list[RemoteBufferHandle] = [] + parent_keepalive: list[TaskArgs] = [] + try: + worker_a = worker.add_remote_worker( + RemoteWorkerSpec( + endpoint=args.machine_a, + platform=args.platform, + runtime=args.runtime, + device_ids=machine_a_devices, + transport="sim", + session_listen_host=args.session_listen_host, + allow_wildcard_session_bind=True, + ) + ) + worker_b = worker.add_remote_worker( + RemoteWorkerSpec( + endpoint=args.machine_b, + platform=args.platform, + runtime=args.runtime, + device_ids=machine_b_devices, + transport="sim", + session_listen_host=args.session_listen_host, + allow_wildcard_session_bind=True, + ) + ) + chip_handle = worker.register(_build_vector_chip_callable(args.platform, args.runtime)) + remote_handle = worker.register(RemoteCallable(REMOTE_ORCH_TARGET), workers=[worker_a, worker_b]) + worker.init() + + tensor_nbytes = ELEMENTS * FLOAT_NBYTES + group_values = { + worker_a: (2.0, 3.0, 4.0, 5.0), + worker_b: (6.0, 7.0, 8.0, 9.0), + } + group_handles: dict[int, list[RemoteBufferHandle]] = {} + output_arrays: dict[int, dict[str, tuple[Any, float]]] = {} + for worker_id, (a0_value, b0_value, a1_value, b1_value) in group_values.items(): + handles = [worker.remote_malloc(worker=worker_id, nbytes=tensor_nbytes) for _ in range(TENSOR_COUNT)] + remote_buffers.extend(handles) + group_handles[worker_id] = handles + initial_arrays = ( + _make_array(a0_value), + _make_array(b0_value), + _make_array(0.0), + _make_array(a1_value), + _make_array(b1_value), + _make_array(0.0), + ) + for handle, array in zip(handles, initial_arrays, strict=True): + worker.remote_copy_to(handle, array, tensor_nbytes) + output_arrays[worker_id] = { + "f0": (_make_array(0.0), _expected(a0_value, b0_value)), + "f1": (_make_array(0.0), _expected(a1_value, b1_value)), + } + + def parent_orch(orch, _args, cfg): + args_a = _make_remote_group_args(group_handles[worker_a], chip_handle.digest) + args_b = _make_remote_group_args(group_handles[worker_b], chip_handle.digest) + parent_keepalive[:] = [args_a, args_b] + orch.submit_next_level(remote_handle, args_a, cfg, worker=worker_a) + orch.submit_next_level(remote_handle, args_b, cfg, worker=worker_b) + + config = CallConfig() + config.aicpu_thread_num = 4 + worker.run(parent_orch, args=None, config=config) + + for worker_id, handles in group_handles.items(): + worker.remote_copy_from(handles[2], output_arrays[worker_id]["f0"][0], tensor_nbytes) + worker.remote_copy_from(handles[5], output_arrays[worker_id]["f1"][0], tensor_nbytes) + + for worker_id, outputs in output_arrays.items(): + for name, (array, expected) in outputs.items(): + max_diff = max(abs(float(array[index]) - expected) for index in range(ELEMENTS)) + print(f"[remote-l4-group] worker={worker_id} output={name} max_diff={max_diff:.3e}") + if max_diff > 1e-4: + raise AssertionError(f"worker {worker_id} {name} golden mismatch: max_diff={max_diff}") + + print( + "remote L4 L3-group NPU smoke passed: " + f"{args.machine_a}[devices={args.machine_a_devices}], " + f"{args.machine_b}[devices={args.machine_b_devices}], " + f"elements={ELEMENTS}" + ) + return 0 + finally: + parent_keepalive.clear() + for handle in reversed(remote_buffers): + try: + worker.remote_free(handle) + except Exception: # noqa: BLE001 + pass + worker.close() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/remote_l4_npu/run_parent_smoke.sh b/tools/remote_l4_npu/run_parent_smoke.sh new file mode 100755 index 0000000000..0e2c3db87c --- /dev/null +++ b/tools/remote_l4_npu/run_parent_smoke.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" + +: "${SIMPLER_REMOTE_L4_NPU_MACHINE_A:?set machine A daemon as HOST:PORT}" +: "${SIMPLER_REMOTE_L4_NPU_MACHINE_B:?set machine B daemon as HOST:PORT}" +: "${SIMPLER_REMOTE_L4_NPU_MACHINE_A_DEVICES:=0,1}" +: "${SIMPLER_REMOTE_L4_NPU_MACHINE_B_DEVICES:=0,1}" +: "${SIMPLER_REMOTE_L4_NPU_PLATFORM:=a2a3}" +: "${SIMPLER_REMOTE_L4_NPU_RUNTIME:=tensormap_and_ringbuffer}" +: "${SIMPLER_REMOTE_L4_NPU_SESSION_LISTEN_HOST:=0.0.0.0}" +: "${SIMPLER_REMOTE_L4_NPU_SESSION_TIMEOUT:=120}" + +cd "${ROOT_DIR}" +source .venv/bin/activate + +exec python -m tools.remote_l4_npu.remote_l4_npu_smoke \ + --machine-a "${SIMPLER_REMOTE_L4_NPU_MACHINE_A}" \ + --machine-b "${SIMPLER_REMOTE_L4_NPU_MACHINE_B}" \ + --machine-a-devices "${SIMPLER_REMOTE_L4_NPU_MACHINE_A_DEVICES}" \ + --machine-b-devices "${SIMPLER_REMOTE_L4_NPU_MACHINE_B_DEVICES}" \ + --platform "${SIMPLER_REMOTE_L4_NPU_PLATFORM}" \ + --runtime "${SIMPLER_REMOTE_L4_NPU_RUNTIME}" \ + --session-listen-host "${SIMPLER_REMOTE_L4_NPU_SESSION_LISTEN_HOST}" \ + --session-timeout "${SIMPLER_REMOTE_L4_NPU_SESSION_TIMEOUT}" diff --git a/tools/remote_l4_npu/start_machine_daemon.sh b/tools/remote_l4_npu/start_machine_daemon.sh new file mode 100755 index 0000000000..981fddec06 --- /dev/null +++ b/tools/remote_l4_npu/start_machine_daemon.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" + +: "${SIMPLER_REMOTE_L4_NPU_ROLE:=machine}" +: "${SIMPLER_REMOTE_L4_NPU_HOST:=0.0.0.0}" +: "${SIMPLER_REMOTE_L4_NPU_PORT:=19073}" + +cd "${ROOT_DIR}" +source .venv/bin/activate + +echo "[remote-l4-npu] role=${SIMPLER_REMOTE_L4_NPU_ROLE}" +echo "[remote-l4-npu] daemon=${SIMPLER_REMOTE_L4_NPU_HOST}:${SIMPLER_REMOTE_L4_NPU_PORT}" + +exec python -m simpler.remote_l3_worker \ + --host "${SIMPLER_REMOTE_L4_NPU_HOST}" \ + --port "${SIMPLER_REMOTE_L4_NPU_PORT}"