diff --git a/dags/multipod/maxtext_e2e_tests.py b/dags/multipod/maxtext_e2e_tests.py index 9b5b9f492..5a1e30c5e 100644 --- a/dags/multipod/maxtext_e2e_tests.py +++ b/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, + ) diff --git a/xlml/utils/github.py b/xlml/utils/github.py index f3ffe373b..530e6ef64 100644 --- a/xlml/utils/github.py +++ b/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", }, diff --git a/xlml/utils/github_test.py b/xlml/utils/github_test.py new file mode 100644 index 000000000..f71505fa3 --- /dev/null +++ b/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()