Change capacity type to reserved to avoid provision overtime

GoogleCloudPlatform · Mar 30, 2024 · 539d5e2 · 539d5e2
1 parent b7b4ff0
commit 539d5e2
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/dags/xlml/solutionsteam_jax_integration.py b/dags/xlml/solutionsteam_jax_integration.py
@@ -47,14 +47,16 @@
 
   pod_latest = task.TpuQueuedResourceTask(
       test_config.JSonnetTpuVmTest.from_jax(
-          "jax-pod-latest-tpu-ubuntu2204-base-func-v2-32-1vm"
+          "jax-pod-latest-tpu-ubuntu2204-base-func-v2-32-1vm",
+          reserved = True,
       ),
       US_CENTRAL1_A,
   ).run()
 
   pod_head = task.TpuQueuedResourceTask(
       test_config.JSonnetTpuVmTest.from_jax(
-          "jax-pod-head-tpu-ubuntu2204-base-func-v2-32-1vm"
+          "jax-pod-head-tpu-ubuntu2204-base-func-v2-32-1vm",
+          reserved = True,
       ),
       US_CENTRAL1_A,
   ).run()