Skip to content

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

Open
RexBearIU wants to merge 4 commits into
mainfrom
maxtext/user/jackyf/refactor-github-utils
Open

refactor(github): convert github utils to TaskFlow @task and chain#306
RexBearIU wants to merge 4 commits into
mainfrom
maxtext/user/jackyf/refactor-github-utils

Conversation

@RexBearIU

@RexBearIU RexBearIU commented Jul 27, 2026

Copy link
Copy Markdown

Description

Refactored xlml/utils/github.py and dags/multipod/maxtext_e2e_tests.py to follow standard Airflow 2.x TaskFlow conventions:

  • Converted validate_git_trigger and fire_github_callback to TaskFlow tasks using @task.
  • Replaced **context with get_current_context().
  • Replaced >> bitshift operators with chain().
  • Standardized DAG parameter names to commit_sha and github_token (with backward-compatible fallbacks).
  • Added unit tests in xlml/utils/github_test.py.

Tests

  1. Pre-commit & Unit Tests:
pre-commit run --all-files
python3 -m unittest xlml/utils/github_test.py
  1. Cloud Composer Integration Tests:
    Uploaded DAG and utility files to test environment jackyf-test and verified task structure and import errors:
gcloud storage cp dags/multipod/maxtext_e2e_tests.py gs://us-east1-jackyf-test-3fc56b08-bucket/dags/multipod/maxtext_e2e_tests.py
gcloud storage cp xlml/utils/github.py gs://us-east1-jackyf-test-3fc56b08-bucket/dags/xlml/utils/github.py
gcloud composer environments run jackyf-test --location us-east1 --project cloud-ml-auto-solutions dags list-import-errors
gcloud composer environments run jackyf-test --location us-east1 --project cloud-ml-auto-solutions tasks list -- maxtext_e2e_tests

Output from Composer:

  • dags list-import-errors: No import errors for maxtext_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

  • I have performed a self-review of my code.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run one-shot tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed.

@RexBearIU
RexBearIU force-pushed the maxtext/user/jackyf/refactor-github-utils branch from d31880b to afb1667 Compare July 27, 2026 09:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant