From 7420ad22aeaf349f1ee81ad212d740546611223f Mon Sep 17 00:00:00 2001 From: Vitaly Terentyev Date: Fri, 21 Aug 2026 10:24:00 +0400 Subject: [PATCH 1/4] Update google-cloud-bigtable --- sdks/python/container/ml/py310/base_image_requirements.txt | 2 +- sdks/python/container/ml/py310/gpu_image_requirements.txt | 2 +- sdks/python/container/ml/py311/base_image_requirements.txt | 2 +- sdks/python/container/ml/py311/gpu_image_requirements.txt | 2 +- sdks/python/container/ml/py312/base_image_requirements.txt | 2 +- sdks/python/container/ml/py312/gpu_image_requirements.txt | 2 +- sdks/python/container/ml/py313/base_image_requirements.txt | 2 +- sdks/python/container/py310/base_image_requirements.txt | 2 +- sdks/python/container/py311/base_image_requirements.txt | 2 +- sdks/python/container/py312/base_image_requirements.txt | 2 +- sdks/python/container/py313/base_image_requirements.txt | 2 +- sdks/python/container/py314/base_image_requirements.txt | 2 +- 12 files changed, 12 insertions(+), 12 deletions(-) diff --git a/sdks/python/container/ml/py310/base_image_requirements.txt b/sdks/python/container/ml/py310/base_image_requirements.txt index 93ce4dcba32a..7edb58d44a63 100644 --- a/sdks/python/container/ml/py310/base_image_requirements.txt +++ b/sdks/python/container/ml/py310/base_image_requirements.txt @@ -70,7 +70,7 @@ google-auth-httplib2==0.2.1 google-cloud-aiplatform==1.164.0 google-cloud-bigquery==3.43.0 google-cloud-bigquery-storage==2.40.0 -google-cloud-bigtable==2.41.0 +google-cloud-bigtable==2.42.0 google-cloud-build==3.38.1 google-cloud-core==2.6.1 google-cloud-dataflow-client==0.13.0 diff --git a/sdks/python/container/ml/py310/gpu_image_requirements.txt b/sdks/python/container/ml/py310/gpu_image_requirements.txt index 79f7005b9ffe..3f74a02528d3 100644 --- a/sdks/python/container/ml/py310/gpu_image_requirements.txt +++ b/sdks/python/container/ml/py310/gpu_image_requirements.txt @@ -90,7 +90,7 @@ google-auth-httplib2==0.2.1 google-cloud-aiplatform==1.164.0 google-cloud-bigquery==3.43.0 google-cloud-bigquery-storage==2.40.0 -google-cloud-bigtable==2.41.0 +google-cloud-bigtable==2.42.0 google-cloud-build==3.38.1 google-cloud-core==2.6.1 google-cloud-dataflow-client==0.13.0 diff --git a/sdks/python/container/ml/py311/base_image_requirements.txt b/sdks/python/container/ml/py311/base_image_requirements.txt index 39bbd6d23af0..e7364e139bd5 100644 --- a/sdks/python/container/ml/py311/base_image_requirements.txt +++ b/sdks/python/container/ml/py311/base_image_requirements.txt @@ -68,7 +68,7 @@ google-auth-httplib2==0.2.1 google-cloud-aiplatform==1.164.0 google-cloud-bigquery==3.43.0 google-cloud-bigquery-storage==2.40.0 -google-cloud-bigtable==2.41.0 +google-cloud-bigtable==2.42.0 google-cloud-build==3.38.1 google-cloud-core==2.6.1 google-cloud-dataflow-client==0.13.0 diff --git a/sdks/python/container/ml/py311/gpu_image_requirements.txt b/sdks/python/container/ml/py311/gpu_image_requirements.txt index 5ab774d0218f..faa238024a7c 100644 --- a/sdks/python/container/ml/py311/gpu_image_requirements.txt +++ b/sdks/python/container/ml/py311/gpu_image_requirements.txt @@ -88,7 +88,7 @@ google-auth-httplib2==0.2.1 google-cloud-aiplatform==1.164.0 google-cloud-bigquery==3.43.0 google-cloud-bigquery-storage==2.40.0 -google-cloud-bigtable==2.41.0 +google-cloud-bigtable==2.42.0 google-cloud-build==3.38.1 google-cloud-core==2.6.1 google-cloud-dataflow-client==0.13.0 diff --git a/sdks/python/container/ml/py312/base_image_requirements.txt b/sdks/python/container/ml/py312/base_image_requirements.txt index 2224b2ed4086..fc268a729b1d 100644 --- a/sdks/python/container/ml/py312/base_image_requirements.txt +++ b/sdks/python/container/ml/py312/base_image_requirements.txt @@ -67,7 +67,7 @@ google-auth-httplib2==0.2.1 google-cloud-aiplatform==1.164.0 google-cloud-bigquery==3.43.0 google-cloud-bigquery-storage==2.40.0 -google-cloud-bigtable==2.41.0 +google-cloud-bigtable==2.42.0 google-cloud-build==3.38.1 google-cloud-core==2.6.1 google-cloud-dataflow-client==0.13.0 diff --git a/sdks/python/container/ml/py312/gpu_image_requirements.txt b/sdks/python/container/ml/py312/gpu_image_requirements.txt index 85987d16d3c5..87498144539e 100644 --- a/sdks/python/container/ml/py312/gpu_image_requirements.txt +++ b/sdks/python/container/ml/py312/gpu_image_requirements.txt @@ -87,7 +87,7 @@ google-auth-httplib2==0.2.1 google-cloud-aiplatform==1.164.0 google-cloud-bigquery==3.43.0 google-cloud-bigquery-storage==2.40.0 -google-cloud-bigtable==2.41.0 +google-cloud-bigtable==2.42.0 google-cloud-build==3.38.1 google-cloud-core==2.6.1 google-cloud-dataflow-client==0.13.0 diff --git a/sdks/python/container/ml/py313/base_image_requirements.txt b/sdks/python/container/ml/py313/base_image_requirements.txt index ecebb0ad1513..96525cc92743 100644 --- a/sdks/python/container/ml/py313/base_image_requirements.txt +++ b/sdks/python/container/ml/py313/base_image_requirements.txt @@ -67,7 +67,7 @@ google-auth-httplib2==0.2.1 google-cloud-aiplatform==1.164.0 google-cloud-bigquery==3.43.0 google-cloud-bigquery-storage==2.40.0 -google-cloud-bigtable==2.41.0 +google-cloud-bigtable==2.42.0 google-cloud-build==3.38.1 google-cloud-core==2.6.1 google-cloud-dataflow-client==0.13.0 diff --git a/sdks/python/container/py310/base_image_requirements.txt b/sdks/python/container/py310/base_image_requirements.txt index 6c9e2d6a6d5a..a322bff46125 100644 --- a/sdks/python/container/py310/base_image_requirements.txt +++ b/sdks/python/container/py310/base_image_requirements.txt @@ -64,7 +64,7 @@ google-auth-httplib2==0.2.1 google-cloud-aiplatform==1.164.0 google-cloud-bigquery==3.43.0 google-cloud-bigquery-storage==2.40.0 -google-cloud-bigtable==2.41.0 +google-cloud-bigtable==2.42.0 google-cloud-build==3.38.1 google-cloud-core==2.6.1 google-cloud-dataflow-client==0.13.0 diff --git a/sdks/python/container/py311/base_image_requirements.txt b/sdks/python/container/py311/base_image_requirements.txt index 5d88cf1fe6a9..161f5e575433 100644 --- a/sdks/python/container/py311/base_image_requirements.txt +++ b/sdks/python/container/py311/base_image_requirements.txt @@ -62,7 +62,7 @@ google-auth-httplib2==0.2.1 google-cloud-aiplatform==1.164.0 google-cloud-bigquery==3.43.0 google-cloud-bigquery-storage==2.40.0 -google-cloud-bigtable==2.41.0 +google-cloud-bigtable==2.42.0 google-cloud-build==3.38.1 google-cloud-core==2.6.1 google-cloud-dataflow-client==0.13.0 diff --git a/sdks/python/container/py312/base_image_requirements.txt b/sdks/python/container/py312/base_image_requirements.txt index c74bf184722b..a16784409a1c 100644 --- a/sdks/python/container/py312/base_image_requirements.txt +++ b/sdks/python/container/py312/base_image_requirements.txt @@ -61,7 +61,7 @@ google-auth-httplib2==0.2.1 google-cloud-aiplatform==1.164.0 google-cloud-bigquery==3.43.0 google-cloud-bigquery-storage==2.40.0 -google-cloud-bigtable==2.41.0 +google-cloud-bigtable==2.42.0 google-cloud-build==3.38.1 google-cloud-core==2.6.1 google-cloud-dataflow-client==0.13.0 diff --git a/sdks/python/container/py313/base_image_requirements.txt b/sdks/python/container/py313/base_image_requirements.txt index bd44f421d3a8..b4b8f9c2e1fa 100644 --- a/sdks/python/container/py313/base_image_requirements.txt +++ b/sdks/python/container/py313/base_image_requirements.txt @@ -61,7 +61,7 @@ google-auth-httplib2==0.2.1 google-cloud-aiplatform==1.164.0 google-cloud-bigquery==3.43.0 google-cloud-bigquery-storage==2.40.0 -google-cloud-bigtable==2.41.0 +google-cloud-bigtable==2.42.0 google-cloud-build==3.38.1 google-cloud-core==2.6.1 google-cloud-dataflow-client==0.13.0 diff --git a/sdks/python/container/py314/base_image_requirements.txt b/sdks/python/container/py314/base_image_requirements.txt index 14407fb49fa4..bdeb1d85c4a9 100644 --- a/sdks/python/container/py314/base_image_requirements.txt +++ b/sdks/python/container/py314/base_image_requirements.txt @@ -61,7 +61,7 @@ google-auth-httplib2==0.2.1 google-cloud-aiplatform==1.164.0 google-cloud-bigquery==3.43.0 google-cloud-bigquery-storage==2.40.0 -google-cloud-bigtable==2.41.0 +google-cloud-bigtable==2.42.0 google-cloud-build==3.38.1 google-cloud-core==2.6.1 google-cloud-dataflow-client==0.13.0 From e626f54690bc80410a9b75934184829f536930b7 Mon Sep 17 00:00:00 2001 From: Jack McCluskey <34928439+jrmccluskey@users.noreply.github.com> Date: Fri, 21 Aug 2026 09:36:35 -0400 Subject: [PATCH 2/4] 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 --- .../trigger_files/beam_PostCommit_Python.json | 2 +- sdks/python/apache_beam/ml/inference/base.py | 28 + .../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(-) create mode 100644 sdks/python/apache_beam/ml/inference/vertex_ai_model_monitoring_v2.py create mode 100644 sdks/python/apache_beam/ml/inference/vertex_ai_model_monitoring_v2_it_test.py create mode 100644 sdks/python/apache_beam/ml/inference/vertex_ai_model_monitoring_v2_test.py diff --git a/.github/trigger_files/beam_PostCommit_Python.json b/.github/trigger_files/beam_PostCommit_Python.json index 89cec619b020..e8079b053aeb 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 84d68ef6c061..875b329f42c4 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -1385,6 +1385,7 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__( 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 @@ def failure_callback(exception: Exception, element: Any): 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 @@ def failure_callback(exception: Exception, element: Any): 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 000000000000..0bd15db515b8 --- /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 000000000000..f09039497527 --- /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 + + +@pytest.mark.it_postcommit +@pytest.mark.vertex_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 000000000000..383e632d8416 --- /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 + + +@dataclasses.dataclass +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())) + + +@pytest.mark.skipif( + 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() + + +@pytest.mark.skipif( + 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() From 3b9e647dec1e56725b21a90103f9cd7b5711c76c Mon Sep 17 00:00:00 2001 From: Elia Liu Date: Sat, 22 Aug 2026 00:24:16 +1000 Subject: [PATCH 3/4] [Prism] Schedule consumers of a self checkpointing source (#39572) * [Prism] Schedule consumers of a self checkpointing source Fixes #39446. An unbounded source SDF that returns a process continuation residual has that residual re-queued as a pending element carrying the input element's event time, MinTimestamp for an Impulse rooted source, so the stage's watermark never advances. updateWatermarks returns no refreshes when the output watermark does not move, and PersistBundle marked only the producing stage as changed, so a consumer that was just handed pending elements was never surfaced to the scheduler. Its elements accumulated forever. checkForQuiescence cannot catch this, because the source stays schedulable, so the job live locks instead of failing fast. PersistBundle now records the consumers that accepted data, and updateWatermarks hands them to the scheduler on the path where the output watermark does not advance. Recording is gated on any bundle having returned a residual, so a pipeline that never self checkpoints keeps its previous bundle scheduling. bundleReady additionally lets a stateful stage with no side inputs run on pending data alone, under the same gate, since statefulStageKind.buildEventTimeBundle takes data at any watermark and gates only timers. Its stillSchedulable now requires buildable work, and a key that supplies nothing no longer consumes the OneKeyPerBundle slot, is not marked in progress, and does not hold the bundle's minimum timestamp. A consumer that reads a side input still waits on the watermark, since side input readiness is derived from it. * Publish all JmsIO records after the pipeline starts The test pre-published part of the records before running the pipeline to work around the starvation this change fixes. All records are now published after the pipeline starts, and the duplicate-tolerant assertion tightens back to an exact count. Also bumps the messaging postcommit trigger file. * Apply suggestion from @Abacn --------- Co-authored-by: Yi Hu --- ...tCommit_Python_Xlang_Messaging_Direct.json | 2 +- .../prism/internal/engine/elementmanager.go | 91 +++++- .../elementmanager_continuation_test.go | 299 ++++++++++++++++++ .../io/external/xlang_jmsio_it_test.py | 13 +- 4 files changed, 381 insertions(+), 24 deletions(-) create mode 100644 sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager_continuation_test.go diff --git a/.github/trigger_files/beam_PostCommit_Python_Xlang_Messaging_Direct.json b/.github/trigger_files/beam_PostCommit_Python_Xlang_Messaging_Direct.json index d6a91b7e2e86..38ae1cf68222 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Xlang_Messaging_Direct.json +++ b/.github/trigger_files/beam_PostCommit_Python_Xlang_Messaging_Direct.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 7 + "modification": 8 } diff --git a/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go b/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go index 2b502c679db9..7ebf174d9ec0 100644 --- a/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go +++ b/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go @@ -235,6 +235,10 @@ type ElementManager struct { livePending atomic.Int64 // An accessible live pending count. DEBUG USE ONLY pendingElements sync.WaitGroup // pendingElements counts all unprocessed elements in a job. Jobs with no pending elements terminate successfully. + // Latched once any bundle returns a residual, after which a stage watermark + // can be pinned indefinitely and can't be relied on to schedule consumers. + sawResidual atomic.Bool + processTimeEvents *stageRefreshQueue // Manages sequence of stage updates when interfacing with processing time. Callers must hold refreshCond.L lock. testStreamHandler *testStreamHandler // Optional test stream handler when a test stream is in the pipeline. } @@ -841,6 +845,17 @@ func reElementResiduals(residuals []Residual, inputInfo PColInfo, rb RunBundle) // input elements, and the committed output elements. func (em *ElementManager) PersistBundle(rb RunBundle, col2Coders map[string]PColInfo, d TentativeData, inputInfo PColInfo, residuals Residuals) { stage := em.stages[rb.StageID] + // Consumers that received data from this bundle, recorded so they can still + // be scheduled when this stage's output watermark is held back. Only needed + // once something self checkpoints, so pipelines that never do keep their + // previous bundle scheduling exactly. + if len(residuals.Data) > 0 { + em.sawResidual.Store(true) + } + var changedConsumers set[string] + if em.sawResidual.Load() { + changedConsumers = set[string]{} + } var seq int for output, data := range d.Raw { info := col2Coders[output] @@ -910,6 +925,9 @@ func (em *ElementManager) PersistBundle(rb RunBundle, col2Coders map[string]PCol count = consumer.AddPending(em, newPending) } em.addPending(count) + if changedConsumers != nil && count > 0 { + changedConsumers.insert(sID) + } } for _, link := range sideConsumers { consumer := em.stages[link.Global] @@ -942,6 +960,12 @@ func (em *ElementManager) PersistBundle(rb RunBundle, col2Coders map[string]PCol // even if a panic occurs during `em.addPending`. This prevents potential deadlocks // if the waitgroup unexpectedly drops below zero due to a runner bug. defer stage.mu.Unlock() + if len(changedConsumers) > 0 { + if stage.consumersWithNewData == nil { + stage.consumersWithNewData = set[string]{} + } + stage.consumersWithNewData.merge(changedConsumers) + } completed := stage.inprogress[rb.BundleID] em.addPending(-len(completed.es)) delete(stage.inprogress, rb.BundleID) @@ -1101,6 +1125,9 @@ func (em *ElementManager) ReturnResiduals(rb RunBundle, firstRsIndex int, inputI stage := em.stages[rb.StageID] stage.splitBundle(rb, firstRsIndex, em) + if len(residuals.Data) > 0 { + em.sawResidual.Store(true) + } unprocessedElements := reElementResiduals(residuals.Data, inputInfo, rb) if len(unprocessedElements) > 0 { slog.Debug("ReturnResiduals: unprocessed elements", "bundle", rb, "count", len(unprocessedElements)) @@ -1211,6 +1238,9 @@ type stageState struct { estimatedOutput mtime.Time // Estimated watermark output from DoFns previousInput mtime.Time // input watermark before the latest watermark refresh + // Consumers handed data by a bundle, pending delivery to the scheduler. + consumersWithNewData set[string] + pending elementHeap // pending input elements for this stage that are to be processesd inprogress map[string]elements // inprogress elements by active bundles, keyed by bundle sideInputs map[LinkID]map[typex.Window][][]byte // side input data for this stage, from {tid, inputID} -> window @@ -1822,12 +1852,6 @@ keysPerBundle: if ss.inprogressKeys.present(k) { continue } - newKeys.insert(k) - // Track the min-timestamp for later watermark handling. - if dnt.elements[0].timestamp < minTs { - minTs = dnt.elements[0].timestamp - } - dataInBundle := false var toProcessForKey []element @@ -1874,22 +1898,51 @@ keysPerBundle: break } } + if len(toProcessForKey) > 0 { + newKeys.insert(k) + // Track the min-timestamp for later watermark handling. Elements pop + // in timestamp order, so the first selected one is the earliest. + if ts := toProcessForKey[0].timestamp; ts < minTs { + minTs = ts + } + } toProcess = append(toProcess, toProcessForKey...) if dnt.elements.Len() == 0 { delete(ss.pendingByKeys, k) } - if OneKeyPerBundle { + // A key that yielded nothing, such as one headed by a timer above the + // watermark, must not consume the single key slot, or the bundle is empty + // and the stage keeps rescheduling on it. + if OneKeyPerBundle && len(toProcessForKey) > 0 { break keysPerBundle } } - // If we're out of data, and timers were not cleared then the watermark is accurate. - stillSchedulable := !(len(ss.pendingByKeys) == 0 && !timerCleared) + // Reschedule only when a later bundle could build something, or a cleared + // timer may have held back the minimum pending timestamp. + stillSchedulable := timerCleared || ss.hasBuildableDataLocked(watermark) return toProcess, minTs, newKeys, holdsInBundle, nil, stillSchedulable, 0 } +// hasBuildableDataLocked reports whether a key that isn't in progress heads its +// heap with data, or with a timer the watermark has reached. Callers hold ss.mu. +func (ss *stageState) hasBuildableDataLocked(watermark mtime.Time) bool { + for k, dnt := range ss.pendingByKeys { + if ss.inprogressKeys.present(k) { + continue + } + if dnt.elements.Len() == 0 { + continue + } + if e := dnt.elements[0]; e.IsData() || e.timestamp <= watermark { + return true + } + } + return false +} + // buildEventTimeBundle for aggregation stages, processes all elements that are within the watermark for completed windows. func (*aggregateStageKind) buildEventTimeBundle(ss *stageState, watermark mtime.Time) (toProcess elementHeap, _ mtime.Time, _ set[string], _ map[mtime.Time]int, panesInBundle []bundlePane, schedulable bool, pendingAdjustment int) { minTs := mtime.MaxTimestamp @@ -2286,9 +2339,13 @@ func (ss *stageState) updateWatermarks(em *ElementManager) set[string] { if minWatermarkHold < newOut { newOut = minWatermarkHold } - // If the newOut is smaller, then don't change downstream watermarks. + // If the newOut is smaller, then don't change downstream watermarks. Any + // consumer that received data still needs scheduling, since an unadvancing + // watermark is otherwise the only thing that would surface it. if newOut <= ss.output { - return nil + refreshes := ss.consumersWithNewData + ss.consumersWithNewData = nil + return refreshes } // If bigger, advance the output watermark @@ -2328,6 +2385,7 @@ func (ss *stageState) updateWatermarks(em *ElementManager) set[string] { // Update this stage's output watermark, and then propagate that to downstream stages refreshes := set[string]{} + ss.consumersWithNewData = nil ss.output = newOut for _, outputCol := range ss.outputIDs { consumers := em.consumers[outputCol] @@ -2435,13 +2493,20 @@ func (ss *stageState) bundleReady(em *ElementManager, emNow mtime.Time) (mtime.T previousInputW := ss.previousInput _, isOrdinaryStage := ss.kind.(*ordinaryStageKind) - if isOrdinaryStage && len(ss.sides) == 0 { + _, isStatefulStage := ss.kind.(*statefulStageKind) + switch { + case isOrdinaryStage && len(ss.sides) == 0: // For ordinary stage with no side inputs, we use whether there are pending elements to determine // whether a bundle is ready or not. if len(ss.pending) == 0 { return mtime.MinTimestamp, false, ptimeEventsReady, injectedReady } - } else if inputW == upstreamW && previousInputW == inputW { + case isStatefulStage && len(ss.sides) == 0 && em.sawResidual.Load() && ss.hasBuildableDataLocked(upstreamW): + // A stateful stage processes pending data at whatever the current + // watermark is, so data alone makes it ready once something is self + // checkpointing. Side input readiness comes from the watermark, so + // stages that read one keep waiting. + case inputW == upstreamW && previousInputW == inputW: // Otherwise, use the progression of watermark to determine the bundle readiness. slog.Debug("bundleReady: unchanged upstream watermark", slog.String("stage", ss.ID), diff --git a/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager_continuation_test.go b/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager_continuation_test.go new file mode 100644 index 000000000000..2c0eac52e5e1 --- /dev/null +++ b/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager_continuation_test.go @@ -0,0 +1,299 @@ +// 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. + +package engine + +import ( + "bytes" + "context" + "fmt" + "io" + "testing" + + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/mtime" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/exec" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex" +) + +// bundleBudget bounds how many bundles a continuation test drives. The source +// never terminates, so these tests stop as soon as the consumer is scheduled and +// use the budget only to declare starvation. +const bundleBudget = 50 + +// continuationInfo is a global window PColInfo, keyed for stateful consumers. +func continuationInfo(t *testing.T, keyed bool) PColInfo { + t.Helper() + readAll := func(r io.Reader) []byte { + b, err := io.ReadAll(r) + if err != nil { + t.Fatalf("error decoding element: %v", err) + } + return b + } + info := PColInfo{ + GlobalID: "continuation_info", + WDec: exec.MakeWindowDecoder(coder.NewGlobalWindow()), + WEnc: exec.MakeWindowEncoder(coder.NewGlobalWindow()), + EDec: readAll, + } + if keyed { + info.KeyDec = readAll + } + return info +} + +// encodeElement produces a global window element at the given event time. +func encodeElement(t *testing.T, info PColInfo, et mtime.Time) []byte { + t.Helper() + var buf bytes.Buffer + if err := exec.EncodeWindowedValueHeader(info.WEnc, []typex.Window{window.GlobalWindow{}}, et, typex.NoFiringPane(), &buf); err != nil { + t.Fatalf("EncodeWindowedValueHeader: %v", err) + } + buf.Write([]byte{3, 65, 66, 67}) // "ABC" + return buf.Bytes() +} + +// TestPersistBundle_ContinuationResidualConsumers covers issue #39446: a source +// whose residual pins its watermark must still get its consumers scheduled. +func TestPersistBundle_ContinuationResidualConsumers(t *testing.T) { + for _, test := range []struct { + name string + keyed bool + stateful bool + }{ + {name: "ordinary consumer"}, + {name: "stateful consumer", keyed: true, stateful: true}, + } { + t.Run(test.name, func(t *testing.T) { + srcInfo := continuationInfo(t, false) + outInfo := continuationInfo(t, test.keyed) + + ctx, cancelFn := context.WithCancelCause(context.Background()) + defer cancelFn(nil) + + em := NewElementManager(Config{}) + em.AddStage("impulse", nil, []string{"src_in"}, nil) + em.AddStage("src", []string{"src_in"}, []string{"sink_in"}, nil) + em.AddStage("sink", []string{"sink_in"}, nil, nil) + if test.stateful { + em.StageStateful("sink", nil) + } + em.Impulse("impulse") + + var i int + ch := em.Bundles(ctx, cancelFn, func() string { + defer func() { i++ }() + return fmt.Sprintf("%v", i) + }) + + src := em.stages["src"] + + // The source emits a record and self checkpoints every round, so it + // always has a residual outstanding and never terminates. + var srcBundles, sinkBundles int + for b := 0; b < bundleBudget && sinkBundles == 0; b++ { + rb, ok := <-ch + if !ok { + t.Fatalf("bundle %d: bundles channel closed early", b) + } + switch rb.StageID { + case "sink": + sinkBundles++ + em.PersistBundle(rb, nil, TentativeData{}, outInfo, Residuals{}) + case "src": + srcBundles++ + td := TentativeData{} + td.WriteData("sink_in", encodeElement(t, outInfo, mtime.Time(100*srcBundles))) + // No reported estimate means MIN_TIMESTAMP, so only the + // arriving data can drive the consumer. + em.PersistBundle(rb, map[string]PColInfo{"sink_in": outInfo}, td, srcInfo, Residuals{ + TransformID: "src", + InputID: "i0", + Data: []Residual{{Element: encodeElement(t, srcInfo, mtime.MinTimestamp)}}, + }) + default: + t.Fatalf("bundle %d: unexpected stage %v", b, rb.StageID) + } + } + + if sinkBundles == 0 { + t.Errorf("consumer stage was never scheduled across %d source bundles, src output watermark = %v; its pending elements are starved", + srcBundles, src.OutputWatermark()) + } + }) + } +} + +// TestPersistBundle_ContinuationResidualTransitive covers a consumer two stages +// below the source. The middle stage returns no residual of its own, yet its +// watermark is still pinned by the source's. +func TestPersistBundle_ContinuationResidualTransitive(t *testing.T) { + info := continuationInfo(t, false) + ctx, cancelFn := context.WithCancelCause(context.Background()) + defer cancelFn(nil) + + em := NewElementManager(Config{}) + em.AddStage("impulse", nil, []string{"src_in"}, nil) + em.AddStage("src", []string{"src_in"}, []string{"mid_in"}, nil) + em.AddStage("mid", []string{"mid_in"}, []string{"sink_in"}, nil) + em.AddStage("sink", []string{"sink_in"}, nil, nil) + em.Impulse("impulse") + + var i int + ch := em.Bundles(ctx, cancelFn, func() string { + defer func() { i++ }() + return fmt.Sprintf("%v", i) + }) + + var srcBundles, midBundles, sinkBundles int + for b := 0; b < bundleBudget && sinkBundles == 0; b++ { + rb, ok := <-ch + if !ok { + t.Fatalf("bundle %d: bundles channel closed early", b) + } + switch rb.StageID { + case "src": + srcBundles++ + td := TentativeData{} + td.WriteData("mid_in", encodeElement(t, info, mtime.Time(100*srcBundles))) + em.PersistBundle(rb, map[string]PColInfo{"mid_in": info}, td, info, Residuals{ + TransformID: "src", + InputID: "i0", + Data: []Residual{{Element: encodeElement(t, info, mtime.MinTimestamp)}}, + }) + case "mid": + midBundles++ + td := TentativeData{} + td.WriteData("sink_in", encodeElement(t, info, mtime.Time(100*midBundles))) + em.PersistBundle(rb, map[string]PColInfo{"sink_in": info}, td, info, Residuals{}) + case "sink": + sinkBundles++ + em.PersistBundle(rb, nil, TentativeData{}, info, Residuals{}) + default: + t.Fatalf("bundle %d: unexpected stage %v", b, rb.StageID) + } + } + + if midBundles == 0 { + t.Error("middle stage was never scheduled") + } + if sinkBundles == 0 { + t.Errorf("stage two below the source was never scheduled across %d source and %d middle bundles; its pending elements are starved", + srcBundles, midBundles) + } +} + +// TestPersistBundle_ContinuationResidualWatermark pins the BundleApplication +// output_watermarks contract: an unreported estimate means MIN_TIMESTAMP. +func TestPersistBundle_ContinuationResidualWatermark(t *testing.T) { + for _, test := range []struct { + name string + report bool + wantHeld bool + }{ + {name: "reported estimate advances the watermark", report: true}, + {name: "no estimate holds the watermark", wantHeld: true}, + } { + t.Run(test.name, func(t *testing.T) { + info := continuationInfo(t, false) + ctx, cancelFn := context.WithCancelCause(context.Background()) + defer cancelFn(nil) + + em := NewElementManager(Config{}) + em.AddStage("impulse", nil, []string{"src_in"}, nil) + em.AddStage("src", []string{"src_in"}, []string{"sink_in"}, nil) + em.AddStage("sink", []string{"sink_in"}, nil, nil) + em.Impulse("impulse") + + var i int + ch := em.Bundles(ctx, cancelFn, func() string { + defer func() { i++ }() + return fmt.Sprintf("%v", i) + }) + + src := em.stages["src"] + + for round := 0; round < 3; round++ { + rb, ok := <-ch + if !ok { + t.Fatalf("round %d: bundles channel closed early", round) + } + residuals := Residuals{ + TransformID: "src", + InputID: "i0", + Data: []Residual{{Element: encodeElement(t, info, mtime.MinTimestamp)}}, + } + if test.report { + residuals.MinOutputWatermarks = map[string]mtime.Time{"sink_in": mtime.Time(1000 * (round + 1))} + } + em.PersistBundle(rb, nil, TentativeData{}, info, residuals) + } + + got := src.OutputWatermark() + if test.wantHeld && got != mtime.MinTimestamp { + t.Errorf("src.OutputWatermark() = %v, want %v: an unreported estimate defaults to MIN_TIMESTAMP", got, mtime.MinTimestamp) + } + if !test.wantHeld && got == mtime.MinTimestamp { + t.Errorf("src.OutputWatermark() = %v, want it to follow the reported estimate", got) + } + }) + } +} + +// TestStatefulBuildEventTimeBundle_OneKeyPerBundle checks that a key holding +// only a timer above the watermark does not consume the single key slot, which +// would build an empty bundle and reschedule the stage on the same key forever. +func TestStatefulBuildEventTimeBundle_OneKeyPerBundle(t *testing.T) { + OneKeyPerBundle = true + t.Cleanup(func() { OneKeyPerBundle = false }) + + // Key iteration order is randomized, so repeat until the timer key leads. + for i := 0; i < 20; i++ { + em := NewElementManager(Config{}) + ss := makeStageState("stateful", []string{"input"}, nil, nil) + ss.kind = &statefulStageKind{} + ss.AddPending(em, []element{{ + window: window.GlobalWindow{}, + timestamp: mtime.MaxTimestamp - 1, + holdTimestamp: mtime.MaxTimestamp - 1, + pane: typex.NoFiringPane(), + transform: "stateful", + family: "timer", + keyBytes: []byte("timerkey"), + sequence: 0, + }, { + window: window.GlobalWindow{}, + timestamp: 10, + pane: typex.NoFiringPane(), + elmBytes: []byte{3, 65, 66, 67}, + keyBytes: []byte("datakey"), + }}) + + toProcess, minTs, newKeys, _, _, _, _ := ss.kind.buildEventTimeBundle(ss, mtime.Time(100)) + if len(toProcess) == 0 { + t.Fatalf("iteration %d: built an empty bundle while a data key was pending", i) + } + // The skipped timer key must not be marked in progress, nor hold the + // bundle's minimum timestamp. + if len(newKeys) != 1 || !newKeys.present("datakey") { + t.Fatalf("iteration %d: newKeys = %v, want only datakey", i, newKeys) + } + if want := mtime.Time(10); minTs != want { + t.Fatalf("iteration %d: minTs = %v, want %v", i, minTs, want) + } + } +} diff --git a/sdks/python/apache_beam/io/external/xlang_jmsio_it_test.py b/sdks/python/apache_beam/io/external/xlang_jmsio_it_test.py index c1922eb26ba6..c3f260973765 100644 --- a/sdks/python/apache_beam/io/external/xlang_jmsio_it_test.py +++ b/sdks/python/apache_beam/io/external/xlang_jmsio_it_test.py @@ -92,7 +92,7 @@ def _run_streaming_test( subscriber_result = {} def publish(): - self.produce(source_queue, remaining_records) + self.produce(source_queue, NUM_RECORDS) stop_event = threading.Event() @@ -110,12 +110,6 @@ def subscribe(): break _LOGGER.info('received %s messages', len(received_messages)) - # TODO(https://github.com/apache/beam/issues/39446): Clean up - # pre-publishing Prism runner issue resolved - initial_records = 10 - remaining_records = NUM_RECORDS - initial_records - self.produce(source_queue, initial_records) - publisher = threading.Thread(target=publish, daemon=True) subscriber = threading.Thread(target=subscribe, daemon=True) @@ -147,7 +141,7 @@ def subscribe(): result = p.run() subscriber.start() try: - subscriber.join(timeout=20) # 1.5 min + subscriber.join(timeout=20) finally: stop_event.set() publisher.join() @@ -160,8 +154,7 @@ def subscribe(): received = subscriber_result.get('received', []) self.assertEqual(len(received), NUM_RECORDS) - # there are identical records - self.assertEqual(len(set(received)), NUM_RECORDS - initial_records) + self.assertEqual(len(set(received)), NUM_RECORDS) class ActiveMQJmsIOTest(_BaseJmsIOTest): From 88b3ee7b488aa8a3e85e8dd04b64349384818d16 Mon Sep 17 00:00:00 2001 From: Vitaly Terentyev Date: Fri, 21 Aug 2026 18:46:38 +0400 Subject: [PATCH 4/4] JmsIO yaml (#39818) * Add JmsIO YAML test * Fix depends * Add IBM MQ --- .../apache_beam/yaml/integration_tests.py | 49 +++++++++++++++ sdks/python/apache_beam/yaml/standard_io.yaml | 58 +++++++++++++++++ .../python/apache_beam/yaml/tests/ibm_mq.yaml | 62 +++++++++++++++++++ sdks/python/apache_beam/yaml/tests/jms.yaml | 56 +++++++++++++++++ sdks/python/apache_beam/yaml/yaml_provider.py | 6 +- sdks/python/build.gradle | 1 + 6 files changed, 230 insertions(+), 2 deletions(-) create mode 100644 sdks/python/apache_beam/yaml/tests/ibm_mq.yaml create mode 100644 sdks/python/apache_beam/yaml/tests/jms.yaml diff --git a/sdks/python/apache_beam/yaml/integration_tests.py b/sdks/python/apache_beam/yaml/integration_tests.py index c6e33f8c6de3..f43292058552 100644 --- a/sdks/python/apache_beam/yaml/integration_tests.py +++ b/sdks/python/apache_beam/yaml/integration_tests.py @@ -804,6 +804,55 @@ def temp_kinesis_localstack(): _LOGGER.info('LocalStack Kinesis fixture stopped.') +@contextlib.contextmanager +def temp_jms_activemq_server(): + """Context manager to provide a temporary ActiveMQ broker for JMS tests.""" + + broker = DockerContainer('apache/activemq-classic:5.18.3').with_exposed_ports( + 61616) + + try: + broker.start() + wait_for_logs(broker, '.*ActiveMQ .* started.*', timeout=30) + + host = broker.get_container_host_ip() + port = broker.get_exposed_port(61616) + + yield { + 'SERVER_URI': f'tcp://{host}:{port}', + 'CONNECTION_FACTORY_CLASS_NAME': 'org.apache.activemq.ActiveMQConnectionFactory', + } + finally: + broker.stop() + + +@contextlib.contextmanager +def temp_ibm_mq_server(): + container = ( + DockerContainer('icr.io/ibm-messaging/mq:9.3.0.25-r1').with_env( + 'LICENSE', 'accept').with_env('MQ_QMGR_NAME', 'QM1').with_env( + 'MQ_APP_PASSWORD', 'admin123').with_exposed_ports(1414)) + + try: + container.start() + wait_for_logs(container, '.*(MQQMNAME|Started queue manager).*', timeout=45) + + host = container.get_container_host_ip() + port = container.get_exposed_port(1414) + + yield { + 'SERVER_URI': f'tcp://{host}:{port}?channel=DEV.APP.SVRCONN&queueManager=QM1', + 'CONNECTION_FACTORY_CLASS_NAME': 'com.ibm.mq.jms.MQConnectionFactory', + 'USERNAME': 'app', + 'PASSWORD': 'admin123', + 'SOURCE_QUEUE': 'DEV.QUEUE.1', + 'SINK_QUEUE': 'DEV.QUEUE.2', + } + + finally: + container.stop() + + @contextlib.contextmanager def temp_kafka_server(): """Context manager to provide a temporary Kafka server for testing. diff --git a/sdks/python/apache_beam/yaml/standard_io.yaml b/sdks/python/apache_beam/yaml/standard_io.yaml index b14db7bbec10..81fd62b05012 100644 --- a/sdks/python/apache_beam/yaml/standard_io.yaml +++ b/sdks/python/apache_beam/yaml/standard_io.yaml @@ -183,6 +183,64 @@ config: gradle_target: 'sdks:java:extensions:schemaio-expansion-service:shadowJar' +# JMS +- type: renaming + transforms: + 'ReadFromJms': 'ReadFromJms' + 'WriteToJms': 'WriteToJms' + config: + mappings: + 'ReadFromJms': + connection_configuration: 'connection_configuration' + queue: 'queue' + topic: 'topic' + max_num_records: 'max_num_records' + max_read_time_seconds: 'max_read_time_seconds' + close_timeout_seconds: 'close_timeout_seconds' + acknowledge_mode: 'acknowledge_mode' + individual_acknowledge_mode_code: 'individual_acknowledge_mode_code' + 'WriteToJms': + connection_configuration: 'connection_configuration' + queue: 'queue' + topic: 'topic' + underlying_provider: + type: beamJar + transforms: + 'ReadFromJms': 'beam:schematransform:org.apache.beam:jms_read:v1' + 'WriteToJms': 'beam:schematransform:org.apache.beam:jms_write:v1' + config: + gradle_target: 'sdks:java:io:messaging-expansion-service:shadowJar' +- type: renaming + transforms: + 'ReadFromIbmMQ': 'ReadFromIbmMQ' + 'WriteToIbmMQ': 'WriteToIbmMQ' + config: + mappings: + 'ReadFromIbmMQ': + connection_configuration: 'connection_configuration' + queue: 'queue' + topic: 'topic' + max_num_records: 'max_num_records' + max_read_time_seconds: 'max_read_time_seconds' + close_timeout_seconds: 'close_timeout_seconds' + acknowledge_mode: 'acknowledge_mode' + individual_acknowledge_mode_code: 'individual_acknowledge_mode_code' + 'WriteToIbmMQ': + connection_configuration: 'connection_configuration' + queue: 'queue' + topic: 'topic' + underlying_provider: + type: beamJar + transforms: + 'ReadFromIbmMQ': 'beam:schematransform:org.apache.beam:jms_read:v1' + 'WriteToIbmMQ': 'beam:schematransform:org.apache.beam:jms_write:v1' + config: + gradle_target: 'sdks:java:io:messaging-expansion-service:shadowJar' + classpath: + - 'com.ibm.mq:com.ibm.mq.allclient:9.3.0.25' + - 'org.json:json:20251224' + + # Debezium - type: renaming transforms: diff --git a/sdks/python/apache_beam/yaml/tests/ibm_mq.yaml b/sdks/python/apache_beam/yaml/tests/ibm_mq.yaml new file mode 100644 index 000000000000..b621e293cc7c --- /dev/null +++ b/sdks/python/apache_beam/yaml/tests/ibm_mq.yaml @@ -0,0 +1,62 @@ +# +# 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. +# + +fixtures: + - name: IBM_MQ + type: apache_beam.yaml.integration_tests.temp_ibm_mq_server + +pipelines: + - pipeline: + type: chain + transforms: + - type: Create + config: + elements: + - payload: message-1 + - payload: message-2 + - payload: message-3 + + - type: WriteToIbmMQ + config: + connection_configuration: + server_uri: "{IBM_MQ[SERVER_URI]}" + connection_factory_class_name: + "{IBM_MQ[CONNECTION_FACTORY_CLASS_NAME]}" + username: "{IBM_MQ[USERNAME]}" + password: "{IBM_MQ[PASSWORD]}" + queue: "{IBM_MQ[SOURCE_QUEUE]}" + + - pipeline: + type: chain + transforms: + - type: ReadFromIbmMQ + config: + connection_configuration: + server_uri: "{IBM_MQ[SERVER_URI]}" + connection_factory_class_name: + "{IBM_MQ[CONNECTION_FACTORY_CLASS_NAME]}" + username: "{IBM_MQ[USERNAME]}" + password: "{IBM_MQ[PASSWORD]}" + queue: "{IBM_MQ[SOURCE_QUEUE]}" + max_num_records: 3 + + - type: AssertEqual + config: + elements: + - payload: message-1 + - payload: message-2 + - payload: message-3 diff --git a/sdks/python/apache_beam/yaml/tests/jms.yaml b/sdks/python/apache_beam/yaml/tests/jms.yaml new file mode 100644 index 000000000000..60070dd2d45e --- /dev/null +++ b/sdks/python/apache_beam/yaml/tests/jms.yaml @@ -0,0 +1,56 @@ +# +# 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. +# + +fixtures: + - name: JMS + type: "apache_beam.yaml.integration_tests.temp_jms_activemq_server" + +pipelines: + - pipeline: + type: chain + transforms: + - type: Create + config: + elements: + - payload: message-1 + - payload: message-2 + - payload: message-3 + + - type: WriteToJms + config: + connection_configuration: + server_uri: "{JMS[SERVER_URI]}" + connection_factory_class_name: "{JMS[CONNECTION_FACTORY_CLASS_NAME]}" + queue: yaml-jms-test + + - pipeline: + type: chain + transforms: + - type: ReadFromJms + config: + connection_configuration: + server_uri: "{JMS[SERVER_URI]}" + connection_factory_class_name: "{JMS[CONNECTION_FACTORY_CLASS_NAME]}" + queue: yaml-jms-test + max_num_records: 3 + + - type: AssertEqual + config: + elements: + - payload: message-1 + - payload: message-2 + - payload: message-3 diff --git a/sdks/python/apache_beam/yaml/yaml_provider.py b/sdks/python/apache_beam/yaml/yaml_provider.py index 324ae0c2e734..6dd7a502eba6 100755 --- a/sdks/python/apache_beam/yaml/yaml_provider.py +++ b/sdks/python/apache_beam/yaml/yaml_provider.py @@ -339,12 +339,14 @@ def beam_jar( managed_replacement=None, appendix=None, version=beam_version, - artifact_id=None): + artifact_id=None, + classpath=None): return ExternalJavaProvider( urns, lambda: subprocess_server.JavaJarServer.path_to_beam_jar( gradle_target=gradle_target, version=version, artifact_id=artifact_id ), - managed_replacement=managed_replacement) + managed_replacement=managed_replacement, + classpath=classpath) @ExternalProvider.register_provider_type('docker') diff --git a/sdks/python/build.gradle b/sdks/python/build.gradle index 1a93983da66d..a15d4719ff7b 100644 --- a/sdks/python/build.gradle +++ b/sdks/python/build.gradle @@ -152,6 +152,7 @@ tasks.register("yamlIntegrationTests") { dependsOn ":sdks:java:io:expansion-service:build" dependsOn ":sdks:java:io:google-cloud-platform:expansion-service:build" dependsOn ":sdks:java:io:debezium:expansion-service:shadowJar" + dependsOn ":sdks:java:io:messaging-expansion-service:shadowJar" doLast { exec {