[AdarshK15]

Add support for TPUs on Slurm (tpu7x-standard-4t)

Summary

Adds support for running Cloud TPU VMs as Slurm compute nodes through the schedmd-slurm-gcp-v6-nodeset module. A TPU slice is declared like any other nodeset by setting machine_type and accelerator_topology, and the cluster brings the nodes up, schedules jobs on them, and tears them down.

TPU VMs are created as a GCE Managed Instance Group attached to a Workload Policy that pins the slice topology, so the VMs come up as one wired TPU slice instead of unrelated VMs.

Supports both static and dynamic (autoscaled) nodes (through schedmd-slurm-gcp-v6-nodeset).

Scope

• The old schedmd-slurm-gcp-v6-nodeset-tpu module is not touched. Existing deployments on it keep working.
• No CPU or GPU behaviour changes. Every new path is behind a TPU check, so a cluster with no TPU nodeset generates the same config as before.
• New config generation is in a new file, conf_v2605.py, used only on Slurm 26.05+. Older versions keep the existing generators.

---

What's supported

Machine type	tpu7x-standard-4t (4 chips per VM)
Slurm version	26.05 and newer
Nodeset modes	Static and dynamic
Provisioning	GCE MIG + Workload Policy

Only machine types ending in -4t are accepted. A TPU machine type that does not end in -4t fails at gcluster deploy with a clear error, before any VMs are created.

Four preconditions added (three in the nodeset module, one in the partition module):

• TPU machine type must end in -4t
• a static TPU nodeset must set a valid accelerator_topology, and node_count_static must be an integer multiple of the slice VM count (total_chips / 4)
• a nodeset cannot mix static and dynamic TPU nodes
• a partition cannot mix static and dynamic TPU nodesets

---

Usage

- id: tpu_nodeset
  source: community/modules/compute/schedmd-slurm-gcp-v6-nodeset
  use: [slurm-network]
  settings:
    machine_type: tpu7x-standard-4t
    accelerator_topology: 2x4x4
    node_count_static: 8
    node_count_dynamic_max: 0
    reservation_name: <RESERVATION_NAME>

srun -N 8 -p tpu --layout=tpu7x=2x4x4 python my_job.py

--layout tells Slurm the slice shape.

---

How provisioning works

1. Chips, Cores, and Slice Sizing:

   • Each tpu7x-standard-4t VM has 4 physical TPU chips (Gres=tpu:tpu7x:4), where each dual-core chip exposes 2 TensorCores (8 VFIO devices /dev/vfio/devices/vfio0..7 per VM).
   • Slurm GRES and GCP topologies (2x4x4 = 32 chips) count physical chips (32 / 4 = 8 VMs per slice) so cli_filter/tpu can validate that -N matches --layout. Setting ConstrainDevices=no in cgroup.conf then allows JAX to open all 8 VFIO TensorCore devices (vfio0..7).
   • Each slice maps 1:1 to a dedicated GCE MIG (targetSizePolicy=BULK) attached to a regional HIGH_THROUGHPUT Workload Policy (acceleratorTopology). If node_count_static spans multiple slices (e.g., 16 VMs with 2x4x4 = 8 VMs/slice), resume.py chunks the nodes into two separate 8-VM MIGs ([0-7] and [8-15]) and registers each slice as a block in cloud_topology.yaml (tpu-block-topo).

2. Static vs. Dynamic Partitions:

   • Static (node_count_static > 0): Nodes get a fixed topology feature (Feature=tpu7x_2x4x4) and PowerDownOnIdle=NO. slurmsync.py powers up all static slices automatically at controller boot and keeps them idle between jobs.
   • Dynamic (node_count_dynamic_max > 0): Nodes start powered down (idle~) with Feature=tpu_dynamic and PowerDownOnIdle=YES. When a job is submitted with --layout=tpu7x=<topology>, job_submit/tpu passes that topology to resume.py via RESUME_FILE, which provisions an on-demand MIG of that shape and deletes it as soon as the job finishes.

Resiliency & Failure Handling

• Whole-Slice Atomicity (resume.py, mig_flex.py):
  • A TPU slice is all-or-nothing. If a static resume request does not cover every node in a slice, the resume is skipped and the slice is forced to POWER_DOWN_FORCE rather than booting a partial slice.
  • Conversely, because a slice cannot be partially deleted, suspending any node in a slice deletes the entire MIG and transitions all peer nodes in that slice to POWER_DOWN_FORCE.
