Remove usage of (leaking) shared handles, which can also accidentally release the last reference holding the model.

[tvalentyn]

RunInference framework uses _ModelStatus objects to store a state of the model, like whether the model is valid or should be reloaded. We need to store one _ModelStatus object per one instance of RunInference. SInce there can be more than one such object, we key them by a tag: {model_tag}_model_status.

To share these objects within the SDK process, we rely on shared.Shared(), which also has a concept of tags and is a tool to share a singleton object across multiple DoFn instances.

However, the usage return shared.Shared().acquire(lambda: _ModelStatus(False), tag=tag) is incorrect since this creates a new shared handle every time get_model_status is called. The subsequent acquire call, creates a new, uninitialized _ModelStatus instance.

This confusion is likely caused by similar Beam APIs: shared vs multiprocess_shared, which use tags slightly differently.

For MultiProcessShared, tag is the key. The same tag from any process gives the same object.

For shared.Shared, it is a version label. The key is the handle's own internal uuid. Acquiring through the same handle with a different tag throws away the old object and builds a new one; the same tag through a new handle also gives a new object.

The Shared handle is generally not expected to be created more than once per stage of the pipeline (recall that this was originally written for Batch pipelines in TFX; with streaming it gets more complicated since we have all stages running at once).

Given that a) RunInference used 2 shared handles, one for the model, one for the ModelStatus, and b) the Shared machinery has a single _SharedMap._keepalive slot for entire process, calling load_model_status invalidated the last existing reference to the model in some cases, causing #40440.

Because we instantiated a new ModelStatus object every time, we didn't store the state of previously loaded models properly, so invalidating a model might have not caused the reload. However, this would only happen if we are NOT sharing the model across processes, so impact should be limited.

fixes: #40440

[github-actions]

Assigning reviewers:

R: @damccorm 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).

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 59.04%. Comparing base (69bdf5f) to head (1d81860).
⚠️ Report is 27 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff            @@
##             master   #40441   +/-   ##
=========================================
  Coverage     59.04%   59.04%           
  Complexity    15624    15624           
=========================================
  Files          2797     2797           
  Lines        280750   280789   +39     
  Branches      12488    12488           
=========================================
+ Hits         165764   165795   +31     
- Misses       108540   108548    +8     
  Partials       6446     6446           

Flag	Coverage Δ
python	79.66% <100.00%> (-0.01%)	⬇️

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.

[damccorm]

Thanks!