[cgraph] Avoid depending on torch CPU module for CPU-only actor#53849
Merged
jjyao merged 2 commits intoray-project:masterfrom Jun 18, 2025
Merged
[cgraph] Avoid depending on torch CPU module for CPU-only actor#53849jjyao merged 2 commits intoray-project:masterfrom
jjyao merged 2 commits intoray-project:masterfrom
Conversation
Signed-off-by: huafengchun <huafengchun@gmail.com>
Contributor
There was a problem hiding this comment.
Pull Request Overview
This PR updates AcceleratorContext to skip importing the torch CPU backend when the device is CPU-only, removing an unnecessary dependency on torch.cpu.
- Conditional import of
torch.<backend>only for non-CPU devices - Eliminates unconditional
torch.cpuimport for CPU-only actors
Comments suppressed due to low confidence (3)
python/ray/experimental/channel/accelerator_context.py:40
- The comment should mention how the CPU case is handled (e.g., that no torch backend is loaded or that
self._torch_modwill beNone). This will clarify expected behavior for maintainers.
# Import the torch backend module (e.g., torch.cuda) if the device is not 'cpu'.
python/ray/experimental/channel/accelerator_context.py:1
- Add a unit test for the CPU-only path to verify that
AcceleratorContextinitializes correctly without importing a torch backend and does not raise errors when accessingself._torch_mod.
def __init__(self, torch_module_name: str, communicator_cls: Type[Communicator])
python/ray/experimental/channel/accelerator_context.py:41
- When
torch_module_nameis "cpu",self._torch_modis never defined, which may lead to an AttributeError later. Consider adding anelseblock to defineself._torch_mod(e.g.,Noneor the basetorchmodule) or refactor usage to handle a missing backend module.
if torch_module_name != "cpu":
Contributor
Author
|
@ruisearch42 @kevin85421 Could you please review this bug fix? Thanks. |
ruisearch42
approved these changes
Jun 17, 2025
Contributor
ruisearch42
left a comment
There was a problem hiding this comment.
Thanks. Let me trigger the tests.
Contributor
Contributor
minerharry
pushed a commit
to minerharry/ray
that referenced
this pull request
Jun 27, 2025
…project#53849) Signed-off-by: huafengchun <huafengchun@gmail.com>
elliot-barn
pushed a commit
that referenced
this pull request
Jul 2, 2025
Signed-off-by: huafengchun <huafengchun@gmail.com> Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>
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
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.
Why are these changes needed?
In cases where the actor involved in the Compiled Graph exclusively uses the CPU, AcceleratorContext unnecessarily loads torch’s CPU module, resulting in an avoidable dependency.
This PR eliminates that dependency by avoiding the torch module loading for CPU-only actors.
Related issue number
Closes #53716
Checks
git commit -s) in this PR.scripts/format.shto lint the changes in this PR.method in Tune, I've added it in
doc/source/tune/api/under thecorresponding
.rstfile.