Skip to content

[Python] Fix max_models_per_worker_hint not being enforced in KeyedModelHandler - #40500

Open
schizophrenicmaniac wants to merge 1 commit into
apache:masterfrom
schizophrenicmaniac:fix/40468-bug-runinference-max-models-per-worker
Open

schizophrenicmaniac wants to merge 1 commit into
apache:masterfrom
schizophrenicmaniac:fix/40468-bug-runinference-max-models-per-worker

Conversation

@schizophrenicmaniac

Copy link
Copy Markdown
Contributor

The max_models_per_worker_hint limit for a KeyedModelHandler with several model handlers was never actually enforced. run_inference created a new threading.Lock() on every call, so acquiring it always succeeded. Each deserialized copy of the handler (one per DoFn instance, plus new copies when bundle processors are recreated) still had the hint set, so every copy added it to the shared _ModelHandlerManager again. In the repro from the issue, five handler copies raised the limit to 5 and all three models stayed loaded with a hint of 1.

This change makes _ModelHandlerManager.increment_max_models take an optional process_id and count each process only once. KeyedModelHandler.run_inference passes os.getpid(). The pid is computed on the caller's side because the manager is proxied from a separate MultiProcessShared process. Calls without a process id work as before.

I added a unit test for the manager and one based on the repro in the issue (five pickled handler copies sharing one manager, hint of 1). The repro test fails on master and passes with this change. I also added a CHANGES.md entry.

Fixes #40468

@codecov

codecov Bot commented Oct 10, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 56.23%. Comparing base (76f6e6d) to head (c00266e).
⚠️ Report is 111 commits behind head on master.

Additional details and impacted files
@@             Coverage Diff              @@
##             master   #40500      +/-   ##
============================================
+ Coverage     56.16%   56.23%   +0.06%     
  Complexity     2288     2288              
============================================
  Files          1124     1124              
  Lines        178336   178362      +26     
  Branches       1489     1489              
============================================
+ Hits         100170   100300     +130     
+ Misses        75645    75541     -104     
  Partials       2521     2521              
Flag Coverage Δ
python 79.68% <100.00%> (+0.13%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@github-actions

Copy link
Copy Markdown
Contributor

Assigning reviewers:

R: @shunping for label python.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

  • stop reviewer notifications - opt out of the automated review tooling
  • remind me after tests pass - tag the comment author after tests pass
  • waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: [RunInference] max_models_per_worker_hint is not enforced

1 participant