This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new d5aed746ad9 Avoid Vertex AI Feature Store system Dag import timeout
(#69823)
d5aed746ad9 is described below
commit d5aed746ad90ed731702ceaf0f87cb43d01a0d28
Author: Vic Wen <[email protected]>
AuthorDate: Sun Oct 18 04:07:03 2026 +0800
Avoid Vertex AI Feature Store system Dag import timeout (#69823)
* Avoid Vertex AI Feature Store system Dag import timeout
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 <[email protected]>
* Pass the Bigtable online store config as a dictionary
Rebasing onto #71238 kept its Bigtable autoscaling config as a typed
FeatureOnlineStore message, which brought back the aiplatform_v1beta1
import into the system Dag. Express it as a dictionary so Dag parsing no
longer imports the package, and document that the Feature Store operators
accept dictionaries.
---------
Signed-off-by: viiccwen <[email protected]>
Co-authored-by: Shahar Epstein <[email protected]>
---
.../google/docs/operators/cloud/vertex_ai.rst | 5 +++
.../google/cloud/hooks/vertex_ai/feature_store.py | 31 ++++++++++-------
.../cloud/operators/vertex_ai/feature_store.py | 17 ++++-----
.../vertex_ai/example_vertex_ai_feature_store.py | 36 +++++++++----------
.../cloud/hooks/vertex_ai/test_feature_store.py | 40 +++++++++++++++++++++-
5 files changed, 88 insertions(+), 41 deletions(-)
diff --git a/providers/google/docs/operators/cloud/vertex_ai.rst
b/providers/google/docs/operators/cloud/vertex_ai.rst
index 6cfd165fe53..e4fd2c3f70e 100644
--- a/providers/google/docs/operators/cloud/vertex_ai.rst
+++ b/providers/google/docs/operators/cloud/vertex_ai.rst
@@ -676,6 +676,11 @@ To create feature online store you can use
:class:`~airflow.providers.google.cloud.operators.vertex_ai.feature_store.CreateFeatureOnlineStoreOperator`.
The operator creation results in :ref:`XCom <concepts:xcom>` under
``return_value`` key.
+The ``feature_online_store``, ``feature_view`` and ``data_key`` parameters of
the Feature Store operators
+accept either the ``google.cloud.aiplatform_v1beta1`` message objects or
equivalent dictionaries.
+Dictionaries keep ``google.cloud.aiplatform_v1beta1`` out of Dag parsing,
which can otherwise
+exceed ``dagbag_import_timeout``.
+
.. exampleinclude::
/../../google/tests/system/google/cloud/vertex_ai/example_vertex_ai_feature_store.py
:language: python
:dedent: 4
diff --git
a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/feature_store.py
b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/feature_store.py
index 17a1b6b7ad1..e5bbc84bda7 100644
---
a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/feature_store.py
+++
b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/feature_store.py
@@ -19,15 +19,9 @@
from __future__ import annotations
from collections.abc import Sequence
-from typing import (
- TYPE_CHECKING,
-)
+from typing import TYPE_CHECKING
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.aiplatform_v1beta1 import (
- FeatureOnlineStoreAdminServiceClient,
- FeatureOnlineStoreServiceClient,
-)
from airflow.providers.common.compat.sdk import AirflowException
from airflow.providers.google.common.consts import CLIENT_INFO
@@ -36,6 +30,10 @@ from airflow.providers.google.common.hooks.base_google
import PROVIDE_PROJECT_ID
if TYPE_CHECKING:
from google.api_core.operation import Operation
from google.api_core.retry import Retry
+ from google.cloud.aiplatform_v1beta1 import (
+ FeatureOnlineStoreAdminServiceClient,
+ FeatureOnlineStoreServiceClient,
+ )
from google.cloud.aiplatform_v1beta1.types import (
FeatureOnlineStore,
FeatureView,
@@ -90,6 +88,8 @@ class FeatureStoreHook(GoogleBaseHook):
If provided and not 'global', the client will be configured to use
the
region-specific API endpoint.
"""
+ from google.cloud.aiplatform_v1beta1 import
FeatureOnlineStoreAdminServiceClient
+
return FeatureOnlineStoreAdminServiceClient(
credentials=self.get_credentials(),
client_info=CLIENT_INFO,
@@ -101,6 +101,8 @@ class FeatureStoreHook(GoogleBaseHook):
location: str | None = None,
custom_endpoint: str | None = None,
) -> FeatureOnlineStoreServiceClient:
+ from google.cloud.aiplatform_v1beta1 import
FeatureOnlineStoreServiceClient
+
return FeatureOnlineStoreServiceClient(
credentials=self.get_credentials(),
client_info=CLIENT_INFO,
@@ -113,7 +115,7 @@ class FeatureStoreHook(GoogleBaseHook):
def create_feature_online_store(
self,
feature_online_store_id: str,
- feature_online_store: FeatureOnlineStore,
+ feature_online_store: FeatureOnlineStore | dict,
project_id: str = PROVIDE_PROJECT_ID,
location: str | None = None,
timeout: float | _MethodDefault = DEFAULT,
@@ -127,7 +129,8 @@ class FeatureStoreHook(GoogleBaseHook):
Feature Online Store aims to serve and manage features data as a part
of VertexAI MLOps.
:param feature_online_store_id: The ID of the online feature store.
- :param feature_online_store: The configuration of the online
repository.
+ :param feature_online_store: The configuration of the online
repository as a
+ FeatureOnlineStore object or dictionary.
:param project_id: The ID of the Google Cloud project that contains the
feature store. If not provided, will attempt to determine from the
environment.
:param location: The Google Cloud region where the feature store is
located
@@ -196,7 +199,7 @@ class FeatureStoreHook(GoogleBaseHook):
def create_feature_view(
self,
feature_view_id: str,
- feature_view: FeatureView,
+ feature_view: FeatureView | dict,
feature_online_store_id: str,
project_id: str = PROVIDE_PROJECT_ID,
run_sync_immediately: bool = False,
@@ -215,7 +218,8 @@ class FeatureStoreHook(GoogleBaseHook):
become the final component of the FeatureView's resource name.
This value may be up to 60 characters, and valid characters are
``[a-z0-9_]``.
The first character cannot be a number. The value must be unique
within a FeatureOnlineStore.
- :param feature_view: The configuration of the FeatureView to create.
+ :param feature_view: The configuration of the FeatureView to create as
a FeatureView object
+ or dictionary.
:param feature_online_store_id: The ID of the online feature store.
:param run_sync_immediately: If set to true, one on demand sync will
be run
immediately, regardless the FeatureView.sync_config.
@@ -327,7 +331,7 @@ class FeatureStoreHook(GoogleBaseHook):
entity_id: str | None = None,
project_id: str = PROVIDE_PROJECT_ID,
endpoint_domain_name: str | None = None,
- data_key: FeatureViewDataKey | None = None,
+ data_key: FeatureViewDataKey | dict | None = None,
data_format: int | None = None,
location: str | None = None,
timeout: float | _MethodDefault = DEFAULT,
@@ -346,7 +350,8 @@ class FeatureStoreHook(GoogleBaseHook):
and default feature store endpoint will be used, based on location
provided.
:param feature_view_id: The FeatureView ID to fetch data from.
:param feature_online_store_id: The ID of the online feature store.
- :param data_key: Optional. The request key to fetch feature values for.
+ :param data_key: Optional. The request key to fetch feature values for
as a
+ FeatureViewDataKey object or dictionary.
:param data_format: Optional. Response data format. If not set,
FeatureViewDataFormat.KEY_VALUE
will be used.
:param project_id: The ID of the Google Cloud project that contains the
diff --git
a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/feature_store.py
b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/feature_store.py
index 64fcdf38b09..0d3aa58d8ce 100644
---
a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/feature_store.py
+++
b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/feature_store.py
@@ -24,7 +24,6 @@ from typing import TYPE_CHECKING, Any
from google.api_core.exceptions import GoogleAPICallError
from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.aiplatform_v1beta1.types import FeatureViewDataFormat
from airflow.providers.common.compat.sdk import AirflowException
from airflow.providers.google.cloud.hooks.vertex_ai.feature_store import
FeatureStoreHook
@@ -128,8 +127,8 @@ class
CreateFeatureOnlineStoreOperator(GoogleCloudBaseOperator, OperationHelper)
This specifies the Google Cloud region where the feature store
resources are located.
:param feature_online_store_id: Required. The ID of the online feature
store that contains
the feature view to be synchronized. This store serves as the online
serving layer.
- :param feature_online_store: FeatureOnlineStore configuration object of
the feature online store
- to be created.
+ :param feature_online_store: FeatureOnlineStore configuration object or
dictionary of the
+ feature online store to be created.
:param gcp_conn_id: The connection ID to use for connecting to Google
Cloud Platform.
Defaults to 'google_cloud_default'.
:param impersonation_chain: Optional service account to impersonate using
short-term
@@ -155,7 +154,7 @@ class
CreateFeatureOnlineStoreOperator(GoogleCloudBaseOperator, OperationHelper)
project_id: str,
location: str,
feature_online_store_id: str,
- feature_online_store: FeatureOnlineStore,
+ feature_online_store: FeatureOnlineStore | dict,
timeout: float | _MethodDefault = DEFAULT,
retry: Retry | _MethodDefault | None = DEFAULT,
metadata: Sequence[tuple[str, str]] = (),
@@ -289,7 +288,7 @@ class CreateFeatureViewOperator(GoogleCloudBaseOperator,
OperationHelper):
of the FeatureView's resource name. This value may be up to 60
characters, and valid characters
are ``[a-z0-9_]``. The first character cannot be a number.
The value must be unique within a FeatureOnlineStore.
- :param feature_view: The configuration of the FeatureView to create.
+ :param feature_view: The FeatureView configuration object or dictionary to
create.
:param feature_online_store_id: The ID of the online feature store.
:param run_sync_immediately: If set to true, one on demand sync will be run
immediately, regardless the FeatureView.sync_config.
@@ -320,7 +319,7 @@ class CreateFeatureViewOperator(GoogleCloudBaseOperator,
OperationHelper):
self,
*,
feature_view_id: str,
- feature_view: FeatureView,
+ feature_view: FeatureView | dict,
feature_online_store_id: str,
run_sync_immediately: bool = False,
project_id: str,
@@ -442,7 +441,7 @@ class FetchFeatureValuesOperator(GoogleCloudBaseOperator,
OperationHelper):
:param entity_id: Simple ID to identify Entity to fetch feature values for.
:param feature_view_id: The FeatureView ID to fetch data from.
:param feature_online_store_id: The ID of the online feature store.
- :param data_key: The request key to fetch feature values for.
+ :param data_key: The FeatureViewDataKey object or dictionary used to fetch
feature values.
:param project_id: Required. The ID of the Google Cloud project that
contains the feature store.
This is used to identify which project's resources to interact with.
:param location: Required. The location of the feature store (e.g.,
'us-central1', 'us-east1').
@@ -476,7 +475,7 @@ class FetchFeatureValuesOperator(GoogleCloudBaseOperator,
OperationHelper):
project_id: str,
location: str,
entity_id: str | None = None,
- data_key: FeatureViewDataKey | None = None,
+ data_key: FeatureViewDataKey | dict | None = None,
timeout: float | _MethodDefault = DEFAULT,
retry: Retry | _MethodDefault | None = DEFAULT,
metadata: Sequence[tuple[str, str]] = (),
@@ -499,6 +498,8 @@ class FetchFeatureValuesOperator(GoogleCloudBaseOperator,
OperationHelper):
def execute(self, context: Context) -> dict[str, Any]:
"""Execute the get feature view sync operation."""
+ from google.cloud.aiplatform_v1beta1.types import FeatureViewDataFormat
+
hook = FeatureStoreHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
diff --git
a/providers/google/tests/system/google/cloud/vertex_ai/example_vertex_ai_feature_store.py
b/providers/google/tests/system/google/cloud/vertex_ai/example_vertex_ai_feature_store.py
index 52e8b649064..bb35a86ca81 100644
---
a/providers/google/tests/system/google/cloud/vertex_ai/example_vertex_ai_feature_store.py
+++
b/providers/google/tests/system/google/cloud/vertex_ai/example_vertex_ai_feature_store.py
@@ -25,8 +25,6 @@ from __future__ import annotations
import os
from datetime import datetime, timedelta
-from google.cloud.aiplatform_v1beta1 import FeatureOnlineStore, FeatureView,
FeatureViewDataKey
-
from airflow import DAG
from airflow.providers.google.cloud.operators.bigquery import (
BigQueryCreateEmptyDatasetOperator,
@@ -127,15 +125,15 @@ with DAG(
project_id=PROJECT_ID,
location=REGION,
feature_online_store_id=FEATURE_ONLINE_STORE_ID,
- feature_online_store=FeatureOnlineStore(
- bigtable=FeatureOnlineStore.Bigtable(
- auto_scaling=FeatureOnlineStore.Bigtable.AutoScaling(
- min_node_count=1,
- max_node_count=3,
- cpu_utilization_target=50,
- )
- )
- ),
+ feature_online_store={
+ "bigtable": {
+ "auto_scaling": {
+ "min_node_count": 1,
+ "max_node_count": 3,
+ "cpu_utilization_target": 50,
+ }
+ }
+ },
)
# [END how_to_cloud_vertex_ai_create_feature_online_store_operator]
@@ -146,13 +144,13 @@ with DAG(
location=REGION,
feature_online_store_id=FEATURE_ONLINE_STORE_ID,
feature_view_id=FEATURE_VIEW_ID,
- feature_view=FeatureView(
- big_query_source=FeatureView.BigQuerySource(
- uri=f"bq://{BQ_VIEW_FQN}",
- entity_id_columns=["entity_id"],
- ),
- sync_config=FeatureView.SyncConfig(cron="TZ=America/Los_Angeles 56
* * * *"),
- ),
+ feature_view={
+ "big_query_source": {
+ "uri": f"bq://{BQ_VIEW_FQN}",
+ "entity_id_columns": ["entity_id"],
+ },
+ "sync_config": {"cron": "TZ=America/Los_Angeles 56 * * * *"},
+ },
)
# [END how_to_cloud_vertex_ai_create_feature_view_store_operator]
@@ -201,7 +199,7 @@ with DAG(
location=REGION,
feature_online_store_id=FEATURE_ONLINE_STORE_ID,
feature_view_id=FEATURE_VIEW_ID,
- data_key=FeatureViewDataKey(FEATURE_VIEW_DATA_KEY),
+ data_key=FEATURE_VIEW_DATA_KEY,
retries=3,
retry_delay=timedelta(minutes=3),
)
diff --git
a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_feature_store.py
b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_feature_store.py
index f6b2a24d911..8b3e92def4e 100644
---
a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_feature_store.py
+++
b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_feature_store.py
@@ -17,8 +17,12 @@
from __future__ import annotations
+import subprocess
+import sys
from unittest import mock
+import pytest
+
from airflow.providers.google.cloud.hooks.vertex_ai.feature_store import
FeatureStoreHook
from unit.google.cloud.utils.base_gcp_mock import (
@@ -27,6 +31,7 @@ from unit.google.cloud.utils.base_gcp_mock import (
BASE_STRING = "airflow.providers.google.common.hooks.base_google.{}"
FEATURE_STORE_STRING =
"airflow.providers.google.cloud.hooks.vertex_ai.feature_store.{}"
+FEATURE_STORE_CLIENT_STRING = "google.cloud.aiplatform_v1beta1.{}"
TEST_GCP_CONN_ID = "test-gcp-conn-id"
TEST_PROJECT_ID = "test-project"
@@ -37,6 +42,22 @@ TEST_FEATURE_VIEW =
f"projects/{TEST_PROJECT_ID}/locations/{TEST_LOCATION}/featu
TEST_FEATURE_VIEW_SYNC_NAME = f"{TEST_FEATURE_VIEW}/featureViewSyncs/sync123"
[email protected](
+ "module",
+ (
+ "airflow.providers.google.cloud.hooks.vertex_ai.feature_store",
+ "airflow.providers.google.cloud.operators.vertex_ai.feature_store",
+ ),
+)
+def test_import_does_not_load_aiplatform_v1beta1(module):
+ code = (
+ f"import sys; import {module}; "
+ "assert 'google.cloud.aiplatform_v1beta1' not in sys.modules, "
+ "sorted(name for name in sys.modules if 'aiplatform_v1beta1' in name)"
+ )
+ subprocess.run([sys.executable, "-c", code], check=True,
capture_output=True, text=True)
+
+
class TestFeatureStoreHook:
def setup_method(self):
with mock.patch(
@@ -45,7 +66,7 @@ class TestFeatureStoreHook:
self.hook = FeatureStoreHook(gcp_conn_id=TEST_GCP_CONN_ID)
@mock.patch(BASE_STRING.format("GoogleBaseHook.get_client_options"))
-
@mock.patch(FEATURE_STORE_STRING.format("FeatureOnlineStoreAdminServiceClient"),
autospec=True)
+
@mock.patch(FEATURE_STORE_CLIENT_STRING.format("FeatureOnlineStoreAdminServiceClient"),
autospec=True)
@mock.patch(BASE_STRING.format("GoogleBaseHook.get_credentials"))
def test_get_feature_online_store_admin_service_client(
self, mock_get_credentials, mock_client, mock_get_client_options
@@ -69,6 +90,23 @@ class TestFeatureStoreHook:
api_endpoint_override =
mock_get_client_options.call_args[1]["api_endpoint_override"]
assert not api_endpoint_override
+ @mock.patch(BASE_STRING.format("GoogleBaseHook.get_client_options"))
+
@mock.patch(FEATURE_STORE_CLIENT_STRING.format("FeatureOnlineStoreServiceClient"),
autospec=True)
+ @mock.patch(BASE_STRING.format("GoogleBaseHook.get_credentials"))
+ def test_get_feature_online_store_service_client(
+ self,
+ mock_get_credentials,
+ mock_client,
+ mock_get_client_options,
+ ):
+
self.hook.get_feature_online_store_service_client(location=TEST_LOCATION)
+
+ mock_client.assert_called_once_with(
+ credentials=mock_get_credentials.return_value,
+ client_info=mock.ANY,
+ client_options=mock_get_client_options.return_value,
+ )
+
@mock.patch(FEATURE_STORE_STRING.format("FeatureStoreHook.get_feature_online_store_admin_service_client"))
def test_get_feature_view_sync(self, mock_client_getter):
mock_client = mock.MagicMock()
