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/.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): 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() 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 { 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