fix: Map v4-8 to v4-16 cluster and clean up invalid command (GoogleCloudPlatform#1079)

- Cluster Mapping: Temporarily maps the v4-8 accelerator to the `XpkClusters.TPU_V4_16_CLUSTER` due to the `TPU_V4_8_MAXTEXT_CLUSTER` lacking Vertex AI Tensorboard configuration. (b/465618653)
- Conditional Naming: A conditional if/elif block is added to dynamically set test_group_name based on slice_num (1 or 2) and accelerator type. This is necessary to prevent test name collisions while multiple test configurations share the temporary v4-16 cluster.
- Code Cleanup: Commented out the invalid `gsutil ls` command within `run_model_cmds`. (b/465619132)

Author: AidenYu1673

dags/multipod/maxtext_profiling_vertex_ai_tensorboard.py

@@ -17,9 +17,9 @@
 """
 import datetime
 from airflow import models
-from dags import composer_env, gcs_bucket
+from dags import gcs_bucket
 from dags.common import test_owner
-from dags.common.vm_resource import TpuVersion, Zone, DockerImage, XpkClusters
+from dags.common.vm_resource import DockerImage, XpkClusters
 from dags.multipod.configs import gke_config
 from dags.multipod.configs.common import SetupMode
 
@@ -48,7 +48,11 @@
   }
   clusters = {
       # accelerator: cluster name
-      "v4-8": XpkClusters.TPU_V4_8_MAXTEXT_CLUSTER,
+      # TODO(b/465618653): Switch back to v4-8 once the issue is resolved.
+      # Temporary use the cluster for v4-16 since the original does not have
+      # the "vertex-tensorboard" compatibility.
+      # "v4-8": XpkClusters.TPU_V4_8_MAXTEXT_CLUSTER
+      "v4-8": XpkClusters.TPU_V4_16_CLUSTER,
       "v4-16": XpkClusters.TPU_V4_16_CLUSTER,
   }
 
@@ -58,17 +62,30 @@
         current_time = datetime.datetime.now()
         current_datetime = current_time.strftime("%Y-%m-%d-%H-%M-%S")
         profiling_in_vertex_ai_tb_cmds = (
-            f"export RUN_NAME=vertex-ai-{mode.value}-{slice_num}x-{accelerator}-{current_datetime}",
-            "python3 -m MaxText.train MaxText/configs/base.yml"
-            f" run_name=$RUN_NAME base_output_directory={base_output_directory}"
-            f" dataset_path={dataset_path} profiler=xplane steps=10",
-            "gsutil ls gs://cloud-ai-platform-*/tensorboard-*/$EXPERIMENT_NAME",
+            (
+                f"export RUN_NAME=vertex-ai-{mode.value}-{slice_num}x"
+                f"-{accelerator}-{current_datetime}"
+            ),
+            (
+                "python3 -m MaxText.train MaxText/configs/base.yml"
+                f" run_name=$RUN_NAME base_output_directory"
+                f"={base_output_directory} dataset_path={dataset_path}"
+                f" profiler=xplane steps=10"
+            ),
+            # TODO(b/465619132): Revert this change once the issue is resolved.
+            # Comment out due to the path is invalid.
+            # "gsutil ls gs://cloud-ai-platform-*/tensorboard-*/
+            # $EXPERIMENT_NAME",
         )
+
         profiling_in_vertex_ai_tb_test = gke_config.get_gke_config(
             num_slices=slice_num,
             cluster=clusters[accelerator],
             time_out_in_min=240,
-            test_name=f"profiling-vertex-ai-tensorboard-{mode.value}",
+            # TODO(b/465618653): Switch back to v4-8 once the issue is resolved.
+            test_name=(
+                f"profiling-vertex-ai-tensorboard-{mode.value}-{accelerator}"
+            ),
             run_model_cmds=profiling_in_vertex_ai_tb_cmds,
             docker_image=image.value,
             test_owner=test_owner.SURBHI_J,