feat: Add properties feature_stats_and_anomalies to FeatureMonitorJob

vertex-sdk-bot · copybara-github · commit 92feb600b72e · 2024-12-10T13:49:52.000-08:00
PiperOrigin-RevId: 704833055
diff --git a/tests/unit/vertexai/feature_store_constants.py b/tests/unit/vertexai/feature_store_constants.py
@@ -421,10 +421,33 @@
 _TEST_FG1_FMJ1_PATH = f"{_TEST_PARENT}/featureGroups/{_TEST_FG1_ID}/featureMonitors/{_TEST_FG1_FM1_ID}/featureMonitorJobs/{_TEST_FG1_FMJ1_ID}"
 _TEST_FG1_FMJ1_DESCRIPTION = "My feature monitor job 1 in feature monitor 1"
 _TEST_FG1_FMJ1_LABELS = {"my_fg1_feature_monitor_job": "fmj1"}
+_TEST_FG1_F1_FEATURE_STATS_AND_ANOMALY = types.feature_monitor.FeatureStatsAndAnomaly(
+    feature_id="my_fg1_f1",
+    distribution_deviation=0.5,
+    drift_detection_threshold=0.4,
+    drift_detected=True,
+    feature_monitor_job_id=_TEST_FG1_FMJ1_ID,
+    feature_monitor_id=_TEST_FG1_FM1_ID,
+)
+_TEST_FG1_F2_FEATURE_STATS_AND_ANOMALY = types.feature_monitor.FeatureStatsAndAnomaly(
+    feature_id="my_fg1_f2",
+    distribution_deviation=0.2,
+    drift_detection_threshold=0.4,
+    drift_detected=False,
+    feature_monitor_job_id=_TEST_FG1_FMJ1_ID,
+    feature_monitor_id=_TEST_FG1_FM1_ID,
+)
+_TEST_FG1_FMJ1_FEATURE_STATS_AND_ANOMALIES = [
+    _TEST_FG1_F1_FEATURE_STATS_AND_ANOMALY,
+    _TEST_FG1_F2_FEATURE_STATS_AND_ANOMALY,
+]
 _TEST_FG1_FMJ1 = types.feature_monitor_job.FeatureMonitorJob(
     name=_TEST_FG1_FMJ1_PATH,
     description=_TEST_FG1_FMJ1_DESCRIPTION,
     labels=_TEST_FG1_FMJ1_LABELS,
+    job_summary=types.feature_monitor_job.FeatureMonitorJob.JobSummary(
+        feature_stats_and_anomalies=_TEST_FG1_FMJ1_FEATURE_STATS_AND_ANOMALIES
+    ),
 )
 _TEST_FG1_FMJ2_ID = "1234567891"
 _TEST_FG1_FMJ2_PATH = f"{_TEST_PARENT}/featureGroups/{_TEST_FG1_ID}/featureMonitors/{_TEST_FG1_FM1_ID}/featureMonitorJobs/{_TEST_FG1_FMJ2_ID}"
diff --git a/tests/unit/vertexai/test_feature_monitor.py b/tests/unit/vertexai/test_feature_monitor.py
@@ -16,7 +16,7 @@
 #
 
 import re
-from typing import Dict, List, Tuple
+from typing import Dict, List, Optional, Tuple
 from unittest.mock import patch
 
 from google.cloud import aiplatform
@@ -34,6 +34,7 @@
     _TEST_FG1_FM1_SCHEDULE_CONFIG,
     _TEST_FG1_FMJ1,
     _TEST_FG1_FMJ1_DESCRIPTION,
+    _TEST_FG1_FMJ1_FEATURE_STATS_AND_ANOMALIES,
     _TEST_FG1_FMJ1_ID,
     _TEST_FG1_FMJ1_LABELS,
     _TEST_FG1_FMJ_LIST,
@@ -127,13 +128,21 @@ def feature_monitor_job_eq(
     location: str,
     description: str,
     labels: Dict[str, str],
+    feature_stats_and_anomalies: Optional[
+        List[types.feature_monitor.FeatureStatsAndAnomaly]
+    ] = None,
 ):
     """Check if a Feature Monitor Job has the appropriate values set."""
     assert feature_monitor_job_to_check.resource_name == resource_name
     assert feature_monitor_job_to_check.project == project
     assert feature_monitor_job_to_check.location == location
     assert feature_monitor_job_to_check.description == description
     assert feature_monitor_job_to_check.labels == labels
+    if feature_stats_and_anomalies:
+        assert (
+            feature_monitor_job_to_check.feature_stats_and_anomalies
+            == feature_stats_and_anomalies
+        )
 
 
 def test_init_with_feature_monitor_id_and_no_fg_id_raises_error():
@@ -247,6 +256,7 @@ def test_init_with_feature_monitor_job_path(get_feature_monitor_job_mock):
         location=_TEST_LOCATION,
         description=_TEST_FG1_FMJ1_DESCRIPTION,
         labels=_TEST_FG1_FMJ1_LABELS,
+        feature_stats_and_anomalies=_TEST_FG1_FMJ1_FEATURE_STATS_AND_ANOMALIES,
     )
 
 
@@ -291,6 +301,7 @@ def test_create_feature_monitor_job(
         location=_TEST_LOCATION,
         description=_TEST_FG1_FMJ1_DESCRIPTION,
         labels=_TEST_FG1_FMJ1_LABELS,
+        feature_stats_and_anomalies=_TEST_FG1_FMJ1_FEATURE_STATS_AND_ANOMALIES,
     )
 
 
@@ -320,6 +331,7 @@ def test_get_feature_monitor_job(
         location=_TEST_LOCATION,
         description=_TEST_FG1_FMJ1_DESCRIPTION,
         labels=_TEST_FG1_FMJ1_LABELS,
+        feature_stats_and_anomalies=_TEST_FG1_FMJ1_FEATURE_STATS_AND_ANOMALIES,
     )
 
 
diff --git a/vertexai/resources/preview/feature_store/feature_monitor.py b/vertexai/resources/preview/feature_store/feature_monitor.py
@@ -189,6 +189,15 @@ def description(self) -> str:
             """The description of the feature monitor."""
             return self._gca_resource.description
 
+        @property
+        def feature_stats_and_anomalies(
+            self,
+        ) -> List[gca_feature_monitor.FeatureStatsAndAnomaly]:
+            """The feature stats and anomaly of the feature monitor job."""
+            if self._gca_resource.job_summary:
+                return self._gca_resource.job_summary.feature_stats_and_anomalies
+            return []
+
     def create_feature_monitor_job(
         self,
         description: Optional[str] = None,
