From 539d5e290371caea36ced60beed60d0da9978089 Mon Sep 17 00:00:00 2001 From: Qinwen Xu Date: Sat, 30 Mar 2024 01:00:30 +0000 Subject: [PATCH] Change capacity type to reserved to avoid provision overtime --- dags/xlml/solutionsteam_jax_integration.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/dags/xlml/solutionsteam_jax_integration.py b/dags/xlml/solutionsteam_jax_integration.py index c479a8340..70ec728e3 100644 --- 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()