From 1d81860a64e578fe685e89a20b9f7fffce3171fb Mon Sep 17 00:00:00 2001
From: Valentyn Tymofieiev <valentyn@google.com>
Date: Tue, 6 Oct 2026 19:04:04 -0700
Subject: [PATCH] Remove usage of a (leaking) shared handles, which can also
 accidentally release the last reference hodling a model in
 _SharedMap._keepalive when all DoFns have been gc'ed.

---
 sdks/python/apache_beam/ml/inference/base.py  |  7 ++-
 .../apache_beam/ml/inference/base_test.py     | 47 +++++++++++++++++++
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py
index 875b329f42c4..8c4b10f8573d 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/sdks/python/apache_beam/ml/inference/base.py
@@ -1812,7 +1812,7 @@ def try_mark_current_model_invalid(self, min_model_life_seconds):
           self._pending_hard_delete.append((
               tag,
               datetime.now() + 2 * timedelta(seconds=min_model_life_seconds)))
-        self._active_tags.remove(tag)
+        self._active_tags.discard(tag)
 
   def get_valid_tag(self, tag: str) -> str:
     """Takes in a proposed valid tag and returns a valid one.
@@ -1874,13 +1874,16 @@ def mark_tags_deleted(self, deleted_tags: set[str]):
         return
 
 
+_model_statuses: dict[str, _ModelStatus] = {}
+
+
 def load_model_status(
     model_tag: str, share_across_processes: bool) -> _ModelStatus:
   tag = f'{model_tag}_model_status'
   if share_across_processes:
     return multi_process_shared.MultiProcessShared(
         lambda: _ModelStatus(True), tag=tag, always_proxy=True).acquire()
-  return shared.Shared().acquire(lambda: _ModelStatus(False), tag=tag)
+  return _model_statuses.setdefault(tag, _ModelStatus(False))
 
 
 class _ProxyLoader:
diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py
index 7bca1fc63385..8cf408db76f4 100644
--- a/sdks/python/apache_beam/ml/inference/base_test.py
+++ b/sdks/python/apache_beam/ml/inference/base_test.py
@@ -16,6 +16,7 @@
 #
 
 """Tests for apache_beam.ml.base."""
+import gc
 import math
 import multiprocessing
 import os
@@ -26,6 +27,7 @@
 import time
 import unittest
 import unittest.mock
+import uuid
 from collections.abc import Iterable
 from collections.abc import Mapping
 from collections.abc import Sequence
@@ -120,6 +122,22 @@ def batch_elements_kwargs(self):
     return {'min_batch_size': 1, 'max_batch_size': 1}
 
 
+_LOAD_COUNTS: dict[str, int] = {}
+
+
+class LoadCountingModelHandler(base.ModelHandler[int, int, FakeModel]):
+  """Counts load_model() calls per key in the module-level _LOAD_COUNTS."""
+  def __init__(self, key: str):
+    self._key = key
+
+  def load_model(self):
+    _LOAD_COUNTS[self._key] = _LOAD_COUNTS.get(self._key, 0) + 1
+    return FakeModel()
+
+  def run_inference(self, batch, model, inference_args=None):
+    return [model.predict(x) for x in batch]
+
+
 class FakeModelHandler(base.ModelHandler[int, int, FakeModel]):
   def __init__(
       self,
@@ -1837,6 +1855,35 @@ def test_model_status_provides_valid_tags(self):
     self.assertTrue(ms.is_valid_tag('tag1_reload_2'))
     self.assertTrue(ms.is_valid_tag('tag2_reload_2'))
 
+  def test_load_model_status_returns_same_status_per_tag(self):
+    tag = 'tag1' + uuid.uuid4().hex
+    status = base.load_model_status(tag, False)
+    self.assertIs(status, base.load_model_status(tag, False))
+    self.assertIsNot(status, base.load_model_status(tag + '_other', False))
+
+    self.assertTrue(status.is_valid_tag('model'))
+    base.load_model_status(tag, False).try_mark_current_model_invalid(0)
+    self.assertFalse(base.load_model_status(tag, False).is_valid_tag('model'))
+
+  def test_new_dofn_instances_reuse_previously_loaded_models(self):
+    # Runners may recreate DoFn instances, e.g. when the SDK
+    # harness evicts an idle bundle processor.
+    model_tag = 'tag2' + uuid.uuid4().hex
+    handler = LoadCountingModelHandler(model_tag)
+    dofn = base._RunInferenceDoFn(
+        handler, clock=FakeClock(), metrics_namespace=None, model_tag=model_tag)
+    serialized_dofn = pickle.dumps(dofn)
+
+    dofn_instance = pickle.loads(serialized_dofn)
+    dofn_instance.setup()
+    del dofn_instance
+    gc.collect()
+
+    dofn_instance = pickle.loads(serialized_dofn)
+    dofn_instance.setup()
+
+    self.assertEqual(1, _LOAD_COUNTS[model_tag])
+
   def test_model_status_provides_valid_garbage_collection(self):
     ms = base._ModelStatus(True)