feat: Separate out GitHub callback job for pre-training and post-training DAG (GoogleCloudPlatform#1325)

Author: RexBearIU

dags/multipod/maxtext_e2e_tests.py

@@ -14,8 +14,8 @@
 
 """
 Parent DAG that triggers MaxText E2E TPU pre-training and post-training child
-DAGs in parallel, waits for both to complete, then fires a GitHub
-repository_dispatch callback with the aggregated result.
+DAGs in parallel, firing separate GitHub repository_dispatch callbacks for
+pre-training and post-training upon completion of each job.
 """
 import datetime
 from airflow import models
@@ -79,19 +79,24 @@
       poke_interval=600,  # check every 10 minutes for child DAG completion
   )
 
-  github_callback = PythonOperator(
-      task_id="fire_github_callback",
+  github_callback_pre_training = PythonOperator(
+      task_id="fire_github_callback_pre_training",
       python_callable=fire_github_callback,
-      trigger_rule=TriggerRule.ALL_SUCCESS,  # Only fire if all upstream tasks succeeded
+      op_kwargs={"test_type": "pre_training"},
+      trigger_rule=TriggerRule.ALL_SUCCESS,
+  )
+
+  github_callback_post_training = PythonOperator(
+      task_id="fire_github_callback_post_training",
+      python_callable=fire_github_callback,
+      op_kwargs={"test_type": "post_training"},
+      trigger_rule=TriggerRule.ALL_SUCCESS,
   )
 
   validate_task = PythonOperator(
       task_id="validate_trigger",
       python_callable=validate_git_trigger,
   )
 
-  (
-      validate_task
-      >> [trigger_pre_training, trigger_post_training]
-      >> github_callback
-  )
+  (validate_task >> trigger_pre_training >> github_callback_pre_training)
+  (validate_task >> trigger_post_training >> github_callback_post_training)

xlml/utils/github.py

@@ -36,11 +36,21 @@ def validate_git_trigger(**context):
     )
 
 
-def fire_github_callback(**context):
+def fire_github_callback(test_type: str | None = None, **context):
   """Fires a GitHub repository_dispatch callback with the DAG run result."""
   params = context["params"]
   dag_run = context["dag_run"]
 
+  client_payload = {
+      "state": "success",
+      "dag_id": dag_run.dag_id,
+      "dag_run_id": dag_run.run_id,
+      "sha": params["maxtext_sha"],
+      "github_run_id": params["github_run_id"],
+  }
+  if test_type:
+    client_payload["test_type"] = test_type
+
   response = requests.post(
       f"https://api.github.com/repos/{params['github_repo']}/dispatches",
       headers={
@@ -50,13 +60,7 @@ def fire_github_callback(**context):
       },
       json={
           "event_type": "airflow-dag-complete",
-          "client_payload": {
-              "state": "success",
-              "dag_id": dag_run.dag_id,
-              "dag_run_id": dag_run.run_id,
-              "sha": params["maxtext_sha"],
-              "github_run_id": params["github_run_id"],
-          },
+          "client_payload": client_payload,
       },
       timeout=30,
   )