From 2bae81814426964aafb2c818aab31816c1edb3c7 Mon Sep 17 00:00:00 2001 From: Abhishek Agrawal Date: Thu, 23 Apr 2026 11:45:54 -0700 Subject: [PATCH] Adds configurations for running p2p and Pathways-based checkpoint manager benchmarks on multi-slice v5litepod-16 TPUs. PiperOrigin-RevId: 904562509 --- .../oss/cloud_run_integration_tests.yaml | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/checkpoint/orbax/checkpoint/_src/testing/oss/cloud_run_integration_tests.yaml b/checkpoint/orbax/checkpoint/_src/testing/oss/cloud_run_integration_tests.yaml index 592234523..1cbc99fb9 100644 --- a/checkpoint/orbax/checkpoint/_src/testing/oss/cloud_run_integration_tests.yaml +++ b/checkpoint/orbax/checkpoint/_src/testing/oss/cloud_run_integration_tests.yaml @@ -15,3 +15,35 @@ tests: verbose: true skip_validation: true enable_pathways: false + - name: p2p_checkpoint_manager_benchmark + cluster_name: orbax-cluster-test-mtc + tpu_type: v5litepod-16 + zone: us-west1-c + gcp_region: us-west1 + gcp_project: orbax-checkpoint + config_file: orbax/checkpoint/_src/testing/benchmarks/configs/p2p_checkpoint_manager_benchmark.yaml + docker_image: gcr.io/orbax-checkpoint/orbax-benchmarks-runner:latest + output_directory: gs://orbax-benchmarks/cloud_runs/ + nodelete_cluster_on_completion: true + ramdisk_directory: /local/test + num_slices: 2 + test_restart_workflow: false + verbose: true + skip_validation: true + enable_pathways: false + - name: pathways_checkpoint_manager_benchmark + cluster_name: orbax-cluster-test-pathways + tpu_type: v5litepod-16 + zone: us-west1-c + gcp_region: us-west1 + gcp_project: orbax-checkpoint + config_file: orbax/checkpoint/_src/testing/benchmarks/configs/checkpoint_manager_benchmark.yaml + docker_image: gcr.io/orbax-checkpoint/orbax-benchmarks-runner:latest + output_directory: gs://orbax-benchmarks/cloud_runs/ + nodelete_cluster_on_completion: true + ramdisk_directory: /local/test + num_slices: 2 + test_restart_workflow: false + verbose: true + skip_validation: true + enable_pathways: true \ No newline at end of file