[Core] Add TPU utility functions to support slice placement groups

[chiayi]

Why are these changes needed?

This PR adds to the utility library for TPU slice placement group scheduling. We generalize the 2 phase approach that the JaxTrainer uses to reserve and schedule the workers on the TPU slice.

Related issue number

#55162

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

[nrghosh]

thanks for the contribution @chiayi - just a reminder to run pre-commit linting to unblock CI checks. ./ci/lint/lint.sh code_format and ./ci/lint/lint.sh pre_commit should do it.

[chiayi]

After building the ray image, I manually ran the slice placement group against v6e 4x4 multihost tpus using the following testcode:

import ray
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
from ray.util.accelerators.tpu import SlicePlacementGroup

ray.init()
slice_handle = SlicePlacementGroup(topology="4x4", accelerator_version="v6e")
slice_pg = slice_handle.placement_group
ray.get(slice_pg.ready(), timeout=10)
            
@ray.remote(num_cpus=0, resources={'TPU': 4})
def spmd_task(world, rank):
    print(f"Current TPU is rank {rank} of {world}")
    tasks = [
        spmd_task.options(
            scheduling_strategy=PlacementGroupSchedulingStrategy(
                placement_group=slice_pg,
            )
        ).remote(world=4, rank=i)
        for i in range(slice_handle.num_workers)]

[chiayi]

@ryanaoleary Please take a look when you get the chance. Thank you!

[ryanaoleary]

cc: @MengjinYan @edoakes

[ryanaoleary]

cc: @andrewsykim @liulehui @matthewdeng for review, for the JaxTrainer we'd call slice_placement_group or multi_slice_placement_group when creating the placement group for the training function. We'll need to add a new argument to the ScalingConfig for num_slices which we'd pass to these util helpers.

[liulehui]

cc: @andrewsykim @liulehui @matthewdeng for review, for the JaxTrainer we'd call slice_placement_group or multi_slice_placement_group when creating the placement group for the training function. We'll need to add a new argument to the ScalingConfig for num_slices which we'd pass to these util helpers.

I feel the for the Ray Train side, the ScalingConfig params list is getting a bit long, e.g. the topology and num_slice is TPU specific parameters, GPU user shouldn't worry about them.

Wondering if we should create something like AcceleratorConfig and start aggregate the terminology like:

GPUAcceleratorConfig(accelerator_type="T4")
GPUAcceleratorConfig(accelerator_type="A100")
TPUAcceleratorConfig(accelerator_type="TPU-V6E", topology="4x4", num_slices = 1)
TPUAcceleratorConfig(accelerator_type="TPU-V6E", topology="4x4", num_slices = 2)

then maybe get rid of the use_gpu and use_tpu.

cc @matthewdeng

[ryanaoleary]

cc: @andrewsykim @liulehui @matthewdeng for review, for the JaxTrainer we'd call slice_placement_group or multi_slice_placement_group when creating the placement group for the training function. We'll need to add a new argument to the ScalingConfig for num_slices which we'd pass to these util helpers.

I feel the for the Ray Train side, the ScalingConfig params list is getting a bit long, e.g. the topology and num_slice is TPU specific parameters, GPU user shouldn't worry about them.

Wondering if we should create something like AcceleratorConfig and start aggregate the terminology like:

GPUAcceleratorConfig(accelerator_type="T4")
GPUAcceleratorConfig(accelerator_type="A100")
TPUAcceleratorConfig(accelerator_type="TPU-V6E", topology="4x4", num_slices = 1)
TPUAcceleratorConfig(accelerator_type="TPU-V6E", topology="4x4", num_slices = 2)

then maybe get rid of the use_gpu and use_tpu.

cc @matthewdeng

That sounds good to me, in the follow-up change to this PR that adds the multi-slice API to the Ray Train side I could add it in a new TPUAcceleratorConfig, and then if we need to add additional fields in the future the TPU specific logic can be contained there.

[edoakes]

do these match the existing accelerator type resource & label names? @ryanaoleary

[ryanaoleary]

All do except for "Accelerator version" which is just the TPU generation:

- Accelerator type -> ray.io/accelerator-type
- Accelerator version -> No exact equivalent label is set, but we can derive it from the other info
- Pod type -> ray.io/tpu-pod-type
- Accelerator topology -> ray.io/tpu-topology

[ryanaoleary]

Oh and yeah in the example the accelerator type names match:

ray/python/ray/util/accelerators/accelerators.py

Line 36 in 4130e4d

GOOGLE_TPU_V6E = "TPU-V6E"

.

[andrewsykim]

I suggest moving SlicePlacementGroup API (and any future scheduling API) to python/ray/util/tpu and keeping ‎python/ray/util/accelerators/tpu.py for small helper functions for getting specific information about TPUs.

[ryanaoleary]

Slightly confused about what the structure should be, there's currently 2 locations with a tpu.py:

• _private/accelerators/tpu.py: This is where helper functions and those called internally by other libraries are (i.e. TPUAcceleratorManager, reserve_tpu_slice, etc.)
  -util/accelerators/tpu.py: This is where the public APIs currently are implemented.

Should the structure be:

• _private/accelerators/tpu.py: internal functions and helpers
• util/accelerators/tpu.py: public APIs but only those like get_current_pod_name which are smaller helpers for getting info about the accelerator
• util/tpu.py: This would be a new file where we implement the APIs related to TPUs and scheduling. I should probably move some functions like reserve_tpu_slice here.

I can change it to the above, my only concern would be that the logic is getting kind of spread out/complicated.

[andrewsykim]

If possible, we should combine everything in util/accelerators/tpu.py into util/tpu.py? Those are all public alpha APIs. If we want to avoid breaking we could have the same functions in util/accelerators/tpu.py call util/tpu.py

_private/accelerators/tpu.py seems fine as it's for internal calls for TPUs as you mentioned

[ryanaoleary]

Done in a23aed9

[edoakes]

LGTM. As next step, should we add a user guide to the Ray docs about how to use Ray w/ TPUs?

[cursor]

Bug: TPU Reservation Logic Inconsistency

The reserve_tpu_slice function creates a placement group to reserve the TPU slice head and returns only the slice name, but the placement group reference is lost. This orphaned placement group will hold resources but won't be tracked. The subsequent placement group created at lines 183-189 doesn't include the head resource reservation that reserve_tpu_slice was supposed to provide, potentially causing scheduling issues. The function should either return both the slice name and placement group reference, or the reservation logic should be integrated directly into _reserve_slice.

 

[ryanaoleary]

This is covered in a TODO comment, I plan to update it in the PR to add multi-slice to the JaxTrainer since it will require editing the call there.

[ryanaoleary]

LGTM. As next step, should we add a user guide to the Ray docs about how to use Ray w/ TPUs?

Yeah that sounds good to me, I can write up a guide on the new API. There's this existing overview of TPUs with KubeRay: https://docs.ray.io/en/latest/cluster/kubernetes/user-guides/tpu.html, and I could add something similar but for the util library / Ray core support under something like Ray Core > User Guides > Advanced Topics.