Fix DAG maxtext_configs_aot_hybridsim_src_path (GoogleCloudPlatform#1075)

This change fixes the DAG `maxtext_configs_aot_hybridsim_src_path` by:
1. Correct the path for in accordance of MaxText directory structure recent changes; requiring `src/` in the path to correctly locate config scripts. e.g.,: `MaxText/configs/...` -> `src/MaxText/configs/...`
2. Added the mandatory arguments to prevent Pydantic validation errors during the train_compile step, the previously optional `DATASET_PATH` and `OUTPUT_PATH` are now set with dummy values.

Author: AidenYu1673

dags/multipod/maxtext_configs_aot_hybridsim.py

@@ -13,39 +13,53 @@
 # limitations under the License.
 
 """
-A DAG to run AOT compilation and HybridSim tests for MaxText model configs on TPU v4, v5e.
+A DAG to run AOT compilation and HybridSim tests
+for MaxText model configs on TPU v4, v5e.
 """
 import datetime
 from airflow import models
 from airflow.utils.task_group import TaskGroup
+
+from xlml.utils import name_format
+from xlml.apis import metric_config
+
 from dags import composer_env
 from dags.common.quarantined_tests import QuarantineTests
 from dags.common import test_owner
-from dags.common.vm_resource import TpuVersion, Zone, DockerImage, XpkClusters, Project
-from dags.multipod.configs import gke_config
-from xlml.utils import name_format
+from dags.common.vm_resource import TpuVersion, DockerImage, XpkClusters
 from dags.multipod.configs import gke_config
-from xlml.apis import metric_config
 
 
 # Run once a day at 1 pm UTC (5 am PST / 6 am PDT)
 SCHEDULED_TIME = "30 2 * * *" if composer_env.is_prod_env() else None
 
 
-def hybridsim_compile_and_run(test_group_id):
-  with TaskGroup(group_id=test_group_id, prefix_group_id=True) as group:
+def hybridsim_compile_and_run(group_id):
+  with TaskGroup(group_id=group_id, prefix_group_id=True) as _:
     gcs_subfolder = f"{test_owner.Team.MULTIPOD.value}/maxtext"
     shared_gcs_location = name_format.generate_gcs_folder_location.override(
-        task_id=f"{test_group_id}_generate_gcs_folder_location"
+        task_id=f"{group_id}_generate_gcs_folder_location"
     )(
         f"{gcs_subfolder}/maxtext_configs_aot_hybridsim/v{tpu.value}",
-        test_group_id,
+        group_id,
     )
 
+    tpu_version_str = (
+        v5e_alt if tpu.value == TpuVersion.V5E.value else tpu.value
+    )
     # Run AOT workload: generate HLO, upload to GCS
     aot_cmd = (
-        'export XLA_FLAGS="--xla_dump_to=/tmp/xla_dump/ --xla_dump_large_constants"',
-        f"bash MaxText/configs/v{v5e_alt if tpu.value == TpuVersion.V5E.value else tpu.value}/{model_size}.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v{v5e_alt if tpu.value == TpuVersion.V5E.value else tpu.value}-{num_cores} M_COMPILE_TOPOLOGY_NUM_SLICES={n}",
+        (
+            'export XLA_FLAGS="--xla_dump_to=/tmp/xla_dump/'
+            ' --xla_dump_large_constants"'
+        ),
+        (
+            f"bash src/MaxText/configs/v{tpu_version_str}/{model_size}.sh"
+            " EXECUTABLE=train_compile"
+            f" M_COMPILE_TOPOLOGY=v{tpu_version_str}-{num_cores}"
+            f" M_COMPILE_TOPOLOGY_NUM_SLICES={n}"
+            f" DATASET_PATH=dummy-dataset OUTPUT_PATH=dummy-output-dir"
+        ),
         "gsutil -m cp -r /tmp/xla_dump/ ${GCS_OUTPUT}",
     )
     maxtext_aot = gke_config.get_gke_config(
@@ -61,7 +75,11 @@ def hybridsim_compile_and_run(test_group_id):
     chip_config = "default" if tpu == TpuVersion.V5E else "megacore"
     hybridsim_cmd = (
         "gsutil cp gs://cloud-hybridsim-prod/run_hybridsim.sh .",
-        f"bash run_hybridsim.sh GCS_XLA_DUMP_PATH=${{GCS_OUTPUT}}xla_dump GCS_OUTPUT_PATH=${{GCS_OUTPUT}}estimated_cost_ns.jsonl CHIP_CONFIG={chip_config} MODULE_NAME_PATTERN=jit_train_step*",
+        (
+            f"bash run_hybridsim.sh GCS_XLA_DUMP_PATH=${{GCS_OUTPUT}}xla_dump"
+            f" GCS_OUTPUT_PATH=${{GCS_OUTPUT}}estimated_cost_ns.jsonl"
+            f" CHIP_CONFIG={chip_config} MODULE_NAME_PATTERN=jit_train_step*"
+        ),
     )
     job_metric_config = metric_config.MetricConfig(
         json_lines=metric_config.JSONLinesConfig(
@@ -72,14 +90,16 @@ def hybridsim_compile_and_run(test_group_id):
     maxtext_hybridsim = gke_config.get_gke_config(
         cluster=cluster,
         time_out_in_min=240,
-        test_name=f"maxtext-{model_size}-{n}xv{tpu.value}-{num_cores}-hybridsim",
+        test_name=(
+            f"maxtext-{model_size}-{n}xv{tpu.value}-{num_cores}-hybridsim"
+        ),
         run_model_cmds=hybridsim_cmd,
         docker_image=DockerImage.CLOUD_HYBRIDSIM_NIGHTLY.value,
         test_owner=test_owner.AIRFLOW,
         user_specified_job_metric_config=job_metric_config,
     ).run(gcs_location=shared_gcs_location)
 
-    shared_gcs_location >> maxtext_aot >> maxtext_hybridsim
+    _ = shared_gcs_location >> maxtext_aot >> maxtext_hybridsim
 
 
 with models.DAG(
@@ -102,7 +122,12 @@ def hybridsim_compile_and_run(test_group_id):
   model_configs = {
       # accelerator: [(model_size, num_cores), ...],
       TpuVersion.V4: [("22b", 128), ("52b", 384)],
-      TpuVersion.V5E: [("16b", 256), ("32b", 256), ("64b", 256), ("128b", 256)],
+      TpuVersion.V5E: [
+          ("16b", 256),
+          ("32b", 256),
+          ("64b", 256),
+          ("128b", 256),
+      ],
   }
   num_slices = [1, 2, 4, 8]
   clusters = {