feat(runtime): L3 DeviceTensor args, reusable prepare() handle, and shared DeviceMemoryHandle#1451
Merged
Merged
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR adds a DeviceMemoryHandle abstraction and DeviceTensor-aware marshalling, updates DistributedCompiledProgram with prepare() to return a reusable DistributedRuntime, refactors the distributed runner for one-shot and prepared execution, updates codegen and docs, and adds unit/integration tests for the new flows. ChangesL3 Distributed Runtime Reuse with Worker-Resident DeviceTensors
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related issues
Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. 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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request implements Distributed (L3+) program reuse via DistributedCompiledProgram.prepare() and adds support for worker-resident DeviceTensor arguments. It introduces a DeviceMemoryHandle base class to unify memory management across runtime levels and refactors the distributed runner for better modularity. Documentation and tests are updated to reflect these changes. Reviewer feedback correctly identified opportunities to improve temporary file handling, reduce coupling with internal APIs, and optimize validation logic.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/en/user/00-getting_started.md`:
- Around line 257-261: The snippet is missing imports causing NameError; add the
required imports at the top so the example is runnable: import torch and from
pypto.runtime import DeviceTensor, then keep the existing lines using
ir.compile(MyDistributedProgram), DeviceTensor(dev_ptr, (1024, 4096),
torch.float16), and compiled(x, weight, out) unchanged so DeviceTensor and
torch.float16 are defined when the snippet executes.
In `@docs/zh-cn/user/00-getting_started.md`:
- Around line 251-255: The snippet uses DeviceTensor and torch.float16 but lacks
imports, causing NameError when copied; update the example so it is runnable by
adding the missing imports (import torch and from pypto.runtime import
DeviceTensor) above the shown lines where ir.compile(MyDistributedProgram),
DeviceTensor(dev_ptr, (1024, 4096), torch.float16) and compiled(...) are used,
so DeviceTensor, torch and related symbols are defined when the snippet is
copied.
In `@python/pypto/runtime/device_tensor.py`:
- Around line 126-128: The alloc_device_tensor function currently coerces
dimensions with int(d) and only rejects negatives, which lets bools, floats,
zeros, and empty shapes slip through; update its validation to mirror
DeviceTensor's contract: do not coerce values–require shape to be a non-empty
iterable, convert to tuple shape_t but validate each dim is an int
(isinstance(d, int) and not isinstance(d, bool)), and ensure each dim is > 0
(reject zeros and negatives); raise TypeError for non-int/bool dims and
ValueError for empty shapes or non-positive dims, and keep the rest of the
allocation logic using shape_t.
In `@python/pypto/runtime/distributed_runner.py`:
- Around line 471-484: Move the setup calls that can fail into the try/finally
so cleanup always runs: keep creating the worker via _construct_worker(...)
where it is, but call _register_callables(...) and w.init() inside the try block
(before calling _dispatch(...)) so that any exception during registration or
init will trigger the finally that calls w.close() and unlinks rootinfo_path;
ensure the order remains _register_callables -> w.init() -> _dispatch and that
the finally still removes rootinfo_path and handles FileNotFoundError.
- Around line 540-568: Make DistributedRuntime.__init__ exception-safe by
wrapping the post-registration setup (after _register_callables returns and
before the constructor returns) in a try/except/finally so partial resources are
cleaned up on error: if any of self._w.init() or self._w._start_hierarchical()
(or subsequent steps) raise, call the worker cleanup method (e.g.
self._w.shutdown() / close / stop whichever exists on the worker) and remove or
close any rootinfo artifacts referenced by self._rootinfo_path, then re-raise
the exception; use the unique symbols _register_callables, self._w,
self._rootinfo_path, self._w.init(), and self._w._start_hierarchical() to locate
where to add the try/finally and cleanup calls.
🪄 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: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 908ca164-d556-4028-bd26-512ac769023f
📒 Files selected for processing (19)
docs/en/user/00-getting_started.mddocs/zh-cn/user/00-getting_started.mdpython/pypto/ir/distributed_compiled_program.pypython/pypto/runtime/__init__.pypython/pypto/runtime/device_memory.pypython/pypto/runtime/device_tensor.pypython/pypto/runtime/distributed_runner.pypython/pypto/runtime/runner.pypython/pypto/runtime/task_interface.pypython/pypto/runtime/tensor_arg.pypython/pypto/runtime/worker.pysrc/codegen/distributed/distributed_codegen.cpptests/st/distributed/test_l3_device_tensor.pytests/st/distributed/test_l3_distributed.pytests/ut/codegen/test_distributed_codegen.pytests/ut/ir/test_distributed_compiled_program.pytests/ut/runtime/test_device_memory.pytests/ut/runtime/test_distributed_runtime.pytests/ut/runtime/test_tensor_arg.py
luohuan19
added a commit
to luohuan19/pypto
that referenced
this pull request
May 21, 2026
- ruff format distributed_runner.py and test_l3_distributed.py - suppress reportAttributeAccessIssue on simpler re-exports (ContinuousTensor/make_tensor_arg/torch_dtype_to_datatype) imported via .task_interface, matching the existing runner.py pattern - use callable MagicMock sentinels for sub_worker_overrides in tests so they satisfy the dict[str, Callable] signature
luohuan19
added a commit
to luohuan19/pypto
that referenced
this pull request
May 21, 2026
- tensor_arg: keep the lazy ``make_tensor_arg`` import as a parenthesized block so the ``# noqa: PLC0415`` stays on the ``from ... import (`` line, surviving ruff isort's re-wrap (fixes the pre-commit ruff failure). - device_tensor: validate ``alloc_device_tensor`` shape up front and without coercion, mirroring DeviceTensor's contract (reject bool/non-int, empty, and non-positive dims) so a wrong logical shape never reaches malloc. - distributed_runner: make the one-shot dispatch path and DistributedRuntime.__init__ exception-safe — construct/register/init now run under cleanup so a setup failure still closes the worker and unlinks the comm rootinfo temp file. - docs: add the missing ``import torch`` / ``DeviceTensor`` imports to the distributed DeviceTensor snippet (en + zh) so it runs as-is.
luohuan19
added a commit
to luohuan19/pypto
that referenced
this pull request
May 21, 2026
…ve-sys#1451 alloc_device_tensor now enforces DeviceTensor's positive-int shape contract, so the pre-malloc rejection message changed from "non-negative" to "positive". Update the test accordingly and add coverage for the newly rejected zero-dim and empty-shape cases (both must fail before malloc).
…/D2H) DistributedCompiledProgram.__call__ now accepts worker-resident DeviceTensor arguments in place of host torch.Tensor, mirroring the L2 CompiledProgram convention: such args skip the H2D/D2H copies and stay caller-managed. Distributed codegen emits a pypto-owned make_tensor_arg (pypto.runtime.tensor_arg) instead of simpler_setup.torch_interop, widening the orchestration conversion to wrap DeviceTensor handles as ContinuousTensor(child_memory=True). The DeviceTensor->ContinuousTensor conversion is extracted into a shared device_tensor_to_continuous helper reused by both the L2 and L3 paths.
Add DistributedCompiledProgram.prepare() returning a reusable DistributedRuntime handle that runs the expensive L3 setup (compile_and_assemble, module load, Worker construction + registration + init/fork) once and dispatches many times on the held Worker. Refactor execute_distributed's setup into shared free helpers (_assemble_chip_callables, _load_orch_entry, _load_sub_worker_fns, _build_chip_bootstrap, _construct_worker, _register_callables, _make_call_config, _dispatch) so the one-shot and reuse paths run identical setup without duplication. The handle exposes device-memory helpers (malloc/copy_to/copy_from/free/ alloc_tensor) for worker-resident DeviceTensor buffers that survive across dispatches; per-call IO uses shared-memory host tensors reused in place. Extract the shared alloc_device_tensor helper into device_tensor.py (reused by L2 Worker.alloc_tensor). Add UT covering the setup-once/dispatch-many contract and per-call validation, rewrite the L3 ST to exercise prepare(), and document prepare() in getting_started (en + zh-cn).
Extract the duplicated alloc_tensor/free_tensor device-tensor surface from Worker (L2) and DistributedRuntime (L3) into a shared DeviceMemoryHandle ABC. The four memory primitives stay abstract; only the genuine per-level differences are hooks: _require_ready (readiness guard) and _prepare_init (host-init upload policy). L2 keeps the defensive contiguous CPU copy; L3 requires shared memory because its upload runs in a forked child. Add DistributedCompiledProgram.prepare(sub_worker_overrides=...) to replace a generated sub-worker placeholder (matched by name) with a caller-supplied callable. Unknown names fail fast via _merge_sub_worker_overrides.
- ruff format distributed_runner.py and test_l3_distributed.py - suppress reportAttributeAccessIssue on simpler re-exports (ContinuousTensor/make_tensor_arg/torch_dtype_to_datatype) imported via .task_interface, matching the existing runner.py pattern - use callable MagicMock sentinels for sub_worker_overrides in tests so they satisfy the dict[str, Callable] signature
- tensor_arg: keep the lazy ``make_tensor_arg`` import as a parenthesized block so the ``# noqa: PLC0415`` stays on the ``from ... import (`` line, surviving ruff isort's re-wrap (fixes the pre-commit ruff failure). - device_tensor: validate ``alloc_device_tensor`` shape up front and without coercion, mirroring DeviceTensor's contract (reject bool/non-int, empty, and non-positive dims) so a wrong logical shape never reaches malloc. - distributed_runner: make the one-shot dispatch path and DistributedRuntime.__init__ exception-safe — construct/register/init now run under cleanup so a setup failure still closes the worker and unlinks the comm rootinfo temp file. - docs: add the missing ``import torch`` / ``DeviceTensor`` imports to the distributed DeviceTensor snippet (en + zh) so it runs as-is.
…ve-sys#1451 alloc_device_tensor now enforces DeviceTensor's positive-int shape contract, so the pre-malloc rejection message changed from "non-negative" to "positive". Update the test accordingly and add coverage for the newly rejected zero-dim and empty-shape cases (both must fail before malloc).
lyfne123
approved these changes
May 25, 2026
This was referenced May 25, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Builds out the L3 distributed runtime so worker-resident device buffers can be reused across dispatches, mirroring the existing L2
CompiledProgramconventions. Three stacked changes:Accept
DeviceTensorargs in L3+ programs —DistributedCompiledProgram.__call__now accepts worker-residentDeviceTensorarguments in place of hosttorch.Tensor. Such args skip the H2D/D2H copies and stay caller-managed, matching L2. Distributed codegen emits a pypto-ownedmake_tensor_argand widens orchestration conversion to wrapDeviceTensorhandles asContinuousTensor(child_memory=True). TheDeviceTensor -> ContinuousTensorconversion is extracted into a shareddevice_tensor_to_continuoushelper reused by both L2 and L3 paths.Reusable
prepare()handle —DistributedCompiledProgram.prepare()returns a reusableDistributedRuntimehandle that runs the expensive L3 setup (compile_and_assemble, module load, Worker construction + registration + init/fork) once and dispatches many times.execute_distributed's setup is refactored into shared free helpers so the one-shot and reuse paths run identical setup without duplication. The handle exposes device-memory helpers (malloc/copy_to/copy_from/free/alloc_tensor) for buffers that survive across dispatches.sub_worker_overrides+ sharedDeviceMemoryHandle— extract the duplicated alloc_tensor/free_tensor device-tensor surface fromWorker(L2) andDistributedRuntime(L3) into a sharedDeviceMemoryHandleABC. The four memory primitives stay abstract; only the genuine per-level differences are hooks (_require_ready,_prepare_init).prepare(sub_worker_overrides=...)replaces a generated sub-worker placeholder (matched by name) with a caller-supplied callable; unknown names fail fast.Testing
getting_starteden + zh-cn documentprepare())New tests:
test_distributed_runtime.py,test_device_memory.py,test_tensor_arg.py,test_distributed_compiled_program.py,test_l3_device_tensor.py; rewritten L3 ST to exerciseprepare().