[Core][TPU] Support physical chip mask translation for multi-device TPU chips

[ryanaoleary]

Description

On GKE, the TPU device plugin explicitly injects TPU_VISIBLE_CHIPS="0,1,2,3" for privileged containers (securityContext.privileged: true), while leaving it unset for non-privileged containers. On dual-device TPUs (TPU7x), this injected physical chip mask causes Ray to clamp the node to 4 logical resources instead of the 8 detected from hardware (or explicitly set with`--resources='{"TPU": 8}', causing startup validation failures.

Previously, TPUAcceleratorManager handled TPU_VISIBLE_CHIPS without translation:

1. get_current_process_visible_accelerator_ids() returned physical chip counts resource accounting to fail when auto-detecting or explicitly requesting all logical devices.
2. set_current_process_visible_accelerator_ids() passed logical device IDs directly into TPU_VISIBLE_CHIPS, setting out-of-bounds physical chip indices for worker processes.

This PR changes:

• get_current_process_visible_accelerator_ids(): Expands physical chip indices to logical
  device IDs (e.g. ["0", "1", "2", "3"] -> ["0", ..., "7"]).
• set_current_process_visible_accelerator_ids(): Maps logical device IDs back to physical
  chip indices (device_id // resource_per_chip) before exporting TPU_VISIBLE_CHIPS
• get_tpu_resource_per_chip(): A single validated accessor for RAY_TPU_RESOURCE_PER_CHIP,
  replacing five copies of int(os.environ.get(...)) in ray/util/tpu.py where a
  non-positive value silently produced zero TPU resources per slice.
• normalize_tpu_accelerator_type(): A single helper for the canonical v{gen} spelling,
  replacing three inlined variants. This also fixes is_valid_tpu_accelerator_topology(),
  which split TPU-V6E on - before parsing the generation and so rejected every
  TPU--prefixed version.

Both masking directions are gated on RAY_TPU_RESOURCE_PER_CHIP (default 1), so
per-device accounting is opt-in and the default host-level behavior is unchanged.

The above change is required to get TPU_VISIBLE_CHIPS set correctly for both legacy SPMD behavior (workloads request all TPU chips on a host) and with frameworks that run a process per TPU device (i.e. PyTorch).

Related issues

Contributes to: #64055

Additional information

There is some overlap in this PR with: #65221 in terms of the RAY_TPU_RESOURCE_PER_CHIP_ENV_VAR change and getter.

[gemini-code-assist]

Code Review

This pull request introduces support for multi-core TPUs in Ray by mapping physical TPU chips to logical device IDs. It adds helper functions to determine TPU resources per chip and updates the process-visible accelerator ID configuration to handle logical-to-physical mapping. The review feedback highlights a potential TypeError when sorting mixed-type physical chip IDs, dead code and a normalization bug in TPU type resolution, and suggests adding test coverage for mixed-type sorting to prevent regressions.

[cursor]

Stale Bugbot comment from a previous run.

[ryanaoleary]

cc: @andrewsykim for review if you get time, this fixes the privileged container issue with TPU

[andrewsykim]

isn't v5p dual chip too?

[ryanaoleary]

v5p is dual-core per chip but they're fused since v5p follows megacore architecture (https://docs.jax.dev/en/latest/pallas/tpu/pipelining.html#tpus-in-megacore-configuration) like v4, when I list the devices on the TPU node I only see 1 per chip

[ryanaoleary]

more context here: #65221 (comment)

[ryanaoleary]

I removed DUAL_DEVICE_TPU_TYPES from the PR actually since we shouldn't be auto-detecting dual device types - the behavior needs to be opt/in or it'll break legacy JAX/SPMD behavior. Users specify RAY_TPU_RESOURCE_PER_CHIP and the total TPU resource on the node, and then Slice PG and the accelerator manager ensure the env vars are mapped correctly, etc.

[andrewsykim]

From AI review:

Problem: In set_current_process_visible_accelerator_ids, logical devices are mapped to physical chips using dev_idx // devices_per_chip:

If Actor A requests TPU: 1 and is assigned logical device 0, 0 // 2 = 0 
→
→ sets TPU_VISIBLE_CHIPS="0", TPU_CHIPS_PER_HOST_BOUNDS="1,1,1".
If Actor B requests TPU: 1 and is assigned logical device 1, 1 // 2 = 0 
→
→ also sets TPU_VISIBLE_CHIPS="0", TPU_CHIPS_PER_HOST_BOUNDS="1,1,1".

WDYT?

[ryanaoleary]

I don't think this was a valid issue - TPU_VISIBLE_CHIPS is chip granular and it's actually intended for actors/logical devices to share a chip in the tpu7x case where each chip has 2 separate PJRT devices and frameworks like TorchTPU will run a process per device.
7d8ccab

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit 1a14ce7. Configure here.}