This is an automated email from the ASF dual-hosted git repository.
tvalentyn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new c8b5acf67ad Scope RunInference model_identifier annotation to
BeamML_RunInference (#40671)
c8b5acf67ad is described below
commit c8b5acf67ad54b051c7146b4e626dbcd305b5bf0
Author: RuiLong J. <[email protected]>
AuthorDate: Thu Oct 22 15:47:18 2026 -0700
Scope RunInference model_identifier annotation to BeamML_RunInference
(#40671)
* Scope RunInference model_identifier annotation to BeamML_RunInference
* Address review: drop VertexModelMonitoringV2 model_identifier shadow, add
TODO(#40715) for threshold transforms
---
sdks/python/apache_beam/ml/inference/base.py | 86 ++++++++++++----------
sdks/python/apache_beam/ml/inference/base_test.py | 19 +++++
.../ml/inference/vertex_ai_model_monitoring_v2.py | 1 -
.../vertex_ai_model_monitoring_v2_test.py | 5 +-
4 files changed, 67 insertions(+), 44 deletions(-)
diff --git a/sdks/python/apache_beam/ml/inference/base.py
b/sdks/python/apache_beam/ml/inference/base.py
index 4c2d9cb4c80..a417bb9560f 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/sdks/python/apache_beam/ml/inference/base.py
@@ -1607,7 +1607,6 @@ class
RunInference(beam.PTransform[beam.PCollection[Union[ExampleT,
'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 @@ class
RunInference(beam.PTransform[beam.PCollection[Union[ExampleT,
) 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 5d6c5bd517c..dbf81035cf7 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 @@ class RunInferenceBaseTest(unittest.TestCase):
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 0bd15db515b..e00a62e3238 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 @@ class VertexModelMonitoringV2(
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 383e632d841..93f2c97f657 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 @@ class
VertexModelMonitoringV2TransformTest(unittest.TestCase):
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 @@ class
VertexModelMonitoringV2TransformTest(unittest.TestCase):
]
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")
