tvalentyn commented on code in PR #40671:
URL: https://github.com/apache/beam/pull/40671#discussion_r4342776734
##########
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:
based on offline discussion sounds like it is probably ok to keep as is.
--
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]
