[fix] Keep inherited Ray lock handles alive in forked children

tests/test_kv_lock.py

@@ -14,9 +14,11 @@
 # limitations under the License.
 
 import asyncio
+import gc
 import importlib
 import threading
 import time
+import weakref
 from unittest.mock import MagicMock
 
 import pytest
@@ -172,6 +174,26 @@ def test_invalid_lock_settings_raise_before_any_actor_starts(fake_lock_actors, m
     fake_lock_actors.options.assert_not_called()
 
 
+def test_fork_child_keeps_inherited_lock_handles_referenced(monkeypatch):
+    class Handle:
+        pass
+
+    owned, managed = Handle(), Handle()
+    for name in ("_state_lock", "_table", "_registry", "_leases", "_renewer_running"):
+        monkeypatch.setattr(kvl, name, getattr(kvl, name))
+    monkeypatch.setattr(kvl, "_owned", [owned])
+    monkeypatch.setattr(kvl, "_managers", {0: managed})
+    monkeypatch.setattr(kvl, "_inherited_after_fork", [])
+    owned_ref, managed_ref = weakref.ref(owned), weakref.ref(managed)
+    del owned, managed
+
+    kvl._reset_after_fork()
+    gc.collect()
+
+    assert kvl._owned == [] and kvl._managers == {}
+    assert owned_ref() is not None and managed_ref() is not None
+
+
 def test_lock_settings_are_ignored_while_the_lock_is_off(monkeypatch):
     monkeypatch.delenv("TQ_ENABLE_KV_GLOBAL_LOCK", raising=False)
     monkeypatch.setenv("TQ_KV_GLOBAL_LOCK_NUM_SHARDS", "0")

transfer_queue/kv_lock.py

@@ -216,6 +216,7 @@ async def async_kv_local_lock(keys: str | list[str], partition_id: str, timeout:
 _managers: dict = {}  # shard -> TransferQueueLockManager handle, looked up on first use
 _owned: list = []  # the lock actors this process's tq.init() created, killed by its tq.close()
 _closes = 0  # this process's tq.close() calls: a lease created before the latest one is given up
+_inherited_after_fork: list = []  # a forked child's copies of the parent's handles, never released
 # Longest one call waits inside a lock actor; waiting longer takes another call that keeps
 # the queue position. The bound is long because every waiter spends one call per window, and
 # those calls share the actor's limited throughput with everyone else's, renewals included.
@@ -585,7 +586,10 @@ def _close_lock_managers() -> None:
 
 def _reset_after_fork() -> None:
     # The child inherits keys, leases and Ray handles that belong to threads and connections it lacks.
+    # Releasing an inherited handle runs the parent's core worker in the child: that crashes the child,
+    # can report the parent's lock actors out of scope, and leaves the parent's GCS calls timing out.
     global _state_lock, _table, _registry, _leases, _renewer_running, _managers, _owned
+    _inherited_after_fork.append((_owned, _managers, _leases))
     _state_lock, _table = threading.Lock(), {}
     _registry, _leases, _renewer_running, _managers, _owned = threading.Condition(), {}, False, {}, []