Skip to content

Add tpu7x support #586

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.

Sign up for GitHub

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

Jump to bottom
Open
wants to merge 1 commit into
base: develop
Choose a base branch
from
Open

Add tpu7x support #586

wants to merge 1 commit into from
+109 −0

Conversation

scaliby
Copy link
Collaborator

@scaliby scaliby commented Aug 11, 2025 •
edited
Loading

Fixes / Features

  • Introduces support for Ironwood TPU's

Testing / Documentation

Testing has been performed in a manual fashion using the tpu7x accelerators. Test commands that suceeded:

python3 xpk.py cluster create \
  --cluster=xpk-test-tpu7x-2 \
  --project=tpu-staging-env-one-vm \
  --tpu-type=tpu7x-2 \
  --zone=us-central2-newvmos \
  --num-nodes=1 \
  --enable-autoprovisioning \
  --custom-cluster-arguments="--cluster-version=1.33.3-gke.1255000" \
  --on-demand \
  --custom-nodepool-arguments="--node-version=1.33.3-gke.1255000"

python3 xpk.py cluster create \
  --cluster=xpk-test-tpu7x-8 \
  --project=tpu-staging-env-one-vm \
  --tpu-type=tpu7x-8 \
  --zone=us-central2-newvmos \
  --num-nodes=1 \
  --enable-autoprovisioning \
  --custom-cluster-arguments="--cluster-version=1.33.3-gke.1255000" \
  --on-demand \
  --custom-nodepool-arguments="--node-version=1.33.3-gke.1255000"
  • [y] Tests pass
  • [y] Appropriate changes to documentation are included in the PR

@scaliby scaliby added the release-features features label Aug 11, 2025
@scaliby scaliby force-pushed the konradkaim/tpu7x-support branch 2 times, most recently from ad98198 to 2deced2 Compare August 11, 2025 12:52
@scaliby scaliby force-pushed the konradkaim/tpu7x-support branch from 2deced2 to 2bda70a Compare August 20, 2025 13:01
@scaliby scaliby changed the title WIP: add tpu7x support Add tpu7x support Aug 21, 2025
@scaliby scaliby force-pushed the konradkaim/tpu7x-support branch from 2bda70a to 86573a3 Compare August 21, 2025 08:25
@scaliby scaliby marked this pull request as ready for review August 21, 2025 08:25
sharabiani
sharabiani reviewed Aug 21, 2025
View reviewed changes
src/xpk/core/system_characteristics.py
@@ -242,6 +242,115 @@ def get_tpu_system_characteristics_map(
'h100-mega-80gb-8',
),
# TPU system characteristics
**get_tpu_system_characteristics_map(
'tpu7x', 2, 'tpu7x', 'tpu7x-standard-1t', ['1x1x1']
Copy link
Collaborator

Choose a reason for hiding this comment

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

Is "2" correct here for the number of cores in tpu7x-standard-1t ?

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Yes, I explicitly tried this configuration and it provisioned cluster correctly.

@sharabiani
Copy link
Collaborator

Fixes / Features

  • Introduces support for Ironwood TPU's

Testing / Documentation

Testing has been performed in a manual fashion using the tpu7x accelerators. Test commands that suceeded:

python3 xpk.py cluster create \
  --cluster=xpk-test-tpu7x-2 \
  --project=tpu-staging-env-one-vm \
  --tpu-type=tpu7x-2 \
  --zone=us-central2-newvmos \
  --num-nodes=1 \
  --enable-autoprovisioning \
  --custom-cluster-arguments="--cluster-version=1.33.3-gke.1255000" \
  --on-demand \
  --custom-nodepool-arguments="--node-version=1.33.3-gke.1255000"

python3 xpk.py cluster create \
  --cluster=xpk-test-tpu7x-8 \
  --project=tpu-staging-env-one-vm \
  --tpu-type=tpu7x-8 \
  --zone=us-central2-newvmos \
  --num-nodes=1 \
  --enable-autoprovisioning \
  --custom-cluster-arguments="--cluster-version=1.33.3-gke.1255000" \
  --on-demand \
  --custom-nodepool-arguments="--node-version=1.33.3-gke.1255000"
  • [y] Tests pass
  • [y] Appropriate changes to documentation are included in the PR

Did you have a chance to test the workloads on the created cluster?

@scaliby
Copy link
Collaborator Author

scaliby commented Aug 21, 2025

Fixes / Features

  • Introduces support for Ironwood TPU's

Testing / Documentation

Testing has been performed in a manual fashion using the tpu7x accelerators. Test commands that suceeded:

python3 xpk.py cluster create \
  --cluster=xpk-test-tpu7x-2 \
  --project=tpu-staging-env-one-vm \
  --tpu-type=tpu7x-2 \
  --zone=us-central2-newvmos \
  --num-nodes=1 \
  --enable-autoprovisioning \
  --custom-cluster-arguments="--cluster-version=1.33.3-gke.1255000" \
  --on-demand \
  --custom-nodepool-arguments="--node-version=1.33.3-gke.1255000"

python3 xpk.py cluster create \
  --cluster=xpk-test-tpu7x-8 \
  --project=tpu-staging-env-one-vm \
  --tpu-type=tpu7x-8 \
  --zone=us-central2-newvmos \
  --num-nodes=1 \
  --enable-autoprovisioning \
  --custom-cluster-arguments="--cluster-version=1.33.3-gke.1255000" \
  --on-demand \
  --custom-nodepool-arguments="--node-version=1.33.3-gke.1255000"
  • [y] Tests pass
  • [y] Appropriate changes to documentation are included in the PR

Did you have a chance to test the workloads on the created cluster?

I tested by deploying a simple echo app.

@scaliby scaliby requested a review from sharabiani August 21, 2025 14:06
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@Obliviour Obliviour Awaiting requested review from Obliviour Obliviour is a code owner

@44past4 44past4 Awaiting requested review from 44past4 44past4 is a code owner

@pawloch00 pawloch00 Awaiting requested review from pawloch00 pawloch00 is a code owner

@BluValor BluValor Awaiting requested review from BluValor BluValor is a code owner

@gcie gcie Awaiting requested review from gcie gcie is a code owner

@RoshaniN RoshaniN Awaiting requested review from RoshaniN RoshaniN is a code owner

@FIoannides FIoannides Awaiting requested review from FIoannides FIoannides is a code owner

@sharabiani sharabiani Awaiting requested review from sharabiani sharabiani is a code owner

At least 1 approving review is required to merge this pull request.

Assignees
No one assigned
Labels
release-features features
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@scaliby @sharabiani