tvalentyn commented on code in PR #40671:
URL: https://github.com/apache/beam/pull/40671#discussion_r4342185727
##########
sdks/python/apache_beam/ml/inference/base.py:
##########
@@ -1715,47 +1714,48 @@ 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)
+ with batched_elements_pcoll.pipeline.transform_annotations(
+ model_identifier=self._model_tag):
Review Comment:
do we want to add this annotation even when not using ModelManager, or
should we downscope this ?
##########
sdks/python/apache_beam/ml/inference/base.py:
##########
@@ -1715,47 +1714,48 @@ 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)
+ 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,
Review Comment:
it appears that if threshold < 1, then we add a few more nodes in the
pipeline graph that will have this annotation, see:
https://github.com/apache/beam/blob/5b6aa703f0ab6a48d92cc63ea89226485cc265f9/sdks/python/apache_beam/transforms/core.py#L2414
.
I don't have an opinion yet how to best account for this.
##########
sdks/python/apache_beam/ml/inference/vertex_ai_model_monitoring_v2_test.py:
##########
@@ -572,7 +572,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"))
Review Comment:
should we remove
https://github.com/apache/beam/blob/5b6aa703f0ab6a48d92cc63ea89226485cc265f9/sdks/python/apache_beam/ml/inference/vertex_ai_model_monitoring_v2.py#L350
?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
