fix: Pin proper torch/xla image for pytorch_multislice DAG (GoogleCloudPlatform#1281)

The main issue was that the image `us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:nightly_3.10_tpuvm_{DATE}` was deprecated and only the latest image is pinned now in the repository.  This is the last image available  --> `nightly_3.10_tpuvm`

Co-authored-by: Camilo Quinones <camilo.quinones@cienet.com>

Author: AidenYu1673

dags/common/vm_resource.py

@@ -338,9 +338,10 @@ class DockerImage(enum.Enum):
       "gcr.io/cloud-tpu-multipod-dev/maxtext-orbax-custom:"
       f"{datetime.datetime.today().strftime('%Y-%m-%d')}"
   )
+  # Pinned to latest available image for torch/xla
   PYTORCH_NIGHTLY = (
       "us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/"
-      f"xla:nightly_3.10_tpuvm_{datetime.datetime.today().strftime('%Y%m%d')}"
+      "xla:nightly_3.10_tpuvm"
   )
   AXLEARN_CUSTOM = (
       # AXLearn's Docker/Artifact Registry bundler uses the run/workload name as

dags/multipod/pytorch.py

@@ -39,8 +39,16 @@
 
   for num_slices, cluster in [(1, v4_8), (2, v4_8), (1, v4_16)]:
     ici_chips = 4 if cluster == v4_8 else 8
+
+    # TODO(b/519376514): Remove the installation of libraries once resolved.
+    #
+    # The original image included all essential libraries, but it
+    # was deleted from the Artifact Registry.
+    # Only an older version remains, which lacks these essential libraries.
+    install_extra_libs = "apt-get update && apt-get install -y libopenblas-dev"
     run_cmds = (
         (
+            f"{install_extra_libs} && "
             "python /pytorch/xla/test/spmd/test_sharding_strategies.py "
             f"--ici_fsdp_parallelism {ici_chips} "
             f"--dcn_data_parallelism {num_slices}"
@@ -61,6 +69,7 @@
           f"export CHKPT_PATH={metric_config.SshEnvVars.GCS_OUTPUT.value}",
           "pip install gcsfs",
           (
+              f"{install_extra_libs} && "
               "python /pytorch/xla/test/spmd/test_xla_distributed_checkpoint.py "
               "EndToEndCheckpointTest.test_multihost_checkpoint"
           ),