• Fast-Fail & User Notification on Invalid or Missing --layout (resume.py & mig_flex.py):
  • When a job is submitted to a dynamic TPU partition without --layout or with an incompatible topology shape (e.g., 1x2x4 on tpu7x-standard-4t), resume.py extracts the human-readable GCE error message, delivers it directly to the active srun/sbatch job via scontrol notify and admincomment before node teardown, cancels the invalid job (scancel), and immediately resets the dynamic TPU nodes (state=down -> state=power_down) back to idle~ instead of hanging until ResumeTimeout.
• Atomic MIG Rollback & Stale Dynamic MIG Cleanup (mig_flex.py):
  • If createInstances fails after the MIG is created (e.g., Spot stockout), _resume_single_tpu_node deletes the MIG (_delete_tpu_mig) before re-raising, so a failed slice does not leave a MIG behind and block subsequent retries.
  • Before creating a new slice on a dynamic TPU nodeset, _resume_single_tpu_node detects and deletes any leftover stale MIG from previously cancelled or interrupted jobs instead of skipping provisioning.
• Compact MIG Descriptions & Deterministic Topology Blocks (mig_flex.py, conf_v2605.py):
  • Slice membership is stored in the MIG description as a compact Slurm hostlist (slurm_nodes:cluster-nodeset-[0-7]) rather than a comma-separated list, to stay under GCE's 2048-character limit. Each slice gets a deterministic block_uuid in cloud_topology.yaml.

---

Key Changes

Nodeset module

File	Why
schedmd-slurm-gcp-v6-nodeset/main.tf	Adds local.is_tpu and local.tpu_slice_vms (slice VM count derived from the 3D topology), used by the preconditions below.
schedmd-slurm-gcp-v6-nodeset/outputs.tf	Relax the existing accelerator_topology precondition, which assumed a GPU machine type and failed for TPUs, and add the three TPU preconditions above.
schedmd-slurm-gcp-v6-partition/main.tf	Extends the existing TPU SuspendTimeout (240s) / ResumeTimeout (600s) defaults to cover GCE TPU nodesets. Previously these applied only to legacy nodeset_tpu. Non-TPU partitions keep 120s / 300s.
schedmd-slurm-gcp-v6-partition/outputs.tf	Precondition: a partition cannot mix static and dynamic TPU nodesets.

Controller scripts

File	Why
conf_v2605.py (new)	Generates the Slurm config for 26.05+. Added conditional TPU config changes so nothing else is affected.
conf.py	Picks the new generator on Slurm 26.05+, and warns if TPU nodesets are configured on an older version.
util.py	Machine type to TPU generation mapping, slice size maths, and the is_tpu_* helpers that gate every new path.
mig_flex.py	Creates the Workload Policy and MIG for a slice, deletes them on suspend, handles partial slices.
resume.py, suspend.py, slurmsync.py	Send TPU nodes down the TPU resume and suspend paths, take the slice topology from the job's --layout when there is one, and clean up orphaned TPU MIGs and workload policies.
slurm_files/main.tf	Ships conf_v2605.py to the nodes.
cleanup_compute.sh	Deletes TPU -wp-* Workload Policies and -copy-* Instance Templates during gcluster destroy.

Examples and docs

File	Why
tpu7x-standard-4t/tpu7x-slurm-blueprint.yaml	Reference blueprint: builds the Slurm TPU image with Packer, then deploys static and dynamic TPU partitions, Filestore /home, login and controller. Also used by the integration test.
tpu7x-slurm-deployment.yaml, README.md	Deployment vars template and user guide.
examples/machine-learning/README.md	Lists the new blueprint.

Tests

File	Why
tests/test_util.py, test_conf.py, test_resume.py	Unit tests for the machine type mapping, slice size maths, generated config and the resume path.
daily-tests/builds/ml-tpu7x-onspot-slurm.yaml, tests/ml-tpu7x-onspot-slurm.yml	New Cloud Build daily integration test. Runs on Spot with a single node 2x2x1 slice to keep it cheap: builds the image, deploys, validates, destroys.
test-validation/test-tpu-slurm.yml	Validation playbook: TPU plugins and GRES in scontrol show config, static partition (tpu-block-topo, PowerDownOnIdle=NO), dynamic partition (PowerDownOnIdle=YES), node GRES and features, VFIO devices visible inside an srun job, on-demand MIG creation on the dynamic partition via --layout, power-down after the job, and a JAX job through sbatch.
test-infra-kueue/configs/kueue-setup.yaml, dummy-device-plugin.yaml	Test-lock entries for the new build, so it does not run concurrently with itself.

---

Verification

Deployed an 8 node 2x4x4 slice (32 chips) of tpu7x-standard-4t.

The MIG and Workload Policy are created from the topology, from the controller
resume.log:

INFO: resume tpu7xslurm-tpunodeset-[0-7]
INFO: Creating TPU Workload Policy: tpu7xslurm-slurmgcp-tpunodeset-wp-2-4-4
INFO: Creating Zonal TPU MIG tpu7xslurm-tpunodeset-.. in zone us-central1-c
INFO: TPU MIG creation of tpu7xslurm-tpunodeset-.. succeeded.
INFO: Resizing TPU MIG tpu7xslurm-tpunodeset-.. to 8 instances
INFO: TPU MIG and Resize Request for tpu7xslurm-tpunodeset-.. succeeded.

The TPU plugins are loaded and the partition uses TPU block topology:

$ sinfo -V
slurm 26.05.0

$ scontrol show config | grep -E "GresTypes|JobSubmitPlugins|CliFilterPlugins"
CliFilterPlugins        = tpu
GresTypes               = tpu
JobSubmitPlugins        = tpu

$ scontrol show partition tpu
PartitionName=tpu
   NodeSets=tpunodeset
   Nodes=tpu7xslurm-tpunodeset-[0-7]
   Topology=tpu-block-topo
   State=UP TotalCPUs=896 TotalNodes=8

Generated config on the controller:

NodeName=tpu7xslurm-tpunodeset-[0-7] State=CLOUD RealMemory=953840 CPUs=112 Gres=tpu:tpu7x:4 Feature=tpu7x_2x4x4

Nodes register with the GRES and the topology feature:

$ scontrol show node tpu7xslurm-tpunodeset-0
NodeName=tpu7xslurm-tpunodeset-0 Arch=x86_64 CoresPerSocket=112
   AvailableFeatures=tpu7x_2x4x4
   ActiveFeatures=tpu7x_2x4x4
   Gres=tpu:tpu7x:4
   NodeAddr=192.168.64.11 NodeHostName=tpu7xslurm-tpunodeset-0 Version=26.05.0
   State=IDLE+CLOUD ThreadsPerCore=1 TmpDisk=0 Weight=1
   Partitions=tpu

A distributed JAX job across the whole slice, from a shared venv on Filestore /home:

$ srun -N 8 -p tpu --gres=tpu:tpu7x:4 --layout=tpu7x=2x4x4 ~/jax-env/bin/python3 -c \
  "import jax; jax.distributed.initialize(); print(f'Rank {jax.process_index()} ({jax.local_devices()[0].device_kind}): {len(jax.devices())} total TPU chips ({len(jax.local_devices())} local)')"
Rank 0 (TPU7x): 64 total TPU chips (8 local)
Rank 1 (TPU7x): 64 total TPU chips (8 local)
...
Rank 7 (TPU7x): 64 total TPU chips (8 local)

All 8 ranks joined one job and each sees the full slice. The numbers printed are JAX devices (TensorCores), 8 per VM across 8 VMs, on 32 physical chips. That matches the tpu7x-64 accelerator type GCE reports.

TPU devices are visible inside the job on every node. This run also shows the warning when --layout is left out:

$ srun -N 8 -p tpu --gres=tpu:tpu7x:4 ls -la /dev/vfio/devices/
srun: cli_filter/tpu: _process_layout: Warning: TPU GRES detected but --layout unspecified.
srun: cli_filter/tpu: _process_layout: Implications: NodesPerBlock limit will not be enforced, and hardware topology may not be optimized.
crw------- 1 root root 238, 0 Sep 19 15:41 vfio0
crw------- 1 root root 238, 1 Sep 19 15:41 vfio1
...
crw------- 1 root root 238, 7 Sep 19 15:41 vfio7

GCE metadata on tpu7xslurm-tpunodeset-1 confirms the VMs are one slice:

In use by:                 tpu7xslurm-tpunodeset-7ab6eae3
accelerator-type:          tpu7x-64
worker-network-endpoints:  tpu7xslurm-tpunodeset-0,...,tpu7xslurm-tpunodeset-7
tpu-env:
  TOPOLOGY: '2x4x4'
  TPU_WORKER_ID: '1'
  CHIPS_PER_HOST_BOUNDS: '2,2,1'

For Static nodeset, we can have multiple slices per partition, but fails if total partition limit is not a multiple of Slice VMs.

tpu7xslurm-slurm-login-001:~$ sinfo
PARTITION AVAIL  TIMELIMIT  NODES  STATE NODELIST
tpu*         up   infinite      8  idle tpu7xslurm-tpunodeset-[0-7] # Two 2x2x4 (4 VM) slices

If static capacity is unavailable for a slice, will reuse the same MIG when slurmsync retries resuming the slice.

2026-09-23 02:15:37,189 INFO: resume t7xspb55d8-tpunodeset-0
2026-09-23 02:15:37,190 INFO: Scaling TPU nodeset: resuming nodes ['t7xspb55d8-tpunodeset-0']
2026-09-23 02:15:37,190 INFO: TPU slice size derived: 1 VMs per slice (4 chips, 4/node)
2026-09-23 02:15:40,579 INFO: TPU node t7xspb55d8-tpunodeset-0 already managed by MIG .../t7xspb55d8-tpunodeset-04eff95b. Skipping.

For Dynamic nodes the topology is taken at runtime during job submission, fails if --layout flag is not specified.

2026-09-23 02:13:40,213 ERROR: Resetting nodes t7xspb55d8-tpudynnodeset-0 to power_down. Reason: GCP TPU
Error: INVALID_FIELD_VALUE: Missing --layout for dynamic TPU nodes t7xspb55d8-tpudynnodeset-0. Jobs on 
dynamic TPU partitions must specify --layout=<tpu_type>=<topology> (e.g., --layout=tpu7x=2x2x2).

When --layout=tpu7x=2x2x2 is provided on tpu-dyn, the slice MIG is created on demand, runs the job, and tears down automatically when idle (PowerDownOnIdle=YES).

Across the 8-node 2x4x4 slice (64 TPU7x TensorCores), we also verified JAX distributed matrix multiplication, 3D ICI pmean / All-Reduce collectives, and 10-step 64-way FSDP MaxText (llama3-8b) pre-training with loss convergence.

---

[gemini-code-assist]

Summary of Changes

Hello, 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 introduces native support for Cloud TPU VMs (specifically the tpu7x-standard-4t family) within Slurm clusters managed by the Cluster Toolkit. By implementing a dedicated configuration generator for Slurm 26.05+ and leveraging GCE Managed Instance Groups combined with Workload Policies, the changes allow for the automated provisioning and teardown of TPU slices. The update includes robust support for both static and dynamic nodes, along with new blueprints and validation tests to ensure cluster stability and performance.

Highlights

• TPU 7x Support: Added native support for Cloud TPU 7x machine types, enabling them to function as Slurm compute nodes.
• Configuration Generation: Implemented a new Slurm configuration generator for version 26.05+ to handle TPU-specific requirements.
• Provisioning Integration: Integrated GCE Managed Instance Groups (MIGs) and Workload Policies to manage the lifecycle of TPU slices.
• Testing and Documentation: Added comprehensive integration tests, reference blueprints, and documentation for TPU 7x deployments.

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize the 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 counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. 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. ↩

[AdarshK15]

/gemini review

[AdarshK15]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces native support for GCE TPU compute nodes (such as TPU 7x) in Slurm 26.05 clusters, updating Terraform modules, Slurm configuration generators, and node scaling scripts (resume/suspend/sync) to handle TPU topologies and workload policies. Feedback on the changes highlights a critical missing import of the math module in util.py that would cause a runtime crash, a case-sensitivity issue when parsing accelerator topologies in Terraform, and potential resource leaks from silently ignoring HTTP 400 errors during TPU MIG deletion.

[LAVEEN]

LGTM

[LAVEEN]

LGTM