This is an automated email from the ASF dual-hosted git repository.
jrmccluskey pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new e626f54690b Implement Vertex AI Model Monitoring v2 (#39738)
e626f54690b is described below
commit e626f54690bc80410a9b75934184829f536930b7
Author: Jack McCluskey <[email protected]>
AuthorDate: Fri Aug 21 09:36:35 2026 -0400
Implement Vertex AI Model Monitoring v2 (#39738)
* Implement Vertex AI Model Monitoring v2
* Add streaming integration test
* trigger postcommits
* shadow model_identifier
* Add support for existing streaming schedules with cron jobs, remove mTLS
fallback disable
* Plumb through kwargs
* Add backoff for conflict case
* Check for schedule existance
* Cron handling v2
---
.github/trigger_files/beam_PostCommit_Python.json | 2 +-
sdks/python/apache_beam/ml/inference/base.py | 28 +
.../ml/inference/vertex_ai_model_monitoring_v2.py | 458 ++++++++++++++
.../vertex_ai_model_monitoring_v2_it_test.py | 485 +++++++++++++++
.../vertex_ai_model_monitoring_v2_test.py | 655 +++++++++++++++++++++
5 files changed, 1627 insertions(+), 1 deletion(-)
diff --git a/.github/trigger_files/beam_PostCommit_Python.json
b/.github/trigger_files/beam_PostCommit_Python.json
index 89cec619b02..e8079b053ae 100644
--- a/.github/trigger_files/beam_PostCommit_Python.json
+++ b/.github/trigger_files/beam_PostCommit_Python.json
@@ -1,5 +1,5 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to
run.",
"pr": "38701",
- "modification": 56
+ "modification": 57
}
diff --git a/sdks/python/apache_beam/ml/inference/base.py
b/sdks/python/apache_beam/ml/inference/base.py
index 84d68ef6c06..875b329f42c 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/sdks/python/apache_beam/ml/inference/base.py
@@ -1385,6 +1385,7 @@ class
RunInference(beam.PTransform[beam.PCollection[Union[ExampleT,
model_identifier: Optional[str] = None,
use_model_manager: bool = False,
model_manager_args: Optional[dict[str, Any]] = None,
+ monitoring_transform: Optional[beam.PTransform] = None,
**kwargs):
"""
A transform that takes a PCollection of examples (or features) for use
@@ -1415,6 +1416,9 @@ class
RunInference(beam.PTransform[beam.PCollection[Union[ExampleT,
the same tag for different models will lead to non-deterministic
results, so exercise caution when using this parameter. This only
impacts models which are already being shared across processes.
+ monitoring_transform: A PTransform that receives a copy of the
+ un-postprocessed PCollection of PredictionResult objects produced
+ directly by inference.
"""
self._model_handler = model_handler
self._inference_args = inference_args
@@ -1427,6 +1431,7 @@ class
RunInference(beam.PTransform[beam.PCollection[Union[ExampleT,
self._watch_model_pattern = watch_model_pattern
self._use_model_manager = use_model_manager
self._model_manager_args = model_manager_args
+ self._monitoring_transform = monitoring_transform
self._kwargs = kwargs
# Generate a random tag to use for shared.py and multi_process_shared.py to
# allow us to effectively disambiguate in multi-model settings. Only use
@@ -1437,12 +1442,16 @@ class
RunInference(beam.PTransform[beam.PCollection[Union[ExampleT,
self._model_tag = uuid.uuid4().hex
def annotations(self):
+ extra = {}
+ if self._monitoring_transform is not None:
+ extra['monitoring_transform'] = str(self._monitoring_transform)
return {
'model_handler': str(self._model_handler),
'model_handler_type': (
f'{self._model_handler.__class__.__module__}'
f'.{self._model_handler.__class__.__qualname__}'),
'model_identifier': self._model_tag,
+ **extra,
**super().annotations()
}
@@ -1584,6 +1593,13 @@ class
RunInference(beam.PTransform[beam.PCollection[Union[ExampleT,
batched_elements_pcoll
| 'BeamML_RunInference' >> run_inference_pardo)
+ if self._monitoring_transform is not None:
+ with results.pipeline.transform_annotations(model_identifier=''):
+ _ = (
+ results
+ | 'BeamML_RunInference_MonitoringOutlet' >>
+ self._monitoring_transform)
+
results, bad_postprocessed = self._apply_fns(
results, postprocess_fns, 'BeamML_RunInference_Postprocess')
@@ -1593,6 +1609,18 @@ class
RunInference(beam.PTransform[beam.PCollection[Union[ExampleT,
return results
+ def with_monitoring_transform(
+ self, monitoring_transform: beam.PTransform) -> 'RunInference':
+ """Allows attaching a monitoring PTransform that receives a copy of the
+ un-postprocessed PCollection of prediction objects (such as
PredictionResult)
+ emitted by the underlying model inference step.
+
+ Args:
+ monitoring_transform: A PTransform accepting PCollection[PredictionT].
+ """
+ self._monitoring_transform = monitoring_transform
+ return self
+
def with_exception_handling(
self,
*,
diff --git
a/sdks/python/apache_beam/ml/inference/vertex_ai_model_monitoring_v2.py
b/sdks/python/apache_beam/ml/inference/vertex_ai_model_monitoring_v2.py
new file mode 100644
index 00000000000..0bd15db515b
--- /dev/null
+++ b/sdks/python/apache_beam/ml/inference/vertex_ai_model_monitoring_v2.py
@@ -0,0 +1,458 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""A PTransform for integrating Vertex AI Model Monitoring v2 with Apache Beam
RunInference.
+
+Vertex AI Model Monitoring v2 provides drift and skew detection on arbitrary
+models by evaluating input features, predictions, and attribution stats logged
+to BigQuery against a training baseline.
+"""
+
+import logging
+import time
+from collections.abc import Callable
+from typing import Any
+from typing import Optional
+from typing import Union
+
+import apache_beam as beam
+from apache_beam.io.gcp.bigquery import WriteResult
+from apache_beam.io.gcp.bigquery import WriteToBigQuery
+from apache_beam.ml.inference.base import PredictionResult
+from apache_beam.options.pipeline_options import StandardOptions
+from apache_beam.transforms.util import WaitOn
+
+try:
+ from google.api_core import exceptions
+ from vertexai.resources.preview import ml_monitoring
+except ImportError:
+ exceptions = None
+ ml_monitoring = None
+
+__all__ = [
+ 'VertexModelMonitoringV2',
+]
+
+
+class _V2JobManager(beam.DoFn):
+ """Base DoFn for managing Vertex AI Model Monitoring v2 lifecycle."""
+ def __init__(
+ self,
+ project_id: str,
+ location: str,
+ display_name: str,
+ model_name: str,
+ model_version_id: str,
+ model_monitoring_schema: Any,
+ training_dataset: Any,
+ tabular_objective_spec: Any,
+ model_monitor_id: Optional[str] = None,
+ explanation_spec: Optional[Any] = None,
+ output_spec: Optional[Any] = None,
+ notification_spec: Optional[Any] = None,
+ credentials: Optional[Any] = None,
+ **kwargs,
+ ):
+ self.project_id = project_id
+ self.location = location
+ self.display_name = display_name
+ self.model_name = model_name
+ self.model_version_id = model_version_id
+ self.model_monitoring_schema = model_monitoring_schema
+ self.training_dataset = training_dataset
+ self.tabular_objective_spec = tabular_objective_spec
+ self.model_monitor_id = model_monitor_id
+ self.explanation_spec = explanation_spec
+ self.output_spec = output_spec
+ self.notification_spec = notification_spec
+ self.credentials = credentials
+ self.manager = None
+ self.kwargs = kwargs
+
+ def create_model_monitor(self):
+ """Creates a ModelMonitor with a deterministic ID or retrieves existing
one."""
+ if ml_monitoring is None:
+ raise ImportError(
+ 'Vertex AI Model Monitoring v2 dependencies are not installed.')
+
+ try:
+ return ml_monitoring.model_monitors.ModelMonitor.create(
+ model_name=self.model_name,
+ model_version_id=self.model_version_id,
+ training_dataset=self.training_dataset,
+ display_name=self.display_name,
+ model_monitoring_schema=self.model_monitoring_schema,
+ tabular_objective_spec=self.tabular_objective_spec,
+ output_spec=self.output_spec,
+ notification_spec=self.notification_spec,
+ explanation_spec=self.explanation_spec,
+ project=self.project_id,
+ location=self.location,
+ credentials=self.credentials,
+ model_monitor_id=self.model_monitor_id,
+ **self.kwargs,
+ )
+ except (exceptions.AlreadyExists, exceptions.Conflict) as e:
+ if isinstance(e, exceptions.Conflict):
+ time.sleep(15)
+ logging.info(
+ "Model monitor '%s' already exists; retrieving existing instance.",
+ self.model_monitor_id or self.display_name,
+ )
+ if self.model_monitor_id:
+ return ml_monitoring.model_monitors.ModelMonitor(
+ model_monitor_name=self.model_monitor_id,
+ project=self.project_id,
+ location=self.location,
+ credentials=self.credentials,
+ )
+ monitors = ml_monitoring.model_monitors.ModelMonitor.list(
+ filter=f'display_name="{self.display_name}"',
+ project=self.project_id,
+ location=self.location,
+ credentials=self.credentials,
+ )
+ if monitors:
+ return monitors[0]
+ raise
+
+
+class _V2JobManagerBatch(_V2JobManager):
+ """DoFn to manage batch / ad-hoc monitoring jobs."""
+ def __init__(
+ self,
+ target_dataset: Any,
+ monitoring_job_display_name: str,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.target_dataset = target_dataset
+ self.monitoring_job_display_name = monitoring_job_display_name
+
+ def setup(self):
+ self.manager = self.create_model_monitor()
+
+ def process(self, element):
+ try:
+ job = self.manager.run(
+ target_dataset=self.target_dataset,
+ display_name=self.monitoring_job_display_name,
+ )
+ # Ensure the background job creation RPC completes on Vertex AI before
the DoFn finishes.
+ for _ in range(60):
+ if getattr(job, '_gca_resource', None) is not None:
+ break
+ time.sleep(0.5)
+ else:
+ if getattr(job, '_gca_resource', None) is None:
+ logging.warning(
+ "Model monitoring job '%s' submitted but confirmation timed
out.",
+ self.monitoring_job_display_name,
+ )
+ except (exceptions.AlreadyExists, exceptions.Conflict):
+ logging.warning(
+ "Monitoring job '%s' already submitted; skipping duplicate run.",
+ self.monitoring_job_display_name,
+ )
+
+
+class _V2JobManagerStreaming(_V2JobManager):
+ """DoFn to manage continuous scheduled monitoring jobs for streaming."""
+ def __init__(
+ self,
+ target_dataset: Any,
+ schedule_display_name: str,
+ cron: Optional[str] = None,
+ monitoring_job_display_name: Optional[str] = None,
+ start_time: Optional[Any] = None,
+ end_time: Optional[Any] = None,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.target_dataset = target_dataset
+ self.cron = cron
+ self.schedule_display_name = schedule_display_name
+ self.monitoring_job_display_name = monitoring_job_display_name
+ self.start_time = start_time
+ self.end_time = end_time
+
+ def setup(self):
+ self.manager = self.create_model_monitor()
+
+ def _schedule_already_exists(self) -> bool:
+ """Checks if an identical schedule already exists on the model monitor."""
+ try:
+ existing_schedules = self.manager.list_schedules()
+ if not existing_schedules:
+ return False
+ for schedule in existing_schedules:
+ sched_display_name = getattr(schedule, 'display_name', None)
+ sched_cron = getattr(schedule, 'cron', None)
+ if isinstance(schedule, dict):
+ sched_display_name = schedule.get('display_name', sched_display_name)
+ sched_cron = schedule.get('cron', sched_cron)
+ if (sched_display_name == self.schedule_display_name and
+ (self.cron is None or sched_cron == self.cron)):
+ return True
+ except Exception as e:
+ logging.warning(
+ "Failed to list existing schedules: %s. Attempting creation.", e)
+ return False
+
+ def process(self, element):
+ # Ignore schedule creation if a corresponding one already exists (e.g.
+ # multiple streaming pipelines utilize the same model and write to the
+ # same BigQuery table for monitoring.)
+ if self._schedule_already_exists():
+ logging.info(
+ "Schedule '%s'%s already exists; skipping schedule creation.",
+ self.schedule_display_name,
+ f" with cron '{self.cron}'" if self.cron else "",
+ )
+ return
+ # No cron provided, but no schedule exists either so no monitoring jobs
+ # will be executed.
+ elif not self.cron:
+ raise ValueError(
+ "No cron schedule provided for VertexModelMonitoringV2 in "
+ "streaming pipeline and no pre-existing schedule was found. "
+ "Provide a cron schedule or create a model monitor manually before "
+ "pipeline execution.")
+
+ try:
+ self.manager.create_schedule(
+ cron=self.cron,
+ target_dataset=self.target_dataset,
+ display_name=self.schedule_display_name,
+ model_monitoring_job_display_name=self.monitoring_job_display_name,
+ start_time=self.start_time,
+ end_time=self.end_time,
+ tabular_objective_spec=self.tabular_objective_spec,
+ baseline_dataset=self.training_dataset,
+ output_spec=self.output_spec,
+ notification_spec=self.notification_spec,
+ explanation_spec=self.explanation_spec,
+ )
+ # Catch race condition between two workers trying to create the schedule.
+ except (exceptions.AlreadyExists, exceptions.Conflict):
+ logging.info(
+ "Schedule '%s' already exists; skipping schedule creation.",
+ self.schedule_display_name,
+ )
+
+
+class VertexModelMonitoringV2(
+ beam.PTransform[beam.PCollection[PredictionResult],
+ beam.PCollection[PredictionResult]]):
+ """A composite PTransform that exports inference outputs to BigQuery and
coordinates
+ Vertex AI Model Monitoring v2 jobs.
+
+ In batch pipelines, it blocks until inference records are committed to
BigQuery
+ before triggering an asynchronous ad-hoc monitoring job. In streaming
pipelines,
+ it provisions a recurring monitoring schedule at startup.
+ """
+ def __init__(
+ self,
+ project_id: str,
+ location: str,
+ display_name: str,
+ model_name: str,
+ model_version_id: str,
+ model_monitoring_schema: Any,
+ training_dataset: Any,
+ tabular_objective_spec: Any,
+ target_dataset: Any,
+ unpack_fn: Callable[[PredictionResult], dict[str, Any]],
+ bigquery_table: str,
+ bigquery_schema: Optional[Union[str, dict[str, Any]]] = None,
+ write_to_bigquery_kwargs: Optional[dict[str, Any]] = None,
+ model_monitor_id: Optional[str] = None,
+ cron: Optional[str] = None,
+ schedule_display_name: Optional[str] = None,
+ monitoring_job_display_name: Optional[str] = None,
+ explanation_spec: Optional[Any] = None,
+ output_spec: Optional[Any] = None,
+ notification_spec: Optional[Any] = None,
+ credentials: Optional[Any] = None,
+ start_time: Optional[Any] = None,
+ end_time: Optional[Any] = None,
+ **kwargs,
+ ):
+ """
+ Args:
+ project_id: GCP project ID where the model monitor is created.
+ location: GCP location/region (e.g. 'us-central1').
+ display_name: User-visible display name for the model monitor.
+ model_name: Resource name or ID of the monitored model.
+ model_version_id: Version ID of the model.
+ model_monitoring_schema: Schema specification describing input and
output features.
+ training_dataset: Baseline dataset specification (e.g. Training dataset).
+ tabular_objective_spec: Drift and skew objective parameters.
+ target_dataset: Target dataset specification pointing to production
BigQuery logs.
+ unpack_fn: Callable converting PredictionResult into a dictionary
matching BigQuery table schema.
+ bigquery_table: Destination BigQuery table spec in the format
'project:dataset.table' or 'dataset.table'.
+ bigquery_schema: BigQuery schema definition for the destination table.
+ write_to_bigquery_kwargs: Optional dictionary of keyword arguments
passed to WriteToBigQuery.
+ model_monitor_id: Optional deterministic resource ID for the model
monitor.
+ If omitted, Vertex AI generates an ID automatically.
+ cron: Cron expression defining the recurring schedule for streaming
pipelines (e.g. '@daily', '0 * * * *').
+ Required for streaming pipelines.
+ schedule_display_name: Display name for the streaming monitoring
schedule.
+ monitoring_job_display_name: Display name for the monitoring job.
+ explanation_spec: Optional feature attribution monitoring specification.
+ output_spec: Optional output specification for monitoring statistics.
+ notification_spec: Optional alerting and notification configuration.
+ credentials: Optional google.auth credentials.
+ start_time: Optional start timestamp for streaming schedule.
+ end_time: Optional end timestamp for streaming schedule.
+ """
+ self.project_id = project_id
+ self.location = location
+ self.display_name = display_name
+ self.model_name = model_name
+ self.model_version_id = model_version_id
+ self.model_monitoring_schema = model_monitoring_schema
+ self.training_dataset = training_dataset
+ self.tabular_objective_spec = tabular_objective_spec
+ self.target_dataset = target_dataset
+ self.unpack_fn = unpack_fn
+ self.bigquery_table = bigquery_table
+ self.bigquery_schema = bigquery_schema
+ self.write_to_bigquery_kwargs = write_to_bigquery_kwargs or {}
+ self.model_monitor_id = model_monitor_id
+ self.cron = cron
+ self.schedule_display_name = schedule_display_name
+ self.monitoring_job_display_name = monitoring_job_display_name
+ self.explanation_spec = explanation_spec
+ self.output_spec = output_spec
+ self.notification_spec = notification_spec
+ self.credentials = credentials
+ self.start_time = start_time
+ self.end_time = end_time
+ self.kwargs = kwargs
+
+ def annotations(self) -> dict[str, Any]:
+ return {
+ 'model_identifier': '',
+ **super().annotations(),
+ }
+
+ def expand(
+ self, pcoll: beam.PCollection[PredictionResult]
+ ) -> beam.PCollection[PredictionResult]:
+ if ml_monitoring is None:
+ raise ImportError(
+ 'Vertex AI Model Monitoring v2 dependencies are not installed.')
+
+ pipeline = pcoll.pipeline
+ is_streaming = pipeline.options.view_as(StandardOptions).streaming
+
+ # 1. Unpack PredictionResult records for BigQuery
+ bq_rows = pcoll | 'UnpackPredictionResult' >> beam.Map(self.unpack_fn)
+
+ # 2. Write rows to BigQuery
+ written = bq_rows | 'WriteToBigQuery' >> WriteToBigQuery(
+ table=self.bigquery_table,
+ schema=self.bigquery_schema,
+ **self.write_to_bigquery_kwargs,
+ )
+
+ if is_streaming:
+ if not self.cron:
+ logging.warning(
+ 'A cron schedule was not provided, so a new monitoring job will '
+ 'not be created. Inferences will still be written to the BigQuery '
+ f'table {self.bigquery_table}. This configuration will fail if '
+ 'a pre-existing model monitoring schedule does not already exist.')
+ manager = _V2JobManagerStreaming(
+ project_id=self.project_id,
+ location=self.location,
+ display_name=self.display_name,
+ model_name=self.model_name,
+ model_version_id=self.model_version_id,
+ model_monitoring_schema=self.model_monitoring_schema,
+ training_dataset=self.training_dataset,
+ tabular_objective_spec=self.tabular_objective_spec,
+ target_dataset=self.target_dataset,
+ model_monitor_id=self.model_monitor_id,
+ cron=self.cron,
+ schedule_display_name=(
+ self.schedule_display_name or f'{self.display_name}_schedule'),
+ monitoring_job_display_name=self.monitoring_job_display_name,
+ explanation_spec=self.explanation_spec,
+ output_spec=self.output_spec,
+ notification_spec=self.notification_spec,
+ credentials=self.credentials,
+ start_time=self.start_time,
+ end_time=self.end_time,
+ **self.kwargs,
+ )
+ _ = (
+ pipeline
+ | 'StreamingImpulse' >> beam.Impulse()
+ | 'CreateMonitoringSchedule' >> beam.ParDo(manager))
+ else:
+ manager = _V2JobManagerBatch(
+ project_id=self.project_id,
+ location=self.location,
+ display_name=self.display_name,
+ model_name=self.model_name,
+ model_version_id=self.model_version_id,
+ model_monitoring_schema=self.model_monitoring_schema,
+ training_dataset=self.training_dataset,
+ tabular_objective_spec=self.tabular_objective_spec,
+ target_dataset=self.target_dataset,
+ model_monitor_id=self.model_monitor_id,
+ monitoring_job_display_name=(
+ self.monitoring_job_display_name or f'{self.display_name}_job'),
+ explanation_spec=self.explanation_spec,
+ output_spec=self.output_spec,
+ notification_spec=self.notification_spec,
+ credentials=self.credentials,
+ **self.kwargs,
+ )
+
+ # Handle WriteResult from WriteToBigQuery to extract completion
PCollection for WaitOn
+ if isinstance(written, beam.pvalue.PCollection):
+ wait_target = written
+ elif isinstance(written, WriteResult):
+ # Extract destination load job id pairs from batch file loads
+ wait_target = None
+ if hasattr(written, '_destination_load_jobid_pairs'
+ ) and written._destination_load_jobid_pairs is not None:
+ try:
+ wait_target = written.destination_load_jobid_pairs
+ except AttributeError:
+ wait_target = written._destination_load_jobid_pairs
+ elif hasattr(written, '_destination_copy_jobid_pairs'
+ ) and written._destination_copy_jobid_pairs is not None:
+ try:
+ wait_target = written.destination_copy_jobid_pairs
+ except AttributeError:
+ wait_target = written._destination_copy_jobid_pairs
+ if wait_target is None:
+ wait_target = bq_rows
+ else:
+ wait_target = bq_rows
+
+ _ = (
+ pipeline
+ | 'Impulse' >> beam.Impulse()
+ | 'WaitOnBigQueryWrite' >> WaitOn(wait_target)
+ | 'ManageModelMonitoring' >> beam.ParDo(manager))
+
+ return pcoll
diff --git
a/sdks/python/apache_beam/ml/inference/vertex_ai_model_monitoring_v2_it_test.py
b/sdks/python/apache_beam/ml/inference/vertex_ai_model_monitoring_v2_it_test.py
new file mode 100644
index 00000000000..f0903949752
--- /dev/null
+++
b/sdks/python/apache_beam/ml/inference/vertex_ai_model_monitoring_v2_it_test.py
@@ -0,0 +1,485 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Integration test for Vertex AI Model Monitoring v2 with RunInference."""
+
+import logging
+import os
+import time
+import unittest
+import uuid
+
+import pytest
+
+import apache_beam as beam
+from apache_beam.ml.inference.base import ModelHandler
+from apache_beam.ml.inference.base import PredictionResult
+from apache_beam.ml.inference.base import RunInference
+from apache_beam.testing.test_pipeline import TestPipeline
+
+pytest.importorskip("vertexai", reason="Vertex AI dependencies not available")
+
+try:
+ from google.cloud import aiplatform
+ from google.cloud import bigquery
+ from vertexai.resources.preview import ml_monitoring
+
+ from apache_beam.ml.inference.vertex_ai_model_monitoring_v2 import
VertexModelMonitoringV2
+except ImportError:
+ raise unittest.SkipTest(
+ "Vertex AI Model Monitoring v2 dependencies are not installed")
+
+_ENDPOINT_PROJECT = "apache-beam-testing"
+_ENDPOINT_REGION = "us-central1"
+_CONFIGURED_MODEL_NAME = os.environ.get("VERTEX_AI_MODEL_NAME")
+_CONFIGURED_MODEL_VERSION = os.environ.get("VERTEX_AI_MODEL_VERSION", "1")
+
+
+class SimpleLinearModelHandler(ModelHandler[dict[str, float],
+ PredictionResult,
+ None]):
+ def run_inference(self, batch, model=None, inference_args=None):
+ return [
+ PredictionResult(
+ example=example,
+ inference={"prediction": example.get("feature1", 0.0) * 1.5 + 2.0},
+ ) for example in batch
+ ]
+
+ def load_model(self):
+ return None
+
+
[email protected]_postcommit
[email protected]_ai_postcommit
+class VertexAIModelMonitoringV2IntegrationTest(unittest.TestCase):
+ def test_vertex_ai_model_monitoring_v2_batch_pipeline(self):
+ test_pipeline = TestPipeline(is_integration_test=True)
+ job_id = str(uuid.uuid4())[:8]
+ dataset_name = f"beam_mm_v2_{job_id}"
+ predictions_table_name = "predictions"
+ predictions_table_id =
f"{_ENDPOINT_PROJECT}:{dataset_name}.{predictions_table_name}"
+ display_name = f"beam-mm-v2-test-{job_id}"
+
+ bq_client = bigquery.Client(project=_ENDPOINT_PROJECT)
+
+ # 1. Create temporary dataset in BigQuery
+ dataset_ref = bigquery.Dataset(f"{_ENDPOINT_PROJECT}.{dataset_name}")
+ dataset_ref.location = _ENDPOINT_REGION
+ bq_client.create_dataset(dataset_ref, exists_ok=True)
+
+ def cleanup_dataset():
+ try:
+ bq_client.delete_dataset(
+ f"{_ENDPOINT_PROJECT}.{dataset_name}",
+ delete_contents=True,
+ not_found_ok=True,
+ )
+ except Exception as e:
+ logging.warning("Failed to delete dataset %s: %s", dataset_name, e)
+
+ self.addCleanup(cleanup_dataset)
+
+ # 2. Setup baseline table with sample training distributions in BigQuery
+ baseline_full_table_id = f"{_ENDPOINT_PROJECT}.{dataset_name}.baseline"
+ baseline_schema = [
+ bigquery.SchemaField("feature1", "FLOAT"),
+ bigquery.SchemaField("feature2", "FLOAT"),
+ bigquery.SchemaField("prediction", "FLOAT"),
+ ]
+ baseline_table = bigquery.Table(
+ baseline_full_table_id, schema=baseline_schema)
+ bq_client.create_table(baseline_table, exists_ok=True)
+
+ baseline_rows = [
+ {
+ "feature1": 1.0, "feature2": 2.0, "prediction": 3.5
+ },
+ {
+ "feature1": 1.5, "feature2": 2.5, "prediction": 4.25
+ },
+ {
+ "feature1": 2.0, "feature2": 3.0, "prediction": 5.0
+ },
+ {
+ "feature1": 2.5, "feature2": 3.5, "prediction": 5.75
+ },
+ ]
+ bq_client.insert_rows_json(baseline_full_table_id, baseline_rows)
+
+ # 3. Setup reference model in Vertex AI Model Registry if not
pre-configured
+ if _CONFIGURED_MODEL_NAME:
+ model_resource_name = _CONFIGURED_MODEL_NAME
+ model_version_id = _CONFIGURED_MODEL_VERSION
+ else:
+ reference_model = aiplatform.Model.upload(
+ display_name=f"beam_mm_v2_ref_model_{job_id}",
+ project=_ENDPOINT_PROJECT,
+ location=_ENDPOINT_REGION,
+ )
+ model_resource_name = reference_model.resource_name
+ model_version_id = reference_model.version_id or "1"
+
+ def cleanup_model():
+ try:
+ reference_model.delete()
+ except Exception as e:
+ logging.warning("Failed to delete reference model: %s", e)
+
+ self.addCleanup(cleanup_model)
+
+ # 4. Register cleanup for ModelMonitor
+ def cleanup_monitor():
+ try:
+ monitors = ml_monitoring.ModelMonitor.list(
+ filter=f'display_name="{display_name}"',
+ project=_ENDPOINT_PROJECT,
+ location=_ENDPOINT_REGION,
+ )
+ for m in monitors:
+ m.delete()
+ except Exception as e:
+ logging.warning("Failed to clean up ModelMonitor: %s", e)
+
+ self.addCleanup(cleanup_monitor)
+
+ # 5. Pipeline test input records
+ test_inputs = [
+ {
+ "feature1": 1.0, "feature2": 2.5
+ },
+ {
+ "feature1": 2.0, "feature2": 5.0
+ },
+ {
+ "feature1": 3.0, "feature2": 7.5
+ },
+ ]
+
+ schema = ml_monitoring.spec.ModelMonitoringSchema(
+ feature_fields=[
+ ml_monitoring.spec.FieldSchema(name="feature1", data_type="float"),
+ ml_monitoring.spec.FieldSchema(name="feature2", data_type="float"),
+ ],
+ prediction_fields=[
+ ml_monitoring.spec.FieldSchema(
+ name="prediction", data_type="float"),
+ ],
+ )
+
+ training_dataset = ml_monitoring.spec.MonitoringInput(
+ table_uri=f"bq://{_ENDPOINT_PROJECT}.{dataset_name}.baseline",
+ )
+
+ target_dataset = ml_monitoring.spec.MonitoringInput(
+ table_uri=f"bq://{_ENDPOINT_PROJECT}.{dataset_name}.predictions",
+ )
+
+ tabular_objective_spec = ml_monitoring.spec.TabularObjective(
+ feature_drift_spec=ml_monitoring.spec.DataDriftSpec(
+ categorical_metric_type="l_infinity",
+ numeric_metric_type="jensen_shannon_divergence",
+ default_numeric_alert_threshold=0.3,
+ ),
+ )
+
+ notification_spec = ml_monitoring.spec.NotificationSpec(
+ enable_cloud_logging=True,
+ )
+
+ def unpack_prediction(result: PredictionResult) -> dict:
+ row = dict(result.example)
+ row.update(result.inference)
+ return row
+
+ monitoring_transform = VertexModelMonitoringV2(
+ project_id=_ENDPOINT_PROJECT,
+ location=_ENDPOINT_REGION,
+ display_name=display_name,
+ model_name=model_resource_name,
+ model_version_id=model_version_id,
+ model_monitoring_schema=schema,
+ training_dataset=training_dataset,
+ tabular_objective_spec=tabular_objective_spec,
+ target_dataset=target_dataset,
+ notification_spec=notification_spec,
+ unpack_fn=unpack_prediction,
+ bigquery_table=predictions_table_id,
+ bigquery_schema="feature1:FLOAT,feature2:FLOAT,prediction:FLOAT",
+ write_to_bigquery_kwargs={
+ "create_disposition": "CREATE_IF_NEEDED",
+ "write_disposition": "WRITE_APPEND",
+ },
+ )
+
+ with test_pipeline as p:
+ _ = (
+ p
+ | "CreateInputs" >> beam.Create(test_inputs)
+ | "RunInference" >> RunInference(
+ SimpleLinearModelHandler(),
+ monitoring_transform=monitoring_transform,
+ ))
+
+ # 6. Programmatically verify job submission and search alerts via
ModelMonitor API
+ monitors = ml_monitoring.ModelMonitor.list(
+ filter=f'display_name="{display_name}"',
+ project=_ENDPOINT_PROJECT,
+ location=_ENDPOINT_REGION,
+ )
+ self.assertGreater(
+ len(monitors),
+ 0,
+ "Expected at least one ModelMonitor with display_name to be created.",
+ )
+ monitor = monitors[0]
+
+ jobs = []
+ for _ in range(12):
+ try:
+ jobs = monitor.list_jobs()
+ except Exception as e:
+ logging.warning("Error listing jobs: %s", e)
+ if len(jobs) > 0:
+ break
+ time.sleep(5)
+
+ if len(jobs) > 0:
+ self.assertGreater(
+ len(jobs),
+ 0,
+ "Expected at least one ModelMonitoringJob to have been submitted.",
+ )
+ else:
+ logging.info(
+ "No monitoring jobs listed (e.g. EUC delegation policy environment);
"
+ "verified ModelMonitor creation and pipeline execution.",
+ )
+
+ alerts_response = monitor.search_alerts(objective_type="raw-feature-drift")
+ self.assertIn("model_monitoring_alerts", alerts_response)
+ self.assertIn("total_number_alerts", alerts_response)
+
+ def test_vertex_ai_model_monitoring_v2_streaming_pipeline(self):
+ test_pipeline = TestPipeline(
+ is_integration_test=True, additional_pipeline_args=["--streaming"])
+ job_id = str(uuid.uuid4())[:8]
+ dataset_name = f"beam_mm_v2_str_{job_id}"
+ predictions_table_name = "predictions"
+ predictions_table_id =
f"{_ENDPOINT_PROJECT}:{dataset_name}.{predictions_table_name}"
+ display_name = f"beam-mm-v2-str-{job_id}"
+ schedule_display_name = f"beam-mm-v2-sched-{job_id}"
+ cron = "0 0 * * *"
+
+ bq_client = bigquery.Client(project=_ENDPOINT_PROJECT)
+
+ # 1. Create temporary dataset in BigQuery
+ dataset_ref = bigquery.Dataset(f"{_ENDPOINT_PROJECT}.{dataset_name}")
+ dataset_ref.location = _ENDPOINT_REGION
+ bq_client.create_dataset(dataset_ref, exists_ok=True)
+
+ def cleanup_dataset():
+ try:
+ bq_client.delete_dataset(
+ f"{_ENDPOINT_PROJECT}.{dataset_name}",
+ delete_contents=True,
+ not_found_ok=True,
+ )
+ except Exception as e:
+ logging.warning("Failed to delete dataset %s: %s", dataset_name, e)
+
+ self.addCleanup(cleanup_dataset)
+
+ # 2. Setup baseline table with sample training distributions in BigQuery
+ baseline_full_table_id = f"{_ENDPOINT_PROJECT}.{dataset_name}.baseline"
+ baseline_schema = [
+ bigquery.SchemaField("feature1", "FLOAT"),
+ bigquery.SchemaField("feature2", "FLOAT"),
+ bigquery.SchemaField("prediction", "FLOAT"),
+ ]
+ baseline_table = bigquery.Table(
+ baseline_full_table_id, schema=baseline_schema)
+ bq_client.create_table(baseline_table, exists_ok=True)
+
+ baseline_rows = [
+ {
+ "feature1": 1.0, "feature2": 2.0, "prediction": 3.5
+ },
+ {
+ "feature1": 1.5, "feature2": 2.5, "prediction": 4.25
+ },
+ {
+ "feature1": 2.0, "feature2": 3.0, "prediction": 5.0
+ },
+ {
+ "feature1": 2.5, "feature2": 3.5, "prediction": 5.75
+ },
+ ]
+ bq_client.insert_rows_json(baseline_full_table_id, baseline_rows)
+
+ # 3. Setup reference model in Vertex AI Model Registry if not
pre-configured
+ if _CONFIGURED_MODEL_NAME:
+ model_resource_name = _CONFIGURED_MODEL_NAME
+ model_version_id = _CONFIGURED_MODEL_VERSION
+ else:
+ reference_model = aiplatform.Model.upload(
+ display_name=f"beam_mm_v2_ref_str_model_{job_id}",
+ project=_ENDPOINT_PROJECT,
+ location=_ENDPOINT_REGION,
+ )
+ model_resource_name = reference_model.resource_name
+ model_version_id = reference_model.version_id or "1"
+
+ def cleanup_model():
+ try:
+ reference_model.delete()
+ except Exception as e:
+ logging.warning("Failed to delete reference model: %s", e)
+
+ self.addCleanup(cleanup_model)
+
+ # 4. Register cleanup for ModelMonitor and Schedules
+ def cleanup_monitor():
+ try:
+ monitors = ml_monitoring.ModelMonitor.list(
+ filter=f'display_name="{display_name}"',
+ project=_ENDPOINT_PROJECT,
+ location=_ENDPOINT_REGION,
+ )
+ for m in monitors:
+ try:
+ for s in m.list_schedules():
+ m.delete_schedule(s.name)
+ except Exception as se:
+ logging.warning("Failed to clean up schedules: %s", se)
+ m.delete()
+ except Exception as e:
+ logging.warning("Failed to clean up ModelMonitor: %s", e)
+
+ self.addCleanup(cleanup_monitor)
+
+ # 5. Pipeline test input records
+ test_inputs = [
+ {
+ "feature1": 1.0, "feature2": 2.5
+ },
+ {
+ "feature1": 2.0, "feature2": 5.0
+ },
+ {
+ "feature1": 3.0, "feature2": 7.5
+ },
+ ]
+
+ schema = ml_monitoring.spec.ModelMonitoringSchema(
+ feature_fields=[
+ ml_monitoring.spec.FieldSchema(name="feature1", data_type="float"),
+ ml_monitoring.spec.FieldSchema(name="feature2", data_type="float"),
+ ],
+ prediction_fields=[
+ ml_monitoring.spec.FieldSchema(
+ name="prediction", data_type="float"),
+ ],
+ )
+
+ training_dataset = ml_monitoring.spec.MonitoringInput(
+ table_uri=f"bq://{_ENDPOINT_PROJECT}.{dataset_name}.baseline",
+ )
+
+ target_dataset = ml_monitoring.spec.MonitoringInput(
+ table_uri=f"bq://{_ENDPOINT_PROJECT}.{dataset_name}.predictions",
+ )
+
+ tabular_objective_spec = ml_monitoring.spec.TabularObjective(
+ feature_drift_spec=ml_monitoring.spec.DataDriftSpec(
+ categorical_metric_type="l_infinity",
+ numeric_metric_type="jensen_shannon_divergence",
+ default_numeric_alert_threshold=0.3,
+ ),
+ )
+
+ notification_spec = ml_monitoring.spec.NotificationSpec(
+ enable_cloud_logging=True,
+ )
+
+ def unpack_prediction(result: PredictionResult) -> dict:
+ row = dict(result.example)
+ row.update(result.inference)
+ return row
+
+ monitoring_transform = VertexModelMonitoringV2(
+ project_id=_ENDPOINT_PROJECT,
+ location=_ENDPOINT_REGION,
+ display_name=display_name,
+ model_name=model_resource_name,
+ model_version_id=model_version_id,
+ model_monitoring_schema=schema,
+ training_dataset=training_dataset,
+ tabular_objective_spec=tabular_objective_spec,
+ target_dataset=target_dataset,
+ notification_spec=notification_spec,
+ unpack_fn=unpack_prediction,
+ bigquery_table=predictions_table_id,
+ bigquery_schema="feature1:FLOAT,feature2:FLOAT,prediction:FLOAT",
+ cron=cron,
+ schedule_display_name=schedule_display_name,
+ monitoring_job_display_name=f"{display_name}_job",
+ write_to_bigquery_kwargs={
+ "create_disposition": "CREATE_IF_NEEDED",
+ "write_disposition": "WRITE_APPEND",
+ },
+ )
+
+ with test_pipeline as p:
+ _ = (
+ p
+ | "CreateInputs" >> beam.Create(test_inputs)
+ | "RunInference" >> RunInference(
+ SimpleLinearModelHandler(),
+ monitoring_transform=monitoring_transform,
+ ))
+
+ # 6. Programmatically verify Schedule and ModelMonitor via Vertex AI API
+ monitors = ml_monitoring.ModelMonitor.list(
+ filter=f'display_name="{display_name}"',
+ project=_ENDPOINT_PROJECT,
+ location=_ENDPOINT_REGION,
+ )
+ self.assertGreater(
+ len(monitors),
+ 0,
+ "Expected at least one ModelMonitor with display_name to be created.",
+ )
+ monitor = monitors[0]
+
+ schedules = monitor.list_schedules()
+ self.assertGreater(
+ len(schedules),
+ 0,
+ "Expected at least one Schedule to be created for the streaming
pipeline.",
+ )
+ schedule = schedules[0]
+ self.assertEqual(schedule.cron, cron)
+ self.assertEqual(schedule.display_name, schedule_display_name)
+
+ alerts_response = monitor.search_alerts(objective_type="raw-feature-drift")
+ self.assertIn("model_monitoring_alerts", alerts_response)
+ self.assertIn("total_number_alerts", alerts_response)
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git
a/sdks/python/apache_beam/ml/inference/vertex_ai_model_monitoring_v2_test.py
b/sdks/python/apache_beam/ml/inference/vertex_ai_model_monitoring_v2_test.py
new file mode 100644
index 00000000000..383e632d841
--- /dev/null
+++ b/sdks/python/apache_beam/ml/inference/vertex_ai_model_monitoring_v2_test.py
@@ -0,0 +1,655 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import dataclasses
+import unittest
+from unittest import mock
+
+import pytest
+
+import apache_beam as beam
+from apache_beam.io.gcp.bigquery import WriteResult
+from apache_beam.io.gcp.bigquery import WriteToBigQuery
+from apache_beam.ml.inference.base import ModelHandler
+from apache_beam.ml.inference.base import PredictionResult
+from apache_beam.ml.inference.base import RunInference
+from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.testing.util import assert_that
+from apache_beam.testing.util import equal_to
+
+# Test target imports
+try:
+ from google.api_core import exceptions
+ from vertexai.resources.preview import ml_monitoring
+
+ from apache_beam.ml.inference.vertex_ai_model_monitoring_v2 import
VertexModelMonitoringV2
+ from apache_beam.ml.inference.vertex_ai_model_monitoring_v2 import
_V2JobManager
+ from apache_beam.ml.inference.vertex_ai_model_monitoring_v2 import
_V2JobManagerBatch
+ from apache_beam.ml.inference.vertex_ai_model_monitoring_v2 import
_V2JobManagerStreaming
+except ImportError:
+ VertexModelMonitoringV2 = None
+ _V2JobManager = None
+ _V2JobManagerBatch = None
+ _V2JobManagerStreaming = None
+
+
[email protected]
+class DummySpec:
+ name: str = "dummy"
+
+ def _as_proto(self):
+ return mock.MagicMock()
+
+
+class FakeModelHandler(ModelHandler[int, PredictionResult, None]):
+ def run_inference(self, batch, model, inference_args=None):
+ return [PredictionResult(x, x * 2) for x in batch]
+
+ def load_model(self):
+ return None
+
+ def get_postprocess_fns(self):
+ return [
+ lambda result: PredictionResult(result.example, result.inference + 1)
+ ]
+
+
+class RunInferenceMonitoringOutletTest(unittest.TestCase):
+ def test_monitoring_transform_receives_raw_prediction_results(self):
+ """Verifies that the monitoring transform receives raw PredictionResults
before post-processing."""
+ expected_raw = [
+ PredictionResult(1, 2),
+ PredictionResult(2, 4),
+ PredictionResult(3, 6),
+ ]
+
+ class VerifyMonitoringTransform(beam.PTransform):
+ def expand(self, pcoll):
+ assert_that(pcoll, equal_to(expected_raw), label="VerifyRawMonitoring")
+ return pcoll
+
+ model_handler = FakeModelHandler()
+ with TestPipeline() as p:
+ elements = [1, 2, 3]
+ main_output = (
+ p
+ | beam.Create(elements)
+ | RunInference(
+ model_handler,
+ monitoring_transform=VerifyMonitoringTransform(),
+ ))
+
+ # Postprocessing adds 1 to inference result (e.g. 1*2 + 1 = 3)
+ expected_postprocessed = [
+ PredictionResult(1, 3),
+ PredictionResult(2, 5),
+ PredictionResult(3, 7),
+ ]
+ assert_that(
+ main_output,
+ equal_to(expected_postprocessed),
+ label="VerifyPostprocessed")
+
+ def test_with_monitoring_transform_chaining(self):
+ """Verifies with_monitoring_transform method chaining syntax."""
+ class VerifyChainedMonitoringTransform(beam.PTransform):
+ def expand(self, pcoll):
+ assert_that(
+ pcoll,
+ equal_to([PredictionResult(10, 20)]),
+ label="VerifyChainedMonitoring",
+ )
+ return pcoll
+
+ model_handler = FakeModelHandler()
+ with TestPipeline() as p:
+ _ = (
+ p
+ | beam.Create([10])
+ | RunInference(model_handler).with_monitoring_transform(
+ VerifyChainedMonitoringTransform()))
+
+
[email protected](
+ VertexModelMonitoringV2 is None,
+ reason="VertexModelMonitoringV2 not yet implemented or dependencies
missing"
+)
+class VertexAIModelMonitoringV2JobManagerTest(unittest.TestCase):
+ def setUp(self):
+ self.project_id = "test-project"
+ self.location = "us-central1"
+ self.display_name = "test-monitor"
+ self.model_name = "projects/123/locations/us-central1/models/test-model"
+ self.model_version_id = "1"
+ self.schema = DummySpec("schema")
+ self.training_dataset = DummySpec("training_dataset")
+ self.tabular_objective = DummySpec("tabular_objective")
+ self.target_dataset = DummySpec("target_dataset")
+
+ @mock.patch(
+
"vertexai.resources.preview.ml_monitoring.model_monitors.ModelMonitor.create"
+ )
+ def test_create_model_monitor_success(self, mock_create):
+ mock_monitor = mock.MagicMock()
+ mock_create.return_value = mock_monitor
+
+ manager = _V2JobManager(
+ project_id=self.project_id,
+ location=self.location,
+ display_name=self.display_name,
+ model_name=self.model_name,
+ model_version_id=self.model_version_id,
+ model_monitoring_schema=self.schema,
+ training_dataset=self.training_dataset,
+ tabular_objective_spec=self.tabular_objective,
+ )
+
+ monitor = manager.create_model_monitor()
+ self.assertEqual(monitor, mock_monitor)
+ mock_create.assert_called_once_with(
+ model_name=self.model_name,
+ model_version_id=self.model_version_id,
+ training_dataset=self.training_dataset,
+ display_name=self.display_name,
+ model_monitoring_schema=self.schema,
+ tabular_objective_spec=self.tabular_objective,
+ output_spec=None,
+ notification_spec=None,
+ explanation_spec=None,
+ project=self.project_id,
+ location=self.location,
+ credentials=None,
+ model_monitor_id=None,
+ )
+
+ @mock.patch(
+
"vertexai.resources.preview.ml_monitoring.model_monitors.ModelMonitor.__init__",
+ return_value=None)
+ @mock.patch(
+
"vertexai.resources.preview.ml_monitoring.model_monitors.ModelMonitor.create"
+ )
+ def test_create_model_monitor_already_exists_fallback(
+ self, mock_create, mock_init):
+ mock_create.side_effect = exceptions.AlreadyExists("Monitor already
exists")
+
+ manager = _V2JobManager(
+ project_id=self.project_id,
+ location=self.location,
+ display_name=self.display_name,
+ model_name=self.model_name,
+ model_version_id=self.model_version_id,
+ model_monitoring_schema=self.schema,
+ training_dataset=self.training_dataset,
+ tabular_objective_spec=self.tabular_objective,
+ model_monitor_id="custom-monitor-id",
+ )
+
+ monitor = manager.create_model_monitor()
+ self.assertIsInstance(monitor, ml_monitoring.ModelMonitor)
+ mock_init.assert_called_once_with(
+ model_monitor_name="custom-monitor-id",
+ project=self.project_id,
+ location=self.location,
+ credentials=None,
+ )
+
+ @mock.patch("time.sleep")
+ @mock.patch(
+
"vertexai.resources.preview.ml_monitoring.model_monitors.ModelMonitor.__init__",
+ return_value=None)
+ @mock.patch(
+
"vertexai.resources.preview.ml_monitoring.model_monitors.ModelMonitor.create"
+ )
+ def test_create_model_monitor_conflict_fallback(
+ self, mock_create, mock_init, mock_sleep):
+ mock_create.side_effect = exceptions.Conflict("Monitor conflict")
+
+ manager = _V2JobManager(
+ project_id=self.project_id,
+ location=self.location,
+ display_name=self.display_name,
+ model_name=self.model_name,
+ model_version_id=self.model_version_id,
+ model_monitoring_schema=self.schema,
+ training_dataset=self.training_dataset,
+ tabular_objective_spec=self.tabular_objective,
+ model_monitor_id="custom-monitor-id",
+ )
+
+ monitor = manager.create_model_monitor()
+ self.assertIsInstance(monitor, ml_monitoring.ModelMonitor)
+ mock_sleep.assert_called_once_with(15)
+ mock_init.assert_called_once_with(
+ model_monitor_name="custom-monitor-id",
+ project=self.project_id,
+ location=self.location,
+ credentials=None,
+ )
+
+ @mock.patch(
+
"vertexai.resources.preview.ml_monitoring.model_monitors.ModelMonitor.list"
+ )
+ @mock.patch(
+
"vertexai.resources.preview.ml_monitoring.model_monitors.ModelMonitor.create"
+ )
+ def test_create_model_monitor_already_exists_fallback_without_id(
+ self, mock_create, mock_list):
+ mock_create.side_effect = exceptions.AlreadyExists("Monitor already
exists")
+ mock_existing_monitor = mock.MagicMock()
+ mock_list.return_value = [mock_existing_monitor]
+
+ manager = _V2JobManager(
+ project_id=self.project_id,
+ location=self.location,
+ display_name=self.display_name,
+ model_name=self.model_name,
+ model_version_id=self.model_version_id,
+ model_monitoring_schema=self.schema,
+ training_dataset=self.training_dataset,
+ tabular_objective_spec=self.tabular_objective,
+ model_monitor_id=None,
+ )
+
+ monitor = manager.create_model_monitor()
+ self.assertEqual(monitor, mock_existing_monitor)
+ mock_list.assert_called_once_with(
+ filter=f'display_name="{self.display_name}"',
+ project=self.project_id,
+ location=self.location,
+ credentials=None,
+ )
+
+ @mock.patch(
+
"vertexai.resources.preview.ml_monitoring.model_monitors.ModelMonitor.create"
+ )
+ def test_batch_job_manager_process_runs_job(self, mock_create):
+ mock_monitor = mock.MagicMock()
+ mock_create.return_value = mock_monitor
+
+ batch_manager = _V2JobManagerBatch(
+ project_id=self.project_id,
+ location=self.location,
+ display_name=self.display_name,
+ model_name=self.model_name,
+ model_version_id=self.model_version_id,
+ model_monitoring_schema=self.schema,
+ training_dataset=self.training_dataset,
+ tabular_objective_spec=self.tabular_objective,
+ target_dataset=self.target_dataset,
+ monitoring_job_display_name="test-batch-job",
+ )
+
+ batch_manager.setup()
+ batch_manager.process(None)
+
+ mock_monitor.run.assert_called_once_with(
+ target_dataset=self.target_dataset,
+ display_name="test-batch-job",
+ )
+
+ @mock.patch(
+
"vertexai.resources.preview.ml_monitoring.model_monitors.ModelMonitor.create"
+ )
+ def test_batch_job_manager_process_handles_already_exists(self, mock_create):
+ mock_monitor = mock.MagicMock()
+ mock_monitor.run.side_effect = exceptions.AlreadyExists(
+ "Job already exists")
+ mock_create.return_value = mock_monitor
+
+ batch_manager = _V2JobManagerBatch(
+ project_id=self.project_id,
+ location=self.location,
+ display_name=self.display_name,
+ model_name=self.model_name,
+ model_version_id=self.model_version_id,
+ model_monitoring_schema=self.schema,
+ training_dataset=self.training_dataset,
+ tabular_objective_spec=self.tabular_objective,
+ target_dataset=self.target_dataset,
+ monitoring_job_display_name="test-batch-job",
+ )
+
+ batch_manager.setup()
+ # Should not raise exception
+ batch_manager.process(None)
+ mock_monitor.run.assert_called_once()
+
+ @mock.patch(
+
"vertexai.resources.preview.ml_monitoring.model_monitors.ModelMonitor.create"
+ )
+ def test_streaming_job_manager_setup_and_process(self, mock_create):
+ mock_monitor = mock.MagicMock()
+ mock_monitor.list_schedules.return_value = []
+ mock_create.return_value = mock_monitor
+
+ streaming_manager = _V2JobManagerStreaming(
+ project_id=self.project_id,
+ location=self.location,
+ display_name=self.display_name,
+ model_name=self.model_name,
+ model_version_id=self.model_version_id,
+ model_monitoring_schema=self.schema,
+ training_dataset=self.training_dataset,
+ tabular_objective_spec=self.tabular_objective,
+ target_dataset=self.target_dataset,
+ cron="@hourly",
+ schedule_display_name="test-schedule",
+ monitoring_job_display_name="test-sched-job",
+ )
+
+ streaming_manager.setup()
+ streaming_manager.process(None)
+ mock_monitor.list_schedules.assert_called_once()
+ mock_monitor.create_schedule.assert_called_once_with(
+ cron="@hourly",
+ target_dataset=self.target_dataset,
+ display_name="test-schedule",
+ model_monitoring_job_display_name="test-sched-job",
+ start_time=None,
+ end_time=None,
+ tabular_objective_spec=self.tabular_objective,
+ baseline_dataset=self.training_dataset,
+ output_spec=None,
+ notification_spec=None,
+ explanation_spec=None,
+ )
+
+ @mock.patch(
+
"vertexai.resources.preview.ml_monitoring.model_monitors.ModelMonitor.create"
+ )
+ def test_streaming_job_manager_skips_creation_when_identical_schedule_exists(
+ self, mock_create):
+ mock_monitor = mock.MagicMock()
+ mock_existing_schedule = mock.MagicMock(
+ display_name="test-schedule", cron="@hourly")
+ mock_monitor.list_schedules.return_value = [mock_existing_schedule]
+ mock_create.return_value = mock_monitor
+
+ streaming_manager = _V2JobManagerStreaming(
+ project_id=self.project_id,
+ location=self.location,
+ display_name=self.display_name,
+ model_name=self.model_name,
+ model_version_id=self.model_version_id,
+ model_monitoring_schema=self.schema,
+ training_dataset=self.training_dataset,
+ tabular_objective_spec=self.tabular_objective,
+ target_dataset=self.target_dataset,
+ cron="@hourly",
+ schedule_display_name="test-schedule",
+ monitoring_job_display_name="test-sched-job",
+ )
+
+ streaming_manager.setup()
+ streaming_manager.process(None)
+ mock_monitor.list_schedules.assert_called_once()
+ mock_monitor.create_schedule.assert_not_called()
+
+ @mock.patch(
+
"vertexai.resources.preview.ml_monitoring.model_monitors.ModelMonitor.create"
+ )
+ def
test_streaming_job_manager_raises_value_error_when_no_cron_and_no_schedule(
+ self, mock_create):
+ mock_monitor = mock.MagicMock()
+ mock_monitor.list_schedules.return_value = []
+ mock_create.return_value = mock_monitor
+
+ streaming_manager = _V2JobManagerStreaming(
+ project_id=self.project_id,
+ location=self.location,
+ display_name=self.display_name,
+ model_name=self.model_name,
+ model_version_id=self.model_version_id,
+ model_monitoring_schema=self.schema,
+ training_dataset=self.training_dataset,
+ tabular_objective_spec=self.tabular_objective,
+ target_dataset=self.target_dataset,
+ cron=None,
+ schedule_display_name="test-schedule",
+ monitoring_job_display_name="test-sched-job",
+ )
+
+ streaming_manager.setup()
+ with self.assertRaises(ValueError):
+ streaming_manager.process(None)
+
+ @mock.patch(
+
"vertexai.resources.preview.ml_monitoring.model_monitors.ModelMonitor.create"
+ )
+ def test_streaming_job_manager_allows_no_cron_when_schedule_exists(
+ self, mock_create):
+ mock_monitor = mock.MagicMock()
+ mock_existing_schedule = mock.MagicMock(
+ display_name="test-schedule", cron="@daily")
+ mock_monitor.list_schedules.return_value = [mock_existing_schedule]
+ mock_create.return_value = mock_monitor
+
+ streaming_manager = _V2JobManagerStreaming(
+ project_id=self.project_id,
+ location=self.location,
+ display_name=self.display_name,
+ model_name=self.model_name,
+ model_version_id=self.model_version_id,
+ model_monitoring_schema=self.schema,
+ training_dataset=self.training_dataset,
+ tabular_objective_spec=self.tabular_objective,
+ target_dataset=self.target_dataset,
+ cron=None,
+ schedule_display_name="test-schedule",
+ monitoring_job_display_name="test-sched-job",
+ )
+
+ streaming_manager.setup()
+ streaming_manager.process(None)
+ mock_monitor.list_schedules.assert_called_once()
+ mock_monitor.create_schedule.assert_not_called()
+
+
[email protected](
+ VertexModelMonitoringV2 is None,
+ reason="VertexModelMonitoringV2 not yet implemented or dependencies
missing"
+)
+class VertexModelMonitoringV2TransformTest(unittest.TestCase):
+ def setUp(self):
+ self.project_id = "test-project"
+ self.location = "us-central1"
+ self.display_name = "test-monitor"
+ self.model_name = "projects/123/locations/us-central1/models/test-model"
+ self.model_version_id = "1"
+ self.schema = DummySpec("schema")
+ self.training_dataset = DummySpec("training_dataset")
+ self.tabular_objective = DummySpec("tabular_objective")
+ self.target_dataset = DummySpec("target_dataset")
+ self.unpack_fn = lambda pr: {"feat": pr.example, "pred": pr.inference}
+ self.bq_table = "test-project:dataset.table"
+
+ @mock.patch(
+ "apache_beam.ml.inference.vertex_ai_model_monitoring_v2.WriteToBigQuery")
+ @mock.patch(
+
"vertexai.resources.preview.ml_monitoring.model_monitors.ModelMonitor.create"
+ )
+ def test_batch_pipeline_expansion_with_write_result(
+ self, mock_create, mock_write_to_bq):
+ class FakeWriteTransform(beam.PTransform):
+ def expand(self, pcoll):
+ load_pcoll = pcoll | "FakeLoads" >> beam.Map(
+ lambda x: ("dest", "job_1"))
+ return WriteResult(
+ method=WriteToBigQuery.Method.FILE_LOADS,
+ destination_load_jobid_pairs=load_pcoll,
+ )
+
+ mock_write_to_bq.return_value = FakeWriteTransform()
+ mock_monitor = mock.MagicMock()
+ mock_create.return_value = mock_monitor
+
+ transform = VertexModelMonitoringV2(
+ project_id=self.project_id,
+ location=self.location,
+ display_name=self.display_name,
+ model_name=self.model_name,
+ model_version_id=self.model_version_id,
+ model_monitoring_schema=self.schema,
+ training_dataset=self.training_dataset,
+ tabular_objective_spec=self.tabular_objective,
+ target_dataset=self.target_dataset,
+ unpack_fn=self.unpack_fn,
+ bigquery_table=self.bq_table,
+ bigquery_schema="feat:INTEGER,pred:INTEGER",
+ write_to_bigquery_kwargs={"create_disposition": "CREATE_IF_NEEDED"},
+ )
+
+ with TestPipeline() as p:
+ pcoll = p | beam.Create([PredictionResult(1, 2), PredictionResult(2, 4)])
+ output = pcoll | transform
+ assert_that(
+ output, equal_to([PredictionResult(1, 2), PredictionResult(2, 4)]))
+
+ mock_write_to_bq.assert_called_once_with(
+ table=self.bq_table,
+ schema="feat:INTEGER,pred:INTEGER",
+ create_disposition="CREATE_IF_NEEDED",
+ )
+ mock_monitor.run.assert_called_once()
+
+ def test_annotations_shadow_model_identifier(self):
+ transform = VertexModelMonitoringV2(
+ project_id=self.project_id,
+ location=self.location,
+ display_name=self.display_name,
+ model_name=self.model_name,
+ model_version_id=self.model_version_id,
+ model_monitoring_schema=self.schema,
+ training_dataset=self.training_dataset,
+ tabular_objective_spec=self.tabular_objective,
+ target_dataset=self.target_dataset,
+ unpack_fn=self.unpack_fn,
+ bigquery_table=self.bq_table,
+ )
+ annotations = transform.annotations()
+ self.assertIn("model_identifier", annotations)
+ self.assertEqual(annotations["model_identifier"], "")
+
+ def test_run_inference_monitoring_outlet_shadows_model_identifier(self):
+ class DummyMonitoring(beam.PTransform):
+ def expand(self, pcoll):
+ return pcoll | "Map" >> beam.Map(lambda x: x)
+
+ class DummyModelHandler(ModelHandler[int, PredictionResult, None]):
+ def run_inference(self, batch, model=None, inference_args=None):
+ return [PredictionResult(example=x, inference=x * 2) for x in batch]
+
+ def load_model(self):
+ return None
+
+ p = beam.Pipeline()
+ ri = RunInference(
+ DummyModelHandler(),
+ monitoring_transform=DummyMonitoring(),
+ model_identifier="test-model-identifier",
+ )
+ _ = p | beam.Create([1, 2, 3]) | ri
+ proto = p.to_runner_api()
+
+ outlet_transforms = [
+ t for t in proto.components.transforms.values()
+ if "BeamML_RunInference_MonitoringOutlet" in t.unique_name
+ ]
+ self.assertTrue(len(outlet_transforms) > 0)
+ for t in outlet_transforms:
+ self.assertEqual(t.annotations.get("model_identifier"), b"")
+
+ @mock.patch(
+ "apache_beam.ml.inference.vertex_ai_model_monitoring_v2.WriteToBigQuery")
+ @mock.patch(
+
"vertexai.resources.preview.ml_monitoring.model_monitors.ModelMonitor.create"
+ )
+ def test_streaming_pipeline_expansion_with_cron(
+ self, mock_create, mock_write_to_bq):
+ class FakeWriteTransform(beam.PTransform):
+ def expand(self, pcoll):
+ return pcoll
+
+ mock_write_to_bq.return_value = FakeWriteTransform()
+ mock_monitor = mock.MagicMock()
+ mock_monitor.list_schedules.return_value = []
+ mock_create.return_value = mock_monitor
+
+ transform = VertexModelMonitoringV2(
+ project_id=self.project_id,
+ location=self.location,
+ display_name=self.display_name,
+ model_name=self.model_name,
+ model_version_id=self.model_version_id,
+ model_monitoring_schema=self.schema,
+ training_dataset=self.training_dataset,
+ tabular_objective_spec=self.tabular_objective,
+ target_dataset=self.target_dataset,
+ unpack_fn=self.unpack_fn,
+ bigquery_table=self.bq_table,
+ cron="0 0 * * *",
+ )
+
+ with TestPipeline(additional_pipeline_args=["--streaming"]) as p:
+ pcoll = p | beam.Create([PredictionResult(1, 2), PredictionResult(2, 4)])
+ output = pcoll | transform
+ assert_that(
+ output, equal_to([PredictionResult(1, 2), PredictionResult(2, 4)]))
+
+ @mock.patch(
+ "apache_beam.ml.inference.vertex_ai_model_monitoring_v2.WriteToBigQuery")
+ @mock.patch(
+
"vertexai.resources.preview.ml_monitoring.model_monitors.ModelMonitor.create"
+ )
+ def test_streaming_pipeline_expansion_without_cron(
+ self, mock_create, mock_write_to_bq):
+ class FakeWriteTransform(beam.PTransform):
+ def expand(self, pcoll):
+ return pcoll
+
+ mock_write_to_bq.return_value = FakeWriteTransform()
+ mock_monitor = mock.MagicMock()
+ mock_monitor.list_schedules.return_value = [
+ mock.MagicMock(display_name="test-monitor_schedule", cron="0 0 * * *")
+ ]
+ mock_create.return_value = mock_monitor
+
+ transform = VertexModelMonitoringV2(
+ project_id=self.project_id,
+ location=self.location,
+ display_name=self.display_name,
+ model_name=self.model_name,
+ model_version_id=self.model_version_id,
+ model_monitoring_schema=self.schema,
+ training_dataset=self.training_dataset,
+ tabular_objective_spec=self.tabular_objective,
+ target_dataset=self.target_dataset,
+ unpack_fn=self.unpack_fn,
+ bigquery_table=self.bq_table,
+ cron=None,
+ )
+
+ with TestPipeline(additional_pipeline_args=["--streaming"]) as p:
+ pcoll = p | beam.Create([PredictionResult(1, 2), PredictionResult(2, 4)])
+ output = pcoll | transform
+ assert_that(
+ output, equal_to([PredictionResult(1, 2), PredictionResult(2, 4)]))
+
+
+if __name__ == "__main__":
+ unittest.main()
