From 13500307f4a662195e12afe6cffa851cd1cffcdf Mon Sep 17 00:00:00 2001 From: Jacky Fang Date: Tue, 28 Jul 2026 02:32:15 +0000 Subject: [PATCH] refactor(github): convert github utils to TaskFlow @task, chain, and decouple dispatches --- dags/multipod/maxtext_e2e_tests.py | 56 ++++++++++++----- xlml/utils/github.py | 63 +++++++++++-------- xlml/utils/github_test.py | 98 ++++++++++++++++++++++++++++++ 3 files changed, 175 insertions(+), 42 deletions(-) create mode 100644 xlml/utils/github_test.py diff --git a/dags/multipod/maxtext_e2e_tests.py b/dags/multipod/maxtext_e2e_tests.py index 9b5b9f492..1cfa9d920 100644 --- a/dags/multipod/maxtext_e2e_tests.py +++ b/dags/multipod/maxtext_e2e_tests.py @@ -18,12 +18,13 @@ pre-training and post-training upon completion of each job. """ 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 +42,7 @@ type="string", description="Build mode: stable or nightly", ), - "maxtext_sha": Param( + "commit_sha": Param( type="string", description="Commit SHA being tested", ), @@ -53,7 +54,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 +80,47 @@ 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, + )( + repo="{{ params.github_repo }}", + token="{{ params.github_token }}", + client_payload={ + "state": "success", + "dag_id": "{{ dag.dag_id }}", + "dag_run_id": "{{ run_id }}", + "sha": "{{ params.commit_sha }}", + "github_run_id": "{{ params.github_run_id }}", + "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, + )( + repo="{{ params.github_repo }}", + token="{{ params.github_token }}", + client_payload={ + "state": "success", + "dag_id": "{{ dag.dag_id }}", + "dag_run_id": "{{ run_id }}", + "sha": "{{ params.commit_sha }}", + "github_run_id": "{{ params.github_run_id }}", + "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, + ) diff --git a/xlml/utils/github.py b/xlml/utils/github.py index f3ffe373b..eae98a1fd 100644 --- a/xlml/utils/github.py +++ b/xlml/utils/github.py @@ -14,54 +14,65 @@ """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 + +@task +def validate_git_trigger() -> None: + """Validates that the DAG run was externally triggered with required GitHub parameters.""" + context = get_current_context() + dag_run = context.get("dag_run") -def validate_git_trigger(**context): - """Validates that the DAG run has required GitHub parameters. + if not dag_run or not dag_run.external_trigger: + raise AirflowFailException( + "This DAG should not be run manually from the Airflow UI." + ) - Prevents manual runs from the Airflow UI by ensuring github_run_id, - github_repo, and github_callback_token are present. - """ - params = context["params"] + 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): - """Fires a GitHub repository_dispatch callback with the DAG run result.""" - params = context["params"] - dag_run = context["dag_run"] +@task +def trigger_github_repository_dispatch( + repo: str, + token: str, + event_type: str = "airflow-dag-complete", + client_payload: dict[str, Any] | None = None, +) -> None: + """Fires a GitHub repository_dispatch event via the GitHub API.""" + 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"], - } - if test_type: - client_payload["test_type"] = test_type + payload = client_payload or {} 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", }, 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 diff --git a/xlml/utils/github_test.py b/xlml/utils/github_test.py new file mode 100644 index 000000000..da6459b3c --- /dev/null +++ b/xlml/utils/github_test.py @@ -0,0 +1,98 @@ +# 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} + 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") + def test_trigger_github_repository_dispatch_success(self, mock_post): + mock_response = mock.MagicMock() + mock_post.return_value = mock_response + + github.trigger_github_repository_dispatch.function( + repo="owner/repo", + token="secret_token", + event_type="custom-event", + client_payload={"state": "success", "sha": "abc1234"}, + ) + + 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": "custom-event", + "client_payload": {"state": "success", "sha": "abc1234"}, + }, + timeout=30, + ) + mock_response.raise_for_status.assert_called_once() + + def test_trigger_github_repository_dispatch_missing_token(self): + with self.assertRaises(AirflowFailException): + github.trigger_github_repository_dispatch.function( + repo="owner/repo", + token="", + ) + + +if __name__ == "__main__": + unittest.main()