Skip to content

Commit 285ec2c

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
feat: Add Create FeatureMonitor function to FeatureGroup in Vertex AI SDK
PiperOrigin-RevId: 698102159
1 parent 41cd5a8 commit 285ec2c

File tree

5 files changed

+270
-9
lines changed

5 files changed

+270
-9
lines changed

tests/unit/vertexai/feature_store_constants.py

+15
Original file line numberDiff line numberDiff line change
@@ -377,4 +377,19 @@
377377
name=_TEST_FG1_FM1_PATH,
378378
description=_TEST_FG1_FM1_DESCRIPTION,
379379
labels=_TEST_FG1_FM1_LABELS,
380+
schedule_config=types.feature_monitor.ScheduleConfig(cron="0 0 * * *"),
381+
feature_selection_config=types.feature_monitor.FeatureSelectionConfig(
382+
feature_configs=[
383+
types.feature_monitor.FeatureSelectionConfig.FeatureConfig(
384+
feature_id="my_fg1_f1",
385+
drift_threshold=0.3,
386+
),
387+
types.feature_monitor.FeatureSelectionConfig.FeatureConfig(
388+
feature_id="my_fg1_f2",
389+
drift_threshold=0.4,
390+
),
391+
]
392+
),
380393
)
394+
_TEST_FG1_FM1_FEATURE_SELECTION_CONFIGS = [("my_fg1_f1", 0.3), ("my_fg1_f2", 0.4)]
395+
_TEST_FG1_FM1_SCHEDULE_CONFIG = "0 0 * * *"

tests/unit/vertexai/test_feature_group.py

+102
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,9 @@
2424
from google.api_core import operation as ga_operation
2525
from google.cloud import aiplatform
2626
from google.cloud.aiplatform import base
27+
from google.cloud.aiplatform_v1beta1.services.feature_registry_service import (
28+
FeatureRegistryServiceClient,
29+
)
2730
from vertexai.resources.preview.feature_store import (
2831
feature_group,
2932
)
@@ -75,8 +78,16 @@
7578
_TEST_FG1_F2_POINT_OF_CONTACT,
7679
_TEST_FG1_F2_VERSION_COLUMN_NAME,
7780
_TEST_FG1_FEATURE_LIST,
81+
_TEST_FG1_FM1,
82+
_TEST_FG1_FM1_ID,
83+
_TEST_FG1_FM1_PATH,
84+
_TEST_FG1_FM1_DESCRIPTION,
85+
_TEST_FG1_FM1_LABELS,
86+
_TEST_FG1_FM1_FEATURE_SELECTION_CONFIGS,
87+
_TEST_FG1_FM1_SCHEDULE_CONFIG,
7888
)
7989
from test_feature import feature_eq
90+
from test_feature_monitor import feature_monitor_eq
8091

8192

8293
pytestmark = pytest.mark.usefixtures("google_auth_mock")
@@ -137,6 +148,18 @@ def create_feature_mock():
137148
yield create_feature_mock
138149

139150

151+
@pytest.fixture
152+
def create_feature_monitor_mock():
153+
with patch.object(
154+
FeatureRegistryServiceClient,
155+
"create_feature_monitor",
156+
) as create_feature_monitor_mock:
157+
create_feature_monitor_lro_mock = mock.Mock(ga_operation.Operation)
158+
create_feature_monitor_lro_mock.result.return_value = _TEST_FG1_FM1
159+
create_feature_monitor_mock.return_value = create_feature_monitor_lro_mock
160+
yield create_feature_monitor_mock
161+
162+
140163
@pytest.fixture
141164
def create_feature_with_version_column_mock():
142165
with patch.object(
@@ -827,3 +850,82 @@ def test_list_features(get_fg_mock, list_features_mock):
827850
point_of_contact=_TEST_FG1_F2_POINT_OF_CONTACT,
828851
version_column_name=_TEST_FG1_F2_VERSION_COLUMN_NAME,
829852
)
853+
854+
855+
@pytest.mark.parametrize("create_request_timeout", [None, 1.0])
856+
def test_create_feature_monitor(
857+
get_fg_mock,
858+
get_feature_monitor_mock,
859+
create_feature_monitor_mock,
860+
fg_logger_mock,
861+
create_request_timeout,
862+
):
863+
aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
864+
865+
fg = FeatureGroup(_TEST_FG1_ID)
866+
feature_monitor = fg.create_feature_monitor(
867+
_TEST_FG1_FM1_ID,
868+
description=_TEST_FG1_FM1_DESCRIPTION,
869+
labels=_TEST_FG1_FM1_LABELS,
870+
schedule_config=_TEST_FG1_FM1_SCHEDULE_CONFIG,
871+
feature_selection_configs=_TEST_FG1_FM1_FEATURE_SELECTION_CONFIGS,
872+
create_request_timeout=create_request_timeout,
873+
)
874+
875+
expected_feature_monitor = types.feature_monitor.FeatureMonitor(
876+
description=_TEST_FG1_FM1_DESCRIPTION,
877+
labels=_TEST_FG1_FM1_LABELS,
878+
schedule_config=types.feature_monitor.ScheduleConfig(
879+
cron=_TEST_FG1_FM1_SCHEDULE_CONFIG
880+
),
881+
feature_selection_config=types.feature_monitor.FeatureSelectionConfig(
882+
feature_configs=[
883+
types.feature_monitor.FeatureSelectionConfig.FeatureConfig(
884+
feature_id="my_fg1_f1", drift_threshold=0.3
885+
),
886+
types.feature_monitor.FeatureSelectionConfig.FeatureConfig(
887+
feature_id="my_fg1_f2", drift_threshold=0.4
888+
),
889+
]
890+
),
891+
)
892+
create_feature_monitor_mock.assert_called_once_with(
893+
parent=_TEST_FG1_PATH,
894+
feature_monitor_id=_TEST_FG1_FM1_ID,
895+
feature_monitor=expected_feature_monitor,
896+
metadata=(),
897+
timeout=create_request_timeout,
898+
)
899+
900+
feature_monitor_eq(
901+
feature_monitor,
902+
name=_TEST_FG1_FM1_ID,
903+
resource_name=_TEST_FG1_FM1_PATH,
904+
project=_TEST_PROJECT,
905+
location=_TEST_LOCATION,
906+
description=_TEST_FG1_FM1_DESCRIPTION,
907+
labels=_TEST_FG1_FM1_LABELS,
908+
schedule_config=_TEST_FG1_FM1_SCHEDULE_CONFIG,
909+
feature_selection_configs=_TEST_FG1_FM1_FEATURE_SELECTION_CONFIGS,
910+
)
911+
912+
fg_logger_mock.assert_has_calls(
913+
[
914+
call("Creating FeatureMonitor"),
915+
call(
916+
f"Create FeatureMonitor backing LRO:"
917+
f" {create_feature_monitor_mock.return_value.operation.name}"
918+
),
919+
call(
920+
"FeatureMonitor created. Resource name:"
921+
" projects/test-project/locations/us-central1/featureGroups/"
922+
"my_fg1/featureMonitors/my_fg1_fm1"
923+
),
924+
call("To use this FeatureMonitor in another session:"),
925+
call(
926+
"feature_monitor = aiplatform.FeatureMonitor("
927+
"'projects/test-project/locations/us-central1/featureGroups/"
928+
"my_fg1/featureMonitors/my_fg1_fm1')"
929+
),
930+
]
931+
)

tests/unit/vertexai/test_feature_monitor.py

+9-1
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
#
1717

1818
import re
19-
from typing import Dict
19+
from typing import Dict, List, Tuple
2020

2121
from google.cloud import aiplatform
2222
from google.cloud.aiplatform import base
@@ -27,6 +27,8 @@
2727
from feature_store_constants import _TEST_FG1_ID
2828
from feature_store_constants import _TEST_LOCATION
2929
from feature_store_constants import _TEST_PROJECT
30+
from feature_store_constants import _TEST_FG1_FM1_SCHEDULE_CONFIG
31+
from feature_store_constants import _TEST_FG1_FM1_FEATURE_SELECTION_CONFIGS
3032
from vertexai.resources.preview import FeatureMonitor
3133
import pytest
3234

@@ -42,6 +44,8 @@ def feature_monitor_eq(
4244
location: str,
4345
description: str,
4446
labels: Dict[str, str],
47+
schedule_config: str,
48+
feature_selection_configs: List[Tuple[str, float]],
4549
):
4650
"""Check if a Feature Monitor has the appropriate values set."""
4751
assert feature_monitor_to_check.name == name
@@ -97,6 +101,8 @@ def test_init_with_feature_monitor_id(get_feature_monitor_mock):
97101
location=_TEST_LOCATION,
98102
description=_TEST_FG1_FM1_DESCRIPTION,
99103
labels=_TEST_FG1_FM1_LABELS,
104+
schedule_config=_TEST_FG1_FM1_SCHEDULE_CONFIG,
105+
feature_selection_configs=_TEST_FG1_FM1_FEATURE_SELECTION_CONFIGS,
100106
)
101107

102108

@@ -118,4 +124,6 @@ def test_init_with_feature_path(get_feature_monitor_mock):
118124
location=_TEST_LOCATION,
119125
description=_TEST_FG1_FM1_DESCRIPTION,
120126
labels=_TEST_FG1_FM1_LABELS,
127+
schedule_config=_TEST_FG1_FM1_SCHEDULE_CONFIG,
128+
feature_selection_configs=_TEST_FG1_FM1_FEATURE_SELECTION_CONFIGS,
121129
)

vertexai/resources/preview/feature_store/feature_group.py

+123-7
Original file line numberDiff line numberDiff line change
@@ -15,20 +15,15 @@
1515
# limitations under the License.
1616
#
1717

18-
from typing import (
19-
Sequence,
20-
Tuple,
21-
Dict,
22-
List,
23-
Optional,
24-
)
18+
from typing import Dict, List, Optional, Sequence, Tuple
2519
from google.auth import credentials as auth_credentials
2620
from google.cloud.aiplatform import base, initializer
2721
from google.cloud.aiplatform import utils
2822
from google.cloud.aiplatform.compat.types import (
2923
feature as gca_feature,
3024
feature_group as gca_feature_group,
3125
io as gca_io,
26+
feature_monitor_v1beta1 as gca_feature_monitor,
3227
)
3328
from vertexai.resources.preview.feature_store.utils import (
3429
FeatureGroupBigQuerySource,
@@ -428,6 +423,127 @@ def get_feature_monitor(
428423
credentials=credentials,
429424
)
430425

426+
def create_feature_monitor(
427+
self,
428+
name: str,
429+
description: Optional[str] = None,
430+
labels: Optional[Dict[str, str]] = None,
431+
schedule_config: Optional[str] = None,
432+
feature_selection_configs: Optional[List[Tuple[str, float]]] = None,
433+
project: Optional[str] = None,
434+
location: Optional[str] = None,
435+
credentials: Optional[auth_credentials.Credentials] = None,
436+
request_metadata: Optional[Sequence[Tuple[str, str]]] = None,
437+
create_request_timeout: Optional[float] = None,
438+
) -> FeatureMonitor:
439+
"""Creates a new feature monitor.
440+
441+
Args:
442+
name: The name of the feature monitor.
443+
description: Description of the feature monitor.
444+
labels:
445+
The labels with user-defined metadata to organize your FeatureMonitors.
446+
Label keys and values can be no longer than 64 characters
447+
(Unicode codepoints), can only contain lowercase letters,
448+
numeric characters, underscores and dashes. International
449+
characters are allowed.
450+
451+
See https://goo.gl/xmQnxf for more information on and examples
452+
of labels. No more than 64 user labels can be associated with
453+
one FeatureMonitor (System labels are excluded)." System reserved label
454+
keys are prefixed with "aiplatform.googleapis.com/" and are
455+
immutable.
456+
schedule_config:
457+
Configures when data is to be monitored for this
458+
FeatureMonitor. At the end of the scheduled time,
459+
the stats and drift are generated for the selected features.
460+
Example format: "TZ=America/New_York 0 9 * * *" (monitors
461+
daily at 9 AM EST).
462+
feature_selection_configs:
463+
List of tuples of feature id and monitoring threshold. If unset,
464+
all features in the feature group will be monitored, and the
465+
default thresholds 0.3 will be used.
466+
project:
467+
Project to create feature in. If unset, the project set in
468+
aiplatform.init will be used.
469+
location:
470+
Location to create feature in. If not set, location set in
471+
aiplatform.init will be used.
472+
credentials:
473+
Custom credentials to use to create this feature. Overrides
474+
credentials set in aiplatform.init.
475+
request_metadata:
476+
Strings which should be sent along with the request as metadata.
477+
create_request_timeout:
478+
The timeout for the create request in seconds.
479+
480+
Returns:
481+
FeatureMonitor - the FeatureMonitor resource object.
482+
"""
483+
484+
gapic_feature_monitor = gca_feature_monitor.FeatureMonitor()
485+
486+
if description:
487+
gapic_feature_monitor.description = description
488+
489+
if labels:
490+
utils.validate_labels(labels)
491+
gapic_feature_monitor.labels = labels
492+
493+
if request_metadata is None:
494+
request_metadata = ()
495+
496+
if schedule_config:
497+
gapic_feature_monitor.schedule_config = gca_feature_monitor.ScheduleConfig(
498+
cron=schedule_config
499+
)
500+
501+
if feature_selection_configs is None:
502+
raise ValueError(
503+
"Please specify feature_configs: features to be monitored and"
504+
" their thresholds."
505+
)
506+
507+
if feature_selection_configs is not None:
508+
gapic_feature_monitor.feature_selection_config.feature_configs = [
509+
gca_feature_monitor.FeatureSelectionConfig.FeatureConfig(
510+
feature_id=feature_id,
511+
drift_threshold=threshold if threshold else 0.3,
512+
)
513+
for feature_id, threshold in feature_selection_configs
514+
]
515+
516+
api_client = self.__class__._instantiate_client(
517+
location=location, credentials=credentials
518+
)
519+
520+
create_feature_monitor_lro = api_client.select_version(
521+
"v1beta1"
522+
).create_feature_monitor(
523+
parent=self.resource_name,
524+
feature_monitor=gapic_feature_monitor,
525+
feature_monitor_id=name,
526+
metadata=request_metadata,
527+
timeout=create_request_timeout,
528+
)
529+
530+
_LOGGER.log_create_with_lro(FeatureMonitor, create_feature_monitor_lro)
531+
532+
created_feature_monitor = create_feature_monitor_lro.result()
533+
534+
_LOGGER.log_create_complete(
535+
FeatureMonitor, created_feature_monitor, "feature_monitor"
536+
)
537+
538+
feature_monitor_obj = FeatureMonitor(
539+
name=created_feature_monitor.name,
540+
project=project,
541+
location=location,
542+
credentials=credentials,
543+
)
544+
545+
return feature_monitor_obj
546+
431547
@property
432548
def source(self) -> FeatureGroupBigQuerySource:
433549
return FeatureGroupBigQuerySource(

vertexai/resources/preview/feature_store/feature_monitor.py

+21-1
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
#
1717

1818
import re
19-
from typing import Optional
19+
from typing import List, Optional, Tuple
2020
from google.auth import credentials as auth_credentials
2121
from google.cloud.aiplatform import base
2222
from google.cloud.aiplatform import utils
@@ -108,3 +108,23 @@ def __init__(
108108
def description(self) -> str:
109109
"""The description of the feature monitor."""
110110
return self._gca_resource.description
111+
112+
@property
113+
def schedule_config(self) -> str:
114+
"""The schedule config of the feature monitor."""
115+
return self._gca_resource.schedule_config.cron
116+
117+
@property
118+
def feature_selection_configs(self) -> List[Tuple[str, float]]:
119+
"""The feature and it's drift threshold configs of the feature monitor."""
120+
configs: List[Tuple[str, float]] = []
121+
for (
122+
feature_config
123+
) in self._gca_resource.feature_selection_config.feature_configs:
124+
configs.append(
125+
(
126+
feature_config.feature_id,
127+
feature_config.threshold if feature_config.threshold else 0.3,
128+
)
129+
)
130+
return configs

0 commit comments

Comments
 (0)