Avoid Vertex AI Feature Store system Dag import timeout#69823
Conversation
|
cc @guan404ming, @jason810496. 🙌 |
shahar1
left a comment
There was a problem hiding this comment.
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] |
There was a problem hiding this comment.
# type: ignore[arg-type] will leak into the usage example in the docs, is it possible to avoid it?
There was a problem hiding this comment.
What if we specify type dict, then remove the ignore flag?
| @@ -27,6 +27,7 @@ | |||
|
|
|||
There was a problem hiding this comment.
Could we prevent regressions by asserting on sys.modules? (to ensure that no one re-adds the problematic import)
There was a problem hiding this comment.
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>
ad137d0 to
4c95ab8
Compare
Thx for extra content, PTAL. |
Prevent the Google Vertex AI Feature Store system Dag from loading
aiplatform_v1beta1while 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:
google.cloud.aiplatform_v1beta1loaded during parsingThe before result exceeds Airflow's default
dagbag_import_timeoutof 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?
{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.