Skip to content

[Bug]: [RunInference] max_models_per_worker_hint is not enforced #40468

Open

Description

@tvalentyn

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

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
appears to be flawed since:

  1. The lock acquisition always succeeds (it's a new lock instance)
  2. 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 = None from 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

Activity

  1. schizophrenicmaniac commented on Oct 10, 2026

    @schizophrenicmaniac
    Contributor

    .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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions