schizophrenicmaniac commented on code in PR #40500:
URL: https://github.com/apache/beam/pull/40500#discussion_r4299312087
##########
sdks/python/apache_beam/ml/inference/base.py:
##########
@@ -620,17 +625,26 @@ def load(self, key: str) -> _ModelLoadStats:
return _ModelLoadStats(
tag, end_time - start_time, memory_after - memory_before)
- def increment_max_models(self, increment: int):
+ def increment_max_models(
+ self, increment: int, process_id: Optional[int] = None):
"""
Increments the number of models that this instance of a
_ModelHandlerManager is able to hold. If it is never called,
no limit is imposed.
Args:
increment: the amount by which we are incrementing the number of models.
+ process_id: the id of the process requesting the increment. If set,
+ only the first increment from each process is applied, since many
+ copies of the model handler can share this manager within a process.
"""
- if self._max_models is None:
- self._max_models = 0
- self._max_models += increment
+ with self._max_models_lock:
+ if process_id is not None:
+ if process_id in self._incremented_process_ids:
+ return
Review Comment:
The hint is meant to be per worker process, but every deserialized copy of
the KeyedModelHandler in a process (one per DoFn instance, plus new ones when
bundle processors get recreated) calls `increment_max_models` on the same
shared manager. The return makes sure only the first call from each process
adds the hint, and later calls from that process are no-ops. Without it, five
copies with a hint of 1 push the limit to 5, which is what the repro in #40468
shows. A different process still adds its own hint, the same way the original
"first time ... from this process" comment intended.
--
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]
