diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 4c2d9cb4c80e..a417bb9560fd 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -1607,7 +1607,6 @@ def annotations(self): 'model_handler_type': ( f'{self._model_handler.__class__.__module__}' f'.{self._model_handler.__class__.__qualname__}'), - 'model_identifier': self._model_tag, **extra, **super().annotations() } @@ -1715,47 +1714,54 @@ def expand( ) if self._model_metadata_pcoll else None).with_resource_hints( **resource_hints) - if self._with_exception_handling: - # On timeouts, report back to the central model metadata - # that the model is invalid - model_tag = self._model_tag - share_across_processes = self._model_handler.share_model_across_processes( - ) - timeout = self._timeout - - def failure_callback(exception: Exception, element: Any): - if type(exception) is not TimeoutError: - return - model_metadata = load_model_status(model_tag, share_across_processes) - model_metadata.try_mark_current_model_invalid(timeout) - logging.warning( - "Inference failed with a timeout, marking the current " + - "model for garbage collection") - - callback = None - if (self._timeout is not None and - self._model_handler.should_garbage_collect_on_timeout()): - callback = failure_callback - results, bad_inference = ( - batched_elements_pcoll - | 'BeamML_RunInference' >> - run_inference_pardo.with_exception_handling( - exc_class=self._exc_class, - use_subprocess=self._use_subprocess, - threshold=self._threshold, - timeout = self._timeout, - on_failure_callback=callback)) - else: - results = ( - batched_elements_pcoll - | 'BeamML_RunInference' >> run_inference_pardo) + # Attach model_identifier only to the inference step so that runners can + # associate it with the ParDo that actually loads and runs the model. + # TODO(https://github.com/apache/beam/issues/40715): When + # with_exception_handling(threshold<1) is used, the threshold-checking + # transforms added by _ExceptionHandlingWrapper also inherit this + # annotation. Refactor to annotate only the inner inference ParDo. + with batched_elements_pcoll.pipeline.transform_annotations( + model_identifier=self._model_tag): + if self._with_exception_handling: + # On timeouts, report back to the central model metadata + # that the model is invalid + model_tag = self._model_tag + share_across_processes = ( + self._model_handler.share_model_across_processes()) + timeout = self._timeout + + def failure_callback(exception: Exception, element: Any): + if type(exception) is not TimeoutError: + return + model_metadata = load_model_status(model_tag, share_across_processes) + model_metadata.try_mark_current_model_invalid(timeout) + logging.warning( + "Inference failed with a timeout, marking the current " + + "model for garbage collection") + + callback = None + if (self._timeout is not None and + self._model_handler.should_garbage_collect_on_timeout()): + callback = failure_callback + results, bad_inference = ( + batched_elements_pcoll + | 'BeamML_RunInference' >> + run_inference_pardo.with_exception_handling( + exc_class=self._exc_class, + use_subprocess=self._use_subprocess, + threshold=self._threshold, + timeout = self._timeout, + on_failure_callback=callback)) + else: + results = ( + 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 + | + 'BeamML_RunInference_MonitoringOutlet' >> self._monitoring_transform) results, bad_postprocessed = self._apply_fns( results, postprocess_fns, 'BeamML_RunInference_Postprocess') diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index 5d6c5bd517c7..dbf81035cf76 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/sdks/python/apache_beam/ml/inference/base_test.py @@ -2219,6 +2219,25 @@ def test_run_inference_watch_file_pattern_keyword_arg_side_input_label(self): FakeModelHandler(), watch_model_pattern='fake/path/*') assert side_input_str in str(result_pcoll.producer.side_inputs[0]) + def test_run_inference_only_attaches_model_identifier_to_inference_step(self): + pipeline = beam.Pipeline() + from apache_beam.ml.inference.utils import WatchFilePattern + file_pattern_side_input = ( + pipeline + | 'WatchFilePattern' >> WatchFilePattern(file_pattern='fake/path/*')) + pcoll = pipeline | 'start' >> beam.Create([1, 2, 3]) + _ = pcoll | base.RunInference( + FakeModelHandler(), + model_metadata_pcoll=file_pattern_side_input, + model_identifier='my_custom_model_id') + proto = pipeline.to_runner_api() + transforms_with_model_id = [ + t.unique_name for t in proto.components.transforms.values() + if t.annotations.get('model_identifier') == b'my_custom_model_id' + ] + self.assertEqual( + transforms_with_model_id, ['RunInference/BeamML_RunInference']) + def test_model_status_provides_valid_tags(self): ms = base._ModelStatus(False) 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 index 0bd15db515b8..e00a62e32381 100644 --- 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 @@ -347,7 +347,6 @@ def __init__( def annotations(self) -> dict[str, Any]: return { - 'model_identifier': '', **super().annotations(), } 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 index 383e632d8416..93f2c97f6575 100644 --- 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 @@ -542,8 +542,7 @@ def test_annotations_shadow_model_identifier(self): bigquery_table=self.bq_table, ) annotations = transform.annotations() - self.assertIn("model_identifier", annotations) - self.assertEqual(annotations["model_identifier"], "") + self.assertNotIn("model_identifier", annotations) def test_run_inference_monitoring_outlet_shadows_model_identifier(self): class DummyMonitoring(beam.PTransform): @@ -572,7 +571,7 @@ def load_model(self): ] self.assertTrue(len(outlet_transforms) > 0) for t in outlet_transforms: - self.assertEqual(t.annotations.get("model_identifier"), b"") + self.assertFalse(t.annotations.get("model_identifier")) @mock.patch( "apache_beam.ml.inference.vertex_ai_model_monitoring_v2.WriteToBigQuery")