Skip to content

[core][GPU Objects] Add related tests for tensordict #54286

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 6 commits into
base: master
Choose a base branch
from

Conversation

Qiaolin-Yu
Copy link
Member

@Qiaolin-Yu Qiaolin-Yu commented Jul 2, 2025

Why are these changes needed?

The current code already supports extracting tensors from tensordict, and using external transport to transfer the tensors. This PR adds some tests to prove this (e.g., checking the extracted tensor stored in the gpu object manager, and testing nested tensordict).

Related issue number

close #52340

Checks

  • I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
  • I've run scripts/format.sh to lint the changes in this PR.
  • I've included any doc changes needed for https://docs.ray.io/en/master/.
    • I've added any new APIs to the API Reference. For example, if I added a
      method in Tune, I've added it in doc/source/tune/api/ under the
      corresponding .rst file.
  • I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
  • Testing Strategy
    • Unit tests
    • Release tests
    • This PR is not tested :(

@Copilot Copilot AI review requested due to automatic review settings July 2, 2025 18:57
Copy link
Contributor

@Copilot Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull Request Overview

This PR adds tests to validate TensorDict support in GPU object transfers using the Gloo backend.

  • Import TensorDict and extend GPUTestActor.double to handle TensorDict inputs.
  • Enhance test_multiple_tensors to include TensorDict items.
  • Introduce three new tests:
    • test_tensordict_transfer
    • test_nested_tensordict
    • test_tensor_extracted_from_tensordict_in_gpu_object_store
Comments suppressed due to low confidence (1)

python/ray/tests/test_gpu_objects_gloo.py:232

  • The code calls sys.exit in the __main__ guard but sys is not imported at the top of the file. Add import sys to avoid a NameError.
if __name__ == "__main__":

@kevin85421 kevin85421 added go add ONLY when ready to merge, run all tests and removed go add ONLY when ready to merge, run all tests labels Jul 2, 2025
@Qiaolin-Yu Qiaolin-Yu added the go add ONLY when ready to merge, run all tests label Jul 3, 2025
@Qiaolin-Yu
Copy link
Member Author

@jjyao , would you mind merging this?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
go add ONLY when ready to merge, run all tests gpu-objects
Projects
None yet
Development

Successfully merging this pull request may close these issues.

[core][gpu-objects] Ability to register custom types for GPU data
3 participants