[Core] Add default TPU Ray node labels by ryanaoleary · Pull Request #55064 · ray-project/ray

ryanaoleary · 2025-07-30T21:06:21Z

Why are these changes needed?

This PR adds useful TPU accelerator information to the node labels of a Ray node when TPU resources are detected. These labels will enable fine-grained control over scheduling on TPU devices and support for SPMD workloads. This information is currently added to the resource configuration of Ray nodes as additional custom resources, but this workaround will be deprecated in favor of using the label_selector API.

I added get_current_node_accelerator_labels as a function to the AcceleratorManager abstract class so that we can extend this to other accelerators in the future.

Related issue number

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 :(

Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>

gemini-code-assist

Summary of Changes

Hello @ryanaoleary, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances Ray's ability to manage and schedule workloads on TPU devices by automatically adding detailed TPU-specific information as labels to Ray nodes. This change aims to provide finer-grained control for scheduling TPU workloads, particularly for SPMD (Single Program, Multiple Data) paradigms, and replaces a previous workaround that relied on custom resources for this information.

Highlights

Enhanced TPU Node Labeling: Ray nodes with detected TPU resources will now automatically include detailed labels such as ray.io/tpu-slice-name, ray.io/tpu-worker-id, ray.io/tpu-topology, and ray.io/tpu-head. This provides richer metadata for scheduling.
Improved Scheduling for SPMD Workloads: The addition of these new TPU-specific labels enables more precise scheduling and resource allocation for multi-host TPU workloads, which is crucial for Single Program, Multiple Data (SPMD) execution paradigms.
Refactored TPU Detection Logic: New methods have been introduced within the TPUAcceleratorManager to accurately detect TPU topology information from both environment variables (for GKE) and GCE metadata.
New Constants and Test Coverage: Corresponding C++ and Python constants for the new TPU labels have been defined across the codebase, and comprehensive unit and integration tests have been added to ensure the correct application and resolution of these labels.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

ryanaoleary · 2025-07-30T21:08:07Z

cc: @andrewsykim

gemini-code-assist

Code Review

This pull request introduces default TPU node labels to provide more fine-grained scheduling control. The changes are well-structured, touching the necessary C++ constants, Python backend logic for label resolution, and adding corresponding tests.

My review focuses on improving code maintainability and clarity. I've identified a few areas for improvement:

Removing unused function parameters to simplify method signatures.
Avoiding calls to private methods from other modules to respect encapsulation.
Fixing a minor typo in a comment.

Overall, the changes are solid and the tests provide good coverage for the new functionality. Addressing the feedback will enhance the long-term quality of the code.

python/ray/_private/resource_and_label_spec.py

src/ray/common/constants.h

Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>

ryanaoleary · 2025-07-30T22:18:52Z

cc: @MengjinYan @edoakes @matthewdeng

python/ray/_private/accelerators/tpu.py

Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> Remove leading underscores and update how topology is retrieved from GCE Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>

ryanaoleary · 2025-07-31T06:38:42Z

Ran a manual test to verify the GCE code path works for retrieving the TPU topology:

Created a TPU vm with 2x2x1 topology and SSHed into it
Ran the following test script:

cat gce-tpu-env-test.py 
import re
import subprocess

def get_tpu_topology_from_metadata() -> str:
    try:
        # Run curl command to fetch tpu-env metadata
        result = subprocess.run(
            [
                "curl",
                "-s",
                "-H", "Metadata-Flavor: Google",
                "http://metadata.google.internal/computeMetadata/v1/instance/attributes/tpu-env"
            ],
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        env_blob = result.stdout
        match = re.search(r"TOPOLOGY:\s*'([^']+)'", env_blob)
        if match:
            return match.group(1)
        else:
            print("TOPOLOGY field not found.")
            return None
    except subprocess.CalledProcessError as e:
        print("Error fetching metadata:", e.stderr)
        return None

if __name__ == "__main__":
    topology = get_tpu_topology_from_metadata()
    print("TPU Topology:", topology)

With result:

python gce-tpu-env-test.py 
TPU Topology: 2x2x1

This matches the logic we use to query GCE in tpu.py for the topology information when it can't be determined from the env var populated by GKE.

Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>

python/ray/_private/accelerators/tpu.py

Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>

edoakes · 2025-08-04T20:04:46Z

auto-merge on. ping me if CI fails and need to re-enable.

can make any changes from offline discussion as follow ups

This PR adds useful TPU accelerator information to the node labels of a Ray node when TPU resources are detected. These labels will enable fine-grained control over scheduling on TPU devices and support for SPMD workloads. This information is currently added to the resource configuration of Ray nodes as additional custom resources, but this workaround will be deprecated in favor of using the `label_selector` API. I added `get_current_node_accelerator_labels` as a function to the `AcceleratorManager` abstract class so that we can extend this to other accelerators in the future. --------- Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> Signed-off-by: Kamil Kaczmarek <kamil@anyscale.com>

This PR adds useful TPU accelerator information to the node labels of a Ray node when TPU resources are detected. These labels will enable fine-grained control over scheduling on TPU devices and support for SPMD workloads. This information is currently added to the resource configuration of Ray nodes as additional custom resources, but this workaround will be deprecated in favor of using the `label_selector` API. I added `get_current_node_accelerator_labels` as a function to the `AcceleratorManager` abstract class so that we can extend this to other accelerators in the future. --------- Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> Signed-off-by: Michael Acar <michael.j.acar@gmail.com>

This PR adds useful TPU accelerator information to the node labels of a Ray node when TPU resources are detected. These labels will enable fine-grained control over scheduling on TPU devices and support for SPMD workloads. This information is currently added to the resource configuration of Ray nodes as additional custom resources, but this workaround will be deprecated in favor of using the `label_selector` API. I added `get_current_node_accelerator_labels` as a function to the `AcceleratorManager` abstract class so that we can extend this to other accelerators in the future. --------- Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>

This PR adds useful TPU accelerator information to the node labels of a Ray node when TPU resources are detected. These labels will enable fine-grained control over scheduling on TPU devices and support for SPMD workloads. This information is currently added to the resource configuration of Ray nodes as additional custom resources, but this workaround will be deprecated in favor of using the `label_selector` API. I added `get_current_node_accelerator_labels` as a function to the `AcceleratorManager` abstract class so that we can extend this to other accelerators in the future. --------- Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> Signed-off-by: sampan <sampan@anyscale.com>

This PR adds useful TPU accelerator information to the node labels of a Ray node when TPU resources are detected. These labels will enable fine-grained control over scheduling on TPU devices and support for SPMD workloads. This information is currently added to the resource configuration of Ray nodes as additional custom resources, but this workaround will be deprecated in favor of using the `label_selector` API. I added `get_current_node_accelerator_labels` as a function to the `AcceleratorManager` abstract class so that we can extend this to other accelerators in the future. --------- Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> Signed-off-by: jugalshah291 <shah.jugal291@gmail.com>

This PR adds useful TPU accelerator information to the node labels of a Ray node when TPU resources are detected. These labels will enable fine-grained control over scheduling on TPU devices and support for SPMD workloads. This information is currently added to the resource configuration of Ray nodes as additional custom resources, but this workaround will be deprecated in favor of using the `label_selector` API. I added `get_current_node_accelerator_labels` as a function to the `AcceleratorManager` abstract class so that we can extend this to other accelerators in the future. --------- Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> Signed-off-by: Douglas Strodtman <douglas@anyscale.com>

Add default TPU Ray node labels

75cde77

Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>

gemini-code-assist bot reviewed Jul 30, 2025

View reviewed changes

ryanaoleary added 2 commits July 30, 2025 22:13

respect encapsulation and add tests

e431508

Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>

Fix comments

68e7a3c

Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>

ryanaoleary marked this pull request as ready for review July 30, 2025 22:17

ryanaoleary requested a review from a team as a code owner July 30, 2025 22:17

edoakes reviewed Jul 30, 2025

View reviewed changes

python/ray/_private/accelerators/tpu.py Outdated Show resolved Hide resolved

python/ray/_private/accelerators/tpu.py Outdated Show resolved Hide resolved

python/ray/_private/accelerators/tpu.py Outdated Show resolved Hide resolved

ray-gardener bot added community-contribution Contributed by the community core Issues that should be addressed in Ray Core labels Jul 31, 2025

Remove leading underscores and update how topology is retrieved from GCE

7f4c5d7

Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> Remove leading underscores and update how topology is retrieved from GCE Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>

Replace tpu head label with tpu pod type

6ba20ff

Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>

edoakes reviewed Jul 31, 2025

View reviewed changes

python/ray/_private/accelerators/tpu.py Show resolved Hide resolved

Add topology regex unit test

123bccf

Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>

ryanaoleary mentioned this pull request Aug 2, 2025

[Train] JaxTrainer Implementation Tracking Issue #55162

Open

ryanaoleary requested a review from edoakes August 4, 2025 18:48

edoakes added the go add ONLY when ready to merge, run all tests label Aug 4, 2025

edoakes approved these changes Aug 4, 2025

View reviewed changes

edoakes enabled auto-merge (squash) August 4, 2025 20:04

edoakes merged commit cd9c9c3 into ray-project:master Aug 4, 2025
7 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Core] Add default TPU Ray node labels#55064

[Core] Add default TPU Ray node labels#55064
edoakes merged 6 commits intoray-project:masterfrom
ryanaoleary:tpu-indexing-labels

ryanaoleary commented Jul 30, 2025 •

edited

Loading

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

ryanaoleary commented Jul 30, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ryanaoleary commented Jul 30, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ryanaoleary commented Jul 31, 2025

Uh oh!

Uh oh!

edoakes commented Aug 4, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

ryanaoleary commented Jul 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why are these changes needed?

Related issue number

Checks

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

ryanaoleary commented Jul 30, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ryanaoleary commented Jul 30, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ryanaoleary commented Jul 31, 2025

Uh oh!

Uh oh!

edoakes commented Aug 4, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

ryanaoleary commented Jul 30, 2025 •

edited

Loading