Repository navigation
feat: add committed_device_memory() API (L1 byte counter + L2/L3 query) - #1592
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds committed device-memory accounting to host and simulation allocators, runtime C APIs, ChipWorker, hierarchical worker control, and Python worker interfaces. Unit and simulation tests cover allocation changes, cleanup, and level-specific behavior. ChangesCommitted device memory
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Worker
participant Orchestrator
participant ChipWorker
participant MemoryAllocator
Worker->>Orchestrator: committed_device_memory(worker_id)
Orchestrator->>ChipWorker: query committed memory
ChipWorker->>MemoryAllocator: committed_bytes()
MemoryAllocator-->>ChipWorker: byte total
ChipWorker-->>Orchestrator: byte total
Orchestrator-->>Worker: byte total
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/simpler/worker.py`:
- Around line 5809-5826: Update committed_device_memory to explicitly raise
NotImplementedError with a message matching “level-3” when a non-level-2 worker
has no chip workers or device IDs, before calling _check_chip_worker_id.
Preserve the existing level-2 behavior and normal level-3 forwarding through
_orch.committed_device_memory for configured chip workers.
In `@tests/ut/py/test_worker/test_committed_memory_sim.py`:
- Around line 75-79: Update the L3 NotImplementedError test around Worker and
committed_device_memory so it uses a child-backed L3 configuration with a worker
ID sourced from the child IDs, ensuring the assertion exercises L3 query
forwarding rather than the _chip_shms worker-ID range check. Preserve the
expected NotImplementedError match.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: c796d20e-72ae-47f2-8cef-45abbc6d7466
📒 Files selected for processing (23)
python/bindings/task_interface.cpppython/bindings/worker_bind.hpython/simpler/orchestrator.pypython/simpler/task_interface.pypython/simpler/worker.pysrc/a2a3/platform/onboard/host/memory_allocator.cppsrc/a5/platform/onboard/host/memory_allocator.cppsrc/common/hierarchical/orchestrator.cppsrc/common/hierarchical/orchestrator.hsrc/common/hierarchical/worker_manager.cppsrc/common/hierarchical/worker_manager.hsrc/common/platform/include/host/memory_allocator.hsrc/common/platform/onboard/host/c_api_shared.cppsrc/common/platform/onboard/host/device_runner_base.hsrc/common/platform/sim/host/c_api_shared.cppsrc/common/platform/sim/host/device_runner_base.hsrc/common/platform/sim/host/memory_allocator.cppsrc/common/worker/chip_worker.cppsrc/common/worker/chip_worker.hsrc/common/worker/pto_runtime_c_api.htests/ut/cpp/CMakeLists.txttests/ut/cpp/common/test_memory_allocator.cpptests/ut/py/test_worker/test_committed_memory_sim.py
|
@coderabbitai review |
✅ Action performedReview finished.
|
fa5c72b to
3f4c177
Compare
Docstring: "excludes profiling buffers" is platform-dependent and wrong for a2a3(Correcting my earlier comment — the runtime call chain I cited was inaccurate; the conclusion holds, but the mechanism and, more importantly, the a2a3-vs-a5 split are below.)
a2a3 — DFX buffers ARE counted. a5 — DFX buffers are NOT counted. So the same docstring is correct for a5 and wrong for a2a3. Please either make the two platforms consistent, or state the exclusion is platform-specific. These allocations also happen at runtime, not just init (a2a3 path). A background thread started in Consequence for the consumer (#1457). |
Coverage scope of
|
| # | Site | Memory | Recommendation |
|---|---|---|---|
| 1 | comm_hccl.cpp (a2a3+a5): aclrtMalloc (CommContext) + aclrtMallocPhysical/aclrtReserveMemAddress (VMM comm windows / peer buffers) |
Can be large under distributed serving | Follow-up. Do not force through mem_alloc_ — the VMM reserve/map path is a different lifecycle than rtMalloc and HCCL windows have RDMA alignment/registration constraints. Add a comm_committed_bytes_ counter in the comm layer and sum it at the query layer. |
| 2 | a5 device_runner.cpp prof_alloc_cb = raw rtMalloc for all DFX collectors |
DFX pools | Should fix (small, ideally in this PR). a2a3 already routes DFX through mem_alloc_; a5 just needs the same capturing-lambda alloc_cb/free_cb. Removes the platform inconsistency and matches the docstring. |
| 3 | load_aicpu_op.cpp DeviceBuf::aclrtMalloc (inner-.so upload staging) |
Transient, RAII-freed at init | Leave out. Counting a short-lived init buffer would make the number spike-then-drop and misrepresent steady state. |
| 4 | AICPU-side aicpu_device_malloc (halMemAlloc, device-tier) |
Device-side dynamic alloc | Out of scope (architectural). The AICPU cannot reach the host MemoryAllocator; tracking would need a device-side counter mirrored back via shared memory. No runtime caller found in trb today. |
| 5 | force_reset_device 64-byte probe (device_runner.cpp) |
64 B, immediately freed | Leave out. Health-check poke, not real footprint. |
Framing recommendation: treat committed_device_memory() as "sum of simpler's committed-HBM counters" rather than "one allocator's byte count". Today that sum has one term (mem_alloc_); #1 adds a comm term and #4 could add a device-side term later. For this PR the minimal asks are: (a) do #2 or at least document the a2a3/a5 asymmetry, and (b) have the docstring state that comm/VMM windows are excluded (already partly there) and that the value is a live snapshot, not a peak — so downstream (#1457) knows it must add comm accounting and headroom itself until the follow-up lands.
Two smaller items still open on the latest revision1. (should fix) The L2 sim test doesn't verify the arena is counted — the one path #1457 actually cares about.
baseline = worker.committed_device_memory()
assert isinstance(baseline, int)
assert baseline >= 0 # arenas are committed during initThe comment claims "arenas are committed during init", but 2. (consider) The L3
if not self._chip_shms:
raise NotImplementedError("committed_device_memory is not yet aggregated across level-3 forked workers")
self._check_chip_worker_id(worker_id)
return int(self._orch.committed_device_memory(worker_id))The exception fires only when there are no forked chip workers (nothing to aggregate). When there are children — the case where aggregation would actually matter — it silently returns a single worker's bytes (the docstring already tells the caller to "sum across worker_ids"). So the wording is backwards relative to the behavior. Suggest something like |
7821911 to
eb6773b
Compare
Downstream runtimes size their caches from `aclrtGetMemInfo`, which can undercount simpler's own device HBM (the rtMalloc-backed ring-heap arenas are invisible to it). Expose simpler's authoritative committed-HBM total so consumers can subtract it explicitly. Closes hw-native-sys#1457. The value is plumbed through all three tiers: - L1 (MemoryAllocator): replace the pointer set with a ptr->size map and a running committed-bytes counter, exposed as `committed_bytes()`. - L2 (per-chip): `DeviceRunnerBase::committed_device_memory()` surfaces the counter over the C ABI (`committed_device_memory_ctx`) and through `ChipWorker`, with a nanobind read-only property. - L3 (hierarchical): a `control_committed_device_memory` control op mirrors `control_malloc` across the endpoint hierarchy (`Orchestrator` + `worker_manager`); Python adds the `_CTRL_COMMITTED_DEVICE_MEMORY` child-handler, and the `Worker` facade reads in-process at level 2 and forwards to the forked chip child at level 3. Scope: counts memory routed through the host `MemoryAllocator` (user tensors + pooled arenas + runtime buffers); excludes HCCL/VMM comm windows. Tests: - tests/ut/cpp/common/test_memory_allocator.cpp — C++ byte-counter unit test - tests/ut/py/test_worker/test_committed_memory_sim.py — Python sim, L2 query and L3 not-implemented guard
…evice_memory counts them (#1602) The a5 profiling collectors (l2_swimlane, args_dump, pmu, scope_stats, dep_gen) allocated their device buffers via raw rtMalloc (prof_alloc_cb / prof_free_cb), bypassing the per-runner MemoryAllocator. As a result their HBM was invisible to committed_device_memory() (added in #1592), while a2a3 already routes DFX through mem_alloc_ and counts it — a platform-inconsistent gap the #1592 docstring ("excludes HCCL/VMM comm windows", implying profiling buffers ARE counted) now describes correctly for both platforms. Replace the raw-rtMalloc callbacks with capturing `alloc_cb` / `free_cb` lambdas that call mem_alloc_.alloc / mem_alloc_.free, matching a2a3's init_pmu / init_dep_gen / init_scope_stats. Both allocators use the same rtMalloc(RT_MEMORY_HBM) underneath, so the device-buffer properties are unchanged; the only effect is that DFX buffers are now tracked and counted. Completes the a5 coverage deferred from #1592 to keep that PR single-purpose. Onboard a5 DFX scene tests (test_pmu / test_dep_gen / test_scope_stats / test_args_dump) exercise the new alloc path.
2 participants
L1: MemoryAllocator tracks committed bytes (ptr->size map + counter). L2: ChipWorker.committed_device_memory property + C ABI committed_device_memory_ctx. L3: control_committed_device_memory control op (mirrors control_malloc across the endpoint hierarchy) + Orchestrator::committed_device_memory + nanobind binding; Python _CTRL_COMMITTED_DEVICE_MEMORY, child-handler, Worker facade L3 branch, orchestrator wrapper. Lets downstream runtimes read simpler's authoritative HBM total (which aclrtGetMemInfo may undercount).
Tests: tests/ut/cpp/common/test_memory_allocator.cpp (C++), tests/ut/py/test_worker/test_committed_memory_sim.py (Python sim).
Closes #1457