Skip to content

Replace airflow.utils.state import with airflow.sdk in Common AI Provider #72410

Description

@jroachgolf84

Description

This is an issue that is reserved for the Airflow Summit "Contributors" Workshop. This is denoted with the label contributors-workshop. Out of respect for the organizers and participants of this workshop, please do not implement a PR that addresses this issue.

If this issue is still open following Airflow Summit, the label will be removed and the issue can be picked up.

Sub-issue of #62063, which investigates provider usage of airflow.utils imports so they can be replaced with Task SDK equivalents and reduce the airflow-core dependency in providers.

Provider common.ai
Import TaskInstanceState from airflow.utils.state
Replacement airflow.sdk.TaskInstanceState
Shape Import-only change in a single file

common.ai requires apache-airflow>=3.0.0 (see providers/common/ai/pyproject.toml), so airflow.sdk is always available and no version guard is needed here, unlike most other providers that still support Airflow 2.11+. airflow.sdk.TaskInstanceState and airflow.utils.state.TaskInstanceState are two independently defined (str, Enum) classes with matching member names and string values, so this is a plain import swap with no behavior change.

What needs to change

Replace from airflow.utils.state import TaskInstanceState with from airflow.sdk import TaskInstanceState in:

  • providers/common/ai/src/airflow/providers/common/ai/plugins/hitl_review.py:83, inside the if AIRFLOW_V_3_1_PLUS: block, used later to build _RUNNING_TI_STATES = frozenset({TaskInstanceState.RUNNING, ...})

No other code in this file needs to change, only the import line.

How to verify it

The existing unit tests should pass unmodified, since the enum values are identical:

breeze testing providers-tests providers/common/ai/tests/unit/common/ai/plugins/test_hitl_review.py

Definition of Done

  1. Update the import in plugins/hitl_review.py.
  2. Confirm no other airflow.utils.state usage remains in providers/common/ai/src (grep -rn "airflow.utils.state" providers/common/ai/src).
  3. No new tests are needed, this is a like-for-like import swap, existing tests should pass unmodified.
  4. This should pass:
breeze testing providers-tests --test-type "Providers[common.ai]"

Drafted-by: Claude Code (Sonnet 5); reviewed and edited by @jroachgolf84 before posting

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions