Refactor: move comm-domain demos to scene tests - #1677
Conversation
- Keep each communication-domain orchestration beside its kernels - Replace hand-built Workers and tolerance loops with cases and goldens - Point direct Worker API readers to the L3 allreduce walkthrough Fixes hw-native-sys#1616
📝 WalkthroughWalkthroughThe pull request converts seven communication-domain demos from standalone Worker and CLI runners to Changesa2a3 SceneTestCase migration
a5 SceneTestCase migration
Routing and example documentation
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 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 |
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
![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
🧹 Nitpick comments (2)
examples/a2a3/tensormap_and_ringbuffer/deferred_notify_demo/test_deferred_notify_demo.py (1)
104-117: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffDuplicated
generate_argsandcompute_goldenacross the migrated demos.The same bodies appear in
examples/a2a3/tensormap_and_ringbuffer/async_notify_demo/test_async_notify_demo.py,examples/a2a3/tensormap_and_ringbuffer/sdma_async_completion_demo/test_sdma_async_completion_demo.py, and the four a5 demos. The linked issue asks for shared helper modules and consolidation of near-duplicate a2a3 and a5 implementations. A small shared helper for per-rank tensor specs and peer-rotation goldens would remove six copies. Keep the orchestration functions next to their kernels as the PR describes.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/a2a3/tensormap_and_ringbuffer/deferred_notify_demo/test_deferred_notify_demo.py` around lines 104 - 117, Extract the repeated generate_args and compute_golden logic into a shared helper module for the migrated tensor-map/ringbuffer demos, including per-rank tensor specification creation and peer-rotation golden computation. Update the a2a3 and a5 demo test classes to reuse those helpers while keeping orchestration functions beside their kernels; remove the duplicated method bodies across the affected demos.examples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/test_sdma_async_completion_demo.py (1)
40-40: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winGuard the standalone path against a missing SDMA workspace.
pytest.mark.skipifonly affects pytest collection.SceneTestCase.run_moduleat line 139 ignores pytest markers. If a user runspython test_sdma_async_completion_demo.py -p a5withoutSIMPLER_ENABLE_PTO_SDMA_WORKSPACE, the case runs and fails inside the device path.
examples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/test_urma_deferred_completion_demo.pyalready solves this:_require_urma_workspace_enabled()runs first in the orchestration function at line 49 and raises a clearRuntimeError. Mirror that pattern here.♻️ Proposed guard
+def _require_sdma_workspace_enabled(): + if not _sdma_workspace_enabled(): + raise RuntimeError( + "sdma_async_completion_demo requires host runtime built with " + f"{_SDMA_WORKSPACE_ENV}=ON; set it before rebuilding simpler." + ) + + def sdma_async_completion_orch_fn(orch, callables, task_args, config): + _require_sdma_workspace_enabled() input_nbytes = N * DTYPE_NBYTESAlso applies to: 73-77
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/test_sdma_async_completion_demo.py` at line 40, Update sdma_async_completion_orch_fn to call a dedicated workspace guard before entering the device execution path, mirroring _require_urma_workspace_enabled() in the URMA demo. Ensure the guard checks SIMPLER_ENABLE_PTO_SDMA_WORKSPACE and raises a clear RuntimeError when the workspace is unavailable, covering standalone SceneTestCase.run_module execution as well as pytest.
🤖 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
`@examples/a2a3/tensormap_and_ringbuffer/async_notify_demo/test_async_notify_demo.py`:
- Around line 91-108: Update generate_args to make each rank’s input tensor
rank-dependent by incorporating rank into the per-element value, such as a
rank-based offset plus i % 251, while preserving the existing dtype and tensor
structure. Leave compute_golden unchanged so it continues deriving expected
outputs from each rank’s own input.
In
`@examples/a5/tensormap_and_ringbuffer/async_notify_demo/test_async_notify_demo.py`:
- Line 10: Update the module docstring in test_async_notify_demo.py to describe
only the notification-counter smoke test, removing the “deferred-completion”
wording while preserving the existing onboard a5 context.
---
Nitpick comments:
In
`@examples/a2a3/tensormap_and_ringbuffer/deferred_notify_demo/test_deferred_notify_demo.py`:
- Around line 104-117: Extract the repeated generate_args and compute_golden
logic into a shared helper module for the migrated tensor-map/ringbuffer demos,
including per-rank tensor specification creation and peer-rotation golden
computation. Update the a2a3 and a5 demo test classes to reuse those helpers
while keeping orchestration functions beside their kernels; remove the
duplicated method bodies across the affected demos.
In
`@examples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/test_sdma_async_completion_demo.py`:
- Line 40: Update sdma_async_completion_orch_fn to call a dedicated workspace
guard before entering the device execution path, mirroring
_require_urma_workspace_enabled() in the URMA demo. Ensure the guard checks
SIMPLER_ENABLE_PTO_SDMA_WORKSPACE and raises a clear RuntimeError when the
workspace is unavailable, covering standalone SceneTestCase.run_module execution
as well as pytest.
🪄 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: 4b3ed449-b3f4-40bd-be00-5d0c9ab9fb40
📒 Files selected for processing (12)
docs/capability-survey.mdexamples/a2a3/tensormap_and_ringbuffer/README.mdexamples/a2a3/tensormap_and_ringbuffer/async_notify_demo/test_async_notify_demo.pyexamples/a2a3/tensormap_and_ringbuffer/deferred_notify_demo/test_deferred_notify_demo.pyexamples/a2a3/tensormap_and_ringbuffer/sdma_async_completion_demo/test_sdma_async_completion_demo.pyexamples/a5/tensormap_and_ringbuffer/README.mdexamples/a5/tensormap_and_ringbuffer/async_notify_demo/test_async_notify_demo.pyexamples/a5/tensormap_and_ringbuffer/deferred_notify_demo/test_deferred_notify_demo.pyexamples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/README.mdexamples/a5/tensormap_and_ringbuffer/sdma_async_completion_demo/test_sdma_async_completion_demo.pyexamples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/README.mdexamples/a5/tensormap_and_ringbuffer/urma_deferred_completion_demo/test_urma_deferred_completion_demo.py
| def generate_args(self, params): | ||
| specs = [] | ||
| for rank in range(NRANKS): | ||
| inp = torch.tensor([float(i % 251) / 10.0 for i in range(N)], dtype=torch.float32) | ||
| specs.extend( | ||
| [ | ||
| Tensor(f"in_{rank}", inp), | ||
| Tensor(f"out_{rank}", torch.zeros(N, dtype=torch.float32)), | ||
| Tensor(f"result_{rank}", torch.zeros(N, dtype=torch.float32)), | ||
| ] | ||
| ) | ||
| return TaskArgsBuilder(*specs) | ||
|
|
||
| def main() -> int: | ||
| parser = argparse.ArgumentParser() | ||
| parser.add_argument("-p", "--platform", default="a2a3") | ||
| parser.add_argument("-d", "--device", default="0-1") | ||
| args = parser.parse_args() | ||
| return run(args.platform, parse_device_range(args.device)) | ||
| def compute_golden(self, args, params): | ||
| for rank in range(NRANKS): | ||
| expected_out = getattr(args, f"in_{rank}") * 2.0 | ||
| getattr(args, f"out_{rank}").copy_(expected_out) | ||
| getattr(args, f"result_{rank}").copy_(expected_out + 1.0) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Inspect the async notify kernels to confirm local vs peer data flow.
fd -t f -p 'a2a3/tensormap_and_ringbuffer/async_notify_demo/kernels' -x sed -n '1,200p' {}Repository: hw-native-sys/simpler
Length of output: 10053
Make the per-rank input rank-dependent.
Each rank computes out = in * 2 only from in_{rank}, so a rank-swap can still pass if every rank uses the same input values. Use a rank-offset input, such as rank * 1000 + (i % 251), so the golden fails when data from the wrong rank is used locally.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@examples/a2a3/tensormap_and_ringbuffer/async_notify_demo/test_async_notify_demo.py`
around lines 91 - 108, Update generate_args to make each rank’s input tensor
rank-dependent by incorporating rank into the per-element value, such as a
rank-based offset plus i % 251, while preserving the existing dtype and tensor
structure. Leave compute_golden unchanged so it continues deriving expected
outputs from each rank’s own input.
| # See LICENSE in the root of the software repository for the full text of the License. | ||
| # ----------------------------------------------------------------------------------------------------------- | ||
| """Notification counter + deferred completion smoke test for onboard a5.""" | ||
| """Notification counter and deferred-completion smoke test for onboard a5.""" |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Correct the module docstring.
This module tests the notification-counter path. The phrase "deferred-completion" describes the deferred_notify_demo and sdma_async_completion_demo scenarios. Remove it here to keep the demo docs accurate.
📝 Proposed docstring fix
-"""Notification counter and deferred-completion smoke test for onboard a5."""
+"""Notification counter smoke test for onboard a5."""📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| """Notification counter and deferred-completion smoke test for onboard a5.""" | |
| """Notification counter smoke test for onboard a5.""" |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@examples/a5/tensormap_and_ringbuffer/async_notify_demo/test_async_notify_demo.py`
at line 10, Update the module docstring in test_async_notify_demo.py to describe
only the notification-counter smoke test, removing the “deferred-completion”
wording while preserving the existing onboard a5 context.
1 participant
Summary
Testing
--rounds 2onboard-arch-precheckcannot identify silicon becausenpu-smiDCMI initialization fails with-8005)Fixes #1616