Skip to content

Avoid Vertex AI Feature Store system Dag import timeout#69823

Open
viiccwen wants to merge 1 commit into
apache:mainfrom
viiccwen:fix-google-feature-store-dag-imports
Open

Avoid Vertex AI Feature Store system Dag import timeout#69823
viiccwen wants to merge 1 commit into
apache:mainfrom
viiccwen:fix-google-feature-store-dag-imports

Conversation

@viiccwen

@viiccwen viiccwen commented Jul 13, 2026

Copy link
Copy Markdown
Contributor

Prevent the Google Vertex AI Feature Store system Dag from loading aiplatform_v1beta1 while Airflow discovers Dags. The package's third-party import chain can exceed the default 30-second Dag import timeout before any task runs.

Defer the Feature Store clients and data-format type until their execution paths need them. Use GAPIC-supported mappings in the system Dag so parsing does not construct protobuf messages, while keeping the existing hook and operator type signatures unchanged.

Update the client mock targets for the deferred imports and cover both client getter paths.

Before/after cold-import comparison in the same Breeze environment:

Before After
System Dag import time 30+ seconds (timeout) 10.278 seconds
google.cloud.aiplatform_v1beta1 loaded during parsing Yes No

The before result exceeds Airflow's default dagbag_import_timeout of 30 seconds. The after result keeps the Vertex AI Feature Store package out of the Dag parsing process entirely.


Was generative AI tooling used to co-author this PR?
  • Yes - Codex 5.6 Sol

  • Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
  • For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
  • When adding dependency, check compliance with the ASF 3rd Party License Policy.
  • For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

@viiccwen viiccwen requested a review from shahar1 as a code owner July 13, 2026 11:21
@boring-cyborg boring-cyborg Bot added area:providers provider:google Google (including GCP) related issues labels Jul 13, 2026
@viiccwen

Copy link
Copy Markdown
Contributor Author

cc @guan404ming, @jason810496. 🙌

@shahar1 shahar1 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for your contribution!
Got a small deja vu from a similar issue I encountered in the past: googleapis/python-aiplatform#1994
Few comments to handle and we're good to go :)

location=REGION,
feature_online_store_id=FEATURE_ONLINE_STORE_ID,
feature_online_store=FeatureOnlineStore(optimized=FeatureOnlineStore.Optimized()),
feature_online_store={"optimized": {}}, # type: ignore[arg-type]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

# type: ignore[arg-type] will leak into the usage example in the docs, is it possible to avoid it?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What if we specify type dict, then remove the ignore flag?

@@ -27,6 +27,7 @@

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we prevent regressions by asserting on sys.modules? (to ensure that no one re-adds the problematic import)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure, I'll reference the testing style like test_accessing_bundle_loads_cadwyn(...).

The Vertex AI client package performs expensive third-party imports during Dag discovery, pushing the system example beyond the default 30-second limit and exposing circular-import failures. Keeping parsing lightweight prevents system Dag scans from failing before tasks can run.

Signed-off-by: viiccwen <vicwen@apache.org>
@viiccwen viiccwen force-pushed the fix-google-feature-store-dag-imports branch from ad137d0 to 4c95ab8 Compare July 14, 2026 03:51
@viiccwen

Copy link
Copy Markdown
Contributor Author

Thanks for your contribution! Got a small deja vu from a similar issue I encountered in the past: googleapis/python-aiplatform#1994 Few comments to handle and we're good to go :)

Thx for extra content, PTAL.

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

Labels

area:providers provider:google Google (including GCP) related issues

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Google Vertex AI Feature Store system Dag exceeds import timeout

2 participants