[schizophrenicmaniac]

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 Report

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

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #40500      +/-   ##
============================================
+ Coverage     56.16%   56.35%   +0.18%     
  Complexity     2288     2288              
============================================
  Files          1124     1125       +1     
  Lines        178336   178494     +158     
  Branches       1489     1489              
============================================
+ Hits         100170   100591     +421     
+ Misses        75645    75382     -263     
  Partials       2521     2521              

Flag	Coverage Δ
python	79.91% <100.00%> (+0.35%)	⬆️

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]

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

[shunping]

Hi @schizophrenicmaniac, thanks for contributing.

Could you explain a bit more on why you take away the locking mechanism to protect accessing _incremented_process_ids or _max_models when calling increment_max_models?

[schizophrenicmaniac]

@shunping The lock there was created fresh on each call (lock = threading.Lock() right before acquire), so it always succeeded and never actually protected anything. You're right that the state needs guarding though, since the shared manager handles calls from different processes on separate threads. I've added a lock on _ModelHandlerManager itself that covers the process id check and the _max_models update.

[shunping]

Ok, the overall changes sounds reasonable to me. However, I am not super familiar with the design of this code, so maybe @tvalentyn can give a second opinon?

[tvalentyn]

thanks.