refactor(github): convert github utils to TaskFlow @task and chain - #306
refactor(github): convert github utils to TaskFlow @task and chain#306RexBearIU wants to merge 1 commit into
Conversation
d31880b to
afb1667
Compare
|
|
||
| def validate_git_trigger(**context): | ||
| @task | ||
| def validate_git_trigger() -> None: |
There was a problem hiding this comment.
how about checking external_trigger instead
from airflow.operators.python import get_current_context
from airflow.models.dagrun import DagRun
context = get_current_context()
assert dag_run is not None, "..."
dag_run: DagRun = context.get("dag_run")
assert dag_run is not None, "..."
assert dag_run.external_trigger == True, "..."
There was a problem hiding this comment.
Updated! validate_git_trigger now checks dag_run.external_trigger to ensure manual runs from the Airflow UI are prevented.
|
|
||
| def fire_github_callback(test_type: str | None = None, **context): | ||
| @task | ||
| def fire_github_callback(test_type: str | None = None) -> None: |
There was a problem hiding this comment.
From a decoupling perspective, repo, token, sha (or commit_sha), github_run_id, and the request body (L75–L78) are tightly coupled with a specific DAG.
Therefore, as a shared utility, this function should be more generalized so that any DAG can trigger the GitHub Dispatches API (rename this function also make more sense) — meaning these parameters need to be passed as arguments.
Alternatively, moving the entire function back into the DAG file as a private helper would be simpler and eliminate any decoupling issues.
There was a problem hiding this comment.
Thanks for the feedback! We updated trigger_github_repository_dispatch in xlml/utils/github.py to be fully decoupled as a shared utility. It accepts standard GitHub Repository Dispatch API arguments (repo, token, event_type, client_payload) directly as explicit inputs. If explicit arguments are omitted, it gracefully falls back to reading Airflow DAG params from context for backward compatibility. dags/multipod/maxtext_e2e_tests.py has been updated to use trigger_github_repository_dispatch as the primary function name.
15a0ce6 to
a445997
Compare
…decouple dispatches
a445997 to
49b52e1
Compare
2 participants
Description
Refactored
xlml/utils/github.pyanddags/multipod/maxtext_e2e_tests.pyto follow standard Airflow 2.x TaskFlow conventions:validate_git_triggerandfire_github_callbackto TaskFlow tasks using@task.**contextwithget_current_context().>>bitshift operators withchain().commit_shaandgithub_token(with backward-compatible fallbacks).xlml/utils/github_test.py.Tests
Uploaded DAG and utility files to test environment
jackyf-testand verified task structure and import errors:Output from Composer:
dags list-import-errors: No import errors formaxtext_e2e_tests.tasks list: Successfully listed all tasks:<Task(_TaskDecorator): fire_github_callback_post_training><Task(_TaskDecorator): fire_github_callback_pre_training><Task(TriggerDagRunOperator): trigger_tpu_post_training><Task(TriggerDagRunOperator): trigger_tpu_pre_training><Task(_TaskDecorator): validate_git_trigger>Checklist