update JAX StS path (#526)

khatwanimohit · web-flow · commit a366662aa8d5 · 2025-01-07T17:00:47.000-08:00
diff --git a/dags/common/vm_resource.py b/dags/common/vm_resource.py
@@ -300,7 +300,7 @@ class DockerImage(enum.Enum):
       f"{datetime.datetime.today().strftime('%Y-%m-%d')}"
   )
   MAXTEXT_TPU_JAX_STABLE_STACK = (
-      "gcr.io/tpu-prod-env-multipod/maxtext_jax_stable_stack_0.4.35:"
+      "gcr.io/tpu-prod-env-multipod/maxtext_jax_stable_stack:"
       f"{datetime.datetime.today().strftime('%Y-%m-%d')}"
   )
   MAXTEXT_TPU_STABLE_STACK_NIGHTLY_JAX = (

Original file line number	Diff line number	Diff line change
`@@ -300,7 +300,7 @@ class DockerImage(enum.Enum):`
`300`	`300`	`f"{datetime.datetime.today().strftime('%Y-%m-%d')}"`
`301`	`301`	`)`
`302`	`302`	`MAXTEXT_TPU_JAX_STABLE_STACK = (`
`303`		`- "gcr.io/tpu-prod-env-multipod/maxtext_jax_stable_stack_0.4.35:"`
	`303`	`+ "gcr.io/tpu-prod-env-multipod/maxtext_jax_stable_stack:"`
`304`	`304`	`f"{datetime.datetime.today().strftime('%Y-%m-%d')}"`
`305`	`305`	`)`
`306`	`306`	`MAXTEXT_TPU_STABLE_STACK_NIGHTLY_JAX = (`