feat: Restrict DAG maxtext_e2e_tests to GitHub schedule events (GoogleCloudPlatform#1315)

This change restricts the DAG `maxtext_e2e_tests` to prevent accidental manual runs from the Airflow UI, as part of a coordinated change with the `maxtext` repository to restrict E2E test triggers.

Author: RexBearIU

dags/multipod/maxtext_e2e_tests.py

@@ -18,12 +18,12 @@
 repository_dispatch callback with the aggregated result.
 """
 import datetime
-import requests
 from airflow import models
 from airflow.models.param import Param
 from airflow.operators.python import PythonOperator
 from airflow.operators.trigger_dagrun import TriggerDagRunOperator
 from airflow.utils.trigger_rule import TriggerRule
+from xlml.utils.github import fire_github_callback, validate_git_trigger
 
 with models.DAG(
     dag_id="maxtext_e2e_tests",
@@ -79,35 +79,19 @@
       poke_interval=600,  # check every 10 minutes for child DAG completion
   )
 
-  def fire_github_callback(**context):
-    params = context["params"]
-    dag_run = context["dag_run"]
-
-    response = requests.post(
-        f"https://api.github.com/repos/{params['github_repo']}/dispatches",
-        headers={
-            "Authorization": f"Bearer {params['github_callback_token']}",
-            "Accept": "application/vnd.github+json",
-            "X-GitHub-Api-Version": "2022-11-28",
-        },
-        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"],
-            },
-        },
-        timeout=30,
-    )
-    response.raise_for_status()
-
   github_callback = PythonOperator(
       task_id="fire_github_callback",
       python_callable=fire_github_callback,
       trigger_rule=TriggerRule.ALL_SUCCESS,  # Only fire if all upstream tasks succeeded
   )
 
-  [trigger_pre_training, trigger_post_training] >> github_callback
+  validate_task = PythonOperator(
+      task_id="validate_trigger",
+      python_callable=validate_git_trigger,
+  )
+
+  (
+      validate_task
+      >> [trigger_pre_training, trigger_post_training]
+      >> github_callback
+  )

xlml/utils/github.py

@@ -0,0 +1,63 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilities for GitHub API integration."""
+
+import requests
+from airflow.exceptions import AirflowFailException
+
+
+def validate_git_trigger(**context):
+  """Validates that the DAG run has required GitHub parameters.
+
+  Prevents manual runs from the Airflow UI by ensuring github_run_id,
+  github_repo, and github_callback_token are present.
+  """
+  params = context["params"]
+  run_id = params.get("github_run_id")
+  repo = params.get("github_repo")
+  token = params.get("github_callback_token")
+
+  if not run_id or not repo or not token:
+    raise AirflowFailException(
+        "Missing required GitHub parameters (run_id, repo, token). "
+        "This DAG should not be run manually from the Airflow UI."
+    )
+
+
+def fire_github_callback(**context):
+  """Fires a GitHub repository_dispatch callback with the DAG run result."""
+  params = context["params"]
+  dag_run = context["dag_run"]
+
+  response = requests.post(
+      f"https://api.github.com/repos/{params['github_repo']}/dispatches",
+      headers={
+          "Authorization": f"Bearer {params['github_callback_token']}",
+          "Accept": "application/vnd.github+json",
+          "X-GitHub-Api-Version": "2022-11-28",
+      },
+      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"],
+          },
+      },
+      timeout=30,
+  )
+  response.raise_for_status()