Repository navigation
[Bug]: [RunInference] max_models_per_worker_hint is not enforced #40468
Copy link
Copy link
Open
Description
What happened?
The intent of max_models_per_worker hint is to limit the number of models that will be loaded per SDK process.
The logic to increment allowed max number of workers
beam/sdks/python/apache_beam/ml/inference/base.py
Lines 852 to 857 in dabcf50
| 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) | |
| self._max_models_per_worker_hint = None | |
- The lock acquisition always succeeds (it's a new lock instance)
- If there is more than 1 process bundle descriptor over the life time of the SDK process, we might have more than 1 copy of the deserialized RunInference DoFn with unpickled ModelHandler instance, which won't persist
self._max_models_per_worker_hint = Nonefrom a prior initialization:
AI repro:
from apache_beam.internal import pickler
from apache_beam.ml.inference import base
class Model:
def predict(self, x):
return x
class Handler(base.ModelHandler):
def load_model(self):
return Model()
def run_inference(self, batch, model, inference_args=None):
return [model.predict(x) for x in batch]
mhs = [base.KeyModelMapping([k], Handler()) for k in ('a', 'b', 'c')]
keyed_handler = base.KeyedModelHandler(mhs, max_models_per_worker_hint=1)
# RunInference shares one _ModelHandlerManager per transform across all
# DoFn instances (and processes) via MultiProcessShared.
manager = keyed_handler.load_model()
# 5 DoFn instances in one process (harness threads, re-created bundle
# processors); each deserializes its own copy of the model handler.
for _ in range(5):
handler_copy = pickler.roundtrip(keyed_handler)
handler_copy.override_metrics('ns')
handler_copy.run_inference([('a', 1), ('b', 2), ('c', 3)], manager)
print('model limit:', manager._max_models) # 5, expected 1
print('models in memory:', len(manager._tag_map)) # 3
Issue Priority
Priority: 2 (default / most bugs should be filed as P2)
Issue Components
- Component: Python SDK
- Component: Java SDK
- Component: Go SDK
- Component: Typescript SDK
- Component: IO connector
- Component: Beam YAML
- Component: Beam examples
- Component: Beam playground
- Component: Beam katas
- Component: Website
- Component: Infrastructure
- Component: Spark Runner
- Component: Flink Runner
- Component: Prism Runner
- Component: Twister2 Runner
- Component: Hazelcast Jet Runner
- Component: Google Cloud Dataflow Runner
Reactions are currently unavailable
Activity
schizophrenicmaniac commentedon Oct 10, 2026
.take-issue
Hi, I'd like to work on this. The throwaway lock in KeyedModelHandler.run_inference lets every deserialized handler copy call increment_max_models, so I'd have _ModelHandlerManager track which process IDs have already added the hint and count each process only once, with a test based on the repro.