diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py index 4d5eb4f74c0a..c9f6898d57bb 100644 --- a/sdks/python/apache_beam/ml/inference/base.py +++ b/sdks/python/apache_beam/ml/inference/base.py @@ -1745,48 +1745,48 @@ def expand( ) if self._model_metadata_pcoll else None).with_resource_hints( **resource_hints) - # 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) + # Attach model_identifier only to the inference ParDo (also when it is + # wrapped by with_exception_handling) so that runners can associate it with + # the step that actually loads and runs the model. + default_annotations = run_inference_pardo.annotations + run_inference_pardo.annotations = lambda: { + **default_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: _ = ( diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py index d4cfce8ec4f4..e48c048915a1 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/sdks/python/apache_beam/ml/inference/base_test.py @@ -2256,6 +2256,31 @@ def test_run_inference_only_attaches_model_identifier_to_inference_step(self): self.assertEqual( transforms_with_model_id, ['RunInference/BeamML_RunInference']) + def test_run_inference_with_exception_handling_only_annotates_inference_step( + self): + pipeline = beam.Pipeline() + pcoll = pipeline | 'start' >> beam.Create([1, 2, 3]) + _ = pcoll | base.RunInference( + FakeModelHandler(), + model_identifier='my_custom_model_id').with_exception_handling( + threshold=0.5) + 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/' + 'ParDo(_RunInferenceDoFn)/ParDo(_RunInferenceDoFn)' + ]) + # The threshold-checking transforms exist but are not annotated. + self.assertTrue( + any( + 'CountBad' in t.unique_name + for t in proto.components.transforms.values())) + def test_model_status_provides_valid_tags(self): ms = base._ModelStatus(False) diff --git a/sdks/python/apache_beam/transforms/core.py b/sdks/python/apache_beam/transforms/core.py index 16af71b5d0ba..66992f386aef 100644 --- a/sdks/python/apache_beam/transforms/core.py +++ b/sdks/python/apache_beam/transforms/core.py @@ -1688,7 +1688,8 @@ def with_exception_handling( on_failure_callback, allow_unsafe_userstate_in_process, self.get_resource_hints(), - self.get_type_hints()) + self.get_type_hints(), + self.annotations()) def with_error_handler(self, error_handler, **exception_handling_kwargs): """An alias for `with_exception_handling(error_handler=error_handler, ...)` @@ -2351,7 +2352,8 @@ def __init__( on_failure_callback, allow_unsafe_userstate_in_process, resource_hints, - pardo_type_hints=None): + pardo_type_hints=None, + pardo_annotations=None): if partial and use_subprocess: raise ValueError('partial and use_subprocess are mutually incompatible.') self._fn = fn @@ -2370,6 +2372,7 @@ def __init__( self._allow_unsafe_userstate_in_process = allow_unsafe_userstate_in_process self._resource_hints = resource_hints self._pardo_type_hints = pardo_type_hints + self._pardo_annotations = pardo_annotations self._extra_tags = None def with_outputs(self, *tags, main=None): @@ -2407,6 +2410,13 @@ def _build_pardo(self, pcoll): **self._kwargs, ) pardo.get_resource_hints().update(self._resource_hints) + if self._pardo_annotations: + # Attach to the inner ParDo only, not to the threshold-checking + # transforms added in _post_process_result. + default_annotations = pardo.annotations + pardo.annotations = lambda: { + **default_annotations(), **self._pardo_annotations + } return pardo def _post_process_result(self, pcoll, result): diff --git a/sdks/python/apache_beam/transforms/core_test.py b/sdks/python/apache_beam/transforms/core_test.py index 32594941433c..38cf2e20de93 100644 --- a/sdks/python/apache_beam/transforms/core_test.py +++ b/sdks/python/apache_beam/transforms/core_test.py @@ -430,6 +430,24 @@ def test_timer_exception_handling(self): assert_that(good, equal_to([0, 1, 2]), 'good') assert_that(bad_elements, equal_to([(1, 5), (1, 10)]), 'bad') + def test_pardo_annotations_only_on_inner_pardo_with_threshold(self): + pardo = beam.ParDo(TestDoFn9()) + default_annotations = pardo.annotations + pardo.annotations = lambda: {**default_annotations(), 'my_key': 'my_value'} + pipeline = beam.Pipeline() + _ = ( + pipeline + | beam.Create([1, 2, 3]) + | 'Handled' >> pardo.with_exception_handling(threshold=0.5)) + transforms = pipeline.to_runner_api().components.transforms.values() + annotated = [ + t.unique_name for t in transforms + if t.annotations.get('my_key') == b'my_value' + ] + self.assertEqual(len(annotated), 1, annotated) + self.assertTrue(annotated[0].startswith('Handled/'), annotated) + self.assertTrue(any('CountBad' in t.unique_name for t in transforms)) + def test_tags_with_exception_handling_then_resource_hint(self): class TagHint(ResourceHint): urn = 'beam:resources:tags:v1'