Skip to content

Add TPU Capacity Validation DAG - #191

Open
chengpinglin wants to merge 7 commits into
devfrom
tpu-obs/user/tony/capacity_01_dipstick
Open

Add TPU Capacity Validation DAG#191
chengpinglin wants to merge 7 commits into
devfrom
tpu-obs/user/tony/capacity_01_dipstick

Conversation

@chengpinglin

Copy link
Copy Markdown

Description

This PR introduces the chips_capacity_validation_dag, an automated workflow designed to validate Cloud TPU capacity and health. The DAG automates the lifecycle of a TPU workload—from infrastructure provisioning to deep metric-based health checks—ensuring that TPU chips are not just allocated, but fully functional and utilized.

Introduces four specialized sensors that poll Cloud Monitoring to verify the state of the hardware:

  • Scheduled Chips: Confirms the system has successfully assigned TPU resources to the VM.

  • Active Chips: Validates that the TPU driver is active and reporting to the monitoring backend.

  • Chip State (Health): Specifically checks for the HEALTHY state label to ensure no hardware defects are present.

  • Utilized Chips: Confirms that the workload (JAX/PyTorch) has successfully initialized the device and is performing computation.

Tests

  • GCP Composer name: tony-test (under GCP project: cloud-ml-auto-solutions)
  • GCP Composer version: 2.13.1

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run one-shot tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed.

@chengpinglin chengpinglin changed the title Tpu obs/user/tony/capacity 01 dipstick Add TPU Capacity Validation DAG Jan 27, 2026
"gs://ml-auto-solutions-dag-configs/tpu_observability/dag_config.yaml"
)

VM_GCS_CONFIG_PATH = "gs://ml-auto-solutions-dag-configs/tpu_observability/dag_config_tony_vm.yaml"

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

is this your personal config?

@alfredyu-cienet
alfredyu-cienet force-pushed the dev branch 2 times, most recently from 0f0e4e2 to 31723ef Compare January 28, 2026 06:23
@chengpinglin
chengpinglin force-pushed the tpu-obs/user/tony/capacity_01_dipstick branch from 206efaa to 91e3a7d Compare January 28, 2026 07:02
@alfredyu-cienet
alfredyu-cienet force-pushed the dev branch 2 times, most recently from c846590 to 395bf7c Compare February 2, 2026 08:04
@chengpinglin
chengpinglin force-pushed the tpu-obs/user/tony/capacity_01_dipstick branch from 91e3a7d to 5528e72 Compare February 3, 2026 05:45
@alfredyu-cienet
alfredyu-cienet force-pushed the dev branch 3 times, most recently from 8eb5556 to acf36a1 Compare May 19, 2026 03:54
@alfredyu-cienet
alfredyu-cienet force-pushed the dev branch 3 times, most recently from 83f8b37 to ea29ef3 Compare May 28, 2026 07:44
@alfredyu-cienet
alfredyu-cienet force-pushed the dev branch 4 times, most recently from 4af284a to a9fa8fb Compare June 8, 2026 08:58
@alfredyu-cienet
alfredyu-cienet force-pushed the dev branch 4 times, most recently from f8824db to 09c9e42 Compare June 17, 2026 08:08
@alfredyu-cienet
alfredyu-cienet force-pushed the dev branch 5 times, most recently from d242c8d to f76e189 Compare July 3, 2026 08:25
@alfredyu-cienet
alfredyu-cienet force-pushed the dev branch 4 times, most recently from a577785 to 71b8d8a Compare July 14, 2026 03:12
@alfredyu-cienet
alfredyu-cienet force-pushed the dev branch 4 times, most recently from b588728 to 0e938b1 Compare July 23, 2026 01:43
@alfredyu-cienet
alfredyu-cienet force-pushed the dev branch 2 times, most recently from 54f9a78 to 8e4ea82 Compare July 27, 2026 07:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants