AMOOOMA commented on code in PR #40671:
URL: https://github.com/apache/beam/pull/40671#discussion_r4342853551
##########
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:
https://github.com/apache/beam/issues/40715 here for tracking
--
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]
