This is an automated email from the ASF dual-hosted git repository.
tvalentyn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 8eda31b281f [Python] Fix max_models_per_worker_hint not being enforced
in KeyedModelHandler (#40500)
8eda31b281f is described below
commit 8eda31b281f292854e83e3f6a3d84218007eacfb
Author: Divyansh Yadav <[email protected]>
AuthorDate: Sat Oct 17 02:47:40 2026 +0530
[Python] Fix max_models_per_worker_hint not being enforced in
KeyedModelHandler (#40500)
* [Python] Fix max_models_per_worker_hint not being enforced in
KeyedModelHandler
* Guard increment_max_models with a lock on the model handler manager
---
CHANGES.md | 1 +
sdks/python/apache_beam/ml/inference/base.py | 31 ++++++++++++++------
sdks/python/apache_beam/ml/inference/base_test.py | 35 +++++++++++++++++++++++
3 files changed, 59 insertions(+), 8 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index 537c3a32447..d9ca6f0b83d 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -104,6 +104,7 @@
* (Go) Fixed the harness leaking Data/State gRPC streams after the worker
stops, and a deadlock when Send returns EOF
([#40260](https://github.com/apache/beam/issues/40260)).
* (Go) Wrap last dataWriter/timerWriter send errors on Close
([#40405](https://github.com/apache/beam/issues/40405)).
* (Go) Fixed a deadlock recreating a data channel while holding the channel
lock ([#40414](https://github.com/apache/beam/issues/40414)).
+* (Python) Fixed `KeyedModelHandler` not enforcing
`max_models_per_worker_hint` when several copies of the model handler run in
the same SDK process ([#40468](https://github.com/apache/beam/issues/40468)).
* (Java) Fixed the declared schema of the error output of the Kafka write
SchemaTransform, which wrapped the error schema a second time and did not match
the rows it emits ([#39760](https://github.com/apache/beam/issues/39760)).
* (Go) Fixed pubsubio importing a `google.golang.org/genproto` package removed
in recent releases, which broke builds of Go modules depending on a current
`genproto` version ([#40018](https://github.com/apache/beam/issues/40018)).
* (Python) `VertexAIFeatureStoreEnrichmentHandler` now logs a warning instead
of failing when the caller lacks permission to read the feature store metadata
([#40442](https://github.com/apache/beam/issues/40442)).
diff --git a/sdks/python/apache_beam/ml/inference/base.py
b/sdks/python/apache_beam/ml/inference/base.py
index 823a4a3fa71..ce014e4b04f 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/sdks/python/apache_beam/ml/inference/base.py
@@ -577,6 +577,11 @@ class _ModelHandlerManager:
model.
"""
self._max_models = None
+ # Process ids that have already contributed to _max_models
+ self._incremented_process_ids: set[int] = set()
+ # Guards _max_models and _incremented_process_ids, since the shared
+ # manager serves calls from several processes on separate threads.
+ self._max_models_lock = threading.Lock()
# Map keys to model handlers
self._mh_map: dict[str, ModelHandler] = mh_map
# Map keys to the last updated model path for that key
@@ -631,17 +636,26 @@ class _ModelHandlerManager:
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
+ self._incremented_process_ids.add(process_id)
+ if self._max_models is None:
+ self._max_models = 0
+ self._max_models += increment
def update_model_handler(self, key: str, model_path: str, previous_key: str):
"""
@@ -867,11 +881,12 @@ class KeyedModelHandler(Generic[KeyT, ExampleT,
PredictionT, ModelT],
self._unkeyed.run_inference(unkeyed_batch, model, inference_args))
# The first time a MultiProcessShared ModelManager is used for inference
- # from this process, we should increment its max model count
+ # from this process, we should increment its max model count. The manager
+ # may live in another process and be shared by several deserialized copies
+ # of this handler, so it dedupes increments by process id.
if self._max_models_per_worker_hint is not None:
- lock = threading.Lock()
- if lock.acquire(blocking=False):
- model.increment_max_models(self._max_models_per_worker_hint)
+ model.increment_max_models(
+ self._max_models_per_worker_hint, process_id=os.getpid())
self._max_models_per_worker_hint = None
batch_by_key = defaultdict(list)
diff --git a/sdks/python/apache_beam/ml/inference/base_test.py
b/sdks/python/apache_beam/ml/inference/base_test.py
index 568bdf7ea12..64f5c3b1525 100644
--- a/sdks/python/apache_beam/ml/inference/base_test.py
+++ b/sdks/python/apache_beam/ml/inference/base_test.py
@@ -1978,6 +1978,41 @@ class RunInferenceBaseTest(unittest.TestCase):
mh3.load_model, tag=tag3).acquire()
self.assertEqual(8, model3.predict(10))
+ def test_model_handler_manager_increments_once_per_process(self):
+ mhs = {
+ 'key1': FakeModelHandler(state=1),
+ 'key2': FakeModelHandler(state=2),
+ 'key3': FakeModelHandler(state=3)
+ }
+ mm = base._ModelHandlerManager(mh_map=mhs)
+ mm.increment_max_models(1, process_id=1)
+ mm.increment_max_models(1, process_id=1)
+ mm.load('key1')
+ mm.load('key2')
+ mm.load('key3')
+ self.assertEqual(['key3'], list(mm._tag_map.keys()))
+
+ mm.increment_max_models(1, process_id=2)
+ mm.load('key1')
+ self.assertEqual(['key3', 'key1'], list(mm._tag_map.keys()))
+
+ def test_keyed_model_handler_max_models_hint_with_handler_copies(self):
+ mhs = [
+ base.KeyModelMapping([k],
+ FakeModelHandler(
+ state=i, multi_process_shared=True))
+ for i, k in enumerate(['a', 'b', 'c'])
+ ]
+ keyed_mh = base.KeyedModelHandler(mhs, max_models_per_worker_hint=1)
+ mm = keyed_mh.load_model()
+ # Each DoFn instance in a process deserializes its own copy of the model
+ # handler, but they all share the same _ModelHandlerManager.
+ for _ in range(5):
+ mh_copy = pickle.loads(pickle.dumps(keyed_mh))
+ mh_copy.override_metrics('test_namespace')
+ list(mh_copy.run_inference([('a', 1), ('b', 2), ('c', 3)], mm))
+ self.assertEqual(1, len(mm._tag_map))
+
def test_run_inference_loads_different_models(self):
mh1 = FakeModelHandler(incrementing=True, min_batch_size=3)
with TestPipeline() as pipeline:
