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

dags/multipod/maxtext_e2e_tests.py

@@ -19,8 +19,8 @@
 """
 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
@@ -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 = fire_github_callback.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 = fire_github_callback.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

@@ -15,19 +15,23 @@
 """Utilities for GitHub API integration."""
 
 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):
+@task
+def validate_git_trigger() -> None:
   """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.
+  github_repo, and github_token are present.
   """
-  params = context["params"]
+  context = get_current_context()
+  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(
@@ -36,25 +40,35 @@ def validate_git_trigger(**context):
     )
 
 
-def fire_github_callback(test_type: str | None = None, **context):
+@task
+def fire_github_callback(test_type: str | 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")
+
+  repo = params.get("github_repo")
+  token = params.get("github_token") or params.get("github_callback_token")
+
+  if not repo or not 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"],
+      "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("sha"),
+      "github_run_id": params.get("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",
+      f"https://api.github.com/repos/{repo}/dispatches",
       headers={
-          "Authorization": f"Bearer {params['github_callback_token']}",
+          "Authorization": f"Bearer {token}",
           "Accept": "application/vnd.github+json",
           "X-GitHub-Api-Version": "2022-11-28",
       },

xlml/utils/github_test.py

@@ -0,0 +1,104 @@
+# 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_context.return_value = {
+        "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_missing_params(self, mock_context):
+    mock_context.return_value = {
+        "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("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()