Add an new DAG for goodput replica resize (GoogleCloudPlatform#1323)

This change adds a new DAG for goodput elastic replica resizing to validate the resilience, recovery, and scaling behaviors of MaxText Pathways Elastic Training. It also confirms that goodput is properly enabled and tracking metrics via three core Airflow tasks/sensors: `check_goodput_logs`, `check_goodput_logname`, and `check_workload_goodput`.

Author: RUEI4341

dags/common/scheduling_helper/scheduling_helper.py

@@ -57,6 +57,7 @@ class DayOfWeek(enum.Enum):
 
 DagIdToTimeout: TypeAlias = dict[str, dt.timedelta]
 DefaultTimeout: dt.timedelta = dt.timedelta(minutes=30)
+# pylint: disable=line-too-long
 REGISTERED_DAGS: dict[str, DagIdToTimeout] = {
     TPU_OBS_MOCK_CLUSTER.name: {
         "gke_node_pool_label_update": DefaultTimeout,
@@ -96,8 +97,10 @@ class DayOfWeek(enum.Enum):
         "pw_elastic_pause_resume": DefaultTimeout,
         "pw_elastic_replica_resize": DefaultTimeout,
         "pw_elastic_goodput": DefaultTimeout,
+        "pw_elastic_goodput_replica": DefaultTimeout,
     },
 }
+# pylint: enable=line-too-long
 
 
 def get_dag_timeout(dag_id: str) -> dt.timedelta:

dags/maxtext_pathways/configs/parameters.py

@@ -134,11 +134,11 @@
         ),
     ),
     "priority": Param(
-        "medium",
+        "high",
         type="string",
         title="Priority",
         description="Priority for the workload",
-        enum=["very high", "high", "medium", "low"],
+        enum=["very-high", "high", "medium", "low"],
     ),
     "max_restarts": Param(
         1,

dags/maxtext_pathways/configs/utils.py

@@ -23,10 +23,16 @@
 from xlml.utils import gke, xpk
 
 
+# TODO(cienet): Replace this with an official one.
+COLOCATED_PYTHON_IMAGE = (
+    "gcr.io/tpu-prod-env-multipod/lidanny_maxtext-colocated-python:latest"
+)
+
+
 def generate_recipe_workload_id(dag_id: str) -> tuple[str, str]:
   """Generate a workload_id following the standard naming convention."""
   time.localtime()
-  timestamp = time.strftime("%Y%m%d%H%M%S", time.localtime())
+  timestamp = time.strftime("%m%d%H%M%S", time.localtime())
   name = f"{dag_id[:10]}-{timestamp[:10]}"
   name = name[:40].replace("_", "-")
 

dags/maxtext_pathways/pw_mcjax_elastic.py

@@ -29,14 +29,20 @@
 from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper
 from dags.maxtext_pathways.configs import parameters as ui_params
 from dags.maxtext_pathways.configs import recipe_config as recipe_cfg
-from dags.maxtext_pathways.configs.utils import get_dag_parameters, generate_install_dependencies_commands, generate_derived_parameters
+from dags.maxtext_pathways.configs.utils import (
+    get_dag_parameters,
+    generate_install_dependencies_commands,
+    generate_derived_parameters,
+    COLOCATED_PYTHON_IMAGE,
+)
 from xlml.utils import kpo, xpk
 
+
 ELASTIC_TYPE = ["Pause-resume", "Replica-resize"]
 elastic_params = ui_params.PARAMETERS.copy()
 elastic_params.update({
     "colocated_python_image": ui_params.Param(
-        "gcr.io/tpu-prod-env-multipod/lidanny_maxtext-colocated-python:latest",
+        COLOCATED_PYTHON_IMAGE,
         type="string",
         title="Colocated Python Image",
         description="Colocated Python image for pathways.",
@@ -266,11 +272,11 @@ def worker_pod_interruption(
       )
 
       # TODO(cienet): Refine the mechanism to chain tasks
-      _ = (
-          wait_for_step
-          >> trigger_interrupt
-          >> wait_for_elastic_attempt
-          >> wait_for_slices_active
+      chain(
+          wait_for_step,
+          trigger_interrupt,
+          wait_for_elastic_attempt,
+          wait_for_slices_active,
       )
 
       if previous_cycle_tail: