refactor(github): convert github utils to TaskFlow @task, chain, and decouple dispatches

Author: RexBearIU

dags/multipod/maxtext_e2e_tests.py

@@ -19,11 +19,11 @@
 """
 import datetime
 from airflow import models
+from airflow.models.baseoperator import chain
 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
+from xlml.utils.github import trigger_github_repository_dispatch, validate_git_trigger
 
 with models.DAG(
     dag_id="maxtext_e2e_tests",
@@ -41,7 +41,7 @@
             type="string",
             description="Build mode: stable or nightly",
         ),
-        "maxtext_sha": Param(
+        "commit_sha": Param(
             type="string",
             description="Commit SHA being tested",
         ),
@@ -53,7 +53,7 @@
             type="string",
             description="GitHub repository in owner/repo format",
         ),
-        "github_callback_token": Param(
+        "github_token": Param(
             type="string",
             description="GitHub PAT used to fire the repository_dispatch callback",
         ),
@@ -79,24 +79,25 @@
       poke_interval=600,  # check every 10 minutes for child DAG completion
   )
 
-  github_callback_pre_training = PythonOperator(
+  github_callback_pre_training = trigger_github_repository_dispatch.override(
       task_id="fire_github_callback_pre_training",
-      python_callable=fire_github_callback,
-      op_kwargs={"test_type": "pre_training"},
       trigger_rule=TriggerRule.ALL_SUCCESS,
-  )
+  )(test_type="pre_training")
 
-  github_callback_post_training = PythonOperator(
+  github_callback_post_training = trigger_github_repository_dispatch.override(
       task_id="fire_github_callback_post_training",
-      python_callable=fire_github_callback,
-      op_kwargs={"test_type": "post_training"},
       trigger_rule=TriggerRule.ALL_SUCCESS,
-  )
+  )(test_type="post_training")
 
-  validate_task = PythonOperator(
-      task_id="validate_trigger",
-      python_callable=validate_git_trigger,
-  )
+  validate_task = validate_git_trigger()
 
-  (validate_task >> trigger_pre_training >> github_callback_pre_training)
-  (validate_task >> trigger_post_training >> github_callback_post_training)
+  chain(
+      validate_task,
+      trigger_pre_training,
+      github_callback_pre_training,
+  )
+  chain(
+      validate_task,
+      trigger_post_training,
+      github_callback_post_training,
+  )

xlml/utils/github.py

@@ -14,54 +14,89 @@
 
 """Utilities for GitHub API integration."""
 
+from typing import Any
 import requests
+from airflow.decorators import task
 from airflow.exceptions import AirflowFailException
+from airflow.operators.python import get_current_context
 
 
-def validate_git_trigger(**context):
-  """Validates that the DAG run has required GitHub parameters.
+@task
+def validate_git_trigger() -> None:
+  """Validates that the DAG run has required GitHub parameters and was externally triggered."""
+  context = get_current_context()
+  dag_run = context.get("dag_run")
 
-  Prevents manual runs from the Airflow UI by ensuring github_run_id,
-  github_repo, and github_callback_token are present.
-  """
-  params = context["params"]
+  if not dag_run or not dag_run.external_trigger:
+    raise AirflowFailException(
+        "This DAG should not be run manually from the Airflow UI."
+    )
+
+  params = context.get("params", {})
   run_id = params.get("github_run_id")
   repo = params.get("github_repo")
-  token = params.get("github_callback_token")
+  token = params.get("github_token") or 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."
+        "Missing required GitHub parameters (run_id, repo, token)."
     )
 
 
-def fire_github_callback(test_type: str | None = None, **context):
+@task
+def trigger_github_repository_dispatch(
+    test_type: str | None = None,
+    repo: str | None = None,
+    token: str | None = None,
+    event_type: str = "airflow-dag-complete",
+    client_payload: dict[str, Any] | None = None,
+) -> None:
   """Fires a GitHub repository_dispatch callback with the DAG run result."""
-  params = context["params"]
-  dag_run = context["dag_run"]
+  context = get_current_context()
+  params = context.get("params", {})
+  dag_run = context.get("dag_run")
+
+  target_repo = repo or params.get("github_repo")
+  target_token = (
+      token or params.get("github_token") or params.get("github_callback_token")
+  )
+
+  if not target_repo or not target_token:
+    raise AirflowFailException(
+        "Missing required GitHub parameters (repo, token) to fire callback."
+    )
 
-  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
+  if client_payload is None:
+    payload: dict[str, Any] = {
+        "state": "success",
+        "dag_id": dag_run.dag_id if dag_run else None,
+        "dag_run_id": dag_run.run_id if dag_run else None,
+        "sha": params.get("commit_sha")
+        or params.get("maxtext_sha")
+        or params.get("sha"),
+        "github_run_id": params.get("github_run_id"),
+    }
+    if test_type:
+      payload["test_type"] = test_type
+  else:
+    payload = dict(client_payload)
+    if test_type and "test_type" not in payload:
+      payload["test_type"] = test_type
 
   response = requests.post(
-      f"https://api.github.com/repos/{params['github_repo']}/dispatches",
+      f"https://api.github.com/repos/{target_repo}/dispatches",
       headers={
-          "Authorization": f"Bearer {params['github_callback_token']}",
+          "Authorization": f"Bearer {target_token}",
           "Accept": "application/vnd.github+json",
           "X-GitHub-Api-Version": "2022-11-28",
       },
       json={
-          "event_type": "airflow-dag-complete",
-          "client_payload": client_payload,
+          "event_type": event_type,
+          "client_payload": payload,
       },
       timeout=30,
   )
   response.raise_for_status()
+
+
+fire_github_callback = trigger_github_repository_dispatch

xlml/utils/github_test.py

@@ -0,0 +1,156 @@
+# 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.
+
+"""Tests for github.py."""
+
+import unittest
+from unittest import mock
+from airflow.exceptions import AirflowFailException
+from xlml.utils import github
+
+
+class GithubTest(unittest.TestCase):
+
+  @mock.patch("xlml.utils.github.get_current_context")
+  def test_validate_git_trigger_success(self, mock_context):
+    mock_dag_run = mock.MagicMock()
+    mock_dag_run.external_trigger = True
+    mock_context.return_value = {
+        "dag_run": mock_dag_run,
+        "params": {
+            "github_run_id": "12345",
+            "github_repo": "owner/repo",
+            "github_token": "secret_token",
+        },
+    }
+    # Should not raise exception
+    github.validate_git_trigger.function()
+
+  @mock.patch("xlml.utils.github.get_current_context")
+  def test_validate_git_trigger_manual_run(self, mock_context):
+    mock_dag_run = mock.MagicMock()
+    mock_dag_run.external_trigger = False
+    mock_context.return_value = {
+        "dag_run": mock_dag_run,
+        "params": {
+            "github_run_id": "12345",
+            "github_repo": "owner/repo",
+            "github_token": "secret_token",
+        },
+    }
+    with self.assertRaises(AirflowFailException):
+      github.validate_git_trigger.function()
+
+  @mock.patch("xlml.utils.github.get_current_context")
+  def test_validate_git_trigger_missing_params(self, mock_context):
+    mock_dag_run = mock.MagicMock()
+    mock_dag_run.external_trigger = True
+    mock_context.return_value = {
+        "dag_run": mock_dag_run,
+        "params": {
+            "github_run_id": "12345",
+            # missing repo and token
+        },
+    }
+    with self.assertRaises(AirflowFailException):
+      github.validate_git_trigger.function()
+
+  @mock.patch("requests.post")
+  @mock.patch("xlml.utils.github.get_current_context")
+  def test_fire_github_callback_success(self, mock_context, mock_post):
+    mock_dag_run = mock.MagicMock()
+    mock_dag_run.dag_id = "test_dag"
+    mock_dag_run.run_id = "manual__1"
+
+    mock_context.return_value = {
+        "dag_run": mock_dag_run,
+        "params": {
+            "github_repo": "owner/repo",
+            "github_run_id": "12345",
+            "github_token": "secret_token",
+            "commit_sha": "abc1234",
+        },
+    }
+
+    mock_response = mock.MagicMock()
+    mock_post.return_value = mock_response
+
+    github.fire_github_callback.function(test_type="pre_training")
+
+    mock_post.assert_called_once_with(
+        "https://api.github.com/repos/owner/repo/dispatches",
+        headers={
+            "Authorization": "Bearer secret_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": "test_dag",
+                "dag_run_id": "manual__1",
+                "sha": "abc1234",
+                "github_run_id": "12345",
+                "test_type": "pre_training",
+            },
+        },
+        timeout=30,
+    )
+    mock_response.raise_for_status.assert_called_once()
+
+  @mock.patch("requests.post")
+  @mock.patch("xlml.utils.github.get_current_context")
+  def test_trigger_github_repository_dispatch_explicit_args(
+      self, mock_context, mock_post
+  ):
+    mock_context.return_value = {}
+    mock_response = mock.MagicMock()
+    mock_post.return_value = mock_response
+
+    github.trigger_github_repository_dispatch.function(
+        repo="custom/repo",
+        token="custom_token",
+        event_type="custom-event",
+        client_payload={"foo": "bar"},
+    )
+
+    mock_post.assert_called_once_with(
+        "https://api.github.com/repos/custom/repo/dispatches",
+        headers={
+            "Authorization": "Bearer custom_token",
+            "Accept": "application/vnd.github+json",
+            "X-GitHub-Api-Version": "2022-11-28",
+        },
+        json={
+            "event_type": "custom-event",
+            "client_payload": {"foo": "bar"},
+        },
+        timeout=30,
+    )
+    mock_response.raise_for_status.assert_called_once()
+
+  @mock.patch("xlml.utils.github.get_current_context")
+  def test_fire_github_callback_missing_token(self, mock_context):
+    mock_context.return_value = {
+        "params": {
+            "github_repo": "owner/repo",
+        }
+    }
+    with self.assertRaises(AirflowFailException):
+      github.fire_github_callback.function()
+
+
+if __name__ == "__main__":
+  unittest.main()