From 79350c8c285ce04c741c5ff8a102310993756405 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Sat, 3 Oct 2026 23:15:58 +0800 Subject: [PATCH 01/42] [feat] Add a process-local kv_local_lock for TransferQueue KV calls Each KV call is safe on its own, but a sequence such as kv_batch_get -> compute -> kv_put is not: concurrent writers lose updates. kv_local_lock / async_kv_local_lock let threads and asyncio tasks of one process serialize such sequences. Semantics: advisory exclusive mutex per (partition_id, key); one key or a list (deduplicated, acquired in sorted order); timeout is a total deadline and raises TimeoutError; nesting raises RuntimeError; sync kv_local_lock inside a running event loop raises RuntimeError. Waiters are served in FIFO order with direct handoff, so threads and tasks on any loop wake each other, and a timed-out or cancelled waiter never keeps the key. The lock is process-local only; it does not exclude other processes or Ray actors. Signed-off-by: OutstanderWang --- tests/test_kv_lock.py | 335 +++++++++++++++++++++++++++++++++++++ transfer_queue/__init__.py | 6 + transfer_queue/kv_lock.py | 187 +++++++++++++++++++++ 3 files changed, 528 insertions(+) create mode 100644 tests/test_kv_lock.py create mode 100644 transfer_queue/kv_lock.py diff --git a/tests/test_kv_lock.py b/tests/test_kv_lock.py new file mode 100644 index 00000000..686b04e6 --- /dev/null +++ b/tests/test_kv_lock.py @@ -0,0 +1,335 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import gc +import importlib +import threading +import time +from types import SimpleNamespace + +import pytest +import ray +import torch +from omegaconf import OmegaConf + +import transfer_queue as tq +from transfer_queue import async_kv_local_lock, kv_local_lock + +# `transfer_queue.kv_lock` resolves to the re-exported function, so fetch the module itself. +kvl = importlib.import_module("transfer_queue.kv_lock") +P = "p" +# Bound every wait so that a lock regression fails the test instead of hanging it. +TIMEOUT = 10.0 +E2E_TIMEOUT = 120.0 + + +@pytest.fixture(autouse=True) +def table_is_empty_afterwards(): + yield + assert kvl._table == {} + + +def n_waiters(key): + with kvl._state_lock: + return len(kvl._table.get((P, key), ())) + + +def wait_until(predicate): + deadline = time.monotonic() + TIMEOUT + while not predicate(): + assert time.monotonic() < deadline, "condition not reached" + time.sleep(0.001) + + +async def async_wait_until(predicate): + deadline = time.monotonic() + TIMEOUT + while not predicate(): + assert time.monotonic() < deadline, "condition not reached" + await asyncio.sleep(0.001) + + +def start_holder(keys, release): + """Hold `keys` on a new thread until `release` is set; return once they are held.""" + entered = threading.Event() + + def run(): + with kv_local_lock(keys, P, timeout=TIMEOUT): + entered.set() + release.wait(TIMEOUT) + + thread = threading.Thread(target=run) + thread.start() + assert entered.wait(TIMEOUT) + return thread + + +def test_threads_and_tasks_on_several_loops_serialize_read_modify_write(): + counter = [0] + + def thread_worker(): + for _ in range(100): + with kv_local_lock("c", P, timeout=TIMEOUT): + value = counter[0] + time.sleep(0) + counter[0] = value + 1 + + async def loop_main(): + async def task_worker(): + for _ in range(100): + async with async_kv_local_lock("c", P, timeout=TIMEOUT): + value = counter[0] + await asyncio.sleep(0) + counter[0] = value + 1 + + await asyncio.gather(*(task_worker() for _ in range(4))) + + threads = [threading.Thread(target=thread_worker) for _ in range(4)] + threads += [threading.Thread(target=asyncio.run, args=(loop_main(),)) for _ in range(2)] + for thread in threads: + thread.start() + asyncio.run(loop_main()) + for thread in threads: + thread.join(TIMEOUT) + assert not thread.is_alive() + assert counter[0] == 4 * 100 + 3 * 4 * 100 + + +def test_thread_and_task_wake_each_other(): + order = [] + + def thread_body(): + with kv_local_lock("k", P, timeout=TIMEOUT): + order.append("thread") + wait_until(lambda: n_waiters("k") == 1) + + async def main(): + thread = threading.Thread(target=thread_body) + async with async_kv_local_lock("k", P, timeout=TIMEOUT): + thread.start() + wait_until(lambda: n_waiters("k") == 1) + order.append("task") + async with async_kv_local_lock("k", P, timeout=TIMEOUT): + order.append("task again") + thread.join(TIMEOUT) + assert not thread.is_alive() + + asyncio.run(main()) + assert order == ["task", "thread", "task again"] + + +def test_waiters_are_served_fifo(): + order, release = [], threading.Event() + holder = start_holder("k", release) + + def waiter(i): + with kv_local_lock("k", P, timeout=TIMEOUT): + order.append(i) + + threads = [threading.Thread(target=waiter, args=(i,)) for i in range(5)] + for i, thread in enumerate(threads): + thread.start() + wait_until(lambda n=i + 1: n_waiters("k") == n) + release.set() + for thread in [holder, *threads]: + thread.join(TIMEOUT) + assert not thread.is_alive() + assert order == list(range(5)) + + +def test_release_hands_off_directly_so_a_newcomer_cannot_barge(): + outcome = [] + + def holder_then_newcomer(): + with kv_local_lock("k", P, timeout=TIMEOUT): + wait_until(lambda: n_waiters("k") == 1) + try: + with kv_local_lock("k", P, timeout=0): + outcome.append("barged") + except TimeoutError: + outcome.append("queued behind the woken waiter") + + thread = threading.Thread(target=holder_then_newcomer) + thread.start() + wait_until(lambda: (P, "k") in kvl._table) + with kv_local_lock("k", P, timeout=TIMEOUT): + thread.join(TIMEOUT) + assert not thread.is_alive() + assert outcome == ["queued behind the woken waiter"] + + +def test_timeout_raises_and_leaves_no_waiter(): + release = threading.Event() + holder = start_holder("k", release) + + async def take(): + async with async_kv_local_lock("k", P, timeout=0.05): + pass + + with pytest.raises(TimeoutError): + with kv_local_lock("k", P, timeout=0.05): + pass + with pytest.raises(TimeoutError): + asyncio.run(take()) + assert n_waiters("k") == 0 + release.set() + holder.join(TIMEOUT) + assert not holder.is_alive() + + +def test_timeout_is_a_total_deadline_and_partial_acquisition_is_released(): + release_a, release_b = threading.Event(), threading.Event() + holders = [start_holder("a", release_a), start_holder("b", release_b)] + threading.Timer(0.3, release_a.set).start() + start = time.monotonic() + with pytest.raises(TimeoutError): + with kv_local_lock(["b", "a"], P, timeout=0.5): + pass + assert time.monotonic() - start < 0.7 # a per-key budget would take 0.3 + 0.5 s + assert (P, "a") not in kvl._table + release_b.set() + for holder in holders: + holder.join(TIMEOUT) + assert not holder.is_alive() + + +def test_grant_racing_a_timeout_counts_as_acquired(monkeypatch): + release = threading.Event() + holder = start_holder("k", release) + + class GrantedThenTimedOut(threading.Event): + def wait(self, timeout=None): + release.set() + holder.join(TIMEOUT) # the holder hands the key to this waiter... + return False # ...but the wait still reports a timeout + + monkeypatch.setattr(kvl, "threading", SimpleNamespace(Event=GrantedThenTimedOut)) + with kv_local_lock("k", P, timeout=1): + assert (P, "k") in kvl._table and n_waiters("k") == 0 + + +def test_cancellation_never_strands_the_key(): + release, order = threading.Event(), [] + holder = start_holder("k", release) + + async def take(tag): + async with async_kv_local_lock("k", P, timeout=TIMEOUT): + order.append(tag) + + async def main(): + waiting = asyncio.create_task(take("cancelled while waiting")) + await async_wait_until(lambda: n_waiters("k") == 1) + waiting.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(waiting, TIMEOUT) + assert n_waiters("k") == 0 + + granted, nxt = asyncio.create_task(take("cancelled after grant")), asyncio.create_task(take("next")) + await async_wait_until(lambda: n_waiters("k") == 2) + release.set() + holder.join(TIMEOUT) # blocks this loop, so `granted` is cancelled before it can run + granted.cancel() + await asyncio.wait_for(asyncio.gather(granted, nxt, return_exceptions=True), TIMEOUT) + assert granted.cancelled() + + asyncio.run(main()) + assert order == ["next"] + + +def test_waiter_on_a_closed_loop_is_skipped(): + release, taken = threading.Event(), [] + holder = start_holder("k", release) + + async def take(): + async with async_kv_local_lock("k", P, timeout=TIMEOUT): + taken.append("abandoned task") + + def thread_waiter(): + with kv_local_lock("k", P, timeout=TIMEOUT): + taken.append("thread") + + loop = asyncio.new_event_loop() + loop.create_task(take()) + loop.run_until_complete(async_wait_until(lambda: n_waiters("k") == 1)) + loop.close() # leaves the waiting task pending on a closed loop + thread = threading.Thread(target=thread_waiter) + thread.start() + wait_until(lambda: n_waiters("k") == 2) + release.set() + for t in (holder, thread): + t.join(TIMEOUT) + assert not t.is_alive() + assert taken == ["thread"] + gc.collect() # finalizing the abandoned task must leave the table alone + + +def test_invalid_and_nested_use_is_rejected(): + with pytest.raises(ValueError): + with kv_local_lock([], P): + pass + with kv_local_lock(["a", "a", "b"], P): + assert set(kvl._table) == {(P, "a"), (P, "b")} + for keys in ("a", "c"): + with pytest.raises(RuntimeError, match="Nested"): + with kv_local_lock(keys, P, timeout=TIMEOUT): + pass + + async def main(): + with pytest.raises(RuntimeError, match="async_kv_local_lock"): + with kv_local_lock("a", P): + pass + async with async_kv_local_lock("a", P): + with pytest.raises(RuntimeError, match="Nested"): + async with async_kv_local_lock("b", P): + pass + + async def child(): + async with async_kv_local_lock("c", P): + pass + + with pytest.raises(RuntimeError, match="Nested"): + await asyncio.create_task(child()) + + asyncio.run(main()) + + +@pytest.fixture +def simple_storage_tq(): + if not ray.is_initialized(): + ray.init(namespace="TestKVLock") + backend = {"storage_backend": "SimpleStorage", "SimpleStorage": {"total_storage_size": 200}} + tq.init(OmegaConf.create({"controller": {"polling_mode": True}, "backend": backend})) + yield + tq.close() + ray.shutdown() + + +def test_kv_local_lock_serializes_simple_storage_read_modify_write(simple_storage_tq): + tq.kv_put("counter", "lock_e2e", fields={"v": torch.tensor([0])}) # also creates the client on this thread + + def worker(): + for _ in range(50): + with tq.kv_local_lock("counter", "lock_e2e", timeout=E2E_TIMEOUT): + value = tq.kv_batch_get("counter", "lock_e2e")["v"] + tq.kv_put("counter", "lock_e2e", fields={"v": value[0] + 1}) + + threads = [threading.Thread(target=worker) for _ in range(8)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(E2E_TIMEOUT) + assert not thread.is_alive() + assert tq.kv_batch_get("counter", "lock_e2e")["v"][0].item() == 400 + tq.kv_clear("counter", "lock_e2e") diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py index 754bb4d8..60ff7091 100644 --- a/transfer_queue/__init__.py +++ b/transfer_queue/__init__.py @@ -37,6 +37,7 @@ load_checkpoint, save_checkpoint, ) +from .kv_lock import async_kv_local_lock, kv_local_lock from .metadata import BatchMeta, KVBatchMeta from .sampler import BaseSampler from .sampler.grpo_group_n_sampler import GRPOGroupNSampler @@ -65,6 +66,11 @@ "async_kv_clear", "KVBatchMeta", ] + + [ + # Process-local KV key locks + "kv_local_lock", + "async_kv_local_lock", + ] + [ # Checkpoint Interface "save_checkpoint", diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py new file mode 100644 index 00000000..1f298929 --- /dev/null +++ b/transfer_queue/kv_lock.py @@ -0,0 +1,187 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Advisory, process-local locks on TransferQueue KV keys. + +``kv_local_lock`` and ``async_kv_local_lock`` hold an exclusive mutex per +``(partition_id, key)``, shared by every thread and asyncio task of this process. They +are advisory: KV calls never check them. They do not exclude other processes or Ray +actors. Waiters on a key are served in FIFO order. Nesting is rejected; pass all keys +to a single call. +""" + +import asyncio +import os +import threading +import time +from collections import deque +from contextlib import asynccontextmanager, contextmanager +from contextvars import ContextVar + +# Not asyncio.Lock: it is bound to one event loop and is not thread-safe, while a process +# may run several loops plus plain threads. A key is held exactly while it is in _table, +# which maps it to its FIFO of waiters. _state_lock guards every change and is never held +# while waiting. +_state_lock = threading.Lock() +_table: dict[tuple[str, str], deque["_Waiter"]] = {} +# Whether this thread or task holds (or is acquiring) a local lock. A future global lock +# reads it to enforce the "global, then local" order. +_holding: ContextVar[bool] = ContextVar("_holding_kv_lock", default=False) + + +class _Waiter: + def __init__(self, wake): + self.wake = wake # raises RuntimeError if the waiter's event loop is closed + self.granted = False + + +def _take_or_queue(name: tuple[str, str], waiter: _Waiter) -> bool: + with _state_lock: + if name in _table: + _table[name].append(waiter) + return False + _table[name] = deque() + return True + + +def _release(name: tuple[str, str]) -> None: + """Hand the key to its first live waiter, or free it. Call under _state_lock.""" + # Direct handoff keeps strict FIFO: no newcomer can take the key between this + # release and the woken waiter running. + waiters = _table[name] + while waiters: + waiter = waiters.popleft() + try: + waiter.wake() + except RuntimeError: # its event loop is closed, so nobody is left to own the key + continue + waiter.granted = True + return + del _table[name] + + +def _settle(name: tuple[str, str], waiter: _Waiter, keep_grant: bool) -> bool: + """Resolve a wait that ended without a wake-up; return whether the caller owns the key.""" + # A grant can race the timeout or cancellation, so `granted` decides, not the + # Event/Future. A raced grant is kept on timeout but passed on when the caller is + # leaving, so a caller that has left never owns the key. + with _state_lock: + if not waiter.granted: + if waiter in _table.get(name, ()): # _release already dropped it if its loop closed + _table[name].remove(waiter) + return False + if not keep_grant: + _release(name) + return keep_grant + + +def _remaining(deadline: float | None) -> float | None: + return None if deadline is None else max(0.0, deadline - time.monotonic()) + + +def _acquire(name: tuple[str, str], deadline: float | None) -> None: + event = threading.Event() + waiter = _Waiter(event.set) + if _take_or_queue(name, waiter): + return + try: + if not event.wait(_remaining(deadline)): + raise TimeoutError(f"Timed out waiting for kv_local_lock on {name}") + except BaseException as e: # timeout or KeyboardInterrupt + if not _settle(name, waiter, keep_grant=isinstance(e, TimeoutError)): + raise + + +async def _acquire_async(name: tuple[str, str], deadline: float | None) -> None: + loop = asyncio.get_running_loop() + future = loop.create_future() + waiter = _Waiter(lambda: loop.call_soon_threadsafe(future.set_result, None)) + if _take_or_queue(name, waiter): + return + try: + await asyncio.wait([future], timeout=_remaining(deadline)) + if not future.done(): + raise TimeoutError(f"Timed out waiting for async_kv_local_lock on {name}") + except BaseException as e: # timeout or cancellation + if not _settle(name, waiter, keep_grant=isinstance(e, TimeoutError)): + raise + + +def _lock_names(keys: str | list[str], partition_id: str) -> list[tuple[str, str]]: + keys = [keys] if isinstance(keys, str) else keys + if not keys: + raise ValueError("kv_local_lock needs at least one key") + if _holding.get(): + raise RuntimeError("Nested kv_local_lock is not supported; pass all keys to a single call") + # A fixed order keeps two overlapping multi-key calls from deadlocking. + return [(partition_id, key) for key in sorted(set(keys))] + + +def _release_all(names: list[tuple[str, str]]) -> None: + with _state_lock: + for name in reversed(names): + _release(name) + + +@contextmanager +def kv_local_lock(keys: str | list[str], partition_id: str, timeout: float | None = None): + """Hold an exclusive process-local lock on ``keys`` in ``partition_id``. + + Blocks the calling thread, so use ``async_kv_local_lock`` inside a running event loop. + Raises ``TimeoutError`` if not all keys are acquired within ``timeout`` seconds + (``None`` waits forever) and ``RuntimeError`` if a local lock is already held. + """ + try: + asyncio.get_running_loop() + except RuntimeError: + pass + else: + raise RuntimeError("kv_local_lock would block the running event loop; use async_kv_local_lock instead") + names, acquired = _lock_names(keys, partition_id), [] + deadline = None if timeout is None else time.monotonic() + timeout + _holding.set(True) + try: + for name in names: + _acquire(name, deadline) + acquired.append(name) + yield + finally: + _release_all(acquired) + _holding.set(False) + + +@asynccontextmanager +async def async_kv_local_lock(keys: str | list[str], partition_id: str, timeout: float | None = None): + """Async version of ``kv_local_lock``. Tasks created inside the block cannot take a local lock.""" + names, acquired = _lock_names(keys, partition_id), [] + deadline = None if timeout is None else time.monotonic() + timeout + _holding.set(True) + try: + for name in names: + await _acquire_async(name, deadline) + acquired.append(name) + yield + finally: + _release_all(acquired) + _holding.set(False) + + +def _reset_after_fork() -> None: + # The child inherits a snapshot that may show keys held by threads it does not have. + global _state_lock, _table + _state_lock, _table = threading.Lock(), {} + + +os.register_at_fork(after_in_child=_reset_after_fork) From eb256069a5bb8e33660c05436e18fa6373d5590c Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 7 Oct 2026 11:51:59 +0800 Subject: [PATCH 02/42] [feat] Add kv_local_locked wrappers to run one KV call under the local lock kv_local_locked(fn, keys, partition_id, *, lock_timeout=None, **kwargs) and its async counterpart hold kv_local_lock on keys while calling fn(keys, partition_id, **kwargs), so keys and partition are written once. fn may be any TQ KV call or a user helper with the same leading arguments, such as a get-compute-put step. Signed-off-by: OutstanderWang --- tests/test_kv_lock.py | 106 ++++++++++++++++++++++++++++++++++--- transfer_queue/__init__.py | 4 +- transfer_queue/kv_lock.py | 29 +++++++++- 3 files changed, 131 insertions(+), 8 deletions(-) diff --git a/tests/test_kv_lock.py b/tests/test_kv_lock.py index 686b04e6..ab057884 100644 --- a/tests/test_kv_lock.py +++ b/tests/test_kv_lock.py @@ -24,9 +24,10 @@ import ray import torch from omegaconf import OmegaConf +from tensordict import TensorDict import transfer_queue as tq -from transfer_queue import async_kv_local_lock, kv_local_lock +from transfer_queue import async_kv_local_lock, async_kv_local_locked, kv_local_lock, kv_local_locked # `transfer_queue.kv_lock` resolves to the re-exported function, so fetch the module itself. kvl = importlib.import_module("transfer_queue.kv_lock") @@ -305,6 +306,85 @@ async def child(): asyncio.run(main()) +def is_held_elsewhere(keys): + """Whether another thread's kv_local_lock on `keys` times out.""" + outcome = [] + + def run(): + try: + with kv_local_lock(keys, P, timeout=0): + outcome.append(False) + except TimeoutError: + outcome.append(True) + + thread = threading.Thread(target=run) + thread.start() + thread.join(TIMEOUT) + assert not thread.is_alive() + return outcome == [True] + + +def test_locked_wrapper_passes_arguments_through_and_holds_the_lock_during_the_call(): + def fn(keys, partition_id, **kwargs): + assert is_held_elsewhere(keys) + return keys, partition_id, kwargs + + assert kv_local_locked(fn, "a", P, lock_timeout=TIMEOUT, x=1) == ("a", P, {"x": 1}) + assert kv_local_locked(fn, ["b", "a"], P) == (["b", "a"], P, {}) + assert not is_held_elsewhere(["a", "b"]) + + def boom(keys, partition_id): + raise ValueError("boom") + + with pytest.raises(ValueError, match="boom"): + kv_local_locked(boom, "a", P) + assert not is_held_elsewhere("a") + + +def test_async_locked_wrapper_holds_the_lock_and_releases_on_error(): + async def fn(keys, partition_id, **kwargs): + assert is_held_elsewhere(keys) + return keys, partition_id, kwargs + + async def boom(keys, partition_id): + raise ValueError("boom") + + async def main(): + assert await async_kv_local_locked(fn, ["a", "b"], P, lock_timeout=TIMEOUT, x=1) == (["a", "b"], P, {"x": 1}) + with pytest.raises(ValueError, match="boom"): + await async_kv_local_locked(boom, "a", P) + + asyncio.run(main()) + assert not is_held_elsewhere(["a", "b"]) + + +def test_locked_wrapper_lock_timeout_skips_the_call(): + calls, release = [], threading.Event() + holder = start_holder("k", release) + + async def async_fn(keys, partition_id): + calls.append(keys) + + with pytest.raises(TimeoutError): + kv_local_locked(lambda keys, partition_id: calls.append(keys), ["j", "k"], P, lock_timeout=0.05) + with pytest.raises(TimeoutError): + asyncio.run(async_kv_local_locked(async_fn, "k", P, lock_timeout=0.05)) + assert calls == [] and n_waiters("k") == 0 + release.set() + holder.join(TIMEOUT) + assert not holder.is_alive() + + +def test_locked_wrappers_reject_the_wrong_kind_of_function(): + async def async_fn(keys, partition_id): + pass + + with pytest.raises(TypeError, match="async_kv_local_locked"): + kv_local_locked(async_fn, "a", P) + with pytest.raises(TypeError, match="coroutine function"): + asyncio.run(async_kv_local_locked(lambda keys, partition_id: None, "a", P)) + + @pytest.fixture def simple_storage_tq(): if not ray.is_initialized(): @@ -319,17 +399,31 @@ def simple_storage_tq(): def test_kv_local_lock_serializes_simple_storage_read_modify_write(simple_storage_tq): tq.kv_put("counter", "lock_e2e", fields={"v": torch.tensor([0])}) # also creates the client on this thread - def worker(): + def increment(key, partition_id): + value = tq.kv_batch_get(key, partition_id)["v"] + tq.kv_put(key, partition_id, fields={"v": value[0] + 1}) + + def context_worker(): for _ in range(50): with tq.kv_local_lock("counter", "lock_e2e", timeout=E2E_TIMEOUT): - value = tq.kv_batch_get("counter", "lock_e2e")["v"] - tq.kv_put("counter", "lock_e2e", fields={"v": value[0] + 1}) + increment("counter", "lock_e2e") + + def wrapper_worker(): + for _ in range(50): + tq.kv_local_locked(increment, "counter", "lock_e2e", lock_timeout=E2E_TIMEOUT) - threads = [threading.Thread(target=worker) for _ in range(8)] + threads = [threading.Thread(target=worker) for worker in [context_worker, wrapper_worker] * 4] for thread in threads: thread.start() for thread in threads: thread.join(E2E_TIMEOUT) assert not thread.is_alive() assert tq.kv_batch_get("counter", "lock_e2e")["v"][0].item() == 400 - tq.kv_clear("counter", "lock_e2e") + + fields = TensorDict({"v": torch.tensor([[1], [2]])}, batch_size=2) + tq.kv_local_locked(tq.kv_batch_put, ["a", "b"], "lock_e2e", fields=fields) + data = tq.kv_local_locked(tq.kv_batch_get, ["a", "b"], "lock_e2e", select_fields="v") + assert [row.tolist() for row in data["v"]] == [[1], [2]] # rows may come back as a nested tensor + data = asyncio.run(tq.async_kv_local_locked(tq.async_kv_batch_get, "b", "lock_e2e", lock_timeout=E2E_TIMEOUT)) + assert [row.tolist() for row in data["v"]] == [[2]] + tq.kv_clear(["counter", "a", "b"], "lock_e2e") diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py index 60ff7091..49947aca 100644 --- a/transfer_queue/__init__.py +++ b/transfer_queue/__init__.py @@ -37,7 +37,7 @@ load_checkpoint, save_checkpoint, ) -from .kv_lock import async_kv_local_lock, kv_local_lock +from .kv_lock import async_kv_local_lock, async_kv_local_locked, kv_local_lock, kv_local_locked from .metadata import BatchMeta, KVBatchMeta from .sampler import BaseSampler from .sampler.grpo_group_n_sampler import GRPOGroupNSampler @@ -70,6 +70,8 @@ # Process-local KV key locks "kv_local_lock", "async_kv_local_lock", + "kv_local_locked", + "async_kv_local_locked", ] + [ # Checkpoint Interface diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 1f298929..76f72a72 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -19,10 +19,12 @@ ``(partition_id, key)``, shared by every thread and asyncio task of this process. They are advisory: KV calls never check them. They do not exclude other processes or Ray actors. Waiters on a key are served in FIFO order. Nesting is rejected; pass all keys -to a single call. +to a single call. ``kv_local_locked`` and ``async_kv_local_locked`` run one KV call +under the lock without repeating its keys and partition. """ import asyncio +import inspect import os import threading import time @@ -178,6 +180,31 @@ async def async_kv_local_lock(keys: str | list[str], partition_id: str, timeout: _holding.set(False) +def kv_local_locked(fn, keys: str | list[str], partition_id: str, *, lock_timeout: float | None = None, **kwargs): + """Call ``fn(keys, partition_id, **kwargs)`` under ``kv_local_lock(keys, partition_id)``. + + ``fn`` is any KV call such as ``tq.kv_batch_put``, or a user helper with the same + leading ``(keys, partition_id)`` arguments, e.g. one doing get -> compute -> put:: + + meta = tq.kv_local_locked(tq.kv_batch_put, ["a", "b"], "train", fields=td, lock_timeout=5) + tq.kv_local_locked(increment, "counter", "train") + """ + if inspect.iscoroutinefunction(fn): + raise TypeError("kv_local_locked would release the lock before the coroutine runs; use async_kv_local_locked") + with kv_local_lock(keys, partition_id, timeout=lock_timeout): + return fn(keys, partition_id, **kwargs) + + +async def async_kv_local_locked( + fn, keys: str | list[str], partition_id: str, *, lock_timeout: float | None = None, **kwargs +): + """Async version of ``kv_local_locked``; ``fn`` must be a coroutine function such as ``tq.async_kv_batch_get``.""" + if not inspect.iscoroutinefunction(fn): + raise TypeError("async_kv_local_locked needs a coroutine function; a sync call would block the event loop") + async with async_kv_local_lock(keys, partition_id, timeout=lock_timeout): + return await fn(keys, partition_id, **kwargs) + + def _reset_after_fork() -> None: # The child inherits a snapshot that may show keys held by threads it does not have. global _state_lock, _table From bd75f37caeb42b41525a99c7271264d56f212896 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Sat, 3 Oct 2026 23:37:27 +0800 Subject: [PATCH 03/42] [fix] Only the process that created TransferQueue tears it down in close() close() always killed the controller, and processes that attach via init() without a conf also hold a controller handle, so close() in any worker tore TransferQueue down for the whole cluster. init() now records ownership when it actually creates the controller actor. In other processes close() only closes the client and drops the handles, so they can attach again later. Storage cleanup was already owner-only. Signed-off-by: OutstanderWang --- tests/e2e/test_close_ownership_e2e.py | 69 +++++++++++++++++++++ tests/test_interface_close.py | 86 +++++++++++++++++++++++++++ transfer_queue/interface.py | 22 ++++--- 3 files changed, 169 insertions(+), 8 deletions(-) create mode 100644 tests/e2e/test_close_ownership_e2e.py create mode 100644 tests/test_interface_close.py diff --git a/tests/e2e/test_close_ownership_e2e.py b/tests/e2e/test_close_ownership_e2e.py new file mode 100644 index 00000000..fc3b3fee --- /dev/null +++ b/tests/e2e/test_close_ownership_e2e.py @@ -0,0 +1,69 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""A Ray worker that attaches to TransferQueue and closes it must not tear it down for the driver.""" + +import time + +import pytest +import ray +import torch +from omegaconf import OmegaConf + +import transfer_queue as tq + +TIMEOUT_S = 60 +CONF = { + "controller": {"polling_mode": True}, + "backend": {"storage_backend": "SimpleStorage", "SimpleStorage": {"total_storage_size": 100}}, +} + + +@ray.remote +def attach_use_and_close() -> int: + tq.init() + tq.kv_put("from_worker", "close_e2e", fields={"v": torch.tensor([7])}) + value = int(tq.kv_batch_get("from_driver", "close_e2e")["v"][0]) + tq.close() + return value + + +@pytest.fixture +def ray_session(): + if not ray.is_initialized(): + ray.init(namespace="TestCloseOwnershipE2E") + yield + ray.shutdown() + + +def test_only_the_creating_process_tears_down(ray_session): + tq.init(OmegaConf.create(CONF)) + tq.kv_put("from_driver", "close_e2e", fields={"v": torch.tensor([3])}) + + assert ray.get(attach_use_and_close.remote(), timeout=TIMEOUT_S) == 3 + + controller = ray.get_actor("TransferQueueController", namespace="transfer_queue") + ray.get(controller.list_partitions.remote(), timeout=TIMEOUT_S) + assert int(tq.kv_batch_get("from_worker", "close_e2e")["v"][0]) == 7 + + tq.close() + deadline = time.monotonic() + TIMEOUT_S + while True: + try: + ray.get_actor("TransferQueueController", namespace="transfer_queue") + except ValueError: + break + assert time.monotonic() < deadline, "controller still alive after the owner's close()" + time.sleep(0.1) diff --git a/tests/test_interface_close.py b/tests/test_interface_close.py new file mode 100644 index 00000000..761c0c01 --- /dev/null +++ b/tests/test_interface_close.py @@ -0,0 +1,86 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for `close()` ownership: only the process that created TransferQueue tears it down.""" + +from unittest.mock import MagicMock + +import pytest +from omegaconf import OmegaConf + +import transfer_queue.interface as iface + +STORED_CONF = OmegaConf.create( + {"controller": {"zmq_info": None}, "backend": {"storage_backend": "SimpleStorage", "SimpleStorage": {}}} +) + + +@pytest.fixture +def fake(monkeypatch): + for name in ("_TQ_CLIENT", "_TQ_STORAGE", "_TQ_CONTROLLER"): + monkeypatch.setattr(iface, name, None) + monkeypatch.setattr(iface, "_TQ_IS_OWNER", False) + + controller = MagicMock(name="controller") + ray = MagicMock(name="ray") + ray.get.return_value = STORED_CONF + controller_cls = MagicMock(name="TransferQueueController") + controller_cls.options.return_value.remote.return_value = controller + monkeypatch.setattr(iface, "ray", ray) + monkeypatch.setattr(iface, "TransferQueueController", controller_cls) + monkeypatch.setattr(iface, "TransferQueueClient", MagicMock(name="TransferQueueClient")) + monkeypatch.setattr(iface, "process_zmq_server_info", MagicMock(return_value=None)) + monkeypatch.setattr(iface, "_maybe_create_tq_storage", lambda conf: conf) + return ray, controller, controller_cls + + +def test_owner_close_kills_controller(fake): + ray, controller, _ = fake + ray.get_actor.side_effect = ValueError("no controller yet") + + iface.init(OmegaConf.create({})) + client = iface._TQ_CLIENT + assert iface._TQ_IS_OWNER + + iface.close() + client.close.assert_called_once() + ray.kill.assert_called_once_with(controller) + assert (iface._TQ_CLIENT, iface._TQ_CONTROLLER, iface._TQ_IS_OWNER) == (None, None, False) + + +@pytest.mark.parametrize("lost_creation_race", [False, True], ids=["existing", "lost_creation_race"]) +def test_attaching_close_keeps_controller_and_can_reattach(fake, lost_creation_race): + ray, controller, controller_cls = fake + if lost_creation_race: + # Another process creates the named actor between our lookup and our creation attempt. + ray.get_actor.side_effect = [ValueError("no controller yet"), controller] + controller_cls.options.return_value.remote.side_effect = ValueError("name already taken") + else: + ray.get_actor.return_value = controller + + iface.init(OmegaConf.create({})) + client = iface._TQ_CLIENT + assert iface._TQ_CONTROLLER is controller and not iface._TQ_IS_OWNER + + iface.close() + client.close.assert_called_once() + ray.kill.assert_not_called() + assert (iface._TQ_CLIENT, iface._TQ_CONTROLLER, iface._TQ_IS_OWNER) == (None, None, False) + + ray.get_actor.side_effect = None + ray.get_actor.return_value = controller + iface.init() + assert iface._TQ_CONTROLLER is controller and iface._TQ_CLIENT is not None + assert not iface._TQ_IS_OWNER diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index eca8762e..15bdf7f4 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -59,6 +59,9 @@ _TQ_CLIENT: Any = None _TQ_STORAGE: Any = None _TQ_CONTROLLER: Any = None +# True only in the process whose init() created the controller. Other processes +# attach to the same named actor, so their close() must not tear it down. +_TQ_IS_OWNER = False # Idle workers and proxies observe shutdown within one second; leave time for # Ray dispatch, in-flight work, and SSD cleanup before forcing termination. @@ -147,7 +150,8 @@ def init(conf: DictConfig | None = None) -> DictConfig | None: It should be called once at the beginning of the program before any data operations. If a controller already exists, reuse it and only initialize the client; - the provided `conf` will be ignored in this case. + the provided `conf` will be ignored in this case. Only the process that + created the controller tears TransferQueue down in `close()`. Args: conf: Optional custom config merged with default `config.yaml`. @@ -196,10 +200,11 @@ def init(conf: DictConfig | None = None) -> DictConfig | None: raise ValueError(f"Could not find sampler {final_conf.controller.sampler}") from None try: - global _TQ_CONTROLLER + global _TQ_CONTROLLER, _TQ_IS_OWNER _TQ_CONTROLLER = TransferQueueController.options( # type: ignore[attr-defined] name="TransferQueueController", namespace="transfer_queue" ).remote(sampler=sampler, polling_mode=final_conf.controller.polling_mode) + _TQ_IS_OWNER = True logger.info("TransferQueueController has been created.") except ValueError: logger.info("Some other rank has initialized TransferQueueController. Try to connect to existing controller.") @@ -241,10 +246,9 @@ def init(conf: DictConfig | None = None) -> DictConfig | None: def close(): """Close the TransferQueue system. - This function cleans up the TransferQueue system, including: - - Closing the client and its associated resources - - Cleaning up distributed storage (only for the process that initialized it) - - Killing the controller actor + Every process closes its client and drops its handles. Only the process whose + `init()` created TransferQueue also cleans up distributed storage and kills the + controller actor; processes that attached leave them running and may `init()` again. Note: This function should be called when the TransferQueue system is no longer needed. @@ -252,6 +256,7 @@ def close(): global _TQ_CLIENT global _TQ_STORAGE global _TQ_CONTROLLER + global _TQ_IS_OWNER try: if _TQ_STORAGE: @@ -311,12 +316,13 @@ def close(): _TQ_CLIENT.close() _TQ_CLIENT = None - if _TQ_CONTROLLER: + if _TQ_CONTROLLER and _TQ_IS_OWNER: try: ray.kill(_TQ_CONTROLLER) except Exception: pass - _TQ_CONTROLLER = None + _TQ_CONTROLLER = None + _TQ_IS_OWNER = False # ==================== Metrics API ==================== From b9b58e9c47398850244fff5059872edddccec57c Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 7 Oct 2026 12:58:03 +0800 Subject: [PATCH 04/42] [feat] Add a cluster-wide kv_global_lock backed by sharded Ray actors kv_global_lock / async_kv_global_lock hold an advisory exclusive lock per (partition_id, key) across the Ray cluster. tq.init() creates controller.num_lock_shards (default 8) named TransferQueueLockManager actors next to the controller, and each key goes to one of them by crc32 of (partition_id, key). Within a shard all keys of a request are granted at once; a request spanning shards visits them in ascending order, so overlapping requests cannot deadlock. Leases are renewed in batches by one daemon thread per process. The yielded lease exposes per-key fencing tokens and check(), which raises LockLostError once the lease may have passed to someone else. Callers that give up while waiting withdraw their acquire by token, so a grant never goes to a caller that left, and a release wakes only the waiters of the freed keys. kv_lock_list() shows holders and waiters. Only the owning process kills the lock actors in close(); other processes release their leases and leave them running. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 316 +++++++++++++++++++++++++++ tests/test_interface_close.py | 17 +- tests/test_kv_lock.py | 63 ++++++ transfer_queue/__init__.py | 17 +- transfer_queue/config.yaml | 2 + transfer_queue/interface.py | 30 ++- transfer_queue/kv_lock.py | 274 +++++++++++++++++++++-- transfer_queue/lock_manager.py | 150 +++++++++++++ 8 files changed, 845 insertions(+), 24 deletions(-) create mode 100644 tests/e2e/test_kv_global_lock_e2e.py create mode 100644 transfer_queue/lock_manager.py diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py new file mode 100644 index 00000000..b8cd8832 --- /dev/null +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -0,0 +1,316 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""kv_global_lock across Ray actors: exclusion, expiry, withdrawal, lost leases, and close() ownership.""" + +import asyncio +import importlib +import time + +import pytest +import ray +import torch +from omegaconf import OmegaConf + +import transfer_queue as tq + +kvl = importlib.import_module("transfer_queue.kv_lock") +P = "global_lock_e2e" +# Bound every wait so that a lock regression fails the test instead of hanging it. +TIMEOUT_S = 60 +CONF = { + "controller": {"polling_mode": True}, + "backend": {"storage_backend": "SimpleStorage", "SimpleStorage": {"total_storage_size": 100}}, +} + + +@ray.remote +class Worker: + def __init__(self): + tq.init() + + def increment(self, n, keys="counter"): + for _ in range(n): + with tq.kv_global_lock(keys, P, timeout=TIMEOUT_S): + value = int(tq.kv_batch_get("counter", P)["v"][0]) + time.sleep(0.001) + tq.kv_put("counter", P, fields={"v": torch.tensor([value + 1])}) + + def cycle(self, keys, n): + for _ in range(n): + with tq.kv_global_lock(keys, P, timeout=TIMEOUT_S): + pass + + def hold(self, key, lease_s=30): + self.held = tq.kv_global_lock(key, P, timeout=TIMEOUT_S, lease_s=lease_s) + return self.held.__enter__().fence[key] + + def release(self): + self.held.__exit__(None, None, None) + + def close(self): + tq.close() + + +@pytest.fixture(scope="module", autouse=True) +def tq_session(): + if not ray.is_initialized(): + ray.init(namespace="TestKVGlobalLockE2E") + tq.init(OmegaConf.create(CONF)) + yield + tq.close() + ray.shutdown() + + +def wait_until(predicate): + deadline = time.monotonic() + TIMEOUT_S + while not predicate(): + assert time.monotonic() < deadline, "condition not reached" + time.sleep(0.05) + + +def holders(): + # Exit sends its release without waiting for it, so wait for holders to drain, not assert. + return {h["key"] for h in tq.kv_lock_list(P)["holders"]} + + +def shard_of(key): + return kvl._shard((P, key), kvl._shard_count()) + + +def manager(key): + return ray.get_actor(f"TransferQueueLockManager_{shard_of(key)}", namespace="transfer_queue") + + +def keys_on_two_shards(prefix): + """Return two keys on different lock actors, the one on the lower shard first.""" + keys = sorted((f"{prefix}{i}" for i in range(50)), key=shard_of) + assert shard_of(keys[0]) < shard_of(keys[-1]) + return keys[0], keys[-1] + + +def counter(): + return int(tq.kv_batch_get("counter", P)["v"][0]) + + +def test_actors_serialize_read_modify_write(): + tq.kv_put("counter", P, fields={"v": torch.tensor([0])}) + workers = [Worker.remote() for _ in range(3)] + ray.get([w.increment.remote(20) for w in workers], timeout=TIMEOUT_S) + assert counter() == 60 + tq.kv_clear("counter", P) + + +def test_opposite_key_orders_on_two_shards_do_not_deadlock(): + a, b = keys_on_two_shards("ab") + tq.kv_put("counter", P, fields={"v": torch.tensor([0])}) + first, second = Worker.remote(), Worker.remote() + ray.get([first.increment.remote(30, [a, b]), second.increment.remote(30, [b, a])], timeout=TIMEOUT_S) + assert counter() == 60 + tq.kv_clear("counter", P) + wait_until(lambda: holders() == set()) + + +def test_multi_shard_timeout_releases_the_earlier_shard(): + early, late = keys_on_two_shards("mt") + worker = Worker.remote() + ray.get(worker.hold.remote(late), timeout=TIMEOUT_S) + with pytest.raises(TimeoutError): + with tq.kv_global_lock([early, late], P, timeout=0.5): + pass + wait_until(lambda: holders() == {late}) + assert tq.kv_lock_list(P)["waiters"] == 0 + with tq.kv_global_lock(early, P, timeout=5): # a leaked grant would hold it for its 30 s lease + pass + ray.get(worker.release.remote(), timeout=TIMEOUT_S) + wait_until(lambda: tq.kv_lock_list(P) == {"holders": [], "waiters": 0}) + + +def test_lease_lost_on_one_shard_raises(): + early, late = keys_on_two_shards("ll") + with pytest.raises(tq.LockLostError): + with tq.kv_global_lock([early, late], P, lease_s=1) as lease: + assert holders() == {early, late} and len(lease.shards) == 2 + ray.get(manager(late).release.remote(lease.token), timeout=TIMEOUT_S) + wait_until(lambda: lease.lost) + with pytest.raises(tq.LockLostError): + lease.check() + wait_until(lambda: holders() == set()) + + +@ray.remote +def lock_without_init(keys): + with tq.kv_global_lock(keys, P, timeout=TIMEOUT_S) as lease: + return list(lease.shards) + + +def test_process_without_init_learns_the_shard_count(): + a, b = keys_on_two_shards("ni") + assert ray.get(lock_without_init.remote([a, b]), timeout=TIMEOUT_S) == [shard_of(a), shard_of(b)] + wait_until(lambda: holders() == set()) + + +def test_timeout_leaves_no_holder_or_waiter(): + worker = Worker.remote() + ray.get(worker.hold.remote("t"), timeout=TIMEOUT_S) + with pytest.raises(TimeoutError): + with tq.kv_global_lock("t", P, timeout=0.3): + pass + assert tq.kv_lock_list(P)["waiters"] == 0 + ray.get(worker.release.remote(), timeout=TIMEOUT_S) + wait_until(lambda: tq.kv_lock_list(P) == {"holders": [], "waiters": 0}) + + +def test_cancelled_waiter_is_withdrawn_and_never_granted(): + worker = Worker.remote() + fence = ray.get(worker.hold.remote("c"), timeout=TIMEOUT_S) + + async def wait_for_lock(): + async with tq.async_kv_global_lock("c", P): + pass + + async def main(): + task = asyncio.create_task(wait_for_lock()) + deadline = time.monotonic() + TIMEOUT_S + while tq.kv_lock_list(P)["waiters"] != 1: + assert time.monotonic() < deadline, "the task never started waiting" + await asyncio.sleep(0.05) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, TIMEOUT_S) + + asyncio.run(main()) + wait_until(lambda: tq.kv_lock_list(P)["waiters"] == 0) + ray.get(worker.release.remote(), timeout=TIMEOUT_S) + # A ghost grant to the cancelled waiter would hold "c" for its whole 30 s lease. + with tq.kv_global_lock("c", P, timeout=5) as lease: + assert lease.fence["c"] == fence + 1 + + +def test_release_before_acquire_withdraws_it(): + lock_manager = manager("w") + ray.get(lock_manager.release.remote("early"), timeout=TIMEOUT_S) + assert ray.get(lock_manager.acquire.remote([(P, "w")], "early", None, 5, {}), timeout=TIMEOUT_S) is None + assert holders() == set() + + +def test_killed_holder_frees_the_key_after_its_lease(): + worker = Worker.remote() + ray.get(worker.hold.remote("k", lease_s=1.5), timeout=TIMEOUT_S) + ray.kill(worker) + start = time.monotonic() + with tq.kv_global_lock("k", P, timeout=TIMEOUT_S): + assert time.monotonic() - start < 10 + + +def test_lease_renews_and_lost_lease_raises(): + with tq.kv_global_lock("r", P, lease_s=2) as lease: + time.sleep(3) # outlives one lease only through renewal + lease.check() + + with pytest.raises(tq.LockLostError): + with tq.kv_global_lock("r", P, lease_s=1) as lease: + with kvl._registry: # stalls the renewer past the lease + time.sleep(1.5) + with pytest.raises(tq.LockLostError): + lease.check() + + with pytest.raises(ValueError, match="body"): # never masked by LockLostError + with tq.kv_global_lock("r", P, lease_s=1) as lease: + ray.get(manager("r").release.remote(lease.token), timeout=TIMEOUT_S) # the server drops it + wait_until(lambda: lease.lost) + raise ValueError("body") + wait_until(lambda: holders() == set()) + + +def test_fences_grow_across_holders(): + with tq.kv_global_lock(["f", "g"], P) as first: + pass + worker = Worker.remote() + fence = ray.get(worker.hold.remote("f"), timeout=TIMEOUT_S) + ray.get(worker.release.remote(), timeout=TIMEOUT_S) + with tq.kv_global_lock("f", P) as last: + pass + assert first.fence["f"] < fence < last.fence["f"] and set(first.fence) == {"f", "g"} + + +def test_nesting_and_ordering_rules(): + with pytest.raises(ValueError): + with tq.kv_global_lock([], P): + pass + with tq.kv_global_lock("n", P): + with pytest.raises(RuntimeError, match="Nested kv_global_lock"): + with tq.kv_global_lock("m", P): + pass + with tq.kv_local_lock("n", P): # global, then local is allowed + pass + with tq.kv_local_lock("n", P): + with pytest.raises(RuntimeError, match="before kv_local_lock"): + with tq.kv_global_lock("n", P): + pass + + async def main(): + with pytest.raises(RuntimeError, match="async_kv_global_lock"): + with tq.kv_global_lock("n", P): + pass + async with tq.async_kv_global_lock("n", P): + with pytest.raises(RuntimeError, match="Nested kv_global_lock"): + async with tq.async_kv_global_lock("m", P): + pass + + asyncio.run(main()) + wait_until(lambda: holders() == set()) + + +def actors_named(prefix): + return [a["name"] for a in ray.util.list_named_actors(all_namespaces=True) if a["name"].startswith(prefix)] + + +def test_non_owner_close_releases_its_locks_and_keeps_the_managers(): + worker = Worker.remote() + ray.get(worker.hold.remote("o"), timeout=TIMEOUT_S) + ray.get(worker.close.remote(), timeout=TIMEOUT_S) + wait_until(lambda: holders() == set()) + with tq.kv_global_lock("o", P, timeout=5): + pass + assert len(actors_named("TransferQueueLockManager_")) == kvl._shard_count() == 8 + + +def test_owner_close_kills_the_managers_and_fails_waiters(): + """Tears TransferQueue down for this module; only the single-shard test runs after it.""" + held = tq.kv_global_lock("z", P) + held.__enter__() + worker = Worker.remote() + waiting = worker.cycle.remote(["z"], 1) + wait_until(lambda: tq.kv_lock_list(P)["waiters"] == 1) + + tq.close() + with pytest.raises(RuntimeError, match="TransferQueueLockManager actor is gone"): + ray.get(waiting, timeout=TIMEOUT_S) + with pytest.raises(tq.LockLostError): + held.__exit__(None, None, None) + + wait_until(lambda: actors_named("TransferQueueLockManager_") == []) + with pytest.raises(RuntimeError, match="call tq.init"): + tq.kv_lock_list() + + +def test_one_lock_shard_still_works(): + wait_until(lambda: actors_named("TransferQueue") == []) + tq.init(OmegaConf.create({**CONF, "controller": {"polling_mode": True, "num_lock_shards": 1}})) + assert actors_named("TransferQueueLockManager_") == ["TransferQueueLockManager_0"] + with tq.kv_global_lock(["a", "b"], P, timeout=5) as lease: + assert holders() == {"a", "b"} and list(lease.shards) == [0] + wait_until(lambda: holders() == set()) diff --git a/tests/test_interface_close.py b/tests/test_interface_close.py index 761c0c01..af3c9453 100644 --- a/tests/test_interface_close.py +++ b/tests/test_interface_close.py @@ -40,6 +40,7 @@ def fake(monkeypatch): controller_cls.options.return_value.remote.return_value = controller monkeypatch.setattr(iface, "ray", ray) monkeypatch.setattr(iface, "TransferQueueController", controller_cls) + monkeypatch.setattr(iface, "TransferQueueLockManager", MagicMock(name="TransferQueueLockManager")) monkeypatch.setattr(iface, "TransferQueueClient", MagicMock(name="TransferQueueClient")) monkeypatch.setattr(iface, "process_zmq_server_info", MagicMock(return_value=None)) monkeypatch.setattr(iface, "_maybe_create_tq_storage", lambda conf: conf) @@ -51,13 +52,23 @@ def test_owner_close_kills_controller(fake): ray.get_actor.side_effect = ValueError("no controller yet") iface.init(OmegaConf.create({})) - client = iface._TQ_CLIENT - assert iface._TQ_IS_OWNER + client, lock_managers = iface._TQ_CLIENT, iface._TQ_LOCK_MANAGERS + assert iface._TQ_IS_OWNER and len(lock_managers) == 8 iface.close() client.close.assert_called_once() - ray.kill.assert_called_once_with(controller) + assert [c.args for c in ray.kill.call_args_list] == [(controller,)] + [(m,) for m in lock_managers] assert (iface._TQ_CLIENT, iface._TQ_CONTROLLER, iface._TQ_IS_OWNER) == (None, None, False) + assert iface._TQ_LOCK_MANAGERS == [] + + +@pytest.mark.parametrize("num_lock_shards", [0, "8"]) +def test_init_rejects_invalid_num_lock_shards(fake, num_lock_shards): + ray, _, controller_cls = fake + ray.get_actor.side_effect = ValueError("no controller yet") + with pytest.raises(ValueError, match="num_lock_shards"): + iface.init(OmegaConf.create({"controller": {"num_lock_shards": num_lock_shards}})) + controller_cls.options.assert_not_called() @pytest.mark.parametrize("lost_creation_race", [False, True], ids=["existing", "lost_creation_race"]) diff --git a/tests/test_kv_lock.py b/tests/test_kv_lock.py index ab057884..958d1077 100644 --- a/tests/test_kv_lock.py +++ b/tests/test_kv_lock.py @@ -16,6 +16,10 @@ import asyncio import gc import importlib +import os +import signal +import subprocess +import sys import threading import time from types import SimpleNamespace @@ -385,6 +389,65 @@ async def async_fn(keys, partition_id): asyncio.run(async_kv_local_locked(lambda keys, partition_id: None, "a", P)) +@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") +@pytest.mark.filterwarnings("ignore:.*fork.*:DeprecationWarning") +def test_forked_child_drops_the_parents_locks_leases_and_managers(monkeypatch): + monkeypatch.setattr(kvl, "_registry", threading.Condition()) + monkeypatch.setattr(kvl, "_leases", {"parent-token": object()}) + monkeypatch.setattr(kvl, "_renewer_running", True) + monkeypatch.setattr(kvl, "_managers", {0: object()}) + release, registry_held = threading.Event(), threading.Event() + + def hold_registry(): + with kvl._registry: + registry_held.set() + release.wait(TIMEOUT) + + registry_holder = threading.Thread(target=hold_registry) + registry_holder.start() + assert registry_held.wait(TIMEOUT) + key_holder = start_holder("a", release) + try: + pid = os.fork() + if pid == 0: + ok = False + try: + with kv_local_lock("a", P, timeout=1): + ok = kvl._leases == {} and kvl._managers == {} and not kvl._renewer_running + ok = ok and kvl._registry.acquire(blocking=False) + finally: + os._exit(0 if ok else 1) + deadline = time.monotonic() + TIMEOUT + while (status := os.waitpid(pid, os.WNOHANG))[0] == 0: + if time.monotonic() > deadline: + os.kill(pid, signal.SIGKILL) + pytest.fail("forked child hung") + time.sleep(0.01) + assert os.waitstatus_to_exitcode(status[1]) == 0 + finally: + release.set() + registry_holder.join() + key_holder.join() + + +def test_global_lock_shards_agree_across_processes_and_spread_keys(): + names = [(P, f"key{i}") for i in range(800)] + shards = [kvl._shard(name, 8) for name in names] + assert all(60 < shards.count(shard) < 140 for shard in range(8)) + # str hashes differ per PYTHONHASHSEED; the shard of a key must not. + script = "import transfer_queue.kv_lock as k; print([k._shard(('p', f'key{i}'), 8) for i in range(800)])" + for seed in ("1", "2"): + out = subprocess.run( + [sys.executable, "-c", script], + env={**os.environ, "PYTHONHASHSEED": seed}, + capture_output=True, + text=True, + timeout=E2E_TIMEOUT, + check=True, + ) + assert out.stdout.strip() == str(shards) + + @pytest.fixture def simple_storage_tq(): if not ray.is_initialized(): diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py index 49947aca..c8dc8660 100644 --- a/transfer_queue/__init__.py +++ b/transfer_queue/__init__.py @@ -37,7 +37,16 @@ load_checkpoint, save_checkpoint, ) -from .kv_lock import async_kv_local_lock, async_kv_local_locked, kv_local_lock, kv_local_locked +from .kv_lock import ( + LockLostError, + async_kv_global_lock, + async_kv_local_lock, + async_kv_local_locked, + kv_global_lock, + kv_local_lock, + kv_local_locked, + kv_lock_list, +) from .metadata import BatchMeta, KVBatchMeta from .sampler import BaseSampler from .sampler.grpo_group_n_sampler import GRPOGroupNSampler @@ -67,11 +76,15 @@ "KVBatchMeta", ] + [ - # Process-local KV key locks + # KV key locks: process-local and cluster-wide "kv_local_lock", "async_kv_local_lock", "kv_local_locked", "async_kv_local_locked", + "kv_global_lock", + "async_kv_global_lock", + "kv_lock_list", + "LockLostError", ] + [ # Checkpoint Interface diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index c4da5fba..a593d206 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -12,6 +12,8 @@ controller: sampler: SequentialSampler # Whether return an empty BatchMeta to prevent request blocking when no enough data is available polling_mode: False + # Number of lock actors that kv_global_lock spreads keys over by hash (>= 1) + num_lock_shards: 8 # ZMQ Server IP & Ports (automatically generated during init) zmq_info: null diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 15bdf7f4..29fc67ef 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -45,6 +45,8 @@ from transfer_queue.client import TransferQueueClient from transfer_queue.controller import TransferQueueController +from transfer_queue.kv_lock import _detach_lock_managers +from transfer_queue.lock_manager import TransferQueueLockManager from transfer_queue.metadata import KVBatchMeta from transfer_queue.sampler import * # noqa: F401 from transfer_queue.sampler import BaseSampler @@ -62,6 +64,7 @@ # True only in the process whose init() created the controller. Other processes # attach to the same named actor, so their close() must not tear it down. _TQ_IS_OWNER = False +_TQ_LOCK_MANAGERS: list = [] # set only in the owner, which kills them in close() # Idle workers and proxies observe shutdown within one second; leave time for # Ray dispatch, in-flight work, and SSD cleanup before forcing termination. @@ -198,14 +201,24 @@ def init(conf: DictConfig | None = None) -> DictConfig | None: sampler = globals()[final_conf.controller.sampler] except KeyError: raise ValueError(f"Could not find sampler {final_conf.controller.sampler}") from None + num_lock_shards = final_conf.controller.num_lock_shards + if not isinstance(num_lock_shards, int) or num_lock_shards < 1: + raise ValueError(f"controller.num_lock_shards must be an integer >= 1, got {num_lock_shards!r}") try: - global _TQ_CONTROLLER, _TQ_IS_OWNER + global _TQ_CONTROLLER, _TQ_IS_OWNER, _TQ_LOCK_MANAGERS _TQ_CONTROLLER = TransferQueueController.options( # type: ignore[attr-defined] name="TransferQueueController", namespace="transfer_queue" ).remote(sampler=sampler, polling_mode=final_conf.controller.polling_mode) _TQ_IS_OWNER = True logger.info("TransferQueueController has been created.") + # Created here rather than on first lock: a non-detached actor dies with its creator. + _TQ_LOCK_MANAGERS = [ + TransferQueueLockManager.options( # type: ignore[attr-defined] + name=f"TransferQueueLockManager_{i}", namespace="transfer_queue", get_if_exists=True + ).remote(num_lock_shards) + for i in range(num_lock_shards) + ] except ValueError: logger.info("Some other rank has initialized TransferQueueController. Try to connect to existing controller.") _init_from_existing() @@ -257,6 +270,7 @@ def close(): global _TQ_STORAGE global _TQ_CONTROLLER global _TQ_IS_OWNER + global _TQ_LOCK_MANAGERS try: if _TQ_STORAGE: @@ -316,12 +330,16 @@ def close(): _TQ_CLIENT.close() _TQ_CLIENT = None - if _TQ_CONTROLLER and _TQ_IS_OWNER: - try: - ray.kill(_TQ_CONTROLLER) - except Exception: - pass + # The owner kills the lock managers, so releasing first would only hand keys to doomed waiters. + _detach_lock_managers(release=not _TQ_IS_OWNER) + if _TQ_IS_OWNER: + for actor in (_TQ_CONTROLLER, *_TQ_LOCK_MANAGERS): + try: + ray.kill(actor) + except Exception: + pass _TQ_CONTROLLER = None + _TQ_LOCK_MANAGERS = [] _TQ_IS_OWNER = False diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 76f72a72..361e9fec 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Advisory, process-local locks on TransferQueue KV keys. +"""Advisory locks on TransferQueue KV keys. ``kv_local_lock`` and ``async_kv_local_lock`` hold an exclusive mutex per ``(partition_id, key)``, shared by every thread and asyncio task of this process. They @@ -21,6 +21,10 @@ actors. Waiters on a key are served in FIFO order. Nesting is rejected; pass all keys to a single call. ``kv_local_locked`` and ``async_kv_local_locked`` run one KV call under the lock without repeating its keys and partition. + +``kv_global_lock`` and ``async_kv_global_lock`` do the same across the Ray cluster through +``TransferQueueLockManager`` actors that split keys by hash, under a lease that renews +automatically. Take a global lock before a local one, never inside it. """ import asyncio @@ -28,10 +32,15 @@ import os import threading import time +import uuid +import zlib from collections import deque from contextlib import asynccontextmanager, contextmanager from contextvars import ContextVar +import ray +from ray.exceptions import RayActorError + # Not asyncio.Lock: it is bound to one event loop and is not thread-safe, while a process # may run several loops plus plain threads. A key is held exactly while it is in _table, # which maps it to its FIFO of waiters. _state_lock guards every change and is never held @@ -121,12 +130,14 @@ async def _acquire_async(name: tuple[str, str], deadline: float | None) -> None: raise -def _lock_names(keys: str | list[str], partition_id: str) -> list[tuple[str, str]]: +def _lock_names( + keys: str | list[str], partition_id: str, held: ContextVar[bool] = _holding, kind: str = "kv_local_lock" +) -> list[tuple[str, str]]: keys = [keys] if isinstance(keys, str) else keys if not keys: - raise ValueError("kv_local_lock needs at least one key") - if _holding.get(): - raise RuntimeError("Nested kv_local_lock is not supported; pass all keys to a single call") + raise ValueError(f"{kind} needs at least one key") + if held.get(): + raise RuntimeError(f"Nested {kind} is not supported; pass all keys to a single call") # A fixed order keeps two overlapping multi-key calls from deadlocking. return [(partition_id, key) for key in sorted(set(keys))] @@ -137,6 +148,14 @@ def _release_all(names: list[tuple[str, str]]) -> None: _release(name) +def _reject_running_loop(kind: str) -> None: + try: + asyncio.get_running_loop() + except RuntimeError: + return + raise RuntimeError(f"{kind} would block the running event loop; use async_{kind} instead") + + @contextmanager def kv_local_lock(keys: str | list[str], partition_id: str, timeout: float | None = None): """Hold an exclusive process-local lock on ``keys`` in ``partition_id``. @@ -145,12 +164,7 @@ def kv_local_lock(keys: str | list[str], partition_id: str, timeout: float | Non Raises ``TimeoutError`` if not all keys are acquired within ``timeout`` seconds (``None`` waits forever) and ``RuntimeError`` if a local lock is already held. """ - try: - asyncio.get_running_loop() - except RuntimeError: - pass - else: - raise RuntimeError("kv_local_lock would block the running event loop; use async_kv_local_lock instead") + _reject_running_loop("kv_local_lock") names, acquired = _lock_names(keys, partition_id), [] deadline = None if timeout is None else time.monotonic() + timeout _holding.set(True) @@ -205,10 +219,244 @@ async def async_kv_local_locked( return await fn(keys, partition_id, **kwargs) +_holding_global: ContextVar[bool] = ContextVar("_holding_kv_global_lock", default=False) +# This process's leases, renewed in batches by one daemon thread that runs while any exist. +_registry = threading.Condition() +_leases: dict[str, "GlobalLease"] = {} +_renewer_running = False +_num_shards: int | None = None # controller.num_lock_shards, asked of shard 0 on first use +_managers: dict = {} # shard -> TransferQueueLockManager handle, looked up on first use +_NOT_INITIALIZED = "TransferQueueLockManager not found; call tq.init() first" +_MANAGER_GONE = "The TransferQueueLockManager actor is gone; did the process that ran tq.init() call tq.close()?" + + +class LockLostError(RuntimeError): + """A kv_global_lock lease ran out or failed to renew, so another holder may own the keys.""" + + +class GlobalLease: + """Yielded by ``kv_global_lock``. ``fence`` maps each key to a fencing token that grows with every grant.""" + + def __init__(self, names: list[tuple[str, str]], shards: dict[int, list[tuple[str, str]]], lease_s: float): + self.names, self.shards, self.lease_s = names, shards, lease_s + self.token = f"{os.getpid()}:{uuid.uuid4().hex}" + self.fence: dict[str, int] = {} + self.deadline = float("inf") # local monotonic time the lease is known to last until, once granted + self.lost = False + self.pending: list[int] = [] # shards that may hold or await this token, so exit must release them + + def check(self) -> None: + """Raise ``LockLostError`` if the lease was lost; call it right before writes in long sections.""" + if self.lost or time.monotonic() >= self.deadline: + raise LockLostError(f"kv_global_lock lease on {self.names} was lost") + + +def _shard_count() -> int: + global _num_shards + if _num_shards is None: + try: + _num_shards = ray.get(_lock_manager(0).num_shards.remote()) + except RayActorError as e: + raise RuntimeError(_MANAGER_GONE) from e + return _num_shards + + +def _shard(name: tuple[str, str], num_shards: int) -> int: + # crc32, not hash(): str hashes are salted per process, and every process must send a + # key to the same lock actor. + return zlib.crc32(f"{name[0]}\0{name[1]}".encode()) % num_shards + + +def _lock_manager(shard: int): + if shard not in _managers: + try: + _managers[shard] = ray.get_actor(f"TransferQueueLockManager_{shard}", namespace="transfer_queue") + except ValueError: + raise RuntimeError(_NOT_INITIALIZED) from None + return _managers[shard] + + +def _new_lease(keys: str | list[str], partition_id: str, lease_s: float) -> GlobalLease: + names = _lock_names(keys, partition_id, _holding_global, "kv_global_lock") + if _holding.get(): + raise RuntimeError("Take kv_global_lock before kv_local_lock, not while holding one") + if lease_s <= 0: + raise ValueError("lease_s must be positive") + num_shards, shards = _shard_count(), {} + for name in names: + shards.setdefault(_shard(name, num_shards), []).append(name) + return GlobalLease(names, dict(sorted(shards.items())), lease_s) + + +def _send_acquire(lease: GlobalLease, shard: int, deadline: float | None): + manager = _lock_manager(shard) + holder = {"node_ip": ray.util.get_node_ip_address(), "pid": os.getpid(), "thread": threading.current_thread().name} + lease.pending.append(shard) + return manager.acquire.remote(lease.shards[shard], lease.token, _remaining(deadline), lease.lease_s, holder) + + +def _take_grant(lease: GlobalLease, shard: int, reply, sent: float) -> None: + if reply is None: + lease.pending.remove(shard) + raise TimeoutError(f"Timed out waiting for kv_global_lock on {lease.names}") + fences, waited = reply + lease.fence.update(fences) + # The grant happened at least `waited` after `sent`, so this never outlives the shard's lease. + lease.deadline = min(lease.deadline, sent + waited + lease.lease_s) + + +def _start_lease(lease: GlobalLease) -> None: + global _renewer_running + with _registry: + _leases[lease.token] = lease + _registry.notify() # a shorter lease_s shortens the renewal period + if not _renewer_running: + _renewer_running = True + threading.Thread(target=_renew_loop, name="kv_global_lock_renewer", daemon=True).start() + + +def _release_lease(lease: GlobalLease) -> None: + with _registry: + _leases.pop(lease.token, None) + shards, lease.pending = lease.pending, [] + for shard in shards: + # Releasing by token also withdraws an acquire that is still waiting or in flight, so + # it is never granted to a caller that left; ray.cancel would miss a grant already made. + try: + _managers[shard].release.remote(lease.token) + except Exception: # close() dropped the handle, or the actor is gone + pass + + +def _renew_loop() -> None: + global _renewer_running + last = time.monotonic() + while True: + with _registry: + while True: + if not _leases: + _renewer_running = False + return + period = min(lease.lease_s for lease in _leases.values()) / 3 + remaining = last + period - time.monotonic() + if remaining <= 0: + break + _registry.wait(remaining) + batch = dict(_leases) + by_shard: dict[int, list[str]] = {} + for token, lease in batch.items(): + for shard in lease.shards: + by_shard.setdefault(shard, []).append(token) + # Count from the send time: the server extends the lease from when the call arrives, + # which is later, so the local deadline never outlives the server's. + last = time.monotonic() + calls, renewed = {}, set() + for shard, tokens in by_shard.items(): + try: + calls[shard] = _managers[shard].renew_many.remote(tokens) + except Exception: # close() dropped the handle + pass + for shard, call in calls.items(): + try: + alive = ray.get(call, timeout=max(0.0, last + 3 * period - time.monotonic())) + except Exception: # a dead lock manager or a timeout + continue + renewed.update((shard, token) for token, ok in alive.items() if ok) + # A lease is lost as soon as any one of its shards fails to renew it. + for token, lease in batch.items(): + if all((shard, token) in renewed for shard in lease.shards): + lease.deadline = last + lease.lease_s + else: + lease.lost = True + + +@contextmanager +def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | None = None, lease_s: float = 30): + """Hold an exclusive cluster-wide lock on ``keys`` in ``partition_id``; yield a ``GlobalLease``. + + The locks live in the ``controller.num_lock_shards`` actors that ``tq.init()`` creates, + each key in the one its hash picks, and are held under a ``lease_s`` lease that a + background thread renews. Keys on several actors are taken one actor at a time, each + all at once, so keys already granted stay held while a later actor's keys are awaited. + Raises ``TimeoutError`` if the keys are not all granted within ``timeout`` seconds + (``None`` waits forever), and ``RuntimeError`` when nested, taken while holding a + ``kv_local_lock``, or called with a running event loop. Exit always releases the lock, + then raises ``LockLostError`` if the lease was lost, unless the body already raised: + an exception from the body is never masked. + """ + _reject_running_loop("kv_global_lock") + lease = _new_lease(keys, partition_id, lease_s) + deadline = None if timeout is None else time.monotonic() + timeout + _holding_global.set(True) + try: + # Every caller visits shards in ascending order, so no two callers can each hold a + # shard the other awaits: multi-shard requests cannot deadlock. + for shard in lease.shards: + sent = time.monotonic() + try: + reply = ray.get(_send_acquire(lease, shard, deadline)) + except RayActorError as e: + raise RuntimeError(_MANAGER_GONE) from e + _take_grant(lease, shard, reply, sent) + _start_lease(lease) + yield lease + lease.check() + finally: + _release_lease(lease) + _holding_global.set(False) + + +@asynccontextmanager +async def async_kv_global_lock( + keys: str | list[str], partition_id: str, timeout: float | None = None, lease_s: float = 30 +): + """Async version of ``kv_global_lock``. Tasks created inside the block cannot take a global lock.""" + lease = _new_lease(keys, partition_id, lease_s) + deadline = None if timeout is None else time.monotonic() + timeout + _holding_global.set(True) + try: + for shard in lease.shards: + sent = time.monotonic() + try: + reply = await _send_acquire(lease, shard, deadline) + except RayActorError as e: + raise RuntimeError(_MANAGER_GONE) from e + _take_grant(lease, shard, reply, sent) + _start_lease(lease) + yield lease + lease.check() + finally: + _release_lease(lease) + _holding_global.set(False) + + +def kv_lock_list(partition_id: str | None = None) -> dict: + """Return ``{"holders": [...], "waiters": n}`` for live global locks, optionally in one partition. + + Each holder has ``partition_id``, ``key``, ``holder`` (node IP, pid, thread), ``held_s`` and + ``lease_remaining_s``. + """ + replies = ray.get([_lock_manager(shard).list_locks.remote(partition_id) for shard in range(_shard_count())]) + return {"holders": [h for r in replies for h in r["holders"]], "waiters": sum(r["waiters"] for r in replies)} + + +def _detach_lock_managers(release: bool) -> None: + """Called by ``tq.close()``: give up this process's global leases and forget the actor handles.""" + global _num_shards, _managers + with _registry: + leases = list(_leases.values()) + for lease in leases: + lease.lost = True + if release: + _release_lease(lease) + _num_shards, _managers = None, {} + + def _reset_after_fork() -> None: - # The child inherits a snapshot that may show keys held by threads it does not have. - global _state_lock, _table + # The child inherits keys, leases and Ray handles that belong to threads and connections it lacks. + global _state_lock, _table, _registry, _leases, _renewer_running, _managers _state_lock, _table = threading.Lock(), {} + _registry, _leases, _renewer_running, _managers = threading.Condition(), {}, False, {} os.register_at_fork(after_in_child=_reset_after_fork) diff --git a/transfer_queue/lock_manager.py b/transfer_queue/lock_manager.py new file mode 100644 index 00000000..b8a7eeee --- /dev/null +++ b/transfer_queue/lock_manager.py @@ -0,0 +1,150 @@ +# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. +# Copyright 2025 The TransferQueue Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Server side of ``kv_global_lock``: async Ray actors, each holding leased locks on the keys that hash to it.""" + +import asyncio +import heapq +import time + +import ray + +# A withdrawal whose acquire never arrives (e.g. the caller died first) is dropped after this. +_WITHDRAWN_TTL_S = 300 + + +# Every waiting acquire occupies a concurrency slot, so the limit sits far above Ray's async +# default of 1000 to keep renew and release from queueing behind waiters. +@ray.remote(num_cpus=0, max_concurrency=10_000) +class TransferQueueLockManager: + """Exclusive leased locks on ``(partition_id, key)``. Runs on one event loop, so no thread locks.""" + + def __init__(self, num_shards: int): + # Kept so that a process that never ran tq.init() can learn how many shards to hash over. + self._num_shards = num_shards + self._holder: dict[tuple[str, str], str] = {} # name -> token + self._leases: dict[str, dict] = {} # token -> names, lease_s, expires_at, granted_at, info + # Min-heap of (expires_at, token). Release and renewal leave stale entries behind, which + # _expire skips, so neither has to search the heap. + self._expiries: list[tuple[float, str]] = [] + self._fences: dict[tuple[str, str], int] = {} # never reset, so a key's fence only grows + self._waiting: dict[str, tuple[list[tuple[str, str]], asyncio.Event]] = {} # token -> names, wakeup + # name -> wakeups of its waiters, as an insertion-ordered dict so they wake oldest first + self._wakeups: dict[tuple[str, str], dict[asyncio.Event, None]] = {} + self._withdrawn: dict[str, float] = {} # token -> when its tombstone expires + + def _free(self, token: str) -> bool: + lease = self._leases.pop(token, None) + if lease is None: + return False + for name in lease["names"]: + del self._holder[name] + for wakeup in self._wakeups.get(name, ()): + wakeup.set() + return True + + def _expire(self, now: float) -> None: + while self._expiries and self._expiries[0][0] <= now: + expires_at, token = heapq.heappop(self._expiries) + lease = self._leases.get(token) + if lease is not None and lease["expires_at"] == expires_at: + self._free(token) + + async def acquire(self, names, token, timeout, lease_s, holder_info): + """Grant all ``names`` at once; return ``({key: fence}, waited_s)``, or None on timeout or withdrawal.""" + start = time.monotonic() + deadline = None if timeout is None else start + timeout + # Freeing a name sets the wakeups of its waiters only, so a release costs O(its waiters) + # instead of waking every waiter on the actor. + wakeup = asyncio.Event() + self._waiting[token] = (names, wakeup) + for name in names: + self._wakeups.setdefault(name, {})[wakeup] = None + try: + while True: + now = time.monotonic() + self._expire(now) + if self._withdrawn.pop(token, None) is not None: + return None + # All-or-nothing: a waiter never holds part of its keys, so overlapping + # requests cannot deadlock whatever order their keys come in. + busy = [self._leases[self._holder[name]]["expires_at"] for name in names if name in self._holder] + if not busy: + for name in names: + self._holder[name] = token + self._fences[name] = self._fences.get(name, 0) + 1 + self._leases[token] = dict( + names=names, lease_s=lease_s, expires_at=now + lease_s, granted_at=now, info=holder_info + ) + heapq.heappush(self._expiries, (now + lease_s, token)) + return {key: self._fences[(pid, key)] for pid, key in names}, now - start + if deadline is not None and now >= deadline: + return None + # Nothing runs between this clear and the wait, so no wakeup can be missed. Waking + # at the earliest blocking expiry lets _expire free it without a background task. + wakeup.clear() + wake_at = min(busy) if deadline is None else min(*busy, deadline) + try: + await asyncio.wait_for(wakeup.wait(), wake_at - now) + except asyncio.TimeoutError: + pass + finally: + del self._waiting[token] + for name in names: + del self._wakeups[name][wakeup] + if not self._wakeups[name]: + del self._wakeups[name] + + def num_shards(self) -> int: + return self._num_shards + + def renew_many(self, tokens: list[str]) -> dict[str, bool]: + """Extend each live lease by its own ``lease_s``; an expired or released one stays lost.""" + now = time.monotonic() + alive = {} + for token in tokens: + lease = self._leases.get(token) + alive[token] = lease is not None and lease["expires_at"] > now + if alive[token]: + lease["expires_at"] = now + lease["lease_s"] + heapq.heappush(self._expiries, (lease["expires_at"], token)) + return alive + + async def release(self, token: str) -> None: + """Free ``token``'s locks, or withdraw its acquire if it is still waiting or has not arrived.""" + if not self._free(token): + now = time.monotonic() + self._withdrawn = {t: until for t, until in self._withdrawn.items() if until > now} + self._withdrawn[token] = now + _WITHDRAWN_TTL_S + if token in self._waiting: + self._waiting[token][1].set() + + def list_locks(self, partition_id: str | None = None) -> dict: + now = time.monotonic() + holders = [] + for (pid, key), token in self._holder.items(): + lease = self._leases[token] + if partition_id in (None, pid) and lease["expires_at"] > now: + holders.append( + dict( + partition_id=pid, + key=key, + holder=lease["info"], + held_s=now - lease["granted_at"], + lease_remaining_s=lease["expires_at"] - now, + ) + ) + waiters = sum(partition_id in (None, names[0][0]) for names, _ in self._waiting.values()) + return {"holders": holders, "waiters": waiters} From 3443e993b9265c784332602029cbb571ef37228a Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Wed, 7 Oct 2026 13:06:08 +0800 Subject: [PATCH 05/42] [feat] Add kv_global_locked wrappers kv_global_locked(fn, keys, partition_id, *, lock_timeout=None, lease_s=30, **kwargs) and its async counterpart mirror kv_local_locked: they hold kv_global_lock on keys while calling fn(keys, partition_id, **kwargs), and reject a function of the wrong kind with TypeError. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 22 ++++++++++++++++++++++ transfer_queue/__init__.py | 4 ++++ transfer_queue/kv_lock.py | 27 ++++++++++++++++++++++++++- 3 files changed, 52 insertions(+), 1 deletion(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index b8cd8832..6efeddff 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -274,6 +274,28 @@ async def main(): wait_until(lambda: holders() == set()) +def test_locked_wrappers_hold_the_lock_during_the_call_and_release_on_error(): + def fn(keys, partition_id, **kwargs): + return holders(), partition_id, kwargs + + async def async_fn(keys, partition_id, **kwargs): + return holders(), partition_id, kwargs + + def boom(keys, partition_id): + raise ValueError("boom") + + assert tq.kv_global_locked(fn, ["u", "v"], P, lock_timeout=5, lease_s=2, x=1) == ({"u", "v"}, P, {"x": 1}) + assert asyncio.run(tq.async_kv_global_locked(async_fn, "u", P, lock_timeout=5)) == ({"u"}, P, {}) + with pytest.raises(ValueError, match="boom"): + tq.kv_global_locked(boom, "u", P) + wait_until(lambda: holders() == set()) + + with pytest.raises(TypeError, match="async_kv_global_locked"): + tq.kv_global_locked(async_fn, "u", P) + with pytest.raises(TypeError, match="coroutine function"): + asyncio.run(tq.async_kv_global_locked(fn, "u", P)) + + def actors_named(prefix): return [a["name"] for a in ray.util.list_named_actors(all_namespaces=True) if a["name"].startswith(prefix)] diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py index c8dc8660..d67afa85 100644 --- a/transfer_queue/__init__.py +++ b/transfer_queue/__init__.py @@ -40,9 +40,11 @@ from .kv_lock import ( LockLostError, async_kv_global_lock, + async_kv_global_locked, async_kv_local_lock, async_kv_local_locked, kv_global_lock, + kv_global_locked, kv_local_lock, kv_local_locked, kv_lock_list, @@ -83,6 +85,8 @@ "async_kv_local_locked", "kv_global_lock", "async_kv_global_lock", + "kv_global_locked", + "async_kv_global_locked", "kv_lock_list", "LockLostError", ] diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 361e9fec..572fcfc9 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -24,7 +24,8 @@ ``kv_global_lock`` and ``async_kv_global_lock`` do the same across the Ray cluster through ``TransferQueueLockManager`` actors that split keys by hash, under a lease that renews -automatically. Take a global lock before a local one, never inside it. +automatically. Take a global lock before a local one, never inside it. ``kv_global_locked`` +and ``async_kv_global_locked`` mirror the local wrappers. """ import asyncio @@ -430,6 +431,30 @@ async def async_kv_global_lock( _holding_global.set(False) +def kv_global_locked( + fn, keys: str | list[str], partition_id: str, *, lock_timeout: float | None = None, lease_s: float = 30, **kwargs +): + """Call ``fn(keys, partition_id, **kwargs)`` under ``kv_global_lock(keys, partition_id)``. + + Same calling convention as ``kv_local_locked``. If the lease was lost by the time ``fn`` + returns, its result is discarded and ``LockLostError`` is raised. + """ + if inspect.iscoroutinefunction(fn): + raise TypeError("kv_global_locked would release the lock before the coroutine runs; use async_kv_global_locked") + with kv_global_lock(keys, partition_id, timeout=lock_timeout, lease_s=lease_s): + return fn(keys, partition_id, **kwargs) + + +async def async_kv_global_locked( + fn, keys: str | list[str], partition_id: str, *, lock_timeout: float | None = None, lease_s: float = 30, **kwargs +): + """Async version of ``kv_global_locked``; ``fn`` must be a coroutine function such as ``tq.async_kv_batch_get``.""" + if not inspect.iscoroutinefunction(fn): + raise TypeError("async_kv_global_locked needs a coroutine function; a sync call would block the event loop") + async with async_kv_global_lock(keys, partition_id, timeout=lock_timeout, lease_s=lease_s): + return await fn(keys, partition_id, **kwargs) + + def kv_lock_list(partition_id: str | None = None) -> dict: """Return ``{"holders": [...], "waiters": n}`` for live global locks, optionally in one partition. From 12fa4e4a1b216b4a2930c56ecdbd457e3b5a082f Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 9 Oct 2026 10:23:15 +0800 Subject: [PATCH 06/42] [feat] Make kv_local_lock process-family local Signed-off-by: OutstanderWang --- tests/test_interface_close.py | 2 +- tests/test_kv_lock.py | 107 ++++++++++++++------------ transfer_queue/__init__.py | 2 +- transfer_queue/interface.py | 10 ++- transfer_queue/kv_lock.py | 133 ++++++++++++++++++++++++--------- transfer_queue/lock_manager.py | 12 ++- 6 files changed, 172 insertions(+), 94 deletions(-) diff --git a/tests/test_interface_close.py b/tests/test_interface_close.py index af3c9453..bff0470d 100644 --- a/tests/test_interface_close.py +++ b/tests/test_interface_close.py @@ -53,7 +53,7 @@ def test_owner_close_kills_controller(fake): iface.init(OmegaConf.create({})) client, lock_managers = iface._TQ_CLIENT, iface._TQ_LOCK_MANAGERS - assert iface._TQ_IS_OWNER and len(lock_managers) == 8 + assert iface._TQ_IS_OWNER and len(lock_managers) == 9 iface.close() client.close.assert_called_once() diff --git a/tests/test_kv_lock.py b/tests/test_kv_lock.py index 958d1077..08e9c681 100644 --- a/tests/test_kv_lock.py +++ b/tests/test_kv_lock.py @@ -17,12 +17,10 @@ import gc import importlib import os -import signal import subprocess import sys import threading import time -from types import SimpleNamespace import pytest import ray @@ -47,6 +45,18 @@ def table_is_empty_afterwards(): assert kvl._table == {} +@pytest.fixture(scope="module", autouse=True) +def node_lock_session(): + """Create the per-node lock actor used by every local-lock test in this module.""" + if not ray.is_initialized(): + ray.init(namespace="TestKVLock") + backend = {"storage_backend": "SimpleStorage", "SimpleStorage": {"total_storage_size": 200}} + tq.init(OmegaConf.create({"controller": {"polling_mode": True}, "backend": backend})) + yield + tq.close() + ray.shutdown() + + def n_waiters(key): with kvl._state_lock: return len(kvl._table.get((P, key), ())) @@ -220,7 +230,7 @@ def wait(self, timeout=None): holder.join(TIMEOUT) # the holder hands the key to this waiter... return False # ...but the wait still reports a timeout - monkeypatch.setattr(kvl, "threading", SimpleNamespace(Event=GrantedThenTimedOut)) + monkeypatch.setattr(kvl.threading, "Event", GrantedThenTimedOut) with kv_local_lock("k", P, timeout=1): assert (P, "k") in kvl._table and n_waiters("k") == 0 @@ -389,45 +399,57 @@ async def async_fn(keys, partition_id): asyncio.run(async_kv_local_locked(lambda keys, partition_id: None, "a", P)) -@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") -@pytest.mark.filterwarnings("ignore:.*fork.*:DeprecationWarning") -def test_forked_child_drops_the_parents_locks_leases_and_managers(monkeypatch): +def test_child_process_uses_the_parents_local_lock_actor(): + script = """ +import ray +import transfer_queue as tq + +ray.init(address="auto", logging_level="ERROR") +result = "acquired" +try: + with tq.kv_local_lock("family", "p", timeout=0.3): + pass +except TimeoutError: + result = "timeout" +finally: + ray.shutdown() +print(result) +""" + + def child_result(): + completed = subprocess.run( + [sys.executable, "-c", script], + env=os.environ.copy(), + capture_output=True, + text=True, + timeout=E2E_TIMEOUT, + check=True, + ) + return completed.stdout.strip() + + assert kvl._PROCESS_FAMILY_ENV in os.environ + with kv_local_lock("family", P, timeout=TIMEOUT): + assert child_result() == "timeout" + assert child_result() == "acquired" + + +def test_fork_reset_drops_the_parents_locks_leases_and_managers(monkeypatch): monkeypatch.setattr(kvl, "_registry", threading.Condition()) monkeypatch.setattr(kvl, "_leases", {"parent-token": object()}) monkeypatch.setattr(kvl, "_renewer_running", True) monkeypatch.setattr(kvl, "_managers", {0: object()}) - release, registry_held = threading.Event(), threading.Event() + monkeypatch.setattr(kvl, "_table", {(P, "a"): object()}) - def hold_registry(): - with kvl._registry: - registry_held.set() - release.wait(TIMEOUT) + old_registry = kvl._registry + kvl._reset_after_fork() - registry_holder = threading.Thread(target=hold_registry) - registry_holder.start() - assert registry_held.wait(TIMEOUT) - key_holder = start_holder("a", release) - try: - pid = os.fork() - if pid == 0: - ok = False - try: - with kv_local_lock("a", P, timeout=1): - ok = kvl._leases == {} and kvl._managers == {} and not kvl._renewer_running - ok = ok and kvl._registry.acquire(blocking=False) - finally: - os._exit(0 if ok else 1) - deadline = time.monotonic() + TIMEOUT - while (status := os.waitpid(pid, os.WNOHANG))[0] == 0: - if time.monotonic() > deadline: - os.kill(pid, signal.SIGKILL) - pytest.fail("forked child hung") - time.sleep(0.01) - assert os.waitstatus_to_exitcode(status[1]) == 0 - finally: - release.set() - registry_holder.join() - key_holder.join() + assert kvl._table == {} + assert kvl._leases == {} + assert kvl._managers == {} + assert not kvl._renewer_running + assert kvl._registry is not old_registry + assert kvl._registry.acquire(blocking=False) + kvl._registry.release() def test_global_lock_shards_agree_across_processes_and_spread_keys(): @@ -448,18 +470,7 @@ def test_global_lock_shards_agree_across_processes_and_spread_keys(): assert out.stdout.strip() == str(shards) -@pytest.fixture -def simple_storage_tq(): - if not ray.is_initialized(): - ray.init(namespace="TestKVLock") - backend = {"storage_backend": "SimpleStorage", "SimpleStorage": {"total_storage_size": 200}} - tq.init(OmegaConf.create({"controller": {"polling_mode": True}, "backend": backend})) - yield - tq.close() - ray.shutdown() - - -def test_kv_local_lock_serializes_simple_storage_read_modify_write(simple_storage_tq): +def test_kv_local_lock_serializes_simple_storage_read_modify_write(): tq.kv_put("counter", "lock_e2e", fields={"v": torch.tensor([0])}) # also creates the client on this thread def increment(key, partition_id): diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py index d67afa85..ad8ae7d1 100644 --- a/transfer_queue/__init__.py +++ b/transfer_queue/__init__.py @@ -78,7 +78,7 @@ "KVBatchMeta", ] + [ - # KV key locks: process-local and cluster-wide + # KV key locks: process-family-local and cluster-wide "kv_local_lock", "async_kv_local_lock", "kv_local_locked", diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 29fc67ef..ea129f90 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -45,7 +45,7 @@ from transfer_queue.client import TransferQueueClient from transfer_queue.controller import TransferQueueController -from transfer_queue.kv_lock import _detach_lock_managers +from transfer_queue.kv_lock import _detach_lock_managers, _process_lock_manager_name from transfer_queue.lock_manager import TransferQueueLockManager from transfer_queue.metadata import KVBatchMeta from transfer_queue.sampler import * # noqa: F401 @@ -212,13 +212,17 @@ def init(conf: DictConfig | None = None) -> DictConfig | None: ).remote(sampler=sampler, polling_mode=final_conf.controller.polling_mode) _TQ_IS_OWNER = True logger.info("TransferQueueController has been created.") - # Created here rather than on first lock: a non-detached actor dies with its creator. - _TQ_LOCK_MANAGERS = [ + # Created here rather than on first lock: non-detached actors die with their creator. + global_lock_managers = [ TransferQueueLockManager.options( # type: ignore[attr-defined] name=f"TransferQueueLockManager_{i}", namespace="transfer_queue", get_if_exists=True ).remote(num_lock_shards) for i in range(num_lock_shards) ] + process_lock_manager = TransferQueueLockManager.options( # type: ignore[attr-defined] + name=_process_lock_manager_name(), namespace="transfer_queue", get_if_exists=True + ).remote(1) + _TQ_LOCK_MANAGERS = [*global_lock_managers, process_lock_manager] except ValueError: logger.info("Some other rank has initialized TransferQueueController. Try to connect to existing controller.") _init_from_existing() diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 572fcfc9..19498fc0 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -16,11 +16,11 @@ """Advisory locks on TransferQueue KV keys. ``kv_local_lock`` and ``async_kv_local_lock`` hold an exclusive mutex per -``(partition_id, key)``, shared by every thread and asyncio task of this process. They -are advisory: KV calls never check them. They do not exclude other processes or Ray -actors. Waiters on a key are served in FIFO order. Nesting is rejected; pass all keys -to a single call. ``kv_local_locked`` and ``async_kv_local_locked`` run one KV call -under the lock without repeating its keys and partition. +``(partition_id, key)`` across one process family: the process that initialized +TransferQueue, its child processes, and their threads and tasks. A process-local +FIFO queue avoids unnecessary actor calls within one process; an inherited family +identifier routes every family member to one Ray actor. Unrelated processes on the +same node use different actors. ``kv_global_lock`` and ``async_kv_global_lock`` do the same across the Ray cluster through ``TransferQueueLockManager`` actors that split keys by hash, under a lease that renews @@ -38,6 +38,7 @@ from collections import deque from contextlib import asynccontextmanager, contextmanager from contextvars import ContextVar +from typing import Any import ray from ray.exceptions import RayActorError @@ -48,9 +49,12 @@ # while waiting. _state_lock = threading.Lock() _table: dict[tuple[str, str], deque["_Waiter"]] = {} -# Whether this thread or task holds (or is acquiring) a local lock. A future global lock +# Whether this thread or task holds (or is acquiring) a family-local lock. A global lock # reads it to enforce the "global, then local" order. _holding: ContextVar[bool] = ContextVar("_holding_kv_lock", default=False) +_PROCESS_LOCAL_SHARD = -1 +_PROCESS_LOCAL_LEASE_S = 30.0 +_PROCESS_FAMILY_ENV = "TQ_KV_LOCK_PROCESS_FAMILY" class _Waiter: @@ -159,38 +163,52 @@ def _reject_running_loop(kind: str) -> None: @contextmanager def kv_local_lock(keys: str | list[str], partition_id: str, timeout: float | None = None): - """Hold an exclusive process-local lock on ``keys`` in ``partition_id``. + """Hold an exclusive process-family lock on ``keys`` in ``partition_id``. - Blocks the calling thread, so use ``async_kv_local_lock`` inside a running event loop. - Raises ``TimeoutError`` if not all keys are acquired within ``timeout`` seconds - (``None`` waits forever) and ``RuntimeError`` if a local lock is already held. + The lock is shared by this process, its child processes, and their threads + and tasks. Blocks the calling thread, so use ``async_kv_local_lock`` inside + a running event loop. Raises ``TimeoutError`` if not all keys are acquired + within ``timeout`` seconds (``None`` waits forever), ``LockLostError`` if + its renewable actor lease is lost, and ``RuntimeError`` when nested. """ _reject_running_loop("kv_local_lock") names, acquired = _lock_names(keys, partition_id), [] deadline = None if timeout is None else time.monotonic() + timeout + lease = None _holding.set(True) try: for name in names: _acquire(name, deadline) acquired.append(name) + lease = _new_process_local_lease(names) + _acquire_lease(lease, deadline) yield + lease.check() finally: + if lease is not None: + _release_lease(lease) _release_all(acquired) _holding.set(False) @asynccontextmanager async def async_kv_local_lock(keys: str | list[str], partition_id: str, timeout: float | None = None): - """Async version of ``kv_local_lock``. Tasks created inside the block cannot take a local lock.""" + """Async version of ``kv_local_lock``. Child tasks cannot take another local lock.""" names, acquired = _lock_names(keys, partition_id), [] deadline = None if timeout is None else time.monotonic() + timeout + lease = None _holding.set(True) try: for name in names: await _acquire_async(name, deadline) acquired.append(name) + lease = _new_process_local_lease(names) + await _acquire_lease_async(lease, deadline) yield + lease.check() finally: + if lease is not None: + _release_lease(lease) _release_all(acquired) _holding.set(False) @@ -226,20 +244,27 @@ async def async_kv_local_locked( _leases: dict[str, "GlobalLease"] = {} _renewer_running = False _num_shards: int | None = None # controller.num_lock_shards, asked of shard 0 on first use -_managers: dict = {} # shard -> TransferQueueLockManager handle, looked up on first use -_NOT_INITIALIZED = "TransferQueueLockManager not found; call tq.init() first" +_managers: dict[int, Any] = {} # shard -> lock manager handle, looked up on first use +_NOT_INITIALIZED = "TransferQueue lock manager not found; call tq.init() first" _MANAGER_GONE = "The TransferQueueLockManager actor is gone; did the process that ran tq.init() call tq.close()?" class LockLostError(RuntimeError): - """A kv_global_lock lease ran out or failed to renew, so another holder may own the keys.""" + """A lock lease ran out or failed to renew, so another holder may own the keys.""" class GlobalLease: """Yielded by ``kv_global_lock``. ``fence`` maps each key to a fencing token that grows with every grant.""" - def __init__(self, names: list[tuple[str, str]], shards: dict[int, list[tuple[str, str]]], lease_s: float): + def __init__( + self, + names: list[tuple[str, str]], + shards: dict[int, list[tuple[str, str]]], + lease_s: float, + kind: str = "kv_global_lock", + ): self.names, self.shards, self.lease_s = names, shards, lease_s + self.kind = kind self.token = f"{os.getpid()}:{uuid.uuid4().hex}" self.fence: dict[str, int] = {} self.deadline = float("inf") # local monotonic time the lease is known to last until, once granted @@ -249,7 +274,21 @@ def __init__(self, names: list[tuple[str, str]], shards: dict[int, list[tuple[st def check(self) -> None: """Raise ``LockLostError`` if the lease was lost; call it right before writes in long sections.""" if self.lost or time.monotonic() >= self.deadline: - raise LockLostError(f"kv_global_lock lease on {self.names} was lost") + raise LockLostError(f"{self.kind} lease on {self.names} was lost") + + +def _process_family_id() -> str: + """Return a process-family id inherited by child processes through the environment.""" + family_id = os.environ.get(_PROCESS_FAMILY_ENV) + if family_id is None: + family_id = f"{os.getpid()}_{uuid.uuid4().hex}" + os.environ[_PROCESS_FAMILY_ENV] = family_id + return family_id + + +def _process_lock_manager_name() -> str: + """Return the Ray actor name shared by one process and its descendants.""" + return f"TransferQueueProcessLockManager_{_process_family_id()}" def _shard_count() -> int: @@ -270,8 +309,11 @@ def _shard(name: tuple[str, str], num_shards: int) -> int: def _lock_manager(shard: int): if shard not in _managers: + if not ray.is_initialized(): + raise RuntimeError(_NOT_INITIALIZED) + name = _process_lock_manager_name() if shard == _PROCESS_LOCAL_SHARD else f"TransferQueueLockManager_{shard}" try: - _managers[shard] = ray.get_actor(f"TransferQueueLockManager_{shard}", namespace="transfer_queue") + _managers[shard] = ray.get_actor(name, namespace="transfer_queue") except ValueError: raise RuntimeError(_NOT_INITIALIZED) from None return _managers[shard] @@ -283,12 +325,18 @@ def _new_lease(keys: str | list[str], partition_id: str, lease_s: float) -> Glob raise RuntimeError("Take kv_global_lock before kv_local_lock, not while holding one") if lease_s <= 0: raise ValueError("lease_s must be positive") - num_shards, shards = _shard_count(), {} + num_shards = _shard_count() + shards: dict[int, list[tuple[str, str]]] = {} for name in names: shards.setdefault(_shard(name, num_shards), []).append(name) return GlobalLease(names, dict(sorted(shards.items())), lease_s) +def _new_process_local_lease(names: list[tuple[str, str]]) -> GlobalLease: + """Create the hidden renewable lease shared by one process family.""" + return GlobalLease(names, {_PROCESS_LOCAL_SHARD: names}, _PROCESS_LOCAL_LEASE_S, "kv_local_lock") + + def _send_acquire(lease: GlobalLease, shard: int, deadline: float | None): manager = _lock_manager(shard) holder = {"node_ip": ray.util.get_node_ip_address(), "pid": os.getpid(), "thread": threading.current_thread().name} @@ -299,7 +347,7 @@ def _send_acquire(lease: GlobalLease, shard: int, deadline: float | None): def _take_grant(lease: GlobalLease, shard: int, reply, sent: float) -> None: if reply is None: lease.pending.remove(shard) - raise TimeoutError(f"Timed out waiting for kv_global_lock on {lease.names}") + raise TimeoutError(f"Timed out waiting for {lease.kind} on {lease.names}") fences, waited = reply lease.fence.update(fences) # The grant happened at least `waited` after `sent`, so this never outlives the shard's lease. @@ -313,7 +361,7 @@ def _start_lease(lease: GlobalLease) -> None: _registry.notify() # a shorter lease_s shortens the renewal period if not _renewer_running: _renewer_running = True - threading.Thread(target=_renew_loop, name="kv_global_lock_renewer", daemon=True).start() + threading.Thread(target=_renew_loop, name="kv_lock_renewer", daemon=True).start() def _release_lease(lease: GlobalLease) -> None: @@ -351,7 +399,8 @@ def _renew_loop() -> None: # Count from the send time: the server extends the lease from when the call arrives, # which is later, so the local deadline never outlives the server's. last = time.monotonic() - calls, renewed = {}, set() + calls: dict[int, Any] = {} + renewed: set[tuple[int, str]] = set() for shard, tokens in by_shard.items(): try: calls[shard] = _managers[shard].renew_many.remote(tokens) @@ -371,6 +420,30 @@ def _renew_loop() -> None: lease.lost = True +def _acquire_lease(lease: GlobalLease, deadline: float | None) -> None: + """Acquire every manager route for a sync lock and start lease renewal.""" + for shard in lease.shards: + sent = time.monotonic() + try: + reply = ray.get(_send_acquire(lease, shard, deadline)) + except RayActorError as e: + raise RuntimeError(_MANAGER_GONE) from e + _take_grant(lease, shard, reply, sent) + _start_lease(lease) + + +async def _acquire_lease_async(lease: GlobalLease, deadline: float | None) -> None: + """Acquire every manager route for an async lock and start lease renewal.""" + for shard in lease.shards: + sent = time.monotonic() + try: + reply = await _send_acquire(lease, shard, deadline) + except RayActorError as e: + raise RuntimeError(_MANAGER_GONE) from e + _take_grant(lease, shard, reply, sent) + _start_lease(lease) + + @contextmanager def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | None = None, lease_s: float = 30): """Hold an exclusive cluster-wide lock on ``keys`` in ``partition_id``; yield a ``GlobalLease``. @@ -392,14 +465,7 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No try: # Every caller visits shards in ascending order, so no two callers can each hold a # shard the other awaits: multi-shard requests cannot deadlock. - for shard in lease.shards: - sent = time.monotonic() - try: - reply = ray.get(_send_acquire(lease, shard, deadline)) - except RayActorError as e: - raise RuntimeError(_MANAGER_GONE) from e - _take_grant(lease, shard, reply, sent) - _start_lease(lease) + _acquire_lease(lease, deadline) yield lease lease.check() finally: @@ -416,14 +482,7 @@ async def async_kv_global_lock( deadline = None if timeout is None else time.monotonic() + timeout _holding_global.set(True) try: - for shard in lease.shards: - sent = time.monotonic() - try: - reply = await _send_acquire(lease, shard, deadline) - except RayActorError as e: - raise RuntimeError(_MANAGER_GONE) from e - _take_grant(lease, shard, reply, sent) - _start_lease(lease) + await _acquire_lease_async(lease, deadline) yield lease lease.check() finally: diff --git a/transfer_queue/lock_manager.py b/transfer_queue/lock_manager.py index b8a7eeee..c9f6274e 100644 --- a/transfer_queue/lock_manager.py +++ b/transfer_queue/lock_manager.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Server side of ``kv_global_lock``: async Ray actors, each holding leased locks on the keys that hash to it.""" +"""Async Ray actor that holds leased global-shard or process-family KV locks.""" import asyncio import heapq @@ -29,7 +29,7 @@ # default of 1000 to keep renew and release from queueing behind waiters. @ray.remote(num_cpus=0, max_concurrency=10_000) class TransferQueueLockManager: - """Exclusive leased locks on ``(partition_id, key)``. Runs on one event loop, so no thread locks.""" + """Exclusive leased locks on ``(partition_id, key)`` managed on one event loop.""" def __init__(self, num_shards: int): # Kept so that a process that never ran tq.init() can learn how many shards to hash over. @@ -108,6 +108,7 @@ async def acquire(self, names, token, timeout, lease_s, holder_info): del self._wakeups[name] def num_shards(self) -> int: + """Return the global shard count recorded by this manager.""" return self._num_shards def renew_many(self, tokens: list[str]) -> dict[str, bool]: @@ -116,8 +117,10 @@ def renew_many(self, tokens: list[str]) -> dict[str, bool]: alive = {} for token in tokens: lease = self._leases.get(token) - alive[token] = lease is not None and lease["expires_at"] > now - if alive[token]: + is_alive = lease is not None and lease["expires_at"] > now + alive[token] = is_alive + if is_alive: + assert lease is not None lease["expires_at"] = now + lease["lease_s"] heapq.heappush(self._expiries, (lease["expires_at"], token)) return alive @@ -132,6 +135,7 @@ async def release(self, token: str) -> None: self._waiting[token][1].set() def list_locks(self, partition_id: str | None = None) -> dict: + """Return current holders and the number of waiters, optionally for one partition.""" now = time.monotonic() holders = [] for (pid, key), token in self._holder.items(): From 0dcbcb02b2dc974ed9dc79a09f307c36bf25e31b Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 9 Oct 2026 16:52:04 +0800 Subject: [PATCH 07/42] [refactor] Keep kv_local_lock process-local The process-family scope made every local lock an actor round trip and left processes that attach to an existing TransferQueue, including Ray workers, without a lock actor. Cross-process exclusion is what kv_global_lock is for, so kv_local_lock goes back to a pure in-process mutex with no Ray calls. Signed-off-by: OutstanderWang --- tests/test_interface_close.py | 2 +- tests/test_kv_lock.py | 107 ++++++++++++--------------- transfer_queue/__init__.py | 2 +- transfer_queue/interface.py | 10 +-- transfer_queue/kv_lock.py | 129 +++++++++------------------------ transfer_queue/lock_manager.py | 10 +-- 6 files changed, 93 insertions(+), 167 deletions(-) diff --git a/tests/test_interface_close.py b/tests/test_interface_close.py index bff0470d..af3c9453 100644 --- a/tests/test_interface_close.py +++ b/tests/test_interface_close.py @@ -53,7 +53,7 @@ def test_owner_close_kills_controller(fake): iface.init(OmegaConf.create({})) client, lock_managers = iface._TQ_CLIENT, iface._TQ_LOCK_MANAGERS - assert iface._TQ_IS_OWNER and len(lock_managers) == 9 + assert iface._TQ_IS_OWNER and len(lock_managers) == 8 iface.close() client.close.assert_called_once() diff --git a/tests/test_kv_lock.py b/tests/test_kv_lock.py index 08e9c681..958d1077 100644 --- a/tests/test_kv_lock.py +++ b/tests/test_kv_lock.py @@ -17,10 +17,12 @@ import gc import importlib import os +import signal import subprocess import sys import threading import time +from types import SimpleNamespace import pytest import ray @@ -45,18 +47,6 @@ def table_is_empty_afterwards(): assert kvl._table == {} -@pytest.fixture(scope="module", autouse=True) -def node_lock_session(): - """Create the per-node lock actor used by every local-lock test in this module.""" - if not ray.is_initialized(): - ray.init(namespace="TestKVLock") - backend = {"storage_backend": "SimpleStorage", "SimpleStorage": {"total_storage_size": 200}} - tq.init(OmegaConf.create({"controller": {"polling_mode": True}, "backend": backend})) - yield - tq.close() - ray.shutdown() - - def n_waiters(key): with kvl._state_lock: return len(kvl._table.get((P, key), ())) @@ -230,7 +220,7 @@ def wait(self, timeout=None): holder.join(TIMEOUT) # the holder hands the key to this waiter... return False # ...but the wait still reports a timeout - monkeypatch.setattr(kvl.threading, "Event", GrantedThenTimedOut) + monkeypatch.setattr(kvl, "threading", SimpleNamespace(Event=GrantedThenTimedOut)) with kv_local_lock("k", P, timeout=1): assert (P, "k") in kvl._table and n_waiters("k") == 0 @@ -399,57 +389,45 @@ async def async_fn(keys, partition_id): asyncio.run(async_kv_local_locked(lambda keys, partition_id: None, "a", P)) -def test_child_process_uses_the_parents_local_lock_actor(): - script = """ -import ray -import transfer_queue as tq - -ray.init(address="auto", logging_level="ERROR") -result = "acquired" -try: - with tq.kv_local_lock("family", "p", timeout=0.3): - pass -except TimeoutError: - result = "timeout" -finally: - ray.shutdown() -print(result) -""" - - def child_result(): - completed = subprocess.run( - [sys.executable, "-c", script], - env=os.environ.copy(), - capture_output=True, - text=True, - timeout=E2E_TIMEOUT, - check=True, - ) - return completed.stdout.strip() - - assert kvl._PROCESS_FAMILY_ENV in os.environ - with kv_local_lock("family", P, timeout=TIMEOUT): - assert child_result() == "timeout" - assert child_result() == "acquired" - - -def test_fork_reset_drops_the_parents_locks_leases_and_managers(monkeypatch): +@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") +@pytest.mark.filterwarnings("ignore:.*fork.*:DeprecationWarning") +def test_forked_child_drops_the_parents_locks_leases_and_managers(monkeypatch): monkeypatch.setattr(kvl, "_registry", threading.Condition()) monkeypatch.setattr(kvl, "_leases", {"parent-token": object()}) monkeypatch.setattr(kvl, "_renewer_running", True) monkeypatch.setattr(kvl, "_managers", {0: object()}) - monkeypatch.setattr(kvl, "_table", {(P, "a"): object()}) + release, registry_held = threading.Event(), threading.Event() - old_registry = kvl._registry - kvl._reset_after_fork() + def hold_registry(): + with kvl._registry: + registry_held.set() + release.wait(TIMEOUT) - assert kvl._table == {} - assert kvl._leases == {} - assert kvl._managers == {} - assert not kvl._renewer_running - assert kvl._registry is not old_registry - assert kvl._registry.acquire(blocking=False) - kvl._registry.release() + registry_holder = threading.Thread(target=hold_registry) + registry_holder.start() + assert registry_held.wait(TIMEOUT) + key_holder = start_holder("a", release) + try: + pid = os.fork() + if pid == 0: + ok = False + try: + with kv_local_lock("a", P, timeout=1): + ok = kvl._leases == {} and kvl._managers == {} and not kvl._renewer_running + ok = ok and kvl._registry.acquire(blocking=False) + finally: + os._exit(0 if ok else 1) + deadline = time.monotonic() + TIMEOUT + while (status := os.waitpid(pid, os.WNOHANG))[0] == 0: + if time.monotonic() > deadline: + os.kill(pid, signal.SIGKILL) + pytest.fail("forked child hung") + time.sleep(0.01) + assert os.waitstatus_to_exitcode(status[1]) == 0 + finally: + release.set() + registry_holder.join() + key_holder.join() def test_global_lock_shards_agree_across_processes_and_spread_keys(): @@ -470,7 +448,18 @@ def test_global_lock_shards_agree_across_processes_and_spread_keys(): assert out.stdout.strip() == str(shards) -def test_kv_local_lock_serializes_simple_storage_read_modify_write(): +@pytest.fixture +def simple_storage_tq(): + if not ray.is_initialized(): + ray.init(namespace="TestKVLock") + backend = {"storage_backend": "SimpleStorage", "SimpleStorage": {"total_storage_size": 200}} + tq.init(OmegaConf.create({"controller": {"polling_mode": True}, "backend": backend})) + yield + tq.close() + ray.shutdown() + + +def test_kv_local_lock_serializes_simple_storage_read_modify_write(simple_storage_tq): tq.kv_put("counter", "lock_e2e", fields={"v": torch.tensor([0])}) # also creates the client on this thread def increment(key, partition_id): diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py index ad8ae7d1..d67afa85 100644 --- a/transfer_queue/__init__.py +++ b/transfer_queue/__init__.py @@ -78,7 +78,7 @@ "KVBatchMeta", ] + [ - # KV key locks: process-family-local and cluster-wide + # KV key locks: process-local and cluster-wide "kv_local_lock", "async_kv_local_lock", "kv_local_locked", diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index ea129f90..29fc67ef 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -45,7 +45,7 @@ from transfer_queue.client import TransferQueueClient from transfer_queue.controller import TransferQueueController -from transfer_queue.kv_lock import _detach_lock_managers, _process_lock_manager_name +from transfer_queue.kv_lock import _detach_lock_managers from transfer_queue.lock_manager import TransferQueueLockManager from transfer_queue.metadata import KVBatchMeta from transfer_queue.sampler import * # noqa: F401 @@ -212,17 +212,13 @@ def init(conf: DictConfig | None = None) -> DictConfig | None: ).remote(sampler=sampler, polling_mode=final_conf.controller.polling_mode) _TQ_IS_OWNER = True logger.info("TransferQueueController has been created.") - # Created here rather than on first lock: non-detached actors die with their creator. - global_lock_managers = [ + # Created here rather than on first lock: a non-detached actor dies with its creator. + _TQ_LOCK_MANAGERS = [ TransferQueueLockManager.options( # type: ignore[attr-defined] name=f"TransferQueueLockManager_{i}", namespace="transfer_queue", get_if_exists=True ).remote(num_lock_shards) for i in range(num_lock_shards) ] - process_lock_manager = TransferQueueLockManager.options( # type: ignore[attr-defined] - name=_process_lock_manager_name(), namespace="transfer_queue", get_if_exists=True - ).remote(1) - _TQ_LOCK_MANAGERS = [*global_lock_managers, process_lock_manager] except ValueError: logger.info("Some other rank has initialized TransferQueueController. Try to connect to existing controller.") _init_from_existing() diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 19498fc0..2844e29b 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -16,11 +16,11 @@ """Advisory locks on TransferQueue KV keys. ``kv_local_lock`` and ``async_kv_local_lock`` hold an exclusive mutex per -``(partition_id, key)`` across one process family: the process that initialized -TransferQueue, its child processes, and their threads and tasks. A process-local -FIFO queue avoids unnecessary actor calls within one process; an inherited family -identifier routes every family member to one Ray actor. Unrelated processes on the -same node use different actors. +``(partition_id, key)``, shared by every thread and asyncio task of this process. They +are advisory: KV calls never check them. They do not exclude other processes or Ray +actors. Waiters on a key are served in FIFO order. Nesting is rejected; pass all keys +to a single call. ``kv_local_locked`` and ``async_kv_local_locked`` run one KV call +under the lock without repeating its keys and partition. ``kv_global_lock`` and ``async_kv_global_lock`` do the same across the Ray cluster through ``TransferQueueLockManager`` actors that split keys by hash, under a lease that renews @@ -38,7 +38,6 @@ from collections import deque from contextlib import asynccontextmanager, contextmanager from contextvars import ContextVar -from typing import Any import ray from ray.exceptions import RayActorError @@ -49,12 +48,9 @@ # while waiting. _state_lock = threading.Lock() _table: dict[tuple[str, str], deque["_Waiter"]] = {} -# Whether this thread or task holds (or is acquiring) a family-local lock. A global lock +# Whether this thread or task holds (or is acquiring) a local lock. A future global lock # reads it to enforce the "global, then local" order. _holding: ContextVar[bool] = ContextVar("_holding_kv_lock", default=False) -_PROCESS_LOCAL_SHARD = -1 -_PROCESS_LOCAL_LEASE_S = 30.0 -_PROCESS_FAMILY_ENV = "TQ_KV_LOCK_PROCESS_FAMILY" class _Waiter: @@ -163,52 +159,38 @@ def _reject_running_loop(kind: str) -> None: @contextmanager def kv_local_lock(keys: str | list[str], partition_id: str, timeout: float | None = None): - """Hold an exclusive process-family lock on ``keys`` in ``partition_id``. + """Hold an exclusive process-local lock on ``keys`` in ``partition_id``. - The lock is shared by this process, its child processes, and their threads - and tasks. Blocks the calling thread, so use ``async_kv_local_lock`` inside - a running event loop. Raises ``TimeoutError`` if not all keys are acquired - within ``timeout`` seconds (``None`` waits forever), ``LockLostError`` if - its renewable actor lease is lost, and ``RuntimeError`` when nested. + Blocks the calling thread, so use ``async_kv_local_lock`` inside a running event loop. + Raises ``TimeoutError`` if not all keys are acquired within ``timeout`` seconds + (``None`` waits forever) and ``RuntimeError`` if a local lock is already held. """ _reject_running_loop("kv_local_lock") names, acquired = _lock_names(keys, partition_id), [] deadline = None if timeout is None else time.monotonic() + timeout - lease = None _holding.set(True) try: for name in names: _acquire(name, deadline) acquired.append(name) - lease = _new_process_local_lease(names) - _acquire_lease(lease, deadline) yield - lease.check() finally: - if lease is not None: - _release_lease(lease) _release_all(acquired) _holding.set(False) @asynccontextmanager async def async_kv_local_lock(keys: str | list[str], partition_id: str, timeout: float | None = None): - """Async version of ``kv_local_lock``. Child tasks cannot take another local lock.""" + """Async version of ``kv_local_lock``. Tasks created inside the block cannot take a local lock.""" names, acquired = _lock_names(keys, partition_id), [] deadline = None if timeout is None else time.monotonic() + timeout - lease = None _holding.set(True) try: for name in names: await _acquire_async(name, deadline) acquired.append(name) - lease = _new_process_local_lease(names) - await _acquire_lease_async(lease, deadline) yield - lease.check() finally: - if lease is not None: - _release_lease(lease) _release_all(acquired) _holding.set(False) @@ -244,27 +226,20 @@ async def async_kv_local_locked( _leases: dict[str, "GlobalLease"] = {} _renewer_running = False _num_shards: int | None = None # controller.num_lock_shards, asked of shard 0 on first use -_managers: dict[int, Any] = {} # shard -> lock manager handle, looked up on first use -_NOT_INITIALIZED = "TransferQueue lock manager not found; call tq.init() first" +_managers: dict = {} # shard -> TransferQueueLockManager handle, looked up on first use +_NOT_INITIALIZED = "TransferQueueLockManager not found; call tq.init() first" _MANAGER_GONE = "The TransferQueueLockManager actor is gone; did the process that ran tq.init() call tq.close()?" class LockLostError(RuntimeError): - """A lock lease ran out or failed to renew, so another holder may own the keys.""" + """A kv_global_lock lease ran out or failed to renew, so another holder may own the keys.""" class GlobalLease: """Yielded by ``kv_global_lock``. ``fence`` maps each key to a fencing token that grows with every grant.""" - def __init__( - self, - names: list[tuple[str, str]], - shards: dict[int, list[tuple[str, str]]], - lease_s: float, - kind: str = "kv_global_lock", - ): + def __init__(self, names: list[tuple[str, str]], shards: dict[int, list[tuple[str, str]]], lease_s: float): self.names, self.shards, self.lease_s = names, shards, lease_s - self.kind = kind self.token = f"{os.getpid()}:{uuid.uuid4().hex}" self.fence: dict[str, int] = {} self.deadline = float("inf") # local monotonic time the lease is known to last until, once granted @@ -274,21 +249,7 @@ def __init__( def check(self) -> None: """Raise ``LockLostError`` if the lease was lost; call it right before writes in long sections.""" if self.lost or time.monotonic() >= self.deadline: - raise LockLostError(f"{self.kind} lease on {self.names} was lost") - - -def _process_family_id() -> str: - """Return a process-family id inherited by child processes through the environment.""" - family_id = os.environ.get(_PROCESS_FAMILY_ENV) - if family_id is None: - family_id = f"{os.getpid()}_{uuid.uuid4().hex}" - os.environ[_PROCESS_FAMILY_ENV] = family_id - return family_id - - -def _process_lock_manager_name() -> str: - """Return the Ray actor name shared by one process and its descendants.""" - return f"TransferQueueProcessLockManager_{_process_family_id()}" + raise LockLostError(f"kv_global_lock lease on {self.names} was lost") def _shard_count() -> int: @@ -309,11 +270,8 @@ def _shard(name: tuple[str, str], num_shards: int) -> int: def _lock_manager(shard: int): if shard not in _managers: - if not ray.is_initialized(): - raise RuntimeError(_NOT_INITIALIZED) - name = _process_lock_manager_name() if shard == _PROCESS_LOCAL_SHARD else f"TransferQueueLockManager_{shard}" try: - _managers[shard] = ray.get_actor(name, namespace="transfer_queue") + _managers[shard] = ray.get_actor(f"TransferQueueLockManager_{shard}", namespace="transfer_queue") except ValueError: raise RuntimeError(_NOT_INITIALIZED) from None return _managers[shard] @@ -332,11 +290,6 @@ def _new_lease(keys: str | list[str], partition_id: str, lease_s: float) -> Glob return GlobalLease(names, dict(sorted(shards.items())), lease_s) -def _new_process_local_lease(names: list[tuple[str, str]]) -> GlobalLease: - """Create the hidden renewable lease shared by one process family.""" - return GlobalLease(names, {_PROCESS_LOCAL_SHARD: names}, _PROCESS_LOCAL_LEASE_S, "kv_local_lock") - - def _send_acquire(lease: GlobalLease, shard: int, deadline: float | None): manager = _lock_manager(shard) holder = {"node_ip": ray.util.get_node_ip_address(), "pid": os.getpid(), "thread": threading.current_thread().name} @@ -347,7 +300,7 @@ def _send_acquire(lease: GlobalLease, shard: int, deadline: float | None): def _take_grant(lease: GlobalLease, shard: int, reply, sent: float) -> None: if reply is None: lease.pending.remove(shard) - raise TimeoutError(f"Timed out waiting for {lease.kind} on {lease.names}") + raise TimeoutError(f"Timed out waiting for kv_global_lock on {lease.names}") fences, waited = reply lease.fence.update(fences) # The grant happened at least `waited` after `sent`, so this never outlives the shard's lease. @@ -361,7 +314,7 @@ def _start_lease(lease: GlobalLease) -> None: _registry.notify() # a shorter lease_s shortens the renewal period if not _renewer_running: _renewer_running = True - threading.Thread(target=_renew_loop, name="kv_lock_renewer", daemon=True).start() + threading.Thread(target=_renew_loop, name="kv_global_lock_renewer", daemon=True).start() def _release_lease(lease: GlobalLease) -> None: @@ -399,7 +352,7 @@ def _renew_loop() -> None: # Count from the send time: the server extends the lease from when the call arrives, # which is later, so the local deadline never outlives the server's. last = time.monotonic() - calls: dict[int, Any] = {} + calls = {} renewed: set[tuple[int, str]] = set() for shard, tokens in by_shard.items(): try: @@ -420,30 +373,6 @@ def _renew_loop() -> None: lease.lost = True -def _acquire_lease(lease: GlobalLease, deadline: float | None) -> None: - """Acquire every manager route for a sync lock and start lease renewal.""" - for shard in lease.shards: - sent = time.monotonic() - try: - reply = ray.get(_send_acquire(lease, shard, deadline)) - except RayActorError as e: - raise RuntimeError(_MANAGER_GONE) from e - _take_grant(lease, shard, reply, sent) - _start_lease(lease) - - -async def _acquire_lease_async(lease: GlobalLease, deadline: float | None) -> None: - """Acquire every manager route for an async lock and start lease renewal.""" - for shard in lease.shards: - sent = time.monotonic() - try: - reply = await _send_acquire(lease, shard, deadline) - except RayActorError as e: - raise RuntimeError(_MANAGER_GONE) from e - _take_grant(lease, shard, reply, sent) - _start_lease(lease) - - @contextmanager def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | None = None, lease_s: float = 30): """Hold an exclusive cluster-wide lock on ``keys`` in ``partition_id``; yield a ``GlobalLease``. @@ -465,7 +394,14 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No try: # Every caller visits shards in ascending order, so no two callers can each hold a # shard the other awaits: multi-shard requests cannot deadlock. - _acquire_lease(lease, deadline) + for shard in lease.shards: + sent = time.monotonic() + try: + reply = ray.get(_send_acquire(lease, shard, deadline)) + except RayActorError as e: + raise RuntimeError(_MANAGER_GONE) from e + _take_grant(lease, shard, reply, sent) + _start_lease(lease) yield lease lease.check() finally: @@ -482,7 +418,14 @@ async def async_kv_global_lock( deadline = None if timeout is None else time.monotonic() + timeout _holding_global.set(True) try: - await _acquire_lease_async(lease, deadline) + for shard in lease.shards: + sent = time.monotonic() + try: + reply = await _send_acquire(lease, shard, deadline) + except RayActorError as e: + raise RuntimeError(_MANAGER_GONE) from e + _take_grant(lease, shard, reply, sent) + _start_lease(lease) yield lease lease.check() finally: diff --git a/transfer_queue/lock_manager.py b/transfer_queue/lock_manager.py index c9f6274e..8bceb882 100644 --- a/transfer_queue/lock_manager.py +++ b/transfer_queue/lock_manager.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Async Ray actor that holds leased global-shard or process-family KV locks.""" +"""Server side of ``kv_global_lock``: async Ray actors, each holding leased locks on the keys that hash to it.""" import asyncio import heapq @@ -29,7 +29,7 @@ # default of 1000 to keep renew and release from queueing behind waiters. @ray.remote(num_cpus=0, max_concurrency=10_000) class TransferQueueLockManager: - """Exclusive leased locks on ``(partition_id, key)`` managed on one event loop.""" + """Exclusive leased locks on ``(partition_id, key)``. Runs on one event loop, so no thread locks.""" def __init__(self, num_shards: int): # Kept so that a process that never ran tq.init() can learn how many shards to hash over. @@ -117,10 +117,8 @@ def renew_many(self, tokens: list[str]) -> dict[str, bool]: alive = {} for token in tokens: lease = self._leases.get(token) - is_alive = lease is not None and lease["expires_at"] > now - alive[token] = is_alive - if is_alive: - assert lease is not None + alive[token] = lease is not None and lease["expires_at"] > now + if lease is not None and alive[token]: lease["expires_at"] = now + lease["lease_s"] heapq.heappush(self._expiries, (lease["expires_at"], token)) return alive From 191327a434ef36fd3406ed0cd8c951cde04e1ee3 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 9 Oct 2026 17:43:31 +0800 Subject: [PATCH 08/42] [feat] Grant kv_global_lock waiters in arrival order A release used to wake every waiter on the freed keys and let whichever ran first take them. A request arriving in that window could barge ahead of older waiters, and a multi-key request could wait forever while single-key requests kept one of its keys busy. Each lock actor now keeps a FIFO queue per key and grants a request only when all its keys are free and it heads each of their queues. Release, expiry and withdrawal hand the keys straight to the next eligible waiter and wake only it, and a timer expires leases on time now that waiters no longer poll. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 51 +++++++++++ transfer_queue/kv_lock.py | 6 +- transfer_queue/lock_manager.py | 128 ++++++++++++++++----------- 3 files changed, 133 insertions(+), 52 deletions(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index 6efeddff..27323ebe 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -25,6 +25,7 @@ from omegaconf import OmegaConf import transfer_queue as tq +from transfer_queue.lock_manager import TransferQueueLockManager kvl = importlib.import_module("transfer_queue.kv_lock") P = "global_lock_e2e" @@ -206,6 +207,56 @@ def test_release_before_acquire_withdraws_it(): assert holders() == set() +@pytest.fixture +def own_manager(): + """A lock actor of its own, so these tests control every request it sees.""" + lock_manager = TransferQueueLockManager.remote(1) + yield lock_manager + ray.kill(lock_manager) + + +def acquire(lock_manager, keys, token, timeout=None, lease_s=30): + return lock_manager.acquire.remote([(P, key) for key in keys], token, timeout, lease_s, {}) + + +def n_waiters(lock_manager): + return ray.get(lock_manager.list_locks.remote(P), timeout=TIMEOUT_S)["waiters"] + + +def test_waiters_are_granted_in_arrival_order(own_manager): + ray.get(acquire(own_manager, ["a"], "holder"), timeout=TIMEOUT_S) + both = acquire(own_manager, ["a", "b"], "both") + wait_until(lambda: n_waiters(own_manager) == 1) + # b is free, but granting it past the older [a, b] request would let single-key traffic starve it. + assert ray.get(acquire(own_manager, ["b"], "newcomer", timeout=0), timeout=TIMEOUT_S) is None + only_b = acquire(own_manager, ["b"], "only_b") + wait_until(lambda: n_waiters(own_manager) == 2) + ray.get(own_manager.release.remote("holder"), timeout=TIMEOUT_S) + assert ray.get(both, timeout=TIMEOUT_S) is not None + assert n_waiters(own_manager) == 1 + ray.get(own_manager.release.remote("both"), timeout=TIMEOUT_S) + assert ray.get(only_b, timeout=TIMEOUT_S) is not None + + +def test_a_waiter_that_gives_up_lets_the_next_one_through(own_manager): + ray.get(acquire(own_manager, ["a"], "holder"), timeout=TIMEOUT_S) + both = acquire(own_manager, ["a", "b"], "both", timeout=0.5) + wait_until(lambda: n_waiters(own_manager) == 1) + only_b = acquire(own_manager, ["b"], "only_b") + assert ray.get(both, timeout=TIMEOUT_S) is None + assert ray.get(only_b, timeout=TIMEOUT_S) is not None + + +def test_expired_leases_pass_the_key_down_the_queue(own_manager): + # Nobody releases: each grant has to come from the actor expiring the lease ahead of it. + ray.get(acquire(own_manager, ["a"], "dead", lease_s=0.5), timeout=TIMEOUT_S) + first = acquire(own_manager, ["a"], "first", lease_s=0.5) + wait_until(lambda: n_waiters(own_manager) == 1) + second = acquire(own_manager, ["a"], "second") + assert ray.get(second, timeout=TIMEOUT_S) is not None + assert ray.get(first, timeout=TIMEOUT_S) is not None + + def test_killed_holder_frees_the_key_after_its_lease(): worker = Worker.remote() ray.get(worker.hold.remote("k", lease_s=1.5), timeout=TIMEOUT_S) diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 2844e29b..6b4e065e 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -379,8 +379,10 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No The locks live in the ``controller.num_lock_shards`` actors that ``tq.init()`` creates, each key in the one its hash picks, and are held under a ``lease_s`` lease that a - background thread renews. Keys on several actors are taken one actor at a time, each - all at once, so keys already granted stay held while a later actor's keys are awaited. + background thread renews. Each actor grants in arrival order: a request waits behind + every earlier one that shares a key with it. Keys on several actors are taken one actor + at a time, each all at once, so keys already granted stay held while a later actor's + keys are awaited. Raises ``TimeoutError`` if the keys are not all granted within ``timeout`` seconds (``None`` waits forever), and ``RuntimeError`` when nested, taken while holding a ``kv_local_lock``, or called with a running event loop. Exit always releases the lock, diff --git a/transfer_queue/lock_manager.py b/transfer_queue/lock_manager.py index 8bceb882..79ee0f21 100644 --- a/transfer_queue/lock_manager.py +++ b/transfer_queue/lock_manager.py @@ -18,6 +18,7 @@ import asyncio import heapq import time +from collections import deque import ray @@ -39,73 +40,96 @@ def __init__(self, num_shards: int): # Min-heap of (expires_at, token). Release and renewal leave stale entries behind, which # _expire skips, so neither has to search the heap. self._expiries: list[tuple[float, str]] = [] + # Fires _expire at the earliest expiry, so a dead holder's keys pass on even while every + # waiter sleeps until it is granted. + self._timer: asyncio.TimerHandle | None = None + self._timer_at = float("inf") self._fences: dict[tuple[str, str], int] = {} # never reset, so a key's fence only grows - self._waiting: dict[str, tuple[list[tuple[str, str]], asyncio.Event]] = {} # token -> names, wakeup - # name -> wakeups of its waiters, as an insertion-ordered dict so they wake oldest first - self._wakeups: dict[tuple[str, str], dict[asyncio.Event, None]] = {} + # Waiting tokens per key in arrival order. A request is granted only when all its keys are + # free and it heads every one of their queues: no newcomer can barge, and single-key + # requests cannot starve an older multi-key one. All queues share one arrival order, so the + # oldest waiter heads all of its queues and the wait can never form a cycle. + self._queues: dict[tuple[str, str], deque[str]] = {} + self._waiting: dict[str, dict] = {} # token -> names, lease_s, info, done (set on grant or withdrawal) self._withdrawn: dict[str, float] = {} # token -> when its tombstone expires + def _grant(self, token: str, names: list[tuple[str, str]], lease_s: float, info: dict) -> None: + now = time.monotonic() + for name in names: + self._holder[name] = token + self._fences[name] = self._fences.get(name, 0) + 1 + self._leases[token] = dict(names=names, lease_s=lease_s, expires_at=now + lease_s, granted_at=now, info=info) + heapq.heappush(self._expiries, (now + lease_s, token)) + if now + lease_s < self._timer_at: + if self._timer is not None: + self._timer.cancel() + self._timer_at = now + lease_s + self._timer = asyncio.get_running_loop().call_later(lease_s, self._expire) + + def _leave_queues(self, token: str) -> dict: + waiter = self._waiting.pop(token) + for name in waiter["names"]: + self._queues[name].remove(token) + if not self._queues[name]: + del self._queues[name] + return waiter + + def _hand_off(self, names: list[tuple[str, str]]) -> None: + """Grant each free name to the first waiter in its queue, if that waiter now heads all of its queues.""" + for name in names: + if name in self._holder or name not in self._queues: + continue + token = self._queues[name][0] + wanted = self._waiting[token]["names"] + if all(n not in self._holder and self._queues[n][0] == token for n in wanted): + waiter = self._leave_queues(token) + self._grant(token, wanted, waiter["lease_s"], waiter["info"]) + waiter["done"].set() + def _free(self, token: str) -> bool: lease = self._leases.pop(token, None) if lease is None: return False for name in lease["names"]: del self._holder[name] - for wakeup in self._wakeups.get(name, ()): - wakeup.set() + self._hand_off(lease["names"]) return True - def _expire(self, now: float) -> None: + def _expire(self) -> None: + now = time.monotonic() while self._expiries and self._expiries[0][0] <= now: expires_at, token = heapq.heappop(self._expiries) lease = self._leases.get(token) if lease is not None and lease["expires_at"] == expires_at: self._free(token) + # Re-arm for the next entry; a stale one only makes the timer fire early. + self._timer, self._timer_at = None, float("inf") + if self._expiries: + self._timer_at = self._expiries[0][0] + self._timer = asyncio.get_running_loop().call_later(max(0.0, self._timer_at - now), self._expire) async def acquire(self, names, token, timeout, lease_s, holder_info): """Grant all ``names`` at once; return ``({key: fence}, waited_s)``, or None on timeout or withdrawal.""" start = time.monotonic() - deadline = None if timeout is None else start + timeout - # Freeing a name sets the wakeups of its waiters only, so a release costs O(its waiters) - # instead of waking every waiter on the actor. - wakeup = asyncio.Event() - self._waiting[token] = (names, wakeup) - for name in names: - self._wakeups.setdefault(name, {})[wakeup] = None - try: - while True: - now = time.monotonic() - self._expire(now) - if self._withdrawn.pop(token, None) is not None: - return None - # All-or-nothing: a waiter never holds part of its keys, so overlapping - # requests cannot deadlock whatever order their keys come in. - busy = [self._leases[self._holder[name]]["expires_at"] for name in names if name in self._holder] - if not busy: - for name in names: - self._holder[name] = token - self._fences[name] = self._fences.get(name, 0) + 1 - self._leases[token] = dict( - names=names, lease_s=lease_s, expires_at=now + lease_s, granted_at=now, info=holder_info - ) - heapq.heappush(self._expiries, (now + lease_s, token)) - return {key: self._fences[(pid, key)] for pid, key in names}, now - start - if deadline is not None and now >= deadline: - return None - # Nothing runs between this clear and the wait, so no wakeup can be missed. Waking - # at the earliest blocking expiry lets _expire free it without a background task. - wakeup.clear() - wake_at = min(busy) if deadline is None else min(*busy, deadline) - try: - await asyncio.wait_for(wakeup.wait(), wake_at - now) - except asyncio.TimeoutError: - pass - finally: - del self._waiting[token] + if self._withdrawn.pop(token, None) is not None: + return None + if not any(name in self._holder or name in self._queues for name in names): + self._grant(token, names, lease_s, holder_info) + else: + done = asyncio.Event() + self._waiting[token] = dict(names=names, lease_s=lease_s, info=holder_info, done=done) for name in names: - del self._wakeups[name][wakeup] - if not self._wakeups[name]: - del self._wakeups[name] + self._queues.setdefault(name, deque()).append(token) + try: + await asyncio.wait_for(done.wait(), timeout) + except asyncio.TimeoutError: + pass + finally: + if token in self._waiting: # timed out or cancelled before a grant + self._hand_off(self._leave_queues(token)["names"]) + if token not in self._leases: # withdrawn, or already released by a caller that left + return None + return {key: self._fences[(pid, key)] for pid, key in names}, self._leases[token]["granted_at"] - start def num_shards(self) -> int: """Return the global shard count recorded by this manager.""" @@ -125,12 +149,16 @@ def renew_many(self, tokens: list[str]) -> dict[str, bool]: async def release(self, token: str) -> None: """Free ``token``'s locks, or withdraw its acquire if it is still waiting or has not arrived.""" - if not self._free(token): + if self._free(token): + return + if token in self._waiting: + waiter = self._leave_queues(token) + self._hand_off(waiter["names"]) + waiter["done"].set() # its acquire finds no lease and returns None + else: now = time.monotonic() self._withdrawn = {t: until for t, until in self._withdrawn.items() if until > now} self._withdrawn[token] = now + _WITHDRAWN_TTL_S - if token in self._waiting: - self._waiting[token][1].set() def list_locks(self, partition_id: str | None = None) -> dict: """Return current holders and the number of waiters, optionally for one partition.""" @@ -148,5 +176,5 @@ def list_locks(self, partition_id: str | None = None) -> dict: lease_remaining_s=lease["expires_at"] - now, ) ) - waiters = sum(partition_id in (None, names[0][0]) for names, _ in self._waiting.values()) + waiters = sum(partition_id in (None, waiter["names"][0][0]) for waiter in self._waiting.values()) return {"holders": holders, "waiters": waiters} From a8b52acf37878eba5a72bfdb66f3fd358df5d87c Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 9 Oct 2026 20:40:02 +0800 Subject: [PATCH 09/42] [refactor] Drop kv_global_lock fencing tokens Nothing on the KV write path checks them, so they gave no protection against a stalled holder, while the per-key counters on the lock actors grew with every key ever locked. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 19 ++++--------------- transfer_queue/kv_lock.py | 17 +++++++---------- transfer_queue/lock_manager.py | 6 ++---- 3 files changed, 13 insertions(+), 29 deletions(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index 27323ebe..db90be65 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -56,7 +56,7 @@ def cycle(self, keys, n): def hold(self, key, lease_s=30): self.held = tq.kv_global_lock(key, P, timeout=TIMEOUT_S, lease_s=lease_s) - return self.held.__enter__().fence[key] + self.held.__enter__() def release(self): self.held.__exit__(None, None, None) @@ -176,7 +176,7 @@ def test_timeout_leaves_no_holder_or_waiter(): def test_cancelled_waiter_is_withdrawn_and_never_granted(): worker = Worker.remote() - fence = ray.get(worker.hold.remote("c"), timeout=TIMEOUT_S) + ray.get(worker.hold.remote("c"), timeout=TIMEOUT_S) async def wait_for_lock(): async with tq.async_kv_global_lock("c", P): @@ -196,8 +196,8 @@ async def main(): wait_until(lambda: tq.kv_lock_list(P)["waiters"] == 0) ray.get(worker.release.remote(), timeout=TIMEOUT_S) # A ghost grant to the cancelled waiter would hold "c" for its whole 30 s lease. - with tq.kv_global_lock("c", P, timeout=5) as lease: - assert lease.fence["c"] == fence + 1 + with tq.kv_global_lock("c", P, timeout=5): + pass def test_release_before_acquire_withdraws_it(): @@ -286,17 +286,6 @@ def test_lease_renews_and_lost_lease_raises(): wait_until(lambda: holders() == set()) -def test_fences_grow_across_holders(): - with tq.kv_global_lock(["f", "g"], P) as first: - pass - worker = Worker.remote() - fence = ray.get(worker.hold.remote("f"), timeout=TIMEOUT_S) - ray.get(worker.release.remote(), timeout=TIMEOUT_S) - with tq.kv_global_lock("f", P) as last: - pass - assert first.fence["f"] < fence < last.fence["f"] and set(first.fence) == {"f", "g"} - - def test_nesting_and_ordering_rules(): with pytest.raises(ValueError): with tq.kv_global_lock([], P): diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 6b4e065e..f0968acb 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -236,12 +236,11 @@ class LockLostError(RuntimeError): class GlobalLease: - """Yielded by ``kv_global_lock``. ``fence`` maps each key to a fencing token that grows with every grant.""" + """Yielded by ``kv_global_lock``; ``check()`` raises ``LockLostError`` once the lease is lost.""" def __init__(self, names: list[tuple[str, str]], shards: dict[int, list[tuple[str, str]]], lease_s: float): self.names, self.shards, self.lease_s = names, shards, lease_s self.token = f"{os.getpid()}:{uuid.uuid4().hex}" - self.fence: dict[str, int] = {} self.deadline = float("inf") # local monotonic time the lease is known to last until, once granted self.lost = False self.pending: list[int] = [] # shards that may hold or await this token, so exit must release them @@ -297,12 +296,10 @@ def _send_acquire(lease: GlobalLease, shard: int, deadline: float | None): return manager.acquire.remote(lease.shards[shard], lease.token, _remaining(deadline), lease.lease_s, holder) -def _take_grant(lease: GlobalLease, shard: int, reply, sent: float) -> None: - if reply is None: +def _take_grant(lease: GlobalLease, shard: int, waited: float | None, sent: float) -> None: + if waited is None: lease.pending.remove(shard) raise TimeoutError(f"Timed out waiting for kv_global_lock on {lease.names}") - fences, waited = reply - lease.fence.update(fences) # The grant happened at least `waited` after `sent`, so this never outlives the shard's lease. lease.deadline = min(lease.deadline, sent + waited + lease.lease_s) @@ -399,10 +396,10 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No for shard in lease.shards: sent = time.monotonic() try: - reply = ray.get(_send_acquire(lease, shard, deadline)) + waited = ray.get(_send_acquire(lease, shard, deadline)) except RayActorError as e: raise RuntimeError(_MANAGER_GONE) from e - _take_grant(lease, shard, reply, sent) + _take_grant(lease, shard, waited, sent) _start_lease(lease) yield lease lease.check() @@ -423,10 +420,10 @@ async def async_kv_global_lock( for shard in lease.shards: sent = time.monotonic() try: - reply = await _send_acquire(lease, shard, deadline) + waited = await _send_acquire(lease, shard, deadline) except RayActorError as e: raise RuntimeError(_MANAGER_GONE) from e - _take_grant(lease, shard, reply, sent) + _take_grant(lease, shard, waited, sent) _start_lease(lease) yield lease lease.check() diff --git a/transfer_queue/lock_manager.py b/transfer_queue/lock_manager.py index 79ee0f21..1fc57792 100644 --- a/transfer_queue/lock_manager.py +++ b/transfer_queue/lock_manager.py @@ -44,7 +44,6 @@ def __init__(self, num_shards: int): # waiter sleeps until it is granted. self._timer: asyncio.TimerHandle | None = None self._timer_at = float("inf") - self._fences: dict[tuple[str, str], int] = {} # never reset, so a key's fence only grows # Waiting tokens per key in arrival order. A request is granted only when all its keys are # free and it heads every one of their queues: no newcomer can barge, and single-key # requests cannot starve an older multi-key one. All queues share one arrival order, so the @@ -57,7 +56,6 @@ def _grant(self, token: str, names: list[tuple[str, str]], lease_s: float, info: now = time.monotonic() for name in names: self._holder[name] = token - self._fences[name] = self._fences.get(name, 0) + 1 self._leases[token] = dict(names=names, lease_s=lease_s, expires_at=now + lease_s, granted_at=now, info=info) heapq.heappush(self._expiries, (now + lease_s, token)) if now + lease_s < self._timer_at: @@ -109,7 +107,7 @@ def _expire(self) -> None: self._timer = asyncio.get_running_loop().call_later(max(0.0, self._timer_at - now), self._expire) async def acquire(self, names, token, timeout, lease_s, holder_info): - """Grant all ``names`` at once; return ``({key: fence}, waited_s)``, or None on timeout or withdrawal.""" + """Grant all ``names`` at once; return the seconds waited, or None on timeout or withdrawal.""" start = time.monotonic() if self._withdrawn.pop(token, None) is not None: return None @@ -129,7 +127,7 @@ async def acquire(self, names, token, timeout, lease_s, holder_info): self._hand_off(self._leave_queues(token)["names"]) if token not in self._leases: # withdrawn, or already released by a caller that left return None - return {key: self._fences[(pid, key)] for pid, key in names}, self._leases[token]["granted_at"] - start + return self._leases[token]["granted_at"] - start def num_shards(self) -> int: """Return the global shard count recorded by this manager.""" From d8909fe652680e5a0701b5465e37f8758936f4ee Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 9 Oct 2026 20:42:46 +0800 Subject: [PATCH 10/42] [refactor] Drop the *_locked wrappers A with-block over kv_local_lock / kv_global_lock is one line longer and keeps the public lock API to the two context managers per scope. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 22 ------- tests/test_kv_lock.py | 99 ++-------------------------- transfer_queue/__init__.py | 8 --- transfer_queue/kv_lock.py | 56 +--------------- 4 files changed, 6 insertions(+), 179 deletions(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index db90be65..6b3106f4 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -314,28 +314,6 @@ async def main(): wait_until(lambda: holders() == set()) -def test_locked_wrappers_hold_the_lock_during_the_call_and_release_on_error(): - def fn(keys, partition_id, **kwargs): - return holders(), partition_id, kwargs - - async def async_fn(keys, partition_id, **kwargs): - return holders(), partition_id, kwargs - - def boom(keys, partition_id): - raise ValueError("boom") - - assert tq.kv_global_locked(fn, ["u", "v"], P, lock_timeout=5, lease_s=2, x=1) == ({"u", "v"}, P, {"x": 1}) - assert asyncio.run(tq.async_kv_global_locked(async_fn, "u", P, lock_timeout=5)) == ({"u"}, P, {}) - with pytest.raises(ValueError, match="boom"): - tq.kv_global_locked(boom, "u", P) - wait_until(lambda: holders() == set()) - - with pytest.raises(TypeError, match="async_kv_global_locked"): - tq.kv_global_locked(async_fn, "u", P) - with pytest.raises(TypeError, match="coroutine function"): - asyncio.run(tq.async_kv_global_locked(fn, "u", P)) - - def actors_named(prefix): return [a["name"] for a in ray.util.list_named_actors(all_namespaces=True) if a["name"].startswith(prefix)] diff --git a/tests/test_kv_lock.py b/tests/test_kv_lock.py index 958d1077..d25f3ae4 100644 --- a/tests/test_kv_lock.py +++ b/tests/test_kv_lock.py @@ -28,10 +28,9 @@ import ray import torch from omegaconf import OmegaConf -from tensordict import TensorDict import transfer_queue as tq -from transfer_queue import async_kv_local_lock, async_kv_local_locked, kv_local_lock, kv_local_locked +from transfer_queue import async_kv_local_lock, kv_local_lock # `transfer_queue.kv_lock` resolves to the re-exported function, so fetch the module itself. kvl = importlib.import_module("transfer_queue.kv_lock") @@ -310,85 +309,6 @@ async def child(): asyncio.run(main()) -def is_held_elsewhere(keys): - """Whether another thread's kv_local_lock on `keys` times out.""" - outcome = [] - - def run(): - try: - with kv_local_lock(keys, P, timeout=0): - outcome.append(False) - except TimeoutError: - outcome.append(True) - - thread = threading.Thread(target=run) - thread.start() - thread.join(TIMEOUT) - assert not thread.is_alive() - return outcome == [True] - - -def test_locked_wrapper_passes_arguments_through_and_holds_the_lock_during_the_call(): - def fn(keys, partition_id, **kwargs): - assert is_held_elsewhere(keys) - return keys, partition_id, kwargs - - assert kv_local_locked(fn, "a", P, lock_timeout=TIMEOUT, x=1) == ("a", P, {"x": 1}) - assert kv_local_locked(fn, ["b", "a"], P) == (["b", "a"], P, {}) - assert not is_held_elsewhere(["a", "b"]) - - def boom(keys, partition_id): - raise ValueError("boom") - - with pytest.raises(ValueError, match="boom"): - kv_local_locked(boom, "a", P) - assert not is_held_elsewhere("a") - - -def test_async_locked_wrapper_holds_the_lock_and_releases_on_error(): - async def fn(keys, partition_id, **kwargs): - assert is_held_elsewhere(keys) - return keys, partition_id, kwargs - - async def boom(keys, partition_id): - raise ValueError("boom") - - async def main(): - assert await async_kv_local_locked(fn, ["a", "b"], P, lock_timeout=TIMEOUT, x=1) == (["a", "b"], P, {"x": 1}) - with pytest.raises(ValueError, match="boom"): - await async_kv_local_locked(boom, "a", P) - - asyncio.run(main()) - assert not is_held_elsewhere(["a", "b"]) - - -def test_locked_wrapper_lock_timeout_skips_the_call(): - calls, release = [], threading.Event() - holder = start_holder("k", release) - - async def async_fn(keys, partition_id): - calls.append(keys) - - with pytest.raises(TimeoutError): - kv_local_locked(lambda keys, partition_id: calls.append(keys), ["j", "k"], P, lock_timeout=0.05) - with pytest.raises(TimeoutError): - asyncio.run(async_kv_local_locked(async_fn, "k", P, lock_timeout=0.05)) - assert calls == [] and n_waiters("k") == 0 - release.set() - holder.join(TIMEOUT) - assert not holder.is_alive() - - -def test_locked_wrappers_reject_the_wrong_kind_of_function(): - async def async_fn(keys, partition_id): - pass - - with pytest.raises(TypeError, match="async_kv_local_locked"): - kv_local_locked(async_fn, "a", P) - with pytest.raises(TypeError, match="coroutine function"): - asyncio.run(async_kv_local_locked(lambda keys, partition_id: None, "a", P)) - - @pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") @pytest.mark.filterwarnings("ignore:.*fork.*:DeprecationWarning") def test_forked_child_drops_the_parents_locks_leases_and_managers(monkeypatch): @@ -466,27 +386,16 @@ def increment(key, partition_id): value = tq.kv_batch_get(key, partition_id)["v"] tq.kv_put(key, partition_id, fields={"v": value[0] + 1}) - def context_worker(): + def worker(): for _ in range(50): with tq.kv_local_lock("counter", "lock_e2e", timeout=E2E_TIMEOUT): increment("counter", "lock_e2e") - def wrapper_worker(): - for _ in range(50): - tq.kv_local_locked(increment, "counter", "lock_e2e", lock_timeout=E2E_TIMEOUT) - - threads = [threading.Thread(target=worker) for worker in [context_worker, wrapper_worker] * 4] + threads = [threading.Thread(target=worker) for _ in range(8)] for thread in threads: thread.start() for thread in threads: thread.join(E2E_TIMEOUT) assert not thread.is_alive() assert tq.kv_batch_get("counter", "lock_e2e")["v"][0].item() == 400 - - fields = TensorDict({"v": torch.tensor([[1], [2]])}, batch_size=2) - tq.kv_local_locked(tq.kv_batch_put, ["a", "b"], "lock_e2e", fields=fields) - data = tq.kv_local_locked(tq.kv_batch_get, ["a", "b"], "lock_e2e", select_fields="v") - assert [row.tolist() for row in data["v"]] == [[1], [2]] # rows may come back as a nested tensor - data = asyncio.run(tq.async_kv_local_locked(tq.async_kv_batch_get, "b", "lock_e2e", lock_timeout=E2E_TIMEOUT)) - assert [row.tolist() for row in data["v"]] == [[2]] - tq.kv_clear(["counter", "a", "b"], "lock_e2e") + tq.kv_clear("counter", "lock_e2e") diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py index d67afa85..ea9faed2 100644 --- a/transfer_queue/__init__.py +++ b/transfer_queue/__init__.py @@ -40,13 +40,9 @@ from .kv_lock import ( LockLostError, async_kv_global_lock, - async_kv_global_locked, async_kv_local_lock, - async_kv_local_locked, kv_global_lock, - kv_global_locked, kv_local_lock, - kv_local_locked, kv_lock_list, ) from .metadata import BatchMeta, KVBatchMeta @@ -81,12 +77,8 @@ # KV key locks: process-local and cluster-wide "kv_local_lock", "async_kv_local_lock", - "kv_local_locked", - "async_kv_local_locked", "kv_global_lock", "async_kv_global_lock", - "kv_global_locked", - "async_kv_global_locked", "kv_lock_list", "LockLostError", ] diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index f0968acb..c3072d55 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -19,17 +19,14 @@ ``(partition_id, key)``, shared by every thread and asyncio task of this process. They are advisory: KV calls never check them. They do not exclude other processes or Ray actors. Waiters on a key are served in FIFO order. Nesting is rejected; pass all keys -to a single call. ``kv_local_locked`` and ``async_kv_local_locked`` run one KV call -under the lock without repeating its keys and partition. +to a single call. ``kv_global_lock`` and ``async_kv_global_lock`` do the same across the Ray cluster through ``TransferQueueLockManager`` actors that split keys by hash, under a lease that renews -automatically. Take a global lock before a local one, never inside it. ``kv_global_locked`` -and ``async_kv_global_locked`` mirror the local wrappers. +automatically. Take a global lock before a local one, never inside it. """ import asyncio -import inspect import os import threading import time @@ -195,31 +192,6 @@ async def async_kv_local_lock(keys: str | list[str], partition_id: str, timeout: _holding.set(False) -def kv_local_locked(fn, keys: str | list[str], partition_id: str, *, lock_timeout: float | None = None, **kwargs): - """Call ``fn(keys, partition_id, **kwargs)`` under ``kv_local_lock(keys, partition_id)``. - - ``fn`` is any KV call such as ``tq.kv_batch_put``, or a user helper with the same - leading ``(keys, partition_id)`` arguments, e.g. one doing get -> compute -> put:: - - meta = tq.kv_local_locked(tq.kv_batch_put, ["a", "b"], "train", fields=td, lock_timeout=5) - tq.kv_local_locked(increment, "counter", "train") - """ - if inspect.iscoroutinefunction(fn): - raise TypeError("kv_local_locked would release the lock before the coroutine runs; use async_kv_local_locked") - with kv_local_lock(keys, partition_id, timeout=lock_timeout): - return fn(keys, partition_id, **kwargs) - - -async def async_kv_local_locked( - fn, keys: str | list[str], partition_id: str, *, lock_timeout: float | None = None, **kwargs -): - """Async version of ``kv_local_locked``; ``fn`` must be a coroutine function such as ``tq.async_kv_batch_get``.""" - if not inspect.iscoroutinefunction(fn): - raise TypeError("async_kv_local_locked needs a coroutine function; a sync call would block the event loop") - async with async_kv_local_lock(keys, partition_id, timeout=lock_timeout): - return await fn(keys, partition_id, **kwargs) - - _holding_global: ContextVar[bool] = ContextVar("_holding_kv_global_lock", default=False) # This process's leases, renewed in batches by one daemon thread that runs while any exist. _registry = threading.Condition() @@ -432,30 +404,6 @@ async def async_kv_global_lock( _holding_global.set(False) -def kv_global_locked( - fn, keys: str | list[str], partition_id: str, *, lock_timeout: float | None = None, lease_s: float = 30, **kwargs -): - """Call ``fn(keys, partition_id, **kwargs)`` under ``kv_global_lock(keys, partition_id)``. - - Same calling convention as ``kv_local_locked``. If the lease was lost by the time ``fn`` - returns, its result is discarded and ``LockLostError`` is raised. - """ - if inspect.iscoroutinefunction(fn): - raise TypeError("kv_global_locked would release the lock before the coroutine runs; use async_kv_global_locked") - with kv_global_lock(keys, partition_id, timeout=lock_timeout, lease_s=lease_s): - return fn(keys, partition_id, **kwargs) - - -async def async_kv_global_locked( - fn, keys: str | list[str], partition_id: str, *, lock_timeout: float | None = None, lease_s: float = 30, **kwargs -): - """Async version of ``kv_global_locked``; ``fn`` must be a coroutine function such as ``tq.async_kv_batch_get``.""" - if not inspect.iscoroutinefunction(fn): - raise TypeError("async_kv_global_locked needs a coroutine function; a sync call would block the event loop") - async with async_kv_global_lock(keys, partition_id, timeout=lock_timeout, lease_s=lease_s): - return await fn(keys, partition_id, **kwargs) - - def kv_lock_list(partition_id: str | None = None) -> dict: """Return ``{"holders": [...], "waiters": n}`` for live global locks, optionally in one partition. From 510a48bc6138c0f19da2c2034e5d3b0690f9016c Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 9 Oct 2026 20:46:40 +0800 Subject: [PATCH 11/42] [feat] Make the kv_global_lock actors opt-in via lock.enabled The lock actors cost num_shards resident processes per TransferQueue, so tq.init() starts them only with lock.enabled: true; the shard count moves from controller.num_lock_shards to lock.num_shards. Their naming, creation, lookup and teardown now all live in kv_lock.py. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 5 ++- tests/test_interface_close.py | 28 +++++++------ tests/test_kv_lock.py | 4 ++ transfer_queue/config.yaml | 9 ++++- transfer_queue/interface.py | 32 +++++---------- transfer_queue/kv_lock.py | 59 ++++++++++++++++++++++------ 6 files changed, 87 insertions(+), 50 deletions(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index 6b3106f4..15b53c18 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -34,6 +34,7 @@ CONF = { "controller": {"polling_mode": True}, "backend": {"storage_backend": "SimpleStorage", "SimpleStorage": {"total_storage_size": 100}}, + "lock": {"enabled": True}, } @@ -343,13 +344,13 @@ def test_owner_close_kills_the_managers_and_fails_waiters(): held.__exit__(None, None, None) wait_until(lambda: actors_named("TransferQueueLockManager_") == []) - with pytest.raises(RuntimeError, match="call tq.init"): + with pytest.raises(RuntimeError, match="lock.enabled"): tq.kv_lock_list() def test_one_lock_shard_still_works(): wait_until(lambda: actors_named("TransferQueue") == []) - tq.init(OmegaConf.create({**CONF, "controller": {"polling_mode": True, "num_lock_shards": 1}})) + tq.init(OmegaConf.create({**CONF, "lock": {"enabled": True, "num_shards": 1}})) assert actors_named("TransferQueueLockManager_") == ["TransferQueueLockManager_0"] with tq.kv_global_lock(["a", "b"], P, timeout=5) as lease: assert holders() == {"a", "b"} and list(lease.shards) == [0] diff --git a/tests/test_interface_close.py b/tests/test_interface_close.py index af3c9453..b7453d1d 100644 --- a/tests/test_interface_close.py +++ b/tests/test_interface_close.py @@ -15,6 +15,7 @@ """Unit tests for `close()` ownership: only the process that created TransferQueue tears it down.""" +import importlib from unittest.mock import MagicMock import pytest @@ -22,6 +23,8 @@ import transfer_queue.interface as iface +kvl = importlib.import_module("transfer_queue.kv_lock") + STORED_CONF = OmegaConf.create( {"controller": {"zmq_info": None}, "backend": {"storage_backend": "SimpleStorage", "SimpleStorage": {}}} ) @@ -32,6 +35,7 @@ def fake(monkeypatch): for name in ("_TQ_CLIENT", "_TQ_STORAGE", "_TQ_CONTROLLER"): monkeypatch.setattr(iface, name, None) monkeypatch.setattr(iface, "_TQ_IS_OWNER", False) + monkeypatch.setattr(kvl, "_owned", []) controller = MagicMock(name="controller") ray = MagicMock(name="ray") @@ -39,35 +43,37 @@ def fake(monkeypatch): controller_cls = MagicMock(name="TransferQueueController") controller_cls.options.return_value.remote.return_value = controller monkeypatch.setattr(iface, "ray", ray) + monkeypatch.setattr(kvl, "ray", ray) monkeypatch.setattr(iface, "TransferQueueController", controller_cls) - monkeypatch.setattr(iface, "TransferQueueLockManager", MagicMock(name="TransferQueueLockManager")) + monkeypatch.setattr(kvl, "TransferQueueLockManager", MagicMock(name="TransferQueueLockManager")) monkeypatch.setattr(iface, "TransferQueueClient", MagicMock(name="TransferQueueClient")) monkeypatch.setattr(iface, "process_zmq_server_info", MagicMock(return_value=None)) monkeypatch.setattr(iface, "_maybe_create_tq_storage", lambda conf: conf) return ray, controller, controller_cls -def test_owner_close_kills_controller(fake): +@pytest.mark.parametrize("lock_enabled", [False, True], ids=["locks_off", "locks_on"]) +def test_owner_close_kills_controller(fake, lock_enabled): ray, controller, _ = fake ray.get_actor.side_effect = ValueError("no controller yet") - iface.init(OmegaConf.create({})) - client, lock_managers = iface._TQ_CLIENT, iface._TQ_LOCK_MANAGERS - assert iface._TQ_IS_OWNER and len(lock_managers) == 8 + iface.init(OmegaConf.create({"lock": {"enabled": lock_enabled}})) + client, lock_managers = iface._TQ_CLIENT, kvl._owned + assert iface._TQ_IS_OWNER and len(lock_managers) == (8 if lock_enabled else 0) iface.close() client.close.assert_called_once() - assert [c.args for c in ray.kill.call_args_list] == [(controller,)] + [(m,) for m in lock_managers] + assert [c.args for c in ray.kill.call_args_list] == [(m,) for m in lock_managers] + [(controller,)] assert (iface._TQ_CLIENT, iface._TQ_CONTROLLER, iface._TQ_IS_OWNER) == (None, None, False) - assert iface._TQ_LOCK_MANAGERS == [] + assert kvl._owned == [] -@pytest.mark.parametrize("num_lock_shards", [0, "8"]) -def test_init_rejects_invalid_num_lock_shards(fake, num_lock_shards): +@pytest.mark.parametrize("num_shards", [0, "8"]) +def test_init_rejects_invalid_lock_num_shards(fake, num_shards): ray, _, controller_cls = fake ray.get_actor.side_effect = ValueError("no controller yet") - with pytest.raises(ValueError, match="num_lock_shards"): - iface.init(OmegaConf.create({"controller": {"num_lock_shards": num_lock_shards}})) + with pytest.raises(ValueError, match="lock.num_shards"): + iface.init(OmegaConf.create({"lock": {"enabled": True, "num_shards": num_shards}})) controller_cls.options.assert_not_called() diff --git a/tests/test_kv_lock.py b/tests/test_kv_lock.py index d25f3ae4..1cc045a3 100644 --- a/tests/test_kv_lock.py +++ b/tests/test_kv_lock.py @@ -399,3 +399,7 @@ def worker(): assert not thread.is_alive() assert tq.kv_batch_get("counter", "lock_e2e")["v"][0].item() == 400 tq.kv_clear("counter", "lock_e2e") + # The default config starts no lock actors: the local lock works, the global one says why not. + with pytest.raises(RuntimeError, match="lock.enabled"): + with tq.kv_global_lock("counter", "lock_e2e"): + pass diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index a593d206..4f91fd26 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -12,11 +12,16 @@ controller: sampler: SequentialSampler # Whether return an empty BatchMeta to prevent request blocking when no enough data is available polling_mode: False - # Number of lock actors that kv_global_lock spreads keys over by hash (>= 1) - num_lock_shards: 8 # ZMQ Server IP & Ports (automatically generated during init) zmq_info: null +# Cluster-wide KV key locks (kv_global_lock). The process-local kv_local_lock needs none of this. +lock: + # Whether tq.init() starts the lock actors; kv_global_lock raises while this is false. + enabled: false + # Number of lock actors that kv_global_lock spreads keys over by hash (>= 1) + num_shards: 8 + backend: # Pluggable storage/transport backend of TransferQueue. Choose from: diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 29fc67ef..4de838d2 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -45,8 +45,7 @@ from transfer_queue.client import TransferQueueClient from transfer_queue.controller import TransferQueueController -from transfer_queue.kv_lock import _detach_lock_managers -from transfer_queue.lock_manager import TransferQueueLockManager +from transfer_queue.kv_lock import _check_lock_conf, _close_lock_managers, _start_lock_managers from transfer_queue.metadata import KVBatchMeta from transfer_queue.sampler import * # noqa: F401 from transfer_queue.sampler import BaseSampler @@ -64,7 +63,6 @@ # True only in the process whose init() created the controller. Other processes # attach to the same named actor, so their close() must not tear it down. _TQ_IS_OWNER = False -_TQ_LOCK_MANAGERS: list = [] # set only in the owner, which kills them in close() # Idle workers and proxies observe shutdown within one second; leave time for # Ray dispatch, in-flight work, and SSD cleanup before forcing termination. @@ -201,24 +199,16 @@ def init(conf: DictConfig | None = None) -> DictConfig | None: sampler = globals()[final_conf.controller.sampler] except KeyError: raise ValueError(f"Could not find sampler {final_conf.controller.sampler}") from None - num_lock_shards = final_conf.controller.num_lock_shards - if not isinstance(num_lock_shards, int) or num_lock_shards < 1: - raise ValueError(f"controller.num_lock_shards must be an integer >= 1, got {num_lock_shards!r}") + _check_lock_conf(final_conf.lock) try: - global _TQ_CONTROLLER, _TQ_IS_OWNER, _TQ_LOCK_MANAGERS + global _TQ_CONTROLLER, _TQ_IS_OWNER _TQ_CONTROLLER = TransferQueueController.options( # type: ignore[attr-defined] name="TransferQueueController", namespace="transfer_queue" ).remote(sampler=sampler, polling_mode=final_conf.controller.polling_mode) _TQ_IS_OWNER = True logger.info("TransferQueueController has been created.") - # Created here rather than on first lock: a non-detached actor dies with its creator. - _TQ_LOCK_MANAGERS = [ - TransferQueueLockManager.options( # type: ignore[attr-defined] - name=f"TransferQueueLockManager_{i}", namespace="transfer_queue", get_if_exists=True - ).remote(num_lock_shards) - for i in range(num_lock_shards) - ] + _start_lock_managers(final_conf.lock) except ValueError: logger.info("Some other rank has initialized TransferQueueController. Try to connect to existing controller.") _init_from_existing() @@ -270,7 +260,6 @@ def close(): global _TQ_STORAGE global _TQ_CONTROLLER global _TQ_IS_OWNER - global _TQ_LOCK_MANAGERS try: if _TQ_STORAGE: @@ -330,16 +319,13 @@ def close(): _TQ_CLIENT.close() _TQ_CLIENT = None - # The owner kills the lock managers, so releasing first would only hand keys to doomed waiters. - _detach_lock_managers(release=not _TQ_IS_OWNER) + _close_lock_managers() if _TQ_IS_OWNER: - for actor in (_TQ_CONTROLLER, *_TQ_LOCK_MANAGERS): - try: - ray.kill(actor) - except Exception: - pass + try: + ray.kill(_TQ_CONTROLLER) + except Exception: + pass _TQ_CONTROLLER = None - _TQ_LOCK_MANAGERS = [] _TQ_IS_OWNER = False diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index c3072d55..a85d7073 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -23,7 +23,8 @@ ``kv_global_lock`` and ``async_kv_global_lock`` do the same across the Ray cluster through ``TransferQueueLockManager`` actors that split keys by hash, under a lease that renews -automatically. Take a global lock before a local one, never inside it. +automatically. ``tq.init()`` starts those actors only with ``lock.enabled: true``. Take a +global lock before a local one, never inside it. """ import asyncio @@ -39,6 +40,8 @@ import ray from ray.exceptions import RayActorError +from transfer_queue.lock_manager import TransferQueueLockManager + # Not asyncio.Lock: it is bound to one event loop and is not thread-safe, while a process # may run several loops plus plain threads. A key is held exactly while it is in _table, # which maps it to its FIFO of waiters. _state_lock guards every change and is never held @@ -197,9 +200,13 @@ async def async_kv_local_lock(keys: str | list[str], partition_id: str, timeout: _registry = threading.Condition() _leases: dict[str, "GlobalLease"] = {} _renewer_running = False -_num_shards: int | None = None # controller.num_lock_shards, asked of shard 0 on first use +_num_shards: int | None = None # lock.num_shards, asked of shard 0 on first use _managers: dict = {} # shard -> TransferQueueLockManager handle, looked up on first use -_NOT_INITIALIZED = "TransferQueueLockManager not found; call tq.init() first" +_owned: list = [] # the lock actors this process's tq.init() created, killed by its tq.close() +_NOT_INITIALIZED = ( + "TransferQueueLockManager not found; kv_global_lock needs `lock.enabled: true` " + "in the config of the tq.init() that starts TransferQueue" +) _MANAGER_GONE = "The TransferQueueLockManager actor is gone; did the process that ran tq.init() call tq.close()?" @@ -239,10 +246,14 @@ def _shard(name: tuple[str, str], num_shards: int) -> int: return zlib.crc32(f"{name[0]}\0{name[1]}".encode()) % num_shards +def _manager_name(shard: int) -> str: + return f"TransferQueueLockManager_{shard}" + + def _lock_manager(shard: int): if shard not in _managers: try: - _managers[shard] = ray.get_actor(f"TransferQueueLockManager_{shard}", namespace="transfer_queue") + _managers[shard] = ray.get_actor(_manager_name(shard), namespace="transfer_queue") except ValueError: raise RuntimeError(_NOT_INITIALIZED) from None return _managers[shard] @@ -346,7 +357,7 @@ def _renew_loop() -> None: def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | None = None, lease_s: float = 30): """Hold an exclusive cluster-wide lock on ``keys`` in ``partition_id``; yield a ``GlobalLease``. - The locks live in the ``controller.num_lock_shards`` actors that ``tq.init()`` creates, + The locks live in the ``lock.num_shards`` actors that ``tq.init()`` creates, each key in the one its hash picks, and are held under a ``lease_s`` lease that a background thread renews. Each actor grants in arrival order: a request waits behind every earlier one that shares a key with it. Keys on several actors are taken one actor @@ -414,23 +425,47 @@ def kv_lock_list(partition_id: str | None = None) -> dict: return {"holders": [h for r in replies for h in r["holders"]], "waiters": sum(r["waiters"] for r in replies)} -def _detach_lock_managers(release: bool) -> None: - """Called by ``tq.close()``: give up this process's global leases and forget the actor handles.""" - global _num_shards, _managers +def _check_lock_conf(lock_conf) -> None: + """Called by ``tq.init()`` before it creates anything, so a bad config leaves nothing behind.""" + if lock_conf.enabled and (not isinstance(lock_conf.num_shards, int) or lock_conf.num_shards < 1): + raise ValueError(f"lock.num_shards must be an integer >= 1, got {lock_conf.num_shards!r}") + + +def _start_lock_managers(lock_conf) -> None: + """Called by the ``tq.init()`` that created the controller: start the lock actors if enabled.""" + global _owned + if lock_conf.enabled: + # Created here rather than on first lock: a non-detached actor dies with its creator. + _owned = [ + TransferQueueLockManager.options( # type: ignore[attr-defined] + name=_manager_name(shard), namespace="transfer_queue", get_if_exists=True + ).remote(lock_conf.num_shards) + for shard in range(lock_conf.num_shards) + ] + + +def _close_lock_managers() -> None: + """Called by ``tq.close()``: give up this process's leases, then kill the actors it created.""" + global _num_shards, _managers, _owned with _registry: leases = list(_leases.values()) for lease in leases: lease.lost = True - if release: + if not _owned: # the owner kills the actors, so releasing would only hand keys to doomed waiters _release_lease(lease) - _num_shards, _managers = None, {} + for actor in _owned: + try: + ray.kill(actor) + except Exception: + pass + _num_shards, _managers, _owned = None, {}, [] def _reset_after_fork() -> None: # The child inherits keys, leases and Ray handles that belong to threads and connections it lacks. - global _state_lock, _table, _registry, _leases, _renewer_running, _managers + global _state_lock, _table, _registry, _leases, _renewer_running, _managers, _owned _state_lock, _table = threading.Lock(), {} - _registry, _leases, _renewer_running, _managers = threading.Condition(), {}, False, {} + _registry, _leases, _renewer_running, _managers, _owned = threading.Condition(), {}, False, {}, [] os.register_at_fork(after_in_child=_reset_after_fork) From 795a451bff7c96ce6bccdbfe307c6b95c2aeb29b Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 9 Oct 2026 20:48:44 +0800 Subject: [PATCH 12/42] [fix] Reject kv_global_lock callers that hashed over a stale shard count A process that cached lock.num_shards before TransferQueue restarted with another value would send keys to the wrong lock actor, where nothing excludes the processes that sent them to the right one. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 13 ++++++++++--- transfer_queue/kv_lock.py | 4 +++- transfer_queue/lock_manager.py | 9 ++++++++- 3 files changed, 21 insertions(+), 5 deletions(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index 15b53c18..8f0cdb5d 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -204,7 +204,7 @@ async def main(): def test_release_before_acquire_withdraws_it(): lock_manager = manager("w") ray.get(lock_manager.release.remote("early"), timeout=TIMEOUT_S) - assert ray.get(lock_manager.acquire.remote([(P, "w")], "early", None, 5, {}), timeout=TIMEOUT_S) is None + assert ray.get(lock_manager.acquire.remote([(P, "w")], "early", None, 5, {}, 8), timeout=TIMEOUT_S) is None assert holders() == set() @@ -216,14 +216,21 @@ def own_manager(): ray.kill(lock_manager) -def acquire(lock_manager, keys, token, timeout=None, lease_s=30): - return lock_manager.acquire.remote([(P, key) for key in keys], token, timeout, lease_s, {}) +def acquire(lock_manager, keys, token, timeout=None, lease_s=30, num_shards=1): + return lock_manager.acquire.remote([(P, key) for key in keys], token, timeout, lease_s, {}, num_shards) def n_waiters(lock_manager): return ray.get(lock_manager.list_locks.remote(P), timeout=TIMEOUT_S)["waiters"] +def test_acquire_rejects_a_caller_hashing_over_another_shard_count(own_manager): + # e.g. a process that attached before TransferQueue restarted with a new lock.num_shards + with pytest.raises(ValueError, match="lock.num_shards is 1"): + ray.get(acquire(own_manager, ["a"], "stale", num_shards=8), timeout=TIMEOUT_S) + assert ray.get(acquire(own_manager, ["a"], "fresh", timeout=0), timeout=TIMEOUT_S) is not None + + def test_waiters_are_granted_in_arrival_order(own_manager): ray.get(acquire(own_manager, ["a"], "holder"), timeout=TIMEOUT_S) both = acquire(own_manager, ["a", "b"], "both") diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index a85d7073..45cae686 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -276,7 +276,9 @@ def _send_acquire(lease: GlobalLease, shard: int, deadline: float | None): manager = _lock_manager(shard) holder = {"node_ip": ray.util.get_node_ip_address(), "pid": os.getpid(), "thread": threading.current_thread().name} lease.pending.append(shard) - return manager.acquire.remote(lease.shards[shard], lease.token, _remaining(deadline), lease.lease_s, holder) + return manager.acquire.remote( + lease.shards[shard], lease.token, _remaining(deadline), lease.lease_s, holder, _shard_count() + ) def _take_grant(lease: GlobalLease, shard: int, waited: float | None, sent: float) -> None: diff --git a/transfer_queue/lock_manager.py b/transfer_queue/lock_manager.py index 1fc57792..2b0f9b56 100644 --- a/transfer_queue/lock_manager.py +++ b/transfer_queue/lock_manager.py @@ -106,8 +106,15 @@ def _expire(self) -> None: self._timer_at = self._expiries[0][0] self._timer = asyncio.get_running_loop().call_later(max(0.0, self._timer_at - now), self._expire) - async def acquire(self, names, token, timeout, lease_s, holder_info): + async def acquire(self, names, token, timeout, lease_s, holder_info, num_shards): """Grant all ``names`` at once; return the seconds waited, or None on timeout or withdrawal.""" + # A caller that hashed over another shard count may send a key to the wrong shard, + # where nothing excludes the callers that sent it to the right one. + if num_shards != self._num_shards: + raise ValueError( + f"kv_global_lock hashed keys over {num_shards} lock shards, but lock.num_shards is " + f"{self._num_shards}; TransferQueue restarted, so call tq.close() and tq.init() again" + ) start = time.monotonic() if self._withdrawn.pop(token, None) is not None: return None From 7bc0ee00fe51a290dcf813a118009a46f78da8c3 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 9 Oct 2026 20:52:50 +0800 Subject: [PATCH 13/42] [fix] Renew each kv_global_lock shard as soon as it is granted Renewal used to start only once every shard was granted, so waiting on a later shard for longer than lease_s let an earlier shard's lease expire and pass its keys on while the caller still meant to enter. Check the lease before entering, too. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 22 +++++++++++++++++ transfer_queue/kv_lock.py | 36 +++++++++++++++------------- 2 files changed, 41 insertions(+), 17 deletions(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index 8f0cdb5d..da6e7163 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -17,6 +17,7 @@ import asyncio import importlib +import threading import time import pytest @@ -201,6 +202,27 @@ async def main(): pass +def test_waiting_on_a_later_shard_keeps_the_earlier_keys_held(): + early, late = keys_on_two_shards("lw") + worker = Worker.remote() + ray.get(worker.hold.remote(late), timeout=TIMEOUT_S) + entered = [] + + def take_both(): + with tq.kv_global_lock([early, late], P, timeout=TIMEOUT_S, lease_s=2) as lease: + entered.append(lease) + + thread = threading.Thread(target=take_both) + thread.start() + wait_until(lambda: holders() == {early, late}) + time.sleep(3) # longer than our lease: unrenewed, the grant on `early` would have expired + assert holders() == {early, late} + ray.get(worker.release.remote(), timeout=TIMEOUT_S) + thread.join(TIMEOUT_S) + assert len(entered) == 1 # entered without LockLostError + wait_until(lambda: holders() == set()) + + def test_release_before_acquire_withdraws_it(): lock_manager = manager("w") ray.get(lock_manager.release.remote("early"), timeout=TIMEOUT_S) diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 45cae686..22c826a7 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -223,6 +223,7 @@ def __init__(self, names: list[tuple[str, str]], shards: dict[int, list[tuple[st self.deadline = float("inf") # local monotonic time the lease is known to last until, once granted self.lost = False self.pending: list[int] = [] # shards that may hold or await this token, so exit must release them + self.granted: list[int] = [] # shards that granted it, so the renewer keeps their leases alive def check(self) -> None: """Raise ``LockLostError`` if the lease was lost; call it right before writes in long sections.""" @@ -282,16 +283,15 @@ def _send_acquire(lease: GlobalLease, shard: int, deadline: float | None): def _take_grant(lease: GlobalLease, shard: int, waited: float | None, sent: float) -> None: + global _renewer_running if waited is None: lease.pending.remove(shard) raise TimeoutError(f"Timed out waiting for kv_global_lock on {lease.names}") # The grant happened at least `waited` after `sent`, so this never outlives the shard's lease. lease.deadline = min(lease.deadline, sent + waited + lease.lease_s) - - -def _start_lease(lease: GlobalLease) -> None: - global _renewer_running + # Renew from the first grant on: waiting on a later shard may outlast an earlier one's lease. with _registry: + lease.granted.append(shard) _leases[lease.token] = lease _registry.notify() # a shorter lease_s shortens the renewal period if not _renewer_running: @@ -326,10 +326,10 @@ def _renew_loop() -> None: if remaining <= 0: break _registry.wait(remaining) - batch = dict(_leases) + batch = {token: (lease, list(lease.granted)) for token, lease in _leases.items()} by_shard: dict[int, list[str]] = {} - for token, lease in batch.items(): - for shard in lease.shards: + for token, (_, shards) in batch.items(): + for shard in shards: by_shard.setdefault(shard, []).append(token) # Count from the send time: the server extends the lease from when the call arrives, # which is later, so the local deadline never outlives the server's. @@ -347,12 +347,13 @@ def _renew_loop() -> None: except Exception: # a dead lock manager or a timeout continue renewed.update((shard, token) for token, ok in alive.items() if ok) - # A lease is lost as soon as any one of its shards fails to renew it. - for token, lease in batch.items(): - if all((shard, token) in renewed for shard in lease.shards): - lease.deadline = last + lease.lease_s - else: + # A lease is lost as soon as any one of its shards fails to renew it. A shard granted + # since the snapshot was not renewed here, so it keeps the deadline _take_grant set. + for token, (lease, shards) in batch.items(): + if not all((shard, token) in renewed for shard in shards): lease.lost = True + elif len(lease.granted) == len(shards): + lease.deadline = last + lease.lease_s @contextmanager @@ -367,9 +368,10 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No keys are awaited. Raises ``TimeoutError`` if the keys are not all granted within ``timeout`` seconds (``None`` waits forever), and ``RuntimeError`` when nested, taken while holding a - ``kv_local_lock``, or called with a running event loop. Exit always releases the lock, - then raises ``LockLostError`` if the lease was lost, unless the body already raised: - an exception from the body is never masked. + ``kv_local_lock``, or called with a running event loop. Raises ``LockLostError`` on + entry if an earlier key's lease was lost while a later one was awaited. Exit always + releases the lock, then raises ``LockLostError`` if the lease was lost, unless the body + already raised: an exception from the body is never masked. """ _reject_running_loop("kv_global_lock") lease = _new_lease(keys, partition_id, lease_s) @@ -385,7 +387,7 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No except RayActorError as e: raise RuntimeError(_MANAGER_GONE) from e _take_grant(lease, shard, waited, sent) - _start_lease(lease) + lease.check() yield lease lease.check() finally: @@ -409,7 +411,7 @@ async def async_kv_global_lock( except RayActorError as e: raise RuntimeError(_MANAGER_GONE) from e _take_grant(lease, shard, waited, sent) - _start_lease(lease) + lease.check() yield lease lease.check() finally: From 420b822f8d10a37e30198b27f044497e9c1b1197 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 9 Oct 2026 21:01:02 +0800 Subject: [PATCH 14/42] [fix] Long-poll kv_global_lock acquires so waiters never pin actor slots An acquire used to wait inside the lock actor until granted, so enough waiters could take every concurrency slot and starve release and renew_many. Each call now waits at most one second and a follow-up wait call keeps the queue position. Every reply is awaited with a bound, so an unresponsive actor cannot outlast the caller's timeout either. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 27 ++++++++--- transfer_queue/kv_lock.py | 68 ++++++++++++++++++---------- transfer_queue/lock_manager.py | 31 +++++++------ 3 files changed, 82 insertions(+), 44 deletions(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index da6e7163..790d4c2d 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -85,7 +85,7 @@ def wait_until(predicate): def holders(): - # Exit sends its release without waiting for it, so wait for holders to drain, not assert. + # Exit sends its release without waiting for it, so wait for holders and waiters to drain. return {h["key"] for h in tq.kv_lock_list(P)["holders"]} @@ -134,7 +134,7 @@ def test_multi_shard_timeout_releases_the_earlier_shard(): with tq.kv_global_lock([early, late], P, timeout=0.5): pass wait_until(lambda: holders() == {late}) - assert tq.kv_lock_list(P)["waiters"] == 0 + wait_until(lambda: tq.kv_lock_list(P)["waiters"] == 0) with tq.kv_global_lock(early, P, timeout=5): # a leaked grant would hold it for its 30 s lease pass ray.get(worker.release.remote(), timeout=TIMEOUT_S) @@ -171,7 +171,7 @@ def test_timeout_leaves_no_holder_or_waiter(): with pytest.raises(TimeoutError): with tq.kv_global_lock("t", P, timeout=0.3): pass - assert tq.kv_lock_list(P)["waiters"] == 0 + wait_until(lambda: tq.kv_lock_list(P)["waiters"] == 0) ray.get(worker.release.remote(), timeout=TIMEOUT_S) wait_until(lambda: tq.kv_lock_list(P) == {"holders": [], "waiters": 0}) @@ -259,6 +259,7 @@ def test_waiters_are_granted_in_arrival_order(own_manager): wait_until(lambda: n_waiters(own_manager) == 1) # b is free, but granting it past the older [a, b] request would let single-key traffic starve it. assert ray.get(acquire(own_manager, ["b"], "newcomer", timeout=0), timeout=TIMEOUT_S) is None + ray.get(own_manager.release.remote("newcomer"), timeout=TIMEOUT_S) # as kv_global_lock does on timeout only_b = acquire(own_manager, ["b"], "only_b") wait_until(lambda: n_waiters(own_manager) == 2) ray.get(own_manager.release.remote("holder"), timeout=TIMEOUT_S) @@ -268,12 +269,24 @@ def test_waiters_are_granted_in_arrival_order(own_manager): assert ray.get(only_b, timeout=TIMEOUT_S) is not None -def test_a_waiter_that_gives_up_lets_the_next_one_through(own_manager): +def test_a_timed_out_call_keeps_its_place_until_withdrawn(own_manager): ray.get(acquire(own_manager, ["a"], "holder"), timeout=TIMEOUT_S) - both = acquire(own_manager, ["a", "b"], "both", timeout=0.5) - wait_until(lambda: n_waiters(own_manager) == 1) + assert ray.get(acquire(own_manager, ["a"], "first", timeout=0.1), timeout=TIMEOUT_S) is None + second = acquire(own_manager, ["a"], "second") + wait_until(lambda: n_waiters(own_manager) == 2) + ray.get(own_manager.release.remote("holder"), timeout=TIMEOUT_S) + assert ray.get(own_manager.wait.remote("first", TIMEOUT_S), timeout=TIMEOUT_S) > 0 + assert n_waiters(own_manager) == 1 + ray.get(own_manager.release.remote("first"), timeout=TIMEOUT_S) + assert ray.get(second, timeout=TIMEOUT_S) is not None + + +def test_a_waiter_that_withdraws_lets_the_next_one_through(own_manager): + ray.get(acquire(own_manager, ["a"], "holder"), timeout=TIMEOUT_S) + assert ray.get(acquire(own_manager, ["a", "b"], "both", timeout=0.1), timeout=TIMEOUT_S) is None only_b = acquire(own_manager, ["b"], "only_b") - assert ray.get(both, timeout=TIMEOUT_S) is None + wait_until(lambda: n_waiters(own_manager) == 2) + ray.get(own_manager.release.remote("both"), timeout=TIMEOUT_S) assert ray.get(only_b, timeout=TIMEOUT_S) is not None diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 22c826a7..2dd52522 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -38,7 +38,7 @@ from contextvars import ContextVar import ray -from ray.exceptions import RayActorError +from ray.exceptions import GetTimeoutError, RayActorError from transfer_queue.lock_manager import TransferQueueLockManager @@ -203,6 +203,9 @@ async def async_kv_local_lock(keys: str | list[str], partition_id: str, timeout: _num_shards: int | None = None # lock.num_shards, asked of shard 0 on first use _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() +# Longest one call waits inside a lock actor; waiting longer takes another call that keeps +# the queue position. Bounding each call keeps waiters from pinning the actor's concurrency slots. +_POLL_S = 1.0 _NOT_INITIALIZED = ( "TransferQueueLockManager not found; kv_global_lock needs `lock.enabled: true` " "in the config of the tq.init() that starts TransferQueue" @@ -273,22 +276,31 @@ def _new_lease(keys: str | list[str], partition_id: str, lease_s: float) -> Glob return GlobalLease(names, dict(sorted(shards.items())), lease_s) -def _send_acquire(lease: GlobalLease, shard: int, deadline: float | None): +def _poll_window(deadline: float | None) -> float: + remaining = _remaining(deadline) + return _POLL_S if remaining is None else min(_POLL_S, remaining) + + +def _send_acquire(lease: GlobalLease, shard: int, window: float): + """Queue for ``shard`` on the first call, then keep waiting there; reply as the actor's ``wait``.""" manager = _lock_manager(shard) + if shard in lease.pending: + return manager.wait.remote(lease.token, window) holder = {"node_ip": ray.util.get_node_ip_address(), "pid": os.getpid(), "thread": threading.current_thread().name} lease.pending.append(shard) - return manager.acquire.remote( - lease.shards[shard], lease.token, _remaining(deadline), lease.lease_s, holder, _shard_count() - ) + return manager.acquire.remote(lease.shards[shard], lease.token, window, lease.lease_s, holder, _shard_count()) -def _take_grant(lease: GlobalLease, shard: int, waited: float | None, sent: float) -> None: +def _take_grant(lease: GlobalLease, shard: int, remaining: float | None, sent: float, deadline: float | None) -> bool: + """Record a reply from ``shard``; return whether it granted, or raise once ``deadline`` passed.""" global _renewer_running - if waited is None: - lease.pending.remove(shard) - raise TimeoutError(f"Timed out waiting for kv_global_lock on {lease.names}") - # The grant happened at least `waited` after `sent`, so this never outlives the shard's lease. - lease.deadline = min(lease.deadline, sent + waited + lease.lease_s) + if remaining is None: + if deadline is not None and time.monotonic() >= deadline: + # Still queued there, so exit's release must withdraw it: the shard stays pending. + raise TimeoutError(f"Timed out waiting for kv_global_lock on {lease.names}") + return False + # The reply left the actor after `sent`, so this never outlives the shard's lease. + lease.deadline = min(lease.deadline, sent + remaining) # Renew from the first grant on: waiting on a later shard may outlast an earlier one's lease. with _registry: lease.granted.append(shard) @@ -297,6 +309,7 @@ def _take_grant(lease: GlobalLease, shard: int, waited: float | None, sent: floa if not _renewer_running: _renewer_running = True threading.Thread(target=_renew_loop, name="kv_global_lock_renewer", daemon=True).start() + return True def _release_lease(lease: GlobalLease) -> None: @@ -381,12 +394,17 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No # Every caller visits shards in ascending order, so no two callers can each hold a # shard the other awaits: multi-shard requests cannot deadlock. for shard in lease.shards: - sent = time.monotonic() - try: - waited = ray.get(_send_acquire(lease, shard, deadline)) - except RayActorError as e: - raise RuntimeError(_MANAGER_GONE) from e - _take_grant(lease, shard, waited, sent) + granted = False + while not granted: + window, sent = _poll_window(deadline), time.monotonic() + try: + # Bounded, so an unresponsive actor cannot outlast the caller's timeout. + remaining = ray.get(_send_acquire(lease, shard, window), timeout=window + _POLL_S) + except GetTimeoutError: + remaining = None + except RayActorError as e: + raise RuntimeError(_MANAGER_GONE) from e + granted = _take_grant(lease, shard, remaining, sent, deadline) lease.check() yield lease lease.check() @@ -405,12 +423,16 @@ async def async_kv_global_lock( _holding_global.set(True) try: for shard in lease.shards: - sent = time.monotonic() - try: - waited = await _send_acquire(lease, shard, deadline) - except RayActorError as e: - raise RuntimeError(_MANAGER_GONE) from e - _take_grant(lease, shard, waited, sent) + granted = False + while not granted: + window, sent = _poll_window(deadline), time.monotonic() + try: + remaining = await asyncio.wait_for(_send_acquire(lease, shard, window), window + _POLL_S) + except asyncio.TimeoutError: + remaining = None + except RayActorError as e: + raise RuntimeError(_MANAGER_GONE) from e + granted = _take_grant(lease, shard, remaining, sent, deadline) lease.check() yield lease lease.check() diff --git a/transfer_queue/lock_manager.py b/transfer_queue/lock_manager.py index 2b0f9b56..bbc745c5 100644 --- a/transfer_queue/lock_manager.py +++ b/transfer_queue/lock_manager.py @@ -26,8 +26,8 @@ _WITHDRAWN_TTL_S = 300 -# Every waiting acquire occupies a concurrency slot, so the limit sits far above Ray's async -# default of 1000 to keep renew and release from queueing behind waiters. +# A waiting acquire or wait call holds a concurrency slot for at most the caller's poll +# window, so even if waiters take every slot, renew and release get one within that window. @ray.remote(num_cpus=0, max_concurrency=10_000) class TransferQueueLockManager: """Exclusive leased locks on ``(partition_id, key)``. Runs on one event loop, so no thread locks.""" @@ -107,7 +107,7 @@ def _expire(self) -> None: self._timer = asyncio.get_running_loop().call_later(max(0.0, self._timer_at - now), self._expire) async def acquire(self, names, token, timeout, lease_s, holder_info, num_shards): - """Grant all ``names`` at once; return the seconds waited, or None on timeout or withdrawal.""" + """Queue ``token`` for all ``names`` at once, then wait for the grant as ``wait`` does.""" # A caller that hashed over another shard count may send a key to the wrong shard, # where nothing excludes the callers that sent it to the right one. if num_shards != self._num_shards: @@ -115,26 +115,29 @@ async def acquire(self, names, token, timeout, lease_s, holder_info, num_shards) f"kv_global_lock hashed keys over {num_shards} lock shards, but lock.num_shards is " f"{self._num_shards}; TransferQueue restarted, so call tq.close() and tq.init() again" ) - start = time.monotonic() if self._withdrawn.pop(token, None) is not None: return None if not any(name in self._holder or name in self._queues for name in names): self._grant(token, names, lease_s, holder_info) else: - done = asyncio.Event() - self._waiting[token] = dict(names=names, lease_s=lease_s, info=holder_info, done=done) + self._waiting[token] = dict(names=names, lease_s=lease_s, info=holder_info, done=asyncio.Event()) for name in names: self._queues.setdefault(name, deque()).append(token) + return await self.wait(token, timeout) + + async def wait(self, token: str, timeout: float) -> float | None: + """Return the seconds left on ``token``'s lease once granted, else None after ``timeout``. + + A waiter that times out keeps its place in the queues; only ``release`` withdraws it. + """ + waiter = self._waiting.get(token) + if waiter is not None: try: - await asyncio.wait_for(done.wait(), timeout) + await asyncio.wait_for(waiter["done"].wait(), timeout) except asyncio.TimeoutError: pass - finally: - if token in self._waiting: # timed out or cancelled before a grant - self._hand_off(self._leave_queues(token)["names"]) - if token not in self._leases: # withdrawn, or already released by a caller that left - return None - return self._leases[token]["granted_at"] - start + lease = self._leases.get(token) + return None if lease is None else lease["expires_at"] - time.monotonic() def num_shards(self) -> int: """Return the global shard count recorded by this manager.""" @@ -159,7 +162,7 @@ async def release(self, token: str) -> None: if token in self._waiting: waiter = self._leave_queues(token) self._hand_off(waiter["names"]) - waiter["done"].set() # its acquire finds no lease and returns None + waiter["done"].set() # its pending call finds no lease and returns None else: now = time.monotonic() self._withdrawn = {t: until for t, until in self._withdrawn.items() if until > now} From f4566fe4c1f225874189a27584ae40ae4b59ec72 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 9 Oct 2026 21:05:29 +0800 Subject: [PATCH 15/42] [fix] Look the lock actors up off the event loop on async first use The first async_kv_global_lock in a process asked shard 0 for the shard count with ray.get and looked each actor up with ray.get_actor, blocking the running loop. All lookups now happen at once in a worker thread. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 22 ++++++++++++++++++++++ transfer_queue/kv_lock.py | 8 +++++++- 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index 790d4c2d..57cc52d8 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -223,6 +223,28 @@ def take_both(): wait_until(lambda: holders() == set()) +def test_async_first_use_looks_the_actors_up_off_the_event_loop(monkeypatch): + monkeypatch.setattr(kvl, "_num_shards", None) + monkeypatch.setattr(kvl, "_managers", {}) + on_loop = [] + for name in ("get", "get_actor"): + real = getattr(ray, name) + + def spy(*args, _real=real, **kwargs): + on_loop.append(threading.current_thread() is threading.main_thread()) + return _real(*args, **kwargs) + + monkeypatch.setattr(ray, name, spy) + + async def main(): + async with tq.async_kv_global_lock("first_use", P, timeout=TIMEOUT_S): + pass + + asyncio.run(main()) + assert on_loop and not any(on_loop) + assert len(kvl._managers) == kvl._num_shards == 8 + + def test_release_before_acquire_withdraws_it(): lock_manager = manager("w") ray.get(lock_manager.release.remote("early"), timeout=TIMEOUT_S) diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 2dd52522..e531f728 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -235,12 +235,16 @@ def check(self) -> None: def _shard_count() -> int: + """Return lock.num_shards, looking every lock actor up on first use.""" global _num_shards if _num_shards is None: try: - _num_shards = ray.get(_lock_manager(0).num_shards.remote()) + num_shards = ray.get(_lock_manager(0).num_shards.remote()) except RayActorError as e: raise RuntimeError(_MANAGER_GONE) from e + for shard in range(1, num_shards): + _lock_manager(shard) + _num_shards = num_shards return _num_shards @@ -418,6 +422,8 @@ async def async_kv_global_lock( keys: str | list[str], partition_id: str, timeout: float | None = None, lease_s: float = 30 ): """Async version of ``kv_global_lock``. Tasks created inside the block cannot take a global lock.""" + if _num_shards is None: # the first use waits on Ray's GCS, so keep it off the event loop + await asyncio.to_thread(_shard_count) lease = _new_lease(keys, partition_id, lease_s) deadline = None if timeout is None else time.monotonic() + timeout _holding_global.set(True) From 224d15bba6eac74f7bda35589529fded191919ed Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 9 Oct 2026 21:09:17 +0800 Subject: [PATCH 16/42] [fix] Let tasks spawned inside a KV lock block lock again after it exits A task copies its creator's context, so a task created inside an async_kv_local_lock block kept seeing the lock as held and later failed with "Nested kv_local_lock" even for unrelated keys. The context now holds a cell that exit clears for every copy. Signed-off-by: OutstanderWang --- tests/test_kv_lock.py | 2 ++ transfer_queue/kv_lock.py | 44 ++++++++++++++++++++++++--------------- 2 files changed, 29 insertions(+), 17 deletions(-) diff --git a/tests/test_kv_lock.py b/tests/test_kv_lock.py index 1cc045a3..a747a6c6 100644 --- a/tests/test_kv_lock.py +++ b/tests/test_kv_lock.py @@ -305,6 +305,8 @@ async def child(): with pytest.raises(RuntimeError, match="Nested"): await asyncio.create_task(child()) + later = asyncio.create_task(child()) + await later # starts after the block exited, so it no longer counts as nested asyncio.run(main()) diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index e531f728..1153ca4e 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -48,9 +48,16 @@ # while waiting. _state_lock = threading.Lock() _table: dict[tuple[str, str], deque["_Waiter"]] = {} -# Whether this thread or task holds (or is acquiring) a local lock. A future global lock -# reads it to enforce the "global, then local" order. -_holding: ContextVar[bool] = ContextVar("_holding_kv_lock", default=False) +# Whether this thread or task holds (or is acquiring) a local lock; the global lock reads it +# to enforce the "global, then local" order. It holds a cell, not a bool: a task created +# inside the block copies the context, and must see the cell cleared once the block exits. +_holding: ContextVar[list[bool] | None] = ContextVar("_holding_kv_lock", default=None) + + +def _mark(held: ContextVar[list[bool] | None]) -> list[bool]: + cell = [True] + held.set(cell) + return cell class _Waiter: @@ -132,12 +139,15 @@ async def _acquire_async(name: tuple[str, str], deadline: float | None) -> None: def _lock_names( - keys: str | list[str], partition_id: str, held: ContextVar[bool] = _holding, kind: str = "kv_local_lock" + keys: str | list[str], + partition_id: str, + held: ContextVar[list[bool] | None] = _holding, + kind: str = "kv_local_lock", ) -> list[tuple[str, str]]: keys = [keys] if isinstance(keys, str) else keys if not keys: raise ValueError(f"{kind} needs at least one key") - if held.get(): + if held.get() == [True]: raise RuntimeError(f"Nested {kind} is not supported; pass all keys to a single call") # A fixed order keeps two overlapping multi-key calls from deadlocking. return [(partition_id, key) for key in sorted(set(keys))] @@ -168,7 +178,7 @@ def kv_local_lock(keys: str | list[str], partition_id: str, timeout: float | Non _reject_running_loop("kv_local_lock") names, acquired = _lock_names(keys, partition_id), [] deadline = None if timeout is None else time.monotonic() + timeout - _holding.set(True) + held = _mark(_holding) try: for name in names: _acquire(name, deadline) @@ -176,15 +186,15 @@ def kv_local_lock(keys: str | list[str], partition_id: str, timeout: float | Non yield finally: _release_all(acquired) - _holding.set(False) + held[0] = False @asynccontextmanager async def async_kv_local_lock(keys: str | list[str], partition_id: str, timeout: float | None = None): - """Async version of ``kv_local_lock``. Tasks created inside the block cannot take a local lock.""" + """Async version of ``kv_local_lock``. Tasks created inside the block cannot take a local lock while it lasts.""" names, acquired = _lock_names(keys, partition_id), [] deadline = None if timeout is None else time.monotonic() + timeout - _holding.set(True) + held = _mark(_holding) try: for name in names: await _acquire_async(name, deadline) @@ -192,10 +202,10 @@ async def async_kv_local_lock(keys: str | list[str], partition_id: str, timeout: yield finally: _release_all(acquired) - _holding.set(False) + held[0] = False -_holding_global: ContextVar[bool] = ContextVar("_holding_kv_global_lock", default=False) +_holding_global: ContextVar[list[bool] | None] = ContextVar("_holding_kv_global_lock", default=None) # This process's leases, renewed in batches by one daemon thread that runs while any exist. _registry = threading.Condition() _leases: dict[str, "GlobalLease"] = {} @@ -269,7 +279,7 @@ def _lock_manager(shard: int): def _new_lease(keys: str | list[str], partition_id: str, lease_s: float) -> GlobalLease: names = _lock_names(keys, partition_id, _holding_global, "kv_global_lock") - if _holding.get(): + if _holding.get() == [True]: raise RuntimeError("Take kv_global_lock before kv_local_lock, not while holding one") if lease_s <= 0: raise ValueError("lease_s must be positive") @@ -393,7 +403,7 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No _reject_running_loop("kv_global_lock") lease = _new_lease(keys, partition_id, lease_s) deadline = None if timeout is None else time.monotonic() + timeout - _holding_global.set(True) + held = _mark(_holding_global) try: # Every caller visits shards in ascending order, so no two callers can each hold a # shard the other awaits: multi-shard requests cannot deadlock. @@ -414,19 +424,19 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No lease.check() finally: _release_lease(lease) - _holding_global.set(False) + held[0] = False @asynccontextmanager async def async_kv_global_lock( keys: str | list[str], partition_id: str, timeout: float | None = None, lease_s: float = 30 ): - """Async version of ``kv_global_lock``. Tasks created inside the block cannot take a global lock.""" + """Async version of ``kv_global_lock``. Tasks created inside the block cannot take a global lock while it lasts.""" if _num_shards is None: # the first use waits on Ray's GCS, so keep it off the event loop await asyncio.to_thread(_shard_count) lease = _new_lease(keys, partition_id, lease_s) deadline = None if timeout is None else time.monotonic() + timeout - _holding_global.set(True) + held = _mark(_holding_global) try: for shard in lease.shards: granted = False @@ -444,7 +454,7 @@ async def async_kv_global_lock( lease.check() finally: _release_lease(lease) - _holding_global.set(False) + held[0] = False def kv_lock_list(partition_id: str | None = None) -> dict: From d1d7a85f091da0fc3c844d17fb66455dd30297d6 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 9 Oct 2026 21:11:03 +0800 Subject: [PATCH 17/42] [perf] Prune expired withdrawal tombstones from the front Every release that withdrew nothing rebuilt the whole tombstone dict. All tombstones share one TTL, so expired ones are always the oldest and can be dropped from the front in amortized constant time. Signed-off-by: OutstanderWang --- transfer_queue/lock_manager.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/transfer_queue/lock_manager.py b/transfer_queue/lock_manager.py index bbc745c5..e4c850f5 100644 --- a/transfer_queue/lock_manager.py +++ b/transfer_queue/lock_manager.py @@ -165,7 +165,12 @@ async def release(self, token: str) -> None: waiter["done"].set() # its pending call finds no lease and returns None else: now = time.monotonic() - self._withdrawn = {t: until for t, until in self._withdrawn.items() if until > now} + # One TTL for all, so insertion order is expiry order: expired tombstones lead. + while self._withdrawn: + oldest, until = next(iter(self._withdrawn.items())) + if until > now: + break + del self._withdrawn[oldest] self._withdrawn[token] = now + _WITHDRAWN_TTL_S def list_locks(self, partition_id: str | None = None) -> dict: From cca8aaee0e1389f84525a353b2ea005a9664078f Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 9 Oct 2026 21:14:09 +0800 Subject: [PATCH 18/42] [feat] Export GlobalLease, the type kv_global_lock yields Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 1 + transfer_queue/__init__.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index 57cc52d8..93949e74 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -417,5 +417,6 @@ def test_one_lock_shard_still_works(): tq.init(OmegaConf.create({**CONF, "lock": {"enabled": True, "num_shards": 1}})) assert actors_named("TransferQueueLockManager_") == ["TransferQueueLockManager_0"] with tq.kv_global_lock(["a", "b"], P, timeout=5) as lease: + assert isinstance(lease, tq.GlobalLease) assert holders() == {"a", "b"} and list(lease.shards) == [0] wait_until(lambda: holders() == set()) diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py index ea9faed2..1d349a27 100644 --- a/transfer_queue/__init__.py +++ b/transfer_queue/__init__.py @@ -38,6 +38,7 @@ save_checkpoint, ) from .kv_lock import ( + GlobalLease, LockLostError, async_kv_global_lock, async_kv_local_lock, @@ -80,6 +81,7 @@ "kv_global_lock", "async_kv_global_lock", "kv_lock_list", + "GlobalLease", "LockLostError", ] + [ From d36d05eeb18269359eb40a8c310552cd1f6156fb Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Fri, 9 Oct 2026 21:20:36 +0800 Subject: [PATCH 19/42] [doc] Spell out what cross-shard hold-and-wait means for kv_global_lock callers Signed-off-by: OutstanderWang --- transfer_queue/kv_lock.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 1153ca4e..5063c5ce 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -390,9 +390,15 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No The locks live in the ``lock.num_shards`` actors that ``tq.init()`` creates, each key in the one its hash picks, and are held under a ``lease_s`` lease that a background thread renews. Each actor grants in arrival order: a request waits behind - every earlier one that shares a key with it. Keys on several actors are taken one actor - at a time, each all at once, so keys already granted stay held while a later actor's - keys are awaited. + every earlier one that shares a key with it. + + Keys on several actors are taken one actor at a time, each all at once, so a request + holds the keys already granted while it waits for a later actor's. The hash decides + which keys share an actor, so whether a request is all-or-nothing is not up to the + caller: with 8 shards, two keys land on different actors 7 times in 8. A request queues + afresh at each later actor, so its waits add up, and the keys it holds meanwhile block + others even if it then times out. Lock the fewest keys you can and keep the block short. + Raises ``TimeoutError`` if the keys are not all granted within ``timeout`` seconds (``None`` waits forever), and ``RuntimeError`` when nested, taken while holding a ``kv_local_lock``, or called with a running event loop. Raises ``LockLostError`` on @@ -461,7 +467,7 @@ def kv_lock_list(partition_id: str | None = None) -> dict: """Return ``{"holders": [...], "waiters": n}`` for live global locks, optionally in one partition. Each holder has ``partition_id``, ``key``, ``holder`` (node IP, pid, thread), ``held_s`` and - ``lease_remaining_s``. + ``lease_remaining_s``. A holder may still be waiting for its keys on other actors. """ replies = ray.get([_lock_manager(shard).list_locks.remote(partition_id) for shard in range(_shard_count())]) return {"holders": [h for r in replies for h in r["holders"]], "waiters": sum(r["waiters"] for r in replies)} From 5d2b1833bc7cfbb63a1d727ce6811e87a77cc157 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Sat, 10 Oct 2026 16:43:02 +0800 Subject: [PATCH 20/42] [fix] Make kv_global_lock acquire idempotent so an unseen expired grant queues afresh A grant that expired between two polls left the token unknown to the manager, and wait() then returned None at once, so the client spun until its timeout. Re-polls now repeat acquire with the same token, which keeps its place in the queues or queues it again; wait() is removed. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 39 +++++++++++++++++++++++++--- transfer_queue/kv_lock.py | 11 ++++---- transfer_queue/lock_manager.py | 33 ++++++++++++----------- 3 files changed, 57 insertions(+), 26 deletions(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index 93949e74..29a39e48 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -223,6 +223,24 @@ def take_both(): wait_until(lambda: holders() == set()) +def test_a_grant_that_expires_while_the_caller_stalls_is_taken_again(monkeypatch): + worker = Worker.remote() + ray.get(worker.hold.remote("stall"), timeout=TIMEOUT_S) + real_take_grant, stalled = kvl._take_grant, [] + + def take_grant(lease, shard, remaining, *args): + if remaining is None and not stalled: # queued behind the worker: free the key, then stall past our lease + stalled.append(True) + ray.get(worker.release.remote(), timeout=TIMEOUT_S) + time.sleep(1) + return real_take_grant(lease, shard, remaining, *args) + + monkeypatch.setattr(kvl, "_take_grant", take_grant) + with tq.kv_global_lock("stall", P, timeout=10, lease_s=0.3): + assert stalled and holders() == {"stall"} + wait_until(lambda: holders() == set()) + + def test_async_first_use_looks_the_actors_up_off_the_event_loop(monkeypatch): monkeypatch.setattr(kvl, "_num_shards", None) monkeypatch.setattr(kvl, "_managers", {}) @@ -248,8 +266,9 @@ async def main(): def test_release_before_acquire_withdraws_it(): lock_manager = manager("w") ray.get(lock_manager.release.remote("early"), timeout=TIMEOUT_S) - assert ray.get(lock_manager.acquire.remote([(P, "w")], "early", None, 5, {}, 8), timeout=TIMEOUT_S) is None - assert holders() == set() + for _ in range(2): # a retried call that arrives late is withdrawn too + assert ray.get(lock_manager.acquire.remote([(P, "w")], "early", None, 5, {}, 8), timeout=TIMEOUT_S) is None + assert holders() == set() and tq.kv_lock_list(P)["waiters"] == 0 @pytest.fixture @@ -297,12 +316,26 @@ def test_a_timed_out_call_keeps_its_place_until_withdrawn(own_manager): second = acquire(own_manager, ["a"], "second") wait_until(lambda: n_waiters(own_manager) == 2) ray.get(own_manager.release.remote("holder"), timeout=TIMEOUT_S) - assert ray.get(own_manager.wait.remote("first", TIMEOUT_S), timeout=TIMEOUT_S) > 0 + assert ray.get(acquire(own_manager, ["a"], "first"), timeout=TIMEOUT_S) > 0 # calling again keeps the place assert n_waiters(own_manager) == 1 ray.get(own_manager.release.remote("first"), timeout=TIMEOUT_S) assert ray.get(second, timeout=TIMEOUT_S) is not None +def test_a_grant_that_expires_before_the_caller_sees_it_queues_afresh(own_manager): + ray.get(acquire(own_manager, ["a"], "holder"), timeout=TIMEOUT_S) + assert ray.get(acquire(own_manager, ["a"], "slow", timeout=0, lease_s=0.3), timeout=TIMEOUT_S) is None + ray.get(own_manager.release.remote("holder"), timeout=TIMEOUT_S) # grants "slow" between its calls + wait_until(lambda: n_waiters(own_manager) == 0) + time.sleep(0.6) # the grant expires unseen + assert ray.get(acquire(own_manager, ["a"], "other"), timeout=TIMEOUT_S) is not None + # Calling again must not spin on an unknown token: it queues behind the current holder. + assert ray.get(acquire(own_manager, ["a"], "slow", timeout=0.1), timeout=TIMEOUT_S) is None + assert n_waiters(own_manager) == 1 + ray.get(own_manager.release.remote("other"), timeout=TIMEOUT_S) + assert ray.get(acquire(own_manager, ["a"], "slow"), timeout=TIMEOUT_S) > 0 + + def test_a_waiter_that_withdraws_lets_the_next_one_through(own_manager): ray.get(acquire(own_manager, ["a"], "holder"), timeout=TIMEOUT_S) assert ray.get(acquire(own_manager, ["a", "b"], "both", timeout=0.1), timeout=TIMEOUT_S) is None diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 5063c5ce..eff8afe8 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -296,13 +296,12 @@ def _poll_window(deadline: float | None) -> float: def _send_acquire(lease: GlobalLease, shard: int, window: float): - """Queue for ``shard`` on the first call, then keep waiting there; reply as the actor's ``wait``.""" - manager = _lock_manager(shard) - if shard in lease.pending: - return manager.wait.remote(lease.token, window) + if shard not in lease.pending: + lease.pending.append(shard) holder = {"node_ip": ray.util.get_node_ip_address(), "pid": os.getpid(), "thread": threading.current_thread().name} - lease.pending.append(shard) - return manager.acquire.remote(lease.shards[shard], lease.token, window, lease.lease_s, holder, _shard_count()) + return _lock_manager(shard).acquire.remote( + lease.shards[shard], lease.token, window, lease.lease_s, holder, _shard_count() + ) def _take_grant(lease: GlobalLease, shard: int, remaining: float | None, sent: float, deadline: float | None) -> bool: diff --git a/transfer_queue/lock_manager.py b/transfer_queue/lock_manager.py index e4c850f5..da16cb25 100644 --- a/transfer_queue/lock_manager.py +++ b/transfer_queue/lock_manager.py @@ -26,7 +26,7 @@ _WITHDRAWN_TTL_S = 300 -# A waiting acquire or wait call holds a concurrency slot for at most the caller's poll +# A waiting acquire call holds a concurrency slot for at most the caller's poll # window, so even if waiters take every slot, renew and release get one within that window. @ray.remote(num_cpus=0, max_concurrency=10_000) class TransferQueueLockManager: @@ -106,8 +106,13 @@ def _expire(self) -> None: self._timer_at = self._expiries[0][0] self._timer = asyncio.get_running_loop().call_later(max(0.0, self._timer_at - now), self._expire) - async def acquire(self, names, token, timeout, lease_s, holder_info, num_shards): - """Queue ``token`` for all ``names`` at once, then wait for the grant as ``wait`` does.""" + async def acquire(self, names, token, timeout, lease_s, holder_info, num_shards) -> float | None: + """Grant all ``names`` to ``token`` at once; return the seconds left on its lease. + + Returns None if still queued after ``timeout``. Calling again with the same token keeps + its place in the queues; only ``release`` withdraws it. A token that is neither queued + nor held (its grant expired before the caller saw it) queues afresh. + """ # A caller that hashed over another shard count may send a key to the wrong shard, # where nothing excludes the callers that sent it to the right one. if num_shards != self._num_shards: @@ -115,21 +120,15 @@ async def acquire(self, names, token, timeout, lease_s, holder_info, num_shards) f"kv_global_lock hashed keys over {num_shards} lock shards, but lock.num_shards is " f"{self._num_shards}; TransferQueue restarted, so call tq.close() and tq.init() again" ) - if self._withdrawn.pop(token, None) is not None: + if token in self._withdrawn: # kept until it expires: a retried call may still arrive return None - if not any(name in self._holder or name in self._queues for name in names): - self._grant(token, names, lease_s, holder_info) - else: - self._waiting[token] = dict(names=names, lease_s=lease_s, info=holder_info, done=asyncio.Event()) - for name in names: - self._queues.setdefault(name, deque()).append(token) - return await self.wait(token, timeout) - - async def wait(self, token: str, timeout: float) -> float | None: - """Return the seconds left on ``token``'s lease once granted, else None after ``timeout``. - - A waiter that times out keeps its place in the queues; only ``release`` withdraws it. - """ + if token not in self._leases and token not in self._waiting: + if not any(name in self._holder or name in self._queues for name in names): + self._grant(token, names, lease_s, holder_info) + else: + self._waiting[token] = dict(names=names, lease_s=lease_s, info=holder_info, done=asyncio.Event()) + for name in names: + self._queues.setdefault(name, deque()).append(token) waiter = self._waiting.get(token) if waiter is not None: try: From 711061300576237bb9b3169437254218e2f120fd Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Sat, 10 Oct 2026 16:47:57 +0800 Subject: [PATCH 21/42] [fix] Skip kv_global_lock waiters that stopped polling when handing a key off A caller that died while queued was still granted the key, which then sat idle for that caller's whole lease, once per dead waiter in line. The manager now tracks each waiter's calls in flight and last call, and hand-off drops one that has not polled for 5 s; a merely stalled caller queues again on its next call. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 12 +++++++++++ transfer_queue/lock_manager.py | 30 +++++++++++++++++++++------- 2 files changed, 35 insertions(+), 7 deletions(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index 29a39e48..3fb70677 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -355,6 +355,18 @@ def test_expired_leases_pass_the_key_down_the_queue(own_manager): assert ray.get(first, timeout=TIMEOUT_S) is not None +def test_a_waiter_that_stopped_polling_is_skipped(own_manager): + ray.get(acquire(own_manager, ["a"], "holder"), timeout=TIMEOUT_S) + # Queued once, then its caller died: granting it would idle the key for its whole lease. + assert ray.get(acquire(own_manager, ["a"], "dead", timeout=0, lease_s=TIMEOUT_S), timeout=TIMEOUT_S) is None + live = acquire(own_manager, ["a"], "live") + wait_until(lambda: n_waiters(own_manager) == 2) + time.sleep(6) # past the manager's abandonment threshold + ray.get(own_manager.release.remote("holder"), timeout=TIMEOUT_S) + assert ray.get(live, timeout=5) is not None + assert n_waiters(own_manager) == 0 + + def test_killed_holder_frees_the_key_after_its_lease(): worker = Worker.remote() ray.get(worker.hold.remote("k", lease_s=1.5), timeout=TIMEOUT_S) diff --git a/transfer_queue/lock_manager.py b/transfer_queue/lock_manager.py index da16cb25..4594422f 100644 --- a/transfer_queue/lock_manager.py +++ b/transfer_queue/lock_manager.py @@ -24,6 +24,10 @@ # A withdrawal whose acquire never arrives (e.g. the caller died first) is dropped after this. _WITHDRAWN_TTL_S = 300 +# A queued caller re-polls at least every second, so a waiter with no call in flight for this +# long has died. Hand-off skips it instead of granting it a lease that would only expire; if it +# was merely stalled, its next call queues it again at the back. +_ABANDONED_S = 5.0 # A waiting acquire call holds a concurrency slot for at most the caller's poll @@ -49,7 +53,8 @@ def __init__(self, num_shards: int): # requests cannot starve an older multi-key one. All queues share one arrival order, so the # oldest waiter heads all of its queues and the wait can never form a cycle. self._queues: dict[tuple[str, str], deque[str]] = {} - self._waiting: dict[str, dict] = {} # token -> names, lease_s, info, done (set on grant or withdrawal) + # token -> names, lease_s, info, done (set on grant or withdrawal), calls in flight, seen (last call end) + self._waiting: dict[str, dict] = {} self._withdrawn: dict[str, float] = {} # token -> when its tombstone expires def _grant(self, token: str, names: list[tuple[str, str]], lease_s: float, info: dict) -> None: @@ -74,14 +79,19 @@ def _leave_queues(self, token: str) -> dict: def _hand_off(self, names: list[tuple[str, str]]) -> None: """Grant each free name to the first waiter in its queue, if that waiter now heads all of its queues.""" - for name in names: + names, now = list(names), time.monotonic() + while names: + name = names.pop() if name in self._holder or name not in self._queues: continue token = self._queues[name][0] - wanted = self._waiting[token]["names"] - if all(n not in self._holder and self._queues[n][0] == token for n in wanted): - waiter = self._leave_queues(token) - self._grant(token, wanted, waiter["lease_s"], waiter["info"]) + waiter = self._waiting[token] + if waiter["calls"] == 0 and now - waiter["seen"] > _ABANDONED_S: + self._leave_queues(token) + names.extend(waiter["names"]) # each of its queues has a new head + elif all(n not in self._holder and self._queues[n][0] == token for n in waiter["names"]): + self._leave_queues(token) + self._grant(token, waiter["names"], waiter["lease_s"], waiter["info"]) waiter["done"].set() def _free(self, token: str) -> bool: @@ -126,15 +136,21 @@ async def acquire(self, names, token, timeout, lease_s, holder_info, num_shards) if not any(name in self._holder or name in self._queues for name in names): self._grant(token, names, lease_s, holder_info) else: - self._waiting[token] = dict(names=names, lease_s=lease_s, info=holder_info, done=asyncio.Event()) + self._waiting[token] = dict( + names=names, lease_s=lease_s, info=holder_info, done=asyncio.Event(), calls=0, seen=time.monotonic() + ) for name in names: self._queues.setdefault(name, deque()).append(token) waiter = self._waiting.get(token) if waiter is not None: + waiter["calls"] += 1 try: await asyncio.wait_for(waiter["done"].wait(), timeout) except asyncio.TimeoutError: pass + finally: + waiter["calls"] -= 1 + waiter["seen"] = time.monotonic() lease = self._leases.get(token) return None if lease is None else lease["expires_at"] - time.monotonic() From fce1c2bdb5f93980e427da1d5e26bea99e943b32 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Sat, 10 Oct 2026 16:54:11 +0800 Subject: [PATCH 22/42] [perf] Poll kv_global_lock waiters every 5 s, capped at a third of the lease One lock actor serves a few thousand calls a second. With 1 s polls, 8,000 idle waiters on a shard cut uncontended lock cycles there from 1,629/s to 1/s, renewals included; with 5 s polls it keeps ~1,100/s. The window is also capped at lease_s / 3, so a renewal that waits for a slot behind every waiter still gets one well within its lease. Signed-off-by: OutstanderWang --- tests/test_kv_lock.py | 6 ++++++ transfer_queue/kv_lock.py | 23 ++++++++++++++--------- transfer_queue/lock_manager.py | 4 ++-- 3 files changed, 22 insertions(+), 11 deletions(-) diff --git a/tests/test_kv_lock.py b/tests/test_kv_lock.py index a747a6c6..3bb480ec 100644 --- a/tests/test_kv_lock.py +++ b/tests/test_kv_lock.py @@ -370,6 +370,12 @@ def test_global_lock_shards_agree_across_processes_and_spread_keys(): assert out.stdout.strip() == str(shards) +def test_global_lock_poll_window_fits_the_lease_and_the_timeout(): + assert kvl._poll_window(None, 30) == kvl._POLL_S + assert kvl._poll_window(None, 0.6) == pytest.approx(0.2) # a renewal waits at most a third of the lease + assert kvl._poll_window(time.monotonic() + 0.05, 30) <= 0.05 + + @pytest.fixture def simple_storage_tq(): if not ray.is_initialized(): diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index eff8afe8..f4082e2c 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -214,8 +214,11 @@ 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() # Longest one call waits inside a lock actor; waiting longer takes another call that keeps -# the queue position. Bounding each call keeps waiters from pinning the actor's concurrency slots. -_POLL_S = 1.0 +# the queue position. Bounding each call keeps waiters from pinning the actor's concurrency +# slots, and the bound is long because one actor serves only a few thousand calls a second: +# at 1 s, 8,000 idle waiters starved every other call on their shard, renewals included. +_POLL_S = 5.0 +_REPLY_SLACK_S = 1.0 # how late a reply may arrive past its window before the caller re-polls _NOT_INITIALIZED = ( "TransferQueueLockManager not found; kv_global_lock needs `lock.enabled: true` " "in the config of the tq.init() that starts TransferQueue" @@ -290,9 +293,11 @@ def _new_lease(keys: str | list[str], partition_id: str, lease_s: float) -> Glob return GlobalLease(names, dict(sorted(shards.items())), lease_s) -def _poll_window(deadline: float | None) -> float: - remaining = _remaining(deadline) - return _POLL_S if remaining is None else min(_POLL_S, remaining) +def _poll_window(deadline: float | None, lease_s: float) -> float: + # Even when waiters take every actor slot, a renewal waits at most one window for a slot, + # so a window up to lease_s / 3 leaves it most of the lease, as the renewer's period does. + window, remaining = min(_POLL_S, lease_s / 3), _remaining(deadline) + return window if remaining is None else min(window, remaining) def _send_acquire(lease: GlobalLease, shard: int, window: float): @@ -415,10 +420,10 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No for shard in lease.shards: granted = False while not granted: - window, sent = _poll_window(deadline), time.monotonic() + window, sent = _poll_window(deadline, lease.lease_s), time.monotonic() try: # Bounded, so an unresponsive actor cannot outlast the caller's timeout. - remaining = ray.get(_send_acquire(lease, shard, window), timeout=window + _POLL_S) + remaining = ray.get(_send_acquire(lease, shard, window), timeout=window + _REPLY_SLACK_S) except GetTimeoutError: remaining = None except RayActorError as e: @@ -446,9 +451,9 @@ async def async_kv_global_lock( for shard in lease.shards: granted = False while not granted: - window, sent = _poll_window(deadline), time.monotonic() + window, sent = _poll_window(deadline, lease.lease_s), time.monotonic() try: - remaining = await asyncio.wait_for(_send_acquire(lease, shard, window), window + _POLL_S) + remaining = await asyncio.wait_for(_send_acquire(lease, shard, window), window + _REPLY_SLACK_S) except asyncio.TimeoutError: remaining = None except RayActorError as e: diff --git a/transfer_queue/lock_manager.py b/transfer_queue/lock_manager.py index 4594422f..7d2b5b8e 100644 --- a/transfer_queue/lock_manager.py +++ b/transfer_queue/lock_manager.py @@ -24,8 +24,8 @@ # A withdrawal whose acquire never arrives (e.g. the caller died first) is dropped after this. _WITHDRAWN_TTL_S = 300 -# A queued caller re-polls at least every second, so a waiter with no call in flight for this -# long has died. Hand-off skips it instead of granting it a lease that would only expire; if it +# A queued caller sends its next call as soon as one returns, so a waiter with no call in +# flight for this long has died. Hand-off skips it instead of granting it a lease that would only expire; if it # was merely stalled, its next call queues it again at the back. _ABANDONED_S = 5.0 From 04d2230d1d4fd5f942c71c50523b4a65232deea0 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Sat, 10 Oct 2026 17:10:35 +0800 Subject: [PATCH 23/42] [fix] Stop the kv_global_lock renewer on close and set lease deadlines under its lock close() marked this process's leases lost but left them registered, so the renewer kept calling dead actors until each block exited; it now empties the registry and wakes the renewer. A grant's deadline is set under the same lock the renewer holds when it moves it. The docstring notes that an unresponsive actor can stretch timeout by about 1 s. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 1 + transfer_queue/kv_lock.py | 21 ++++++++++++--------- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index 3fb70677..e57cfe2e 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -447,6 +447,7 @@ def test_owner_close_kills_the_managers_and_fails_waiters(): wait_until(lambda: tq.kv_lock_list(P)["waiters"] == 1) tq.close() + wait_until(lambda: not kvl._renewer_running) # even though the block has not exited with pytest.raises(RuntimeError, match="TransferQueueLockManager actor is gone"): ray.get(waiting, timeout=TIMEOUT_S) with pytest.raises(tq.LockLostError): diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index f4082e2c..df449c02 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -317,10 +317,10 @@ def _take_grant(lease: GlobalLease, shard: int, remaining: float | None, sent: f # Still queued there, so exit's release must withdraw it: the shard stays pending. raise TimeoutError(f"Timed out waiting for kv_global_lock on {lease.names}") return False - # The reply left the actor after `sent`, so this never outlives the shard's lease. - lease.deadline = min(lease.deadline, sent + remaining) # Renew from the first grant on: waiting on a later shard may outlast an earlier one's lease. - with _registry: + with _registry: # the renewer also moves the deadline + # The reply left the actor after `sent`, so this never outlives the shard's lease. + lease.deadline = min(lease.deadline, sent + remaining) lease.granted.append(shard) _leases[lease.token] = lease _registry.notify() # a shorter lease_s shortens the renewal period @@ -404,11 +404,12 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No others even if it then times out. Lock the fewest keys you can and keep the block short. Raises ``TimeoutError`` if the keys are not all granted within ``timeout`` seconds - (``None`` waits forever), and ``RuntimeError`` when nested, taken while holding a - ``kv_local_lock``, or called with a running event loop. Raises ``LockLostError`` on - entry if an earlier key's lease was lost while a later one was awaited. Exit always - releases the lock, then raises ``LockLostError`` if the lease was lost, unless the body - already raised: an exception from the body is never masked. + (``None`` waits forever); an unresponsive lock actor can stretch that by up to about + 1 s, the time allowed for its last reply. Raises ``RuntimeError`` when nested, taken + while holding a ``kv_local_lock``, or called with a running event loop. Raises + ``LockLostError`` on entry if an earlier key's lease was lost while a later one was + awaited. Exit always releases the lock, then raises ``LockLostError`` if the lease was + lost, unless the body already raised: an exception from the body is never masked. """ _reject_running_loop("kv_global_lock") lease = _new_lease(keys, partition_id, lease_s) @@ -499,8 +500,10 @@ def _start_lock_managers(lock_conf) -> None: def _close_lock_managers() -> None: """Called by ``tq.close()``: give up this process's leases, then kill the actors it created.""" global _num_shards, _managers, _owned - with _registry: + with _registry: # emptied so the renewer stops now, not when the blocks holding them exit leases = list(_leases.values()) + _leases.clear() + _registry.notify() for lease in leases: lease.lost = True if not _owned: # the owner kills the actors, so releasing would only hand keys to doomed waiters From ef84850e8e361aeaae2504d295798a0101178628 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Sat, 10 Oct 2026 17:35:11 +0800 Subject: [PATCH 24/42] [feat] Make the lock actors' max_concurrency configurable and document the limits lock.max_concurrency (default 10000) replaces the hard-coded value and is validated like lock.num_shards. The kv_global_lock docstring states how many requests can wait at once, which setting raises the total or a hot key's share, and what happens past the limit. Signed-off-by: OutstanderWang --- tests/test_interface_close.py | 19 +++++++++++++++---- transfer_queue/config.yaml | 2 ++ transfer_queue/kv_lock.py | 18 +++++++++++++++--- transfer_queue/lock_manager.py | 7 ++++--- 4 files changed, 36 insertions(+), 10 deletions(-) diff --git a/tests/test_interface_close.py b/tests/test_interface_close.py index b7453d1d..37f6d7e6 100644 --- a/tests/test_interface_close.py +++ b/tests/test_interface_close.py @@ -68,15 +68,26 @@ def test_owner_close_kills_controller(fake, lock_enabled): assert kvl._owned == [] -@pytest.mark.parametrize("num_shards", [0, "8"]) -def test_init_rejects_invalid_lock_num_shards(fake, num_shards): +@pytest.mark.parametrize( + "field, value", [("num_shards", 0), ("num_shards", "8"), ("max_concurrency", 0), ("max_concurrency", 1.5)] +) +def test_init_rejects_invalid_lock_sizes(fake, field, value): ray, _, controller_cls = fake ray.get_actor.side_effect = ValueError("no controller yet") - with pytest.raises(ValueError, match="lock.num_shards"): - iface.init(OmegaConf.create({"lock": {"enabled": True, "num_shards": num_shards}})) + with pytest.raises(ValueError, match=f"lock.{field}"): + iface.init(OmegaConf.create({"lock": {"enabled": True, field: value}})) controller_cls.options.assert_not_called() +def test_init_gives_every_lock_actor_the_configured_max_concurrency(fake): + ray, _, _ = fake + ray.get_actor.side_effect = ValueError("no controller yet") + iface.init(OmegaConf.create({"lock": {"enabled": True, "num_shards": 3, "max_concurrency": 64}})) + calls = kvl.TransferQueueLockManager.options.call_args_list + assert [c.kwargs["max_concurrency"] for c in calls] == [64, 64, 64] + iface.close() + + @pytest.mark.parametrize("lost_creation_race", [False, True], ids=["existing", "lost_creation_race"]) def test_attaching_close_keeps_controller_and_can_reattach(fake, lost_creation_race): ray, controller, controller_cls = fake diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index 4f91fd26..d9dba90b 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -21,6 +21,8 @@ lock: enabled: false # Number of lock actors that kv_global_lock spreads keys over by hash (>= 1) num_shards: 8 + # Calls each lock actor serves at once, waiting acquires included (>= 1); see kv_global_lock + max_concurrency: 10000 backend: diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index df449c02..cad853ef 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -403,6 +403,13 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No afresh at each later actor, so its waits add up, and the keys it holds meanwhile block others even if it then times out. Lock the fewest keys you can and keep the block short. + Each actor serves ``lock.max_concurrency`` calls at once, and a waiting request takes one + of them on the actor it is waiting at. With keys hashing evenly, about ``num_shards * + max_concurrency`` requests can wait at once, so raise ``lock.num_shards`` for more in total; + all waiters on one key share one actor, so raise ``lock.max_concurrency`` for a hot key. + Past the limit, calls queue in Ray: a waiting call frees its slot within its poll window + (5 s, or ``lease_s / 3`` if shorter), so renewal and release are delayed, never stuck. + Raises ``TimeoutError`` if the keys are not all granted within ``timeout`` seconds (``None`` waits forever); an unresponsive lock actor can stretch that by up to about 1 s, the time allowed for its last reply. Raises ``RuntimeError`` when nested, taken @@ -480,8 +487,10 @@ def kv_lock_list(partition_id: str | None = None) -> dict: def _check_lock_conf(lock_conf) -> None: """Called by ``tq.init()`` before it creates anything, so a bad config leaves nothing behind.""" - if lock_conf.enabled and (not isinstance(lock_conf.num_shards, int) or lock_conf.num_shards < 1): - raise ValueError(f"lock.num_shards must be an integer >= 1, got {lock_conf.num_shards!r}") + for field in ("num_shards", "max_concurrency"): + value = lock_conf[field] + if lock_conf.enabled and (not isinstance(value, int) or value < 1): + raise ValueError(f"lock.{field} must be an integer >= 1, got {value!r}") def _start_lock_managers(lock_conf) -> None: @@ -491,7 +500,10 @@ def _start_lock_managers(lock_conf) -> None: # Created here rather than on first lock: a non-detached actor dies with its creator. _owned = [ TransferQueueLockManager.options( # type: ignore[attr-defined] - name=_manager_name(shard), namespace="transfer_queue", get_if_exists=True + name=_manager_name(shard), + namespace="transfer_queue", + get_if_exists=True, + max_concurrency=lock_conf.max_concurrency, ).remote(lock_conf.num_shards) for shard in range(lock_conf.num_shards) ] diff --git a/transfer_queue/lock_manager.py b/transfer_queue/lock_manager.py index 7d2b5b8e..5eb136fe 100644 --- a/transfer_queue/lock_manager.py +++ b/transfer_queue/lock_manager.py @@ -30,9 +30,10 @@ _ABANDONED_S = 5.0 -# A waiting acquire call holds a concurrency slot for at most the caller's poll -# window, so even if waiters take every slot, renew and release get one within that window. -@ray.remote(num_cpus=0, max_concurrency=10_000) +# tq.init() sets max_concurrency from lock.max_concurrency. A waiting acquire call holds a +# slot for at most the caller's poll window, so even if waiters take every slot, renew and +# release get one within that window. +@ray.remote(num_cpus=0) class TransferQueueLockManager: """Exclusive leased locks on ``(partition_id, key)``. Runs on one event loop, so no thread locks.""" From 175b6af4b717030a3e3dc6077448d1fdb4a0b03a Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Sun, 11 Oct 2026 22:03:36 +0800 Subject: [PATCH 25/42] [test] Trim kv lock tests to one case per behavior Drop cases another test already covers: the single-key read-modify-write (the two-shard opposite-order test asserts the same counter), the single-shard timeout (a subset of the multi-shard one), the client-level unseen-grant test (the actor-level one checks the same fix), the thread/task wake test (the mixed read-modify-write test hangs without it), the SimpleStorage local-lock test, the single-shard re-init test, and test_close_ownership_e2e.py, whose checks move into the non-owner and owner close tests. Pin the shard hash instead of comparing subprocesses, and shorten the later-shard renewal wait. Signed-off-by: OutstanderWang --- tests/e2e/test_close_ownership_e2e.py | 69 --------------------- tests/e2e/test_kv_global_lock_e2e.py | 59 +++--------------- tests/test_kv_lock.py | 86 ++------------------------- 3 files changed, 11 insertions(+), 203 deletions(-) delete mode 100644 tests/e2e/test_close_ownership_e2e.py diff --git a/tests/e2e/test_close_ownership_e2e.py b/tests/e2e/test_close_ownership_e2e.py deleted file mode 100644 index fc3b3fee..00000000 --- a/tests/e2e/test_close_ownership_e2e.py +++ /dev/null @@ -1,69 +0,0 @@ -# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. -# Copyright 2025 The TransferQueue Team -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""A Ray worker that attaches to TransferQueue and closes it must not tear it down for the driver.""" - -import time - -import pytest -import ray -import torch -from omegaconf import OmegaConf - -import transfer_queue as tq - -TIMEOUT_S = 60 -CONF = { - "controller": {"polling_mode": True}, - "backend": {"storage_backend": "SimpleStorage", "SimpleStorage": {"total_storage_size": 100}}, -} - - -@ray.remote -def attach_use_and_close() -> int: - tq.init() - tq.kv_put("from_worker", "close_e2e", fields={"v": torch.tensor([7])}) - value = int(tq.kv_batch_get("from_driver", "close_e2e")["v"][0]) - tq.close() - return value - - -@pytest.fixture -def ray_session(): - if not ray.is_initialized(): - ray.init(namespace="TestCloseOwnershipE2E") - yield - ray.shutdown() - - -def test_only_the_creating_process_tears_down(ray_session): - tq.init(OmegaConf.create(CONF)) - tq.kv_put("from_driver", "close_e2e", fields={"v": torch.tensor([3])}) - - assert ray.get(attach_use_and_close.remote(), timeout=TIMEOUT_S) == 3 - - controller = ray.get_actor("TransferQueueController", namespace="transfer_queue") - ray.get(controller.list_partitions.remote(), timeout=TIMEOUT_S) - assert int(tq.kv_batch_get("from_worker", "close_e2e")["v"][0]) == 7 - - tq.close() - deadline = time.monotonic() + TIMEOUT_S - while True: - try: - ray.get_actor("TransferQueueController", namespace="transfer_queue") - except ValueError: - break - assert time.monotonic() < deadline, "controller still alive after the owner's close()" - time.sleep(0.1) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index e57cfe2e..efcc7ec3 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -44,7 +44,7 @@ class Worker: def __init__(self): tq.init() - def increment(self, n, keys="counter"): + def increment(self, n, keys): for _ in range(n): with tq.kv_global_lock(keys, P, timeout=TIMEOUT_S): value = int(tq.kv_batch_get("counter", P)["v"][0]) @@ -108,14 +108,6 @@ def counter(): return int(tq.kv_batch_get("counter", P)["v"][0]) -def test_actors_serialize_read_modify_write(): - tq.kv_put("counter", P, fields={"v": torch.tensor([0])}) - workers = [Worker.remote() for _ in range(3)] - ray.get([w.increment.remote(20) for w in workers], timeout=TIMEOUT_S) - assert counter() == 60 - tq.kv_clear("counter", P) - - def test_opposite_key_orders_on_two_shards_do_not_deadlock(): a, b = keys_on_two_shards("ab") tq.kv_put("counter", P, fields={"v": torch.tensor([0])}) @@ -165,17 +157,6 @@ def test_process_without_init_learns_the_shard_count(): wait_until(lambda: holders() == set()) -def test_timeout_leaves_no_holder_or_waiter(): - worker = Worker.remote() - ray.get(worker.hold.remote("t"), timeout=TIMEOUT_S) - with pytest.raises(TimeoutError): - with tq.kv_global_lock("t", P, timeout=0.3): - pass - wait_until(lambda: tq.kv_lock_list(P)["waiters"] == 0) - ray.get(worker.release.remote(), timeout=TIMEOUT_S) - wait_until(lambda: tq.kv_lock_list(P) == {"holders": [], "waiters": 0}) - - def test_cancelled_waiter_is_withdrawn_and_never_granted(): worker = Worker.remote() ray.get(worker.hold.remote("c"), timeout=TIMEOUT_S) @@ -209,13 +190,13 @@ def test_waiting_on_a_later_shard_keeps_the_earlier_keys_held(): entered = [] def take_both(): - with tq.kv_global_lock([early, late], P, timeout=TIMEOUT_S, lease_s=2) as lease: + with tq.kv_global_lock([early, late], P, timeout=TIMEOUT_S, lease_s=1) as lease: entered.append(lease) thread = threading.Thread(target=take_both) thread.start() wait_until(lambda: holders() == {early, late}) - time.sleep(3) # longer than our lease: unrenewed, the grant on `early` would have expired + time.sleep(1.5) # longer than our lease: unrenewed, the grant on `early` would have expired assert holders() == {early, late} ray.get(worker.release.remote(), timeout=TIMEOUT_S) thread.join(TIMEOUT_S) @@ -223,24 +204,6 @@ def take_both(): wait_until(lambda: holders() == set()) -def test_a_grant_that_expires_while_the_caller_stalls_is_taken_again(monkeypatch): - worker = Worker.remote() - ray.get(worker.hold.remote("stall"), timeout=TIMEOUT_S) - real_take_grant, stalled = kvl._take_grant, [] - - def take_grant(lease, shard, remaining, *args): - if remaining is None and not stalled: # queued behind the worker: free the key, then stall past our lease - stalled.append(True) - ray.get(worker.release.remote(), timeout=TIMEOUT_S) - time.sleep(1) - return real_take_grant(lease, shard, remaining, *args) - - monkeypatch.setattr(kvl, "_take_grant", take_grant) - with tq.kv_global_lock("stall", P, timeout=10, lease_s=0.3): - assert stalled and holders() == {"stall"} - wait_until(lambda: holders() == set()) - - def test_async_first_use_looks_the_actors_up_off_the_event_loop(monkeypatch): monkeypatch.setattr(kvl, "_num_shards", None) monkeypatch.setattr(kvl, "_managers", {}) @@ -378,6 +341,7 @@ def test_killed_holder_frees_the_key_after_its_lease(): def test_lease_renews_and_lost_lease_raises(): with tq.kv_global_lock("r", P, lease_s=2) as lease: + assert isinstance(lease, tq.GlobalLease) time.sleep(3) # outlives one lease only through renewal lease.check() @@ -436,10 +400,11 @@ def test_non_owner_close_releases_its_locks_and_keeps_the_managers(): with tq.kv_global_lock("o", P, timeout=5): pass assert len(actors_named("TransferQueueLockManager_")) == kvl._shard_count() == 8 + assert actors_named("TransferQueueController") == ["TransferQueueController"] def test_owner_close_kills_the_managers_and_fails_waiters(): - """Tears TransferQueue down for this module; only the single-shard test runs after it.""" + """Tears TransferQueue down for this module, so it runs last.""" held = tq.kv_global_lock("z", P) held.__enter__() worker = Worker.remote() @@ -453,16 +418,6 @@ def test_owner_close_kills_the_managers_and_fails_waiters(): with pytest.raises(tq.LockLostError): held.__exit__(None, None, None) - wait_until(lambda: actors_named("TransferQueueLockManager_") == []) + wait_until(lambda: actors_named("TransferQueue") == []) # the controller and storage go too with pytest.raises(RuntimeError, match="lock.enabled"): tq.kv_lock_list() - - -def test_one_lock_shard_still_works(): - wait_until(lambda: actors_named("TransferQueue") == []) - tq.init(OmegaConf.create({**CONF, "lock": {"enabled": True, "num_shards": 1}})) - assert actors_named("TransferQueueLockManager_") == ["TransferQueueLockManager_0"] - with tq.kv_global_lock(["a", "b"], P, timeout=5) as lease: - assert isinstance(lease, tq.GlobalLease) - assert holders() == {"a", "b"} and list(lease.shards) == [0] - wait_until(lambda: holders() == set()) diff --git a/tests/test_kv_lock.py b/tests/test_kv_lock.py index 3bb480ec..7cfbbe29 100644 --- a/tests/test_kv_lock.py +++ b/tests/test_kv_lock.py @@ -18,18 +18,12 @@ import importlib import os import signal -import subprocess -import sys import threading import time from types import SimpleNamespace import pytest -import ray -import torch -from omegaconf import OmegaConf -import transfer_queue as tq from transfer_queue import async_kv_local_lock, kv_local_lock # `transfer_queue.kv_lock` resolves to the re-exported function, so fetch the module itself. @@ -37,7 +31,6 @@ P = "p" # Bound every wait so that a lock regression fails the test instead of hanging it. TIMEOUT = 10.0 -E2E_TIMEOUT = 120.0 @pytest.fixture(autouse=True) @@ -111,29 +104,6 @@ async def task_worker(): assert counter[0] == 4 * 100 + 3 * 4 * 100 -def test_thread_and_task_wake_each_other(): - order = [] - - def thread_body(): - with kv_local_lock("k", P, timeout=TIMEOUT): - order.append("thread") - wait_until(lambda: n_waiters("k") == 1) - - async def main(): - thread = threading.Thread(target=thread_body) - async with async_kv_local_lock("k", P, timeout=TIMEOUT): - thread.start() - wait_until(lambda: n_waiters("k") == 1) - order.append("task") - async with async_kv_local_lock("k", P, timeout=TIMEOUT): - order.append("task again") - thread.join(TIMEOUT) - assert not thread.is_alive() - - asyncio.run(main()) - assert order == ["task", "thread", "task again"] - - def test_waiters_are_served_fifo(): order, release = [], threading.Event() holder = start_holder("k", release) @@ -352,62 +322,14 @@ def hold_registry(): key_holder.join() -def test_global_lock_shards_agree_across_processes_and_spread_keys(): - names = [(P, f"key{i}") for i in range(800)] - shards = [kvl._shard(name, 8) for name in names] +def test_global_lock_shards_are_pinned_and_spread_keys(): + # Pinned rather than hash(): every process, whatever its PYTHONHASHSEED, must pick the same actor. + assert [kvl._shard((P, f"key{i}"), 8) for i in range(4)] == [0, 6, 4, 2] + shards = [kvl._shard((P, f"key{i}"), 8) for i in range(800)] assert all(60 < shards.count(shard) < 140 for shard in range(8)) - # str hashes differ per PYTHONHASHSEED; the shard of a key must not. - script = "import transfer_queue.kv_lock as k; print([k._shard(('p', f'key{i}'), 8) for i in range(800)])" - for seed in ("1", "2"): - out = subprocess.run( - [sys.executable, "-c", script], - env={**os.environ, "PYTHONHASHSEED": seed}, - capture_output=True, - text=True, - timeout=E2E_TIMEOUT, - check=True, - ) - assert out.stdout.strip() == str(shards) def test_global_lock_poll_window_fits_the_lease_and_the_timeout(): assert kvl._poll_window(None, 30) == kvl._POLL_S assert kvl._poll_window(None, 0.6) == pytest.approx(0.2) # a renewal waits at most a third of the lease assert kvl._poll_window(time.monotonic() + 0.05, 30) <= 0.05 - - -@pytest.fixture -def simple_storage_tq(): - if not ray.is_initialized(): - ray.init(namespace="TestKVLock") - backend = {"storage_backend": "SimpleStorage", "SimpleStorage": {"total_storage_size": 200}} - tq.init(OmegaConf.create({"controller": {"polling_mode": True}, "backend": backend})) - yield - tq.close() - ray.shutdown() - - -def test_kv_local_lock_serializes_simple_storage_read_modify_write(simple_storage_tq): - tq.kv_put("counter", "lock_e2e", fields={"v": torch.tensor([0])}) # also creates the client on this thread - - def increment(key, partition_id): - value = tq.kv_batch_get(key, partition_id)["v"] - tq.kv_put(key, partition_id, fields={"v": value[0] + 1}) - - def worker(): - for _ in range(50): - with tq.kv_local_lock("counter", "lock_e2e", timeout=E2E_TIMEOUT): - increment("counter", "lock_e2e") - - threads = [threading.Thread(target=worker) for _ in range(8)] - for thread in threads: - thread.start() - for thread in threads: - thread.join(E2E_TIMEOUT) - assert not thread.is_alive() - assert tq.kv_batch_get("counter", "lock_e2e")["v"][0].item() == 400 - tq.kv_clear("counter", "lock_e2e") - # The default config starts no lock actors: the local lock works, the global one says why not. - with pytest.raises(RuntimeError, match="lock.enabled"): - with tq.kv_global_lock("counter", "lock_e2e"): - pass From efe17ead4a684e0a7640877cf7dc7c9effc5c9ea Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Sun, 11 Oct 2026 22:03:52 +0800 Subject: [PATCH 26/42] [doc] Correct the kv_global_lock capacity note: past the limit, leases can be lost Ray serves queued calls in arrival order, so once waiting acquires take every slot of a lock actor, renewal and release wait behind all of them, not for one poll window. Say that holders can lose their leases there and drop the comments that claimed otherwise. Signed-off-by: OutstanderWang --- transfer_queue/kv_lock.py | 8 ++++---- transfer_queue/lock_manager.py | 4 ++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index cad853ef..c65cdf99 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -294,8 +294,7 @@ def _new_lease(keys: str | list[str], partition_id: str, lease_s: float) -> Glob def _poll_window(deadline: float | None, lease_s: float) -> float: - # Even when waiters take every actor slot, a renewal waits at most one window for a slot, - # so a window up to lease_s / 3 leaves it most of the lease, as the renewer's period does. + # Up to lease_s / 3, so a grant made between two polls is seen with most of its lease left. window, remaining = min(_POLL_S, lease_s / 3), _remaining(deadline) return window if remaining is None else min(window, remaining) @@ -407,8 +406,9 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No of them on the actor it is waiting at. With keys hashing evenly, about ``num_shards * max_concurrency`` requests can wait at once, so raise ``lock.num_shards`` for more in total; all waiters on one key share one actor, so raise ``lock.max_concurrency`` for a hot key. - Past the limit, calls queue in Ray: a waiting call frees its slot within its poll window - (5 s, or ``lease_s / 3`` if shorter), so renewal and release are delayed, never stuck. + Past the limit, calls queue in Ray behind the waiting ones, renewal and release included, + so holders on that actor can lose their leases and waiters can time out: keep the waiters + on one actor below ``max_concurrency``. Raises ``TimeoutError`` if the keys are not all granted within ``timeout`` seconds (``None`` waits forever); an unresponsive lock actor can stretch that by up to about diff --git a/transfer_queue/lock_manager.py b/transfer_queue/lock_manager.py index 5eb136fe..2248242e 100644 --- a/transfer_queue/lock_manager.py +++ b/transfer_queue/lock_manager.py @@ -31,8 +31,8 @@ # tq.init() sets max_concurrency from lock.max_concurrency. A waiting acquire call holds a -# slot for at most the caller's poll window, so even if waiters take every slot, renew and -# release get one within that window. +# slot for up to the caller's poll window, and renew and release queue behind every call +# ahead of them, so more waiters than slots can starve them. @ray.remote(num_cpus=0) class TransferQueueLockManager: """Exclusive leased locks on ``(partition_id, key)``. Runs on one event loop, so no thread locks.""" From 58104759766f673b8c0c3acb5db0ec30a3ab496d Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Sun, 11 Oct 2026 22:05:48 +0800 Subject: [PATCH 27/42] [fix] Keep one kv_global_lock acquire in flight per shard When a reply came later than its window plus slack, the caller sent a new acquire while the old one stayed queued in the actor and later took a slot, so a busy actor's backlog kept growing and waiters starved even after the holder released. Wait on the pending call again instead, and send the next one only after it returns. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 27 +++++++++++++++++++++++++++ transfer_queue/kv_lock.py | 20 +++++++++++++------- 2 files changed, 40 insertions(+), 7 deletions(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index efcc7ec3..786b7bc7 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -204,6 +204,33 @@ def take_both(): wait_until(lambda: holders() == set()) +@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) +def test_a_late_reply_is_awaited_instead_of_sent_again(monkeypatch, use_async): + lock_manager = manager("late") + ray.get(lock_manager.acquire.remote([(P, "late")], "holder", 0, 30, {}, 8), timeout=TIMEOUT_S) + real_send, sends = kvl._send_acquire, [] + + def send_acquire(lease, shard, window): + # The actor keeps the call past the caller's wait, as a call queued behind others would be. + sends.append(window) + return real_send(lease, shard, TIMEOUT_S) + + async def take_async(): + async with tq.async_kv_global_lock("late", P, timeout=2.5, lease_s=1.5): + pass + + monkeypatch.setattr(kvl, "_send_acquire", send_acquire) + with pytest.raises(TimeoutError): # 0.5 s polls, each awaited for up to 1.5 s + if use_async: + asyncio.run(take_async()) + else: + with tq.kv_global_lock("late", P, timeout=2.5, lease_s=1.5): + pass + assert len(sends) == 1 + ray.get(lock_manager.release.remote("holder"), timeout=TIMEOUT_S) + wait_until(lambda: tq.kv_lock_list(P) == {"holders": [], "waiters": 0}) + + def test_async_first_use_looks_the_actors_up_off_the_event_loop(monkeypatch): monkeypatch.setattr(kvl, "_num_shards", None) monkeypatch.setattr(kvl, "_managers", {}) diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index c65cdf99..3eb264e1 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -218,7 +218,7 @@ async def async_kv_local_lock(keys: str | list[str], partition_id: str, timeout: # slots, and the bound is long because one actor serves only a few thousand calls a second: # at 1 s, 8,000 idle waiters starved every other call on their shard, renewals included. _POLL_S = 5.0 -_REPLY_SLACK_S = 1.0 # how late a reply may arrive past its window before the caller re-polls +_REPLY_SLACK_S = 1.0 # how late a reply may arrive past its window before the caller checks its timeout _NOT_INITIALIZED = ( "TransferQueueLockManager not found; kv_global_lock needs `lock.enabled: true` " "in the config of the tq.init() that starts TransferQueue" @@ -426,12 +426,15 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No # Every caller visits shards in ascending order, so no two callers can each hold a # shard the other awaits: multi-shard requests cannot deadlock. for shard in lease.shards: - granted = False + call, granted = None, False while not granted: - window, sent = _poll_window(deadline, lease.lease_s), time.monotonic() + window = _poll_window(deadline, lease.lease_s) + if call is None: # a late reply is awaited again, not re-sent, so calls never pile up + call, sent = _send_acquire(lease, shard, window), time.monotonic() try: # Bounded, so an unresponsive actor cannot outlast the caller's timeout. - remaining = ray.get(_send_acquire(lease, shard, window), timeout=window + _REPLY_SLACK_S) + remaining = ray.get(call, timeout=window + _REPLY_SLACK_S) + call = None except GetTimeoutError: remaining = None except RayActorError as e: @@ -457,11 +460,14 @@ async def async_kv_global_lock( held = _mark(_holding_global) try: for shard in lease.shards: - granted = False + call, granted = None, False while not granted: - window, sent = _poll_window(deadline, lease.lease_s), time.monotonic() + window = _poll_window(deadline, lease.lease_s) + if call is None: + call, sent = asyncio.ensure_future(_send_acquire(lease, shard, window)), time.monotonic() try: - remaining = await asyncio.wait_for(_send_acquire(lease, shard, window), window + _REPLY_SLACK_S) + remaining = await asyncio.wait_for(asyncio.shield(call), window + _REPLY_SLACK_S) + call = None except asyncio.TimeoutError: remaining = None except RayActorError as e: From b256abe96f34f51ec681d462b7ed80dd58954161 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Sun, 11 Oct 2026 22:09:52 +0800 Subject: [PATCH 28/42] [fix] Reserve lock actor slots for renewal and release Waiting acquires could take every slot of a lock actor, and Ray serves queued calls in arrival order, so renewals waited behind them and holders lost their leases. lock.max_concurrency now caps the acquire calls an actor lets wait; tq.init() gives it _RESERVED_SLOTS more for calls that never wait. Past the cap an acquire keeps its queue place and returns at once, and the caller sleeps out the rest of its window before polling again. Such a caller has no call in flight for up to a window, so a waiter now counts as abandoned after 10 s instead of 5 s. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 25 +++++++++++++++++++---- tests/test_interface_close.py | 7 ++++--- transfer_queue/config.yaml | 3 ++- transfer_queue/kv_lock.py | 28 ++++++++++++++++---------- transfer_queue/lock_manager.py | 30 ++++++++++++++++++---------- 5 files changed, 63 insertions(+), 30 deletions(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index 786b7bc7..1b2c06c8 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -26,7 +26,7 @@ from omegaconf import OmegaConf import transfer_queue as tq -from transfer_queue.lock_manager import TransferQueueLockManager +from transfer_queue.lock_manager import _ABANDONED_S, TransferQueueLockManager kvl = importlib.import_module("transfer_queue.kv_lock") P = "global_lock_e2e" @@ -263,8 +263,11 @@ def test_release_before_acquire_withdraws_it(): @pytest.fixture def own_manager(): - """A lock actor of its own, so these tests control every request it sees.""" - lock_manager = TransferQueueLockManager.remote(1) + """A lock actor of its own, so these tests control every request it sees. + + Sized as tq.init() sizes one for ``lock.max_concurrency: 2``. + """ + lock_manager = TransferQueueLockManager.options(max_concurrency=2 + kvl._RESERVED_SLOTS).remote(1, 2) yield lock_manager ray.kill(lock_manager) @@ -351,12 +354,26 @@ def test_a_waiter_that_stopped_polling_is_skipped(own_manager): assert ray.get(acquire(own_manager, ["a"], "dead", timeout=0, lease_s=TIMEOUT_S), timeout=TIMEOUT_S) is None live = acquire(own_manager, ["a"], "live") wait_until(lambda: n_waiters(own_manager) == 2) - time.sleep(6) # past the manager's abandonment threshold + time.sleep(_ABANDONED_S + 1) ray.get(own_manager.release.remote("holder"), timeout=TIMEOUT_S) assert ray.get(live, timeout=5) is not None assert n_waiters(own_manager) == 0 +def test_renewal_and_release_run_while_waiters_fill_the_actor(own_manager): + ray.get(acquire(own_manager, ["a"], "holder"), timeout=TIMEOUT_S) + first, second = acquire(own_manager, ["a"], "first"), acquire(own_manager, ["a"], "second") + wait_until(lambda: n_waiters(own_manager) == 2) + # Both waiting slots are taken: a third waiter is answered at once and keeps its place. + assert ray.get(acquire(own_manager, ["a"], "third"), timeout=1) is None + assert ray.get(own_manager.renew_many.remote(["holder"]), timeout=1) == {"holder": True} + for token, call in (("holder", first), ("first", second)): + ray.get(own_manager.release.remote(token), timeout=1) + assert ray.get(call, timeout=TIMEOUT_S) is not None + ray.get(own_manager.release.remote("second"), timeout=1) + assert ray.get(acquire(own_manager, ["a"], "third"), timeout=1) > 0 + + def test_killed_holder_frees_the_key_after_its_lease(): worker = Worker.remote() ray.get(worker.hold.remote("k", lease_s=1.5), timeout=TIMEOUT_S) diff --git a/tests/test_interface_close.py b/tests/test_interface_close.py index 37f6d7e6..29843101 100644 --- a/tests/test_interface_close.py +++ b/tests/test_interface_close.py @@ -79,12 +79,13 @@ def test_init_rejects_invalid_lock_sizes(fake, field, value): controller_cls.options.assert_not_called() -def test_init_gives_every_lock_actor_the_configured_max_concurrency(fake): +def test_init_sizes_every_lock_actor_from_max_concurrency(fake): ray, _, _ = fake ray.get_actor.side_effect = ValueError("no controller yet") iface.init(OmegaConf.create({"lock": {"enabled": True, "num_shards": 3, "max_concurrency": 64}})) - calls = kvl.TransferQueueLockManager.options.call_args_list - assert [c.kwargs["max_concurrency"] for c in calls] == [64, 64, 64] + options = kvl.TransferQueueLockManager.options + assert [c.kwargs["max_concurrency"] for c in options.call_args_list] == [64 + kvl._RESERVED_SLOTS] * 3 + assert [c.args for c in options.return_value.remote.call_args_list] == [(3, 64)] * 3 iface.close() diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index d9dba90b..8a1f0f05 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -21,7 +21,8 @@ lock: enabled: false # Number of lock actors that kv_global_lock spreads keys over by hash (>= 1) num_shards: 8 - # Calls each lock actor serves at once, waiting acquires included (>= 1); see kv_global_lock + # Acquire calls each lock actor lets wait at once (>= 1); renewal and release always have + # spare slots, and a request past this polls instead. See kv_global_lock max_concurrency: 10000 diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 3eb264e1..4e2d337a 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -40,7 +40,7 @@ import ray from ray.exceptions import GetTimeoutError, RayActorError -from transfer_queue.lock_manager import TransferQueueLockManager +from transfer_queue.lock_manager import _RESERVED_SLOTS, TransferQueueLockManager # Not asyncio.Lock: it is bound to one event loop and is not thread-safe, while a process # may run several loops plus plain threads. A key is held exactly while it is in _table, @@ -294,7 +294,8 @@ def _new_lease(keys: str | list[str], partition_id: str, lease_s: float) -> Glob def _poll_window(deadline: float | None, lease_s: float) -> float: - # Up to lease_s / 3, so a grant made between two polls is seen with most of its lease left. + # Up to lease_s / 3, so a grant made between two polls, e.g. to a caller a full actor turned + # away, is seen with most of its lease left. window, remaining = min(_POLL_S, lease_s / 3), _remaining(deadline) return window if remaining is None else min(window, remaining) @@ -402,13 +403,14 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No afresh at each later actor, so its waits add up, and the keys it holds meanwhile block others even if it then times out. Lock the fewest keys you can and keep the block short. - Each actor serves ``lock.max_concurrency`` calls at once, and a waiting request takes one - of them on the actor it is waiting at. With keys hashing evenly, about ``num_shards * - max_concurrency`` requests can wait at once, so raise ``lock.num_shards`` for more in total; - all waiters on one key share one actor, so raise ``lock.max_concurrency`` for a hot key. - Past the limit, calls queue in Ray behind the waiting ones, renewal and release included, - so holders on that actor can lose their leases and waiters can time out: keep the waiters - on one actor below ``max_concurrency``. + A waiting request keeps one call waiting on the actor it is waiting at, and each actor lets + at most ``lock.max_concurrency`` calls wait at once; renewal and release have slots of + their own, so waiting requests cannot hold them up. Past the limit, a request keeps its + place in the queue but polls once per window (5 s, or ``lease_s / 3`` if shorter), so it + may see its grant up to a window late. With keys hashing evenly, about ``num_shards * + max_concurrency`` requests can wait without that delay, so raise ``lock.num_shards`` for + more in total; all waiters on one key share one actor, so raise ``lock.max_concurrency`` + for a hot key. Raises ``TimeoutError`` if the keys are not all granted within ``timeout`` seconds (``None`` waits forever); an unresponsive lock actor can stretch that by up to about @@ -440,6 +442,8 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No except RayActorError as e: raise RuntimeError(_MANAGER_GONE) from e granted = _take_grant(lease, shard, remaining, sent, deadline) + if not granted: # a reply before the window ends means the actor was full + time.sleep(max(0.0, sent + window - time.monotonic())) lease.check() yield lease lease.check() @@ -473,6 +477,8 @@ async def async_kv_global_lock( except RayActorError as e: raise RuntimeError(_MANAGER_GONE) from e granted = _take_grant(lease, shard, remaining, sent, deadline) + if not granted: + await asyncio.sleep(max(0.0, sent + window - time.monotonic())) lease.check() yield lease lease.check() @@ -509,8 +515,8 @@ def _start_lock_managers(lock_conf) -> None: name=_manager_name(shard), namespace="transfer_queue", get_if_exists=True, - max_concurrency=lock_conf.max_concurrency, - ).remote(lock_conf.num_shards) + max_concurrency=lock_conf.max_concurrency + _RESERVED_SLOTS, + ).remote(lock_conf.num_shards, lock_conf.max_concurrency) for shard in range(lock_conf.num_shards) ] diff --git a/transfer_queue/lock_manager.py b/transfer_queue/lock_manager.py index 2248242e..bc926712 100644 --- a/transfer_queue/lock_manager.py +++ b/transfer_queue/lock_manager.py @@ -24,22 +24,25 @@ # A withdrawal whose acquire never arrives (e.g. the caller died first) is dropped after this. _WITHDRAWN_TTL_S = 300 -# A queued caller sends its next call as soon as one returns, so a waiter with no call in -# flight for this long has died. Hand-off skips it instead of granting it a lease that would only expire; if it -# was merely stalled, its next call queues it again at the back. -_ABANDONED_S = 5.0 +# A queued caller calls again at least once per poll window (5 s at most), even when a full +# actor turned it away, so a waiter with no call in flight for twice that has died. Hand-off +# skips it instead of granting it a lease that would only expire; if it was merely stalled, +# its next call queues it again at the back. +_ABANDONED_S = 10.0 +# Slots beyond lock.max_concurrency for calls that never wait: renew, release, list_locks and +# acquires turned away. Parked acquires never take them, so renewal and release always run. +_RESERVED_SLOTS = 16 -# tq.init() sets max_concurrency from lock.max_concurrency. A waiting acquire call holds a -# slot for up to the caller's poll window, and renew and release queue behind every call -# ahead of them, so more waiters than slots can starve them. +# tq.init() runs each actor with lock.max_concurrency + _RESERVED_SLOTS slots. @ray.remote(num_cpus=0) class TransferQueueLockManager: """Exclusive leased locks on ``(partition_id, key)``. Runs on one event loop, so no thread locks.""" - def __init__(self, num_shards: int): + def __init__(self, num_shards: int, max_parked: int): # Kept so that a process that never ran tq.init() can learn how many shards to hash over. self._num_shards = num_shards + self._max_parked, self._parked = max_parked, 0 # acquire calls waiting inside the actor self._holder: dict[tuple[str, str], str] = {} # name -> token self._leases: dict[str, dict] = {} # token -> names, lease_s, expires_at, granted_at, info # Min-heap of (expires_at, token). Release and renewal leave stale entries behind, which @@ -120,9 +123,10 @@ def _expire(self) -> None: async def acquire(self, names, token, timeout, lease_s, holder_info, num_shards) -> float | None: """Grant all ``names`` to ``token`` at once; return the seconds left on its lease. - Returns None if still queued after ``timeout``. Calling again with the same token keeps - its place in the queues; only ``release`` withdraws it. A token that is neither queued - nor held (its grant expired before the caller saw it) queues afresh. + Returns None if still queued after ``timeout``, or at once if ``max_parked`` calls are + already waiting. Calling again with the same token keeps its place in the queues; only + ``release`` withdraws it. A token that is neither queued nor held (its grant expired + before the caller saw it) queues afresh. """ # A caller that hashed over another shard count may send a key to the wrong shard, # where nothing excludes the callers that sent it to the right one. @@ -144,13 +148,17 @@ async def acquire(self, names, token, timeout, lease_s, holder_info, num_shards) self._queues.setdefault(name, deque()).append(token) waiter = self._waiting.get(token) if waiter is not None: + if self._parked >= self._max_parked: + timeout = 0 # keep the reserved slots free: the caller polls again after its window waiter["calls"] += 1 + self._parked += 1 try: await asyncio.wait_for(waiter["done"].wait(), timeout) except asyncio.TimeoutError: pass finally: waiter["calls"] -= 1 + self._parked -= 1 waiter["seen"] = time.monotonic() lease = self._leases.get(token) return None if lease is None else lease["expires_at"] - time.monotonic() From f61aded93c2e8a787b8d9ff89410677ef97063bb Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Sun, 11 Oct 2026 22:51:52 +0800 Subject: [PATCH 29/42] [test] Queue the slot test's waiters one at a time Ray may run two calls to an async actor out of order, so the second waiter was sometimes queued first and granted on the holder's release, leaving the first to wait out a 30 s lease. Queue them in turn, as the arrival-order test does. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index 1b2c06c8..c3bcd6ce 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -362,7 +362,9 @@ def test_a_waiter_that_stopped_polling_is_skipped(own_manager): def test_renewal_and_release_run_while_waiters_fill_the_actor(own_manager): ray.get(acquire(own_manager, ["a"], "holder"), timeout=TIMEOUT_S) - first, second = acquire(own_manager, ["a"], "first"), acquire(own_manager, ["a"], "second") + first = acquire(own_manager, ["a"], "first") + wait_until(lambda: n_waiters(own_manager) == 1) # an async actor may run two calls out of order + second = acquire(own_manager, ["a"], "second") wait_until(lambda: n_waiters(own_manager) == 2) # Both waiting slots are taken: a third waiter is answered at once and keeps its place. assert ray.get(acquire(own_manager, ["a"], "third"), timeout=1) is None From 9261e6ae07e79bd6da4d7723d6b51c21c2f83d18 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Sun, 11 Oct 2026 22:51:59 +0800 Subject: [PATCH 30/42] [test] Keep only the core kv lock tests Keep ten end-to-end tests for kv_global_lock (exclusion without deadlock, multi-shard timeout cleanup, queue order, unseen expired grants, dead waiters, reserved slots, late replies, killed holders, lease renewal and loss, and the owner's close) and three for kv_local_lock (exclusion across threads and loops, FIFO order, and timeout cleanup). Drop the rest, including tests/test_interface_close.py. The dead-waiter test now backdates the waiter's last call instead of sleeping past the threshold. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 205 ++------------------------ tests/test_interface_close.py | 115 --------------- tests/test_kv_lock.py | 207 --------------------------- 3 files changed, 8 insertions(+), 519 deletions(-) delete mode 100644 tests/test_interface_close.py diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index c3bcd6ce..4fa80bb6 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -13,11 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""kv_global_lock across Ray actors: exclusion, expiry, withdrawal, lost leases, and close() ownership.""" +"""kv_global_lock across Ray actors: exclusion, timeouts, queue order, leases, and close().""" -import asyncio import importlib -import threading import time import pytest @@ -26,7 +24,7 @@ from omegaconf import OmegaConf import transfer_queue as tq -from transfer_queue.lock_manager import _ABANDONED_S, TransferQueueLockManager +from transfer_queue.lock_manager import TransferQueueLockManager kvl = importlib.import_module("transfer_queue.kv_lock") P = "global_lock_e2e" @@ -63,9 +61,6 @@ def hold(self, key, lease_s=30): def release(self): self.held.__exit__(None, None, None) - def close(self): - tq.close() - @pytest.fixture(scope="module", autouse=True) def tq_session(): @@ -133,79 +128,7 @@ def test_multi_shard_timeout_releases_the_earlier_shard(): wait_until(lambda: tq.kv_lock_list(P) == {"holders": [], "waiters": 0}) -def test_lease_lost_on_one_shard_raises(): - early, late = keys_on_two_shards("ll") - with pytest.raises(tq.LockLostError): - with tq.kv_global_lock([early, late], P, lease_s=1) as lease: - assert holders() == {early, late} and len(lease.shards) == 2 - ray.get(manager(late).release.remote(lease.token), timeout=TIMEOUT_S) - wait_until(lambda: lease.lost) - with pytest.raises(tq.LockLostError): - lease.check() - wait_until(lambda: holders() == set()) - - -@ray.remote -def lock_without_init(keys): - with tq.kv_global_lock(keys, P, timeout=TIMEOUT_S) as lease: - return list(lease.shards) - - -def test_process_without_init_learns_the_shard_count(): - a, b = keys_on_two_shards("ni") - assert ray.get(lock_without_init.remote([a, b]), timeout=TIMEOUT_S) == [shard_of(a), shard_of(b)] - wait_until(lambda: holders() == set()) - - -def test_cancelled_waiter_is_withdrawn_and_never_granted(): - worker = Worker.remote() - ray.get(worker.hold.remote("c"), timeout=TIMEOUT_S) - - async def wait_for_lock(): - async with tq.async_kv_global_lock("c", P): - pass - - async def main(): - task = asyncio.create_task(wait_for_lock()) - deadline = time.monotonic() + TIMEOUT_S - while tq.kv_lock_list(P)["waiters"] != 1: - assert time.monotonic() < deadline, "the task never started waiting" - await asyncio.sleep(0.05) - task.cancel() - with pytest.raises(asyncio.CancelledError): - await asyncio.wait_for(task, TIMEOUT_S) - - asyncio.run(main()) - wait_until(lambda: tq.kv_lock_list(P)["waiters"] == 0) - ray.get(worker.release.remote(), timeout=TIMEOUT_S) - # A ghost grant to the cancelled waiter would hold "c" for its whole 30 s lease. - with tq.kv_global_lock("c", P, timeout=5): - pass - - -def test_waiting_on_a_later_shard_keeps_the_earlier_keys_held(): - early, late = keys_on_two_shards("lw") - worker = Worker.remote() - ray.get(worker.hold.remote(late), timeout=TIMEOUT_S) - entered = [] - - def take_both(): - with tq.kv_global_lock([early, late], P, timeout=TIMEOUT_S, lease_s=1) as lease: - entered.append(lease) - - thread = threading.Thread(target=take_both) - thread.start() - wait_until(lambda: holders() == {early, late}) - time.sleep(1.5) # longer than our lease: unrenewed, the grant on `early` would have expired - assert holders() == {early, late} - ray.get(worker.release.remote(), timeout=TIMEOUT_S) - thread.join(TIMEOUT_S) - assert len(entered) == 1 # entered without LockLostError - wait_until(lambda: holders() == set()) - - -@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) -def test_a_late_reply_is_awaited_instead_of_sent_again(monkeypatch, use_async): +def test_a_late_reply_is_awaited_instead_of_sent_again(monkeypatch): lock_manager = manager("late") ray.get(lock_manager.acquire.remote([(P, "late")], "holder", 0, 30, {}, 8), timeout=TIMEOUT_S) real_send, sends = kvl._send_acquire, [] @@ -215,52 +138,15 @@ def send_acquire(lease, shard, window): sends.append(window) return real_send(lease, shard, TIMEOUT_S) - async def take_async(): - async with tq.async_kv_global_lock("late", P, timeout=2.5, lease_s=1.5): - pass - monkeypatch.setattr(kvl, "_send_acquire", send_acquire) with pytest.raises(TimeoutError): # 0.5 s polls, each awaited for up to 1.5 s - if use_async: - asyncio.run(take_async()) - else: - with tq.kv_global_lock("late", P, timeout=2.5, lease_s=1.5): - pass + with tq.kv_global_lock("late", P, timeout=2.5, lease_s=1.5): + pass assert len(sends) == 1 ray.get(lock_manager.release.remote("holder"), timeout=TIMEOUT_S) wait_until(lambda: tq.kv_lock_list(P) == {"holders": [], "waiters": 0}) -def test_async_first_use_looks_the_actors_up_off_the_event_loop(monkeypatch): - monkeypatch.setattr(kvl, "_num_shards", None) - monkeypatch.setattr(kvl, "_managers", {}) - on_loop = [] - for name in ("get", "get_actor"): - real = getattr(ray, name) - - def spy(*args, _real=real, **kwargs): - on_loop.append(threading.current_thread() is threading.main_thread()) - return _real(*args, **kwargs) - - monkeypatch.setattr(ray, name, spy) - - async def main(): - async with tq.async_kv_global_lock("first_use", P, timeout=TIMEOUT_S): - pass - - asyncio.run(main()) - assert on_loop and not any(on_loop) - assert len(kvl._managers) == kvl._num_shards == 8 - - -def test_release_before_acquire_withdraws_it(): - lock_manager = manager("w") - ray.get(lock_manager.release.remote("early"), timeout=TIMEOUT_S) - for _ in range(2): # a retried call that arrives late is withdrawn too - assert ray.get(lock_manager.acquire.remote([(P, "w")], "early", None, 5, {}, 8), timeout=TIMEOUT_S) is None - assert holders() == set() and tq.kv_lock_list(P)["waiters"] == 0 - - @pytest.fixture def own_manager(): """A lock actor of its own, so these tests control every request it sees. @@ -280,13 +166,6 @@ def n_waiters(lock_manager): return ray.get(lock_manager.list_locks.remote(P), timeout=TIMEOUT_S)["waiters"] -def test_acquire_rejects_a_caller_hashing_over_another_shard_count(own_manager): - # e.g. a process that attached before TransferQueue restarted with a new lock.num_shards - with pytest.raises(ValueError, match="lock.num_shards is 1"): - ray.get(acquire(own_manager, ["a"], "stale", num_shards=8), timeout=TIMEOUT_S) - assert ray.get(acquire(own_manager, ["a"], "fresh", timeout=0), timeout=TIMEOUT_S) is not None - - def test_waiters_are_granted_in_arrival_order(own_manager): ray.get(acquire(own_manager, ["a"], "holder"), timeout=TIMEOUT_S) both = acquire(own_manager, ["a", "b"], "both") @@ -303,18 +182,6 @@ def test_waiters_are_granted_in_arrival_order(own_manager): assert ray.get(only_b, timeout=TIMEOUT_S) is not None -def test_a_timed_out_call_keeps_its_place_until_withdrawn(own_manager): - ray.get(acquire(own_manager, ["a"], "holder"), timeout=TIMEOUT_S) - assert ray.get(acquire(own_manager, ["a"], "first", timeout=0.1), timeout=TIMEOUT_S) is None - second = acquire(own_manager, ["a"], "second") - wait_until(lambda: n_waiters(own_manager) == 2) - ray.get(own_manager.release.remote("holder"), timeout=TIMEOUT_S) - assert ray.get(acquire(own_manager, ["a"], "first"), timeout=TIMEOUT_S) > 0 # calling again keeps the place - assert n_waiters(own_manager) == 1 - ray.get(own_manager.release.remote("first"), timeout=TIMEOUT_S) - assert ray.get(second, timeout=TIMEOUT_S) is not None - - def test_a_grant_that_expires_before_the_caller_sees_it_queues_afresh(own_manager): ray.get(acquire(own_manager, ["a"], "holder"), timeout=TIMEOUT_S) assert ray.get(acquire(own_manager, ["a"], "slow", timeout=0, lease_s=0.3), timeout=TIMEOUT_S) is None @@ -329,32 +196,15 @@ def test_a_grant_that_expires_before_the_caller_sees_it_queues_afresh(own_manage assert ray.get(acquire(own_manager, ["a"], "slow"), timeout=TIMEOUT_S) > 0 -def test_a_waiter_that_withdraws_lets_the_next_one_through(own_manager): - ray.get(acquire(own_manager, ["a"], "holder"), timeout=TIMEOUT_S) - assert ray.get(acquire(own_manager, ["a", "b"], "both", timeout=0.1), timeout=TIMEOUT_S) is None - only_b = acquire(own_manager, ["b"], "only_b") - wait_until(lambda: n_waiters(own_manager) == 2) - ray.get(own_manager.release.remote("both"), timeout=TIMEOUT_S) - assert ray.get(only_b, timeout=TIMEOUT_S) is not None - - -def test_expired_leases_pass_the_key_down_the_queue(own_manager): - # Nobody releases: each grant has to come from the actor expiring the lease ahead of it. - ray.get(acquire(own_manager, ["a"], "dead", lease_s=0.5), timeout=TIMEOUT_S) - first = acquire(own_manager, ["a"], "first", lease_s=0.5) - wait_until(lambda: n_waiters(own_manager) == 1) - second = acquire(own_manager, ["a"], "second") - assert ray.get(second, timeout=TIMEOUT_S) is not None - assert ray.get(first, timeout=TIMEOUT_S) is not None - - def test_a_waiter_that_stopped_polling_is_skipped(own_manager): ray.get(acquire(own_manager, ["a"], "holder"), timeout=TIMEOUT_S) # Queued once, then its caller died: granting it would idle the key for its whole lease. assert ray.get(acquire(own_manager, ["a"], "dead", timeout=0, lease_s=TIMEOUT_S), timeout=TIMEOUT_S) is None live = acquire(own_manager, ["a"], "live") wait_until(lambda: n_waiters(own_manager) == 2) - time.sleep(_ABANDONED_S + 1) + # Date its last call back past the abandonment threshold instead of sleeping through it. + backdate = own_manager.__ray_call__.remote(lambda self: self._waiting["dead"].update(seen=float("-inf"))) + ray.get(backdate, timeout=TIMEOUT_S) ray.get(own_manager.release.remote("holder"), timeout=TIMEOUT_S) assert ray.get(live, timeout=5) is not None assert n_waiters(own_manager) == 0 @@ -406,49 +256,10 @@ def test_lease_renews_and_lost_lease_raises(): wait_until(lambda: holders() == set()) -def test_nesting_and_ordering_rules(): - with pytest.raises(ValueError): - with tq.kv_global_lock([], P): - pass - with tq.kv_global_lock("n", P): - with pytest.raises(RuntimeError, match="Nested kv_global_lock"): - with tq.kv_global_lock("m", P): - pass - with tq.kv_local_lock("n", P): # global, then local is allowed - pass - with tq.kv_local_lock("n", P): - with pytest.raises(RuntimeError, match="before kv_local_lock"): - with tq.kv_global_lock("n", P): - pass - - async def main(): - with pytest.raises(RuntimeError, match="async_kv_global_lock"): - with tq.kv_global_lock("n", P): - pass - async with tq.async_kv_global_lock("n", P): - with pytest.raises(RuntimeError, match="Nested kv_global_lock"): - async with tq.async_kv_global_lock("m", P): - pass - - asyncio.run(main()) - wait_until(lambda: holders() == set()) - - def actors_named(prefix): return [a["name"] for a in ray.util.list_named_actors(all_namespaces=True) if a["name"].startswith(prefix)] -def test_non_owner_close_releases_its_locks_and_keeps_the_managers(): - worker = Worker.remote() - ray.get(worker.hold.remote("o"), timeout=TIMEOUT_S) - ray.get(worker.close.remote(), timeout=TIMEOUT_S) - wait_until(lambda: holders() == set()) - with tq.kv_global_lock("o", P, timeout=5): - pass - assert len(actors_named("TransferQueueLockManager_")) == kvl._shard_count() == 8 - assert actors_named("TransferQueueController") == ["TransferQueueController"] - - def test_owner_close_kills_the_managers_and_fails_waiters(): """Tears TransferQueue down for this module, so it runs last.""" held = tq.kv_global_lock("z", P) diff --git a/tests/test_interface_close.py b/tests/test_interface_close.py deleted file mode 100644 index 29843101..00000000 --- a/tests/test_interface_close.py +++ /dev/null @@ -1,115 +0,0 @@ -# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved. -# Copyright 2025 The TransferQueue Team -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Unit tests for `close()` ownership: only the process that created TransferQueue tears it down.""" - -import importlib -from unittest.mock import MagicMock - -import pytest -from omegaconf import OmegaConf - -import transfer_queue.interface as iface - -kvl = importlib.import_module("transfer_queue.kv_lock") - -STORED_CONF = OmegaConf.create( - {"controller": {"zmq_info": None}, "backend": {"storage_backend": "SimpleStorage", "SimpleStorage": {}}} -) - - -@pytest.fixture -def fake(monkeypatch): - for name in ("_TQ_CLIENT", "_TQ_STORAGE", "_TQ_CONTROLLER"): - monkeypatch.setattr(iface, name, None) - monkeypatch.setattr(iface, "_TQ_IS_OWNER", False) - monkeypatch.setattr(kvl, "_owned", []) - - controller = MagicMock(name="controller") - ray = MagicMock(name="ray") - ray.get.return_value = STORED_CONF - controller_cls = MagicMock(name="TransferQueueController") - controller_cls.options.return_value.remote.return_value = controller - monkeypatch.setattr(iface, "ray", ray) - monkeypatch.setattr(kvl, "ray", ray) - monkeypatch.setattr(iface, "TransferQueueController", controller_cls) - monkeypatch.setattr(kvl, "TransferQueueLockManager", MagicMock(name="TransferQueueLockManager")) - monkeypatch.setattr(iface, "TransferQueueClient", MagicMock(name="TransferQueueClient")) - monkeypatch.setattr(iface, "process_zmq_server_info", MagicMock(return_value=None)) - monkeypatch.setattr(iface, "_maybe_create_tq_storage", lambda conf: conf) - return ray, controller, controller_cls - - -@pytest.mark.parametrize("lock_enabled", [False, True], ids=["locks_off", "locks_on"]) -def test_owner_close_kills_controller(fake, lock_enabled): - ray, controller, _ = fake - ray.get_actor.side_effect = ValueError("no controller yet") - - iface.init(OmegaConf.create({"lock": {"enabled": lock_enabled}})) - client, lock_managers = iface._TQ_CLIENT, kvl._owned - assert iface._TQ_IS_OWNER and len(lock_managers) == (8 if lock_enabled else 0) - - iface.close() - client.close.assert_called_once() - assert [c.args for c in ray.kill.call_args_list] == [(m,) for m in lock_managers] + [(controller,)] - assert (iface._TQ_CLIENT, iface._TQ_CONTROLLER, iface._TQ_IS_OWNER) == (None, None, False) - assert kvl._owned == [] - - -@pytest.mark.parametrize( - "field, value", [("num_shards", 0), ("num_shards", "8"), ("max_concurrency", 0), ("max_concurrency", 1.5)] -) -def test_init_rejects_invalid_lock_sizes(fake, field, value): - ray, _, controller_cls = fake - ray.get_actor.side_effect = ValueError("no controller yet") - with pytest.raises(ValueError, match=f"lock.{field}"): - iface.init(OmegaConf.create({"lock": {"enabled": True, field: value}})) - controller_cls.options.assert_not_called() - - -def test_init_sizes_every_lock_actor_from_max_concurrency(fake): - ray, _, _ = fake - ray.get_actor.side_effect = ValueError("no controller yet") - iface.init(OmegaConf.create({"lock": {"enabled": True, "num_shards": 3, "max_concurrency": 64}})) - options = kvl.TransferQueueLockManager.options - assert [c.kwargs["max_concurrency"] for c in options.call_args_list] == [64 + kvl._RESERVED_SLOTS] * 3 - assert [c.args for c in options.return_value.remote.call_args_list] == [(3, 64)] * 3 - iface.close() - - -@pytest.mark.parametrize("lost_creation_race", [False, True], ids=["existing", "lost_creation_race"]) -def test_attaching_close_keeps_controller_and_can_reattach(fake, lost_creation_race): - ray, controller, controller_cls = fake - if lost_creation_race: - # Another process creates the named actor between our lookup and our creation attempt. - ray.get_actor.side_effect = [ValueError("no controller yet"), controller] - controller_cls.options.return_value.remote.side_effect = ValueError("name already taken") - else: - ray.get_actor.return_value = controller - - iface.init(OmegaConf.create({})) - client = iface._TQ_CLIENT - assert iface._TQ_CONTROLLER is controller and not iface._TQ_IS_OWNER - - iface.close() - client.close.assert_called_once() - ray.kill.assert_not_called() - assert (iface._TQ_CLIENT, iface._TQ_CONTROLLER, iface._TQ_IS_OWNER) == (None, None, False) - - ray.get_actor.side_effect = None - ray.get_actor.return_value = controller - iface.init() - assert iface._TQ_CONTROLLER is controller and iface._TQ_CLIENT is not None - assert not iface._TQ_IS_OWNER diff --git a/tests/test_kv_lock.py b/tests/test_kv_lock.py index 7cfbbe29..d9ed2ae9 100644 --- a/tests/test_kv_lock.py +++ b/tests/test_kv_lock.py @@ -14,13 +14,9 @@ # limitations under the License. import asyncio -import gc import importlib -import os -import signal import threading import time -from types import SimpleNamespace import pytest @@ -51,13 +47,6 @@ def wait_until(predicate): time.sleep(0.001) -async def async_wait_until(predicate): - deadline = time.monotonic() + TIMEOUT - while not predicate(): - assert time.monotonic() < deadline, "condition not reached" - await asyncio.sleep(0.001) - - def start_holder(keys, release): """Hold `keys` on a new thread until `release` is set; return once they are held.""" entered = threading.Event() @@ -123,46 +112,6 @@ def waiter(i): assert order == list(range(5)) -def test_release_hands_off_directly_so_a_newcomer_cannot_barge(): - outcome = [] - - def holder_then_newcomer(): - with kv_local_lock("k", P, timeout=TIMEOUT): - wait_until(lambda: n_waiters("k") == 1) - try: - with kv_local_lock("k", P, timeout=0): - outcome.append("barged") - except TimeoutError: - outcome.append("queued behind the woken waiter") - - thread = threading.Thread(target=holder_then_newcomer) - thread.start() - wait_until(lambda: (P, "k") in kvl._table) - with kv_local_lock("k", P, timeout=TIMEOUT): - thread.join(TIMEOUT) - assert not thread.is_alive() - assert outcome == ["queued behind the woken waiter"] - - -def test_timeout_raises_and_leaves_no_waiter(): - release = threading.Event() - holder = start_holder("k", release) - - async def take(): - async with async_kv_local_lock("k", P, timeout=0.05): - pass - - with pytest.raises(TimeoutError): - with kv_local_lock("k", P, timeout=0.05): - pass - with pytest.raises(TimeoutError): - asyncio.run(take()) - assert n_waiters("k") == 0 - release.set() - holder.join(TIMEOUT) - assert not holder.is_alive() - - def test_timeout_is_a_total_deadline_and_partial_acquisition_is_released(): release_a, release_b = threading.Event(), threading.Event() holders = [start_holder("a", release_a), start_holder("b", release_b)] @@ -177,159 +126,3 @@ def test_timeout_is_a_total_deadline_and_partial_acquisition_is_released(): for holder in holders: holder.join(TIMEOUT) assert not holder.is_alive() - - -def test_grant_racing_a_timeout_counts_as_acquired(monkeypatch): - release = threading.Event() - holder = start_holder("k", release) - - class GrantedThenTimedOut(threading.Event): - def wait(self, timeout=None): - release.set() - holder.join(TIMEOUT) # the holder hands the key to this waiter... - return False # ...but the wait still reports a timeout - - monkeypatch.setattr(kvl, "threading", SimpleNamespace(Event=GrantedThenTimedOut)) - with kv_local_lock("k", P, timeout=1): - assert (P, "k") in kvl._table and n_waiters("k") == 0 - - -def test_cancellation_never_strands_the_key(): - release, order = threading.Event(), [] - holder = start_holder("k", release) - - async def take(tag): - async with async_kv_local_lock("k", P, timeout=TIMEOUT): - order.append(tag) - - async def main(): - waiting = asyncio.create_task(take("cancelled while waiting")) - await async_wait_until(lambda: n_waiters("k") == 1) - waiting.cancel() - with pytest.raises(asyncio.CancelledError): - await asyncio.wait_for(waiting, TIMEOUT) - assert n_waiters("k") == 0 - - granted, nxt = asyncio.create_task(take("cancelled after grant")), asyncio.create_task(take("next")) - await async_wait_until(lambda: n_waiters("k") == 2) - release.set() - holder.join(TIMEOUT) # blocks this loop, so `granted` is cancelled before it can run - granted.cancel() - await asyncio.wait_for(asyncio.gather(granted, nxt, return_exceptions=True), TIMEOUT) - assert granted.cancelled() - - asyncio.run(main()) - assert order == ["next"] - - -def test_waiter_on_a_closed_loop_is_skipped(): - release, taken = threading.Event(), [] - holder = start_holder("k", release) - - async def take(): - async with async_kv_local_lock("k", P, timeout=TIMEOUT): - taken.append("abandoned task") - - def thread_waiter(): - with kv_local_lock("k", P, timeout=TIMEOUT): - taken.append("thread") - - loop = asyncio.new_event_loop() - loop.create_task(take()) - loop.run_until_complete(async_wait_until(lambda: n_waiters("k") == 1)) - loop.close() # leaves the waiting task pending on a closed loop - thread = threading.Thread(target=thread_waiter) - thread.start() - wait_until(lambda: n_waiters("k") == 2) - release.set() - for t in (holder, thread): - t.join(TIMEOUT) - assert not t.is_alive() - assert taken == ["thread"] - gc.collect() # finalizing the abandoned task must leave the table alone - - -def test_invalid_and_nested_use_is_rejected(): - with pytest.raises(ValueError): - with kv_local_lock([], P): - pass - with kv_local_lock(["a", "a", "b"], P): - assert set(kvl._table) == {(P, "a"), (P, "b")} - for keys in ("a", "c"): - with pytest.raises(RuntimeError, match="Nested"): - with kv_local_lock(keys, P, timeout=TIMEOUT): - pass - - async def main(): - with pytest.raises(RuntimeError, match="async_kv_local_lock"): - with kv_local_lock("a", P): - pass - async with async_kv_local_lock("a", P): - with pytest.raises(RuntimeError, match="Nested"): - async with async_kv_local_lock("b", P): - pass - - async def child(): - async with async_kv_local_lock("c", P): - pass - - with pytest.raises(RuntimeError, match="Nested"): - await asyncio.create_task(child()) - later = asyncio.create_task(child()) - await later # starts after the block exited, so it no longer counts as nested - - asyncio.run(main()) - - -@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") -@pytest.mark.filterwarnings("ignore:.*fork.*:DeprecationWarning") -def test_forked_child_drops_the_parents_locks_leases_and_managers(monkeypatch): - monkeypatch.setattr(kvl, "_registry", threading.Condition()) - monkeypatch.setattr(kvl, "_leases", {"parent-token": object()}) - monkeypatch.setattr(kvl, "_renewer_running", True) - monkeypatch.setattr(kvl, "_managers", {0: object()}) - release, registry_held = threading.Event(), threading.Event() - - def hold_registry(): - with kvl._registry: - registry_held.set() - release.wait(TIMEOUT) - - registry_holder = threading.Thread(target=hold_registry) - registry_holder.start() - assert registry_held.wait(TIMEOUT) - key_holder = start_holder("a", release) - try: - pid = os.fork() - if pid == 0: - ok = False - try: - with kv_local_lock("a", P, timeout=1): - ok = kvl._leases == {} and kvl._managers == {} and not kvl._renewer_running - ok = ok and kvl._registry.acquire(blocking=False) - finally: - os._exit(0 if ok else 1) - deadline = time.monotonic() + TIMEOUT - while (status := os.waitpid(pid, os.WNOHANG))[0] == 0: - if time.monotonic() > deadline: - os.kill(pid, signal.SIGKILL) - pytest.fail("forked child hung") - time.sleep(0.01) - assert os.waitstatus_to_exitcode(status[1]) == 0 - finally: - release.set() - registry_holder.join() - key_holder.join() - - -def test_global_lock_shards_are_pinned_and_spread_keys(): - # Pinned rather than hash(): every process, whatever its PYTHONHASHSEED, must pick the same actor. - assert [kvl._shard((P, f"key{i}"), 8) for i in range(4)] == [0, 6, 4, 2] - shards = [kvl._shard((P, f"key{i}"), 8) for i in range(800)] - assert all(60 < shards.count(shard) < 140 for shard in range(8)) - - -def test_global_lock_poll_window_fits_the_lease_and_the_timeout(): - assert kvl._poll_window(None, 30) == kvl._POLL_S - assert kvl._poll_window(None, 0.6) == pytest.approx(0.2) # a renewal waits at most a third of the lease - assert kvl._poll_window(time.monotonic() + 0.05, 30) <= 0.05 From 2b8ffba06fa46e65be2c03a748b7d0d10c36ce0a Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 12 Oct 2026 08:40:42 +0800 Subject: [PATCH 31/42] [doc] Correct kv_global_lock docs that overstate its guarantees Signed-off-by: OutstanderWang --- transfer_queue/config.yaml | 10 +++++--- transfer_queue/interface.py | 3 ++- transfer_queue/kv_lock.py | 43 +++++++++++++++++----------------- transfer_queue/lock_manager.py | 27 +++++++++++---------- 4 files changed, 46 insertions(+), 37 deletions(-) diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index 8a1f0f05..db3cb78c 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -19,10 +19,14 @@ controller: lock: # Whether tq.init() starts the lock actors; kv_global_lock raises while this is false. enabled: false - # Number of lock actors that kv_global_lock spreads keys over by hash (>= 1) + # Number of lock actors that kv_global_lock spreads keys over by hash (>= 1). With keys + # hashing evenly, about num_shards * max_concurrency requests can wait without polling delay. num_shards: 8 - # Acquire calls each lock actor lets wait at once (>= 1); renewal and release always have - # spare slots, and a request past this polls instead. See kv_global_lock + # Acquire calls each lock actor lets wait at once (>= 1). Past this, a request keeps its queue + # place but polls once per window (5 s, or lease_s / 3 if shorter), so it may see its grant up + # to a window late; all waiters on one key share one actor. 16 extra slots keep parked calls + # from taking the slots renewal and release need, but every waiter still costs one call per + # window: past the actor's call throughput, renewals are delayed and leases can be lost. max_concurrency: 10000 diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 4de838d2..86ba94ee 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -251,7 +251,8 @@ def close(): Every process closes its client and drops its handles. Only the process whose `init()` created TransferQueue also cleans up distributed storage and kills the - controller actor; processes that attached leave them running and may `init()` again. + controller and lock actors; processes that attached leave them running, release + their own `kv_global_lock` leases, and may `init()` again. Note: This function should be called when the TransferQueue system is no longer needed. diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 4e2d337a..085589ea 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -214,9 +214,8 @@ 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() # Longest one call waits inside a lock actor; waiting longer takes another call that keeps -# the queue position. Bounding each call keeps waiters from pinning the actor's concurrency -# slots, and the bound is long because one actor serves only a few thousand calls a second: -# at 1 s, 8,000 idle waiters starved every other call on their shard, renewals included. +# 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. _POLL_S = 5.0 _REPLY_SLACK_S = 1.0 # how late a reply may arrive past its window before the caller checks its timeout _NOT_INITIALIZED = ( @@ -231,7 +230,12 @@ class LockLostError(RuntimeError): class GlobalLease: - """Yielded by ``kv_global_lock``; ``check()`` raises ``LockLostError`` once the lease is lost.""" + """Yielded by ``kv_global_lock``; ``check()`` raises ``LockLostError`` once the lease is lost. + + Call ``check()`` rather than reading ``lost``, which stays False when the lease merely ran + out locally. ``check()`` only narrows the window: KV calls take no fencing token, so a + write that outlasts the remaining lease can still land after another holder took the keys. + """ def __init__(self, names: list[tuple[str, str]], shards: dict[int, list[tuple[str, str]]], lease_s: float): self.names, self.shards, self.lease_s = names, shards, lease_s @@ -335,8 +339,8 @@ def _release_lease(lease: GlobalLease) -> None: _leases.pop(lease.token, None) shards, lease.pending = lease.pending, [] for shard in shards: - # Releasing by token also withdraws an acquire that is still waiting or in flight, so - # it is never granted to a caller that left; ray.cancel would miss a grant already made. + # Releasing by token also withdraws an acquire that is still waiting or in flight, so a + # caller that left is never left holding a grant; ray.cancel would miss one already made. try: _managers[shard].release.remote(lease.token) except Exception: # close() dropped the handle, or the actor is gone @@ -394,7 +398,8 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No The locks live in the ``lock.num_shards`` actors that ``tq.init()`` creates, each key in the one its hash picks, and are held under a ``lease_s`` lease that a background thread renews. Each actor grants in arrival order: a request waits behind - every earlier one that shares a key with it. + every earlier one that shares a key with it, except one with no call in flight for 10 s: + that waiter is skipped as dead, and if it was alive, its next call queues it at the back. Keys on several actors are taken one actor at a time, each all at once, so a request holds the keys already granted while it waits for a later actor's. The hash decides @@ -402,23 +407,18 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No caller: with 8 shards, two keys land on different actors 7 times in 8. A request queues afresh at each later actor, so its waits add up, and the keys it holds meanwhile block others even if it then times out. Lock the fewest keys you can and keep the block short. - - A waiting request keeps one call waiting on the actor it is waiting at, and each actor lets - at most ``lock.max_concurrency`` calls wait at once; renewal and release have slots of - their own, so waiting requests cannot hold them up. Past the limit, a request keeps its - place in the queue but polls once per window (5 s, or ``lease_s / 3`` if shorter), so it - may see its grant up to a window late. With keys hashing evenly, about ``num_shards * - max_concurrency`` requests can wait without that delay, so raise ``lock.num_shards`` for - more in total; all waiters on one key share one actor, so raise ``lock.max_concurrency`` - for a hot key. + The ``lock`` section of ``config.yaml`` describes capacity limits. Raises ``TimeoutError`` if the keys are not all granted within ``timeout`` seconds - (``None`` waits forever); an unresponsive lock actor can stretch that by up to about - 1 s, the time allowed for its last reply. Raises ``RuntimeError`` when nested, taken - while holding a ``kv_local_lock``, or called with a running event loop. Raises + (``None`` waits forever). The wait can overrun ``timeout`` by up to 1 s, the time allowed + for the last reply, and the first call in a process also looks up the lock actors, which + ``timeout`` does not bound. Raises ``RuntimeError`` if the lock actors were never started + (``lock.enabled: false``) or are gone, when nested, taken while holding a + ``kv_local_lock``, or called with a running event loop; ``ValueError`` for no keys, + ``lease_s <= 0``, or a ``lock.num_shards`` that changed since this process read it. Raises ``LockLostError`` on entry if an earlier key's lease was lost while a later one was awaited. Exit always releases the lock, then raises ``LockLostError`` if the lease was - lost, unless the body already raised: an exception from the body is never masked. + lost, unless the body already raised. """ _reject_running_loop("kv_global_lock") lease = _new_lease(keys, partition_id, lease_s) @@ -491,7 +491,8 @@ def kv_lock_list(partition_id: str | None = None) -> dict: """Return ``{"holders": [...], "waiters": n}`` for live global locks, optionally in one partition. Each holder has ``partition_id``, ``key``, ``holder`` (node IP, pid, thread), ``held_s`` and - ``lease_remaining_s``. A holder may still be waiting for its keys on other actors. + ``lease_remaining_s``. A holder may still be waiting for its keys on other actors, and one + that just left its block may still be listed briefly: exit sends its release without waiting. """ replies = ray.get([_lock_manager(shard).list_locks.remote(partition_id) for shard in range(_shard_count())]) return {"holders": [h for r in replies for h in r["holders"]], "waiters": sum(r["waiters"] for r in replies)} diff --git a/transfer_queue/lock_manager.py b/transfer_queue/lock_manager.py index bc926712..b91e4061 100644 --- a/transfer_queue/lock_manager.py +++ b/transfer_queue/lock_manager.py @@ -24,13 +24,14 @@ # A withdrawal whose acquire never arrives (e.g. the caller died first) is dropped after this. _WITHDRAWN_TTL_S = 300 -# A queued caller calls again at least once per poll window (5 s at most), even when a full -# actor turned it away, so a waiter with no call in flight for twice that has died. Hand-off -# skips it instead of granting it a lease that would only expire; if it was merely stalled, -# its next call queues it again at the back. +# A live caller calls again once per poll window (5 s at most), so a waiter with no call in +# flight for twice that is taken for dead, and hand-off skips it rather than grant it an idle +# lease. A live caller that stalled, or whose next call sat queued in Ray, is skipped too and +# queues again at the back on its next call. _ABANDONED_S = 10.0 -# Slots beyond lock.max_concurrency for calls that never wait: renew, release, list_locks and -# acquires turned away. Parked acquires never take them, so renewal and release always run. +# Slots beyond lock.max_concurrency, shared with no priority by calls that never wait: renew, +# release, list_locks and acquires turned away. They only keep parked acquires from taking +# every slot; past the actor's call throughput, renewals still run late and leases can be lost. _RESERVED_SLOTS = 16 @@ -53,9 +54,10 @@ def __init__(self, num_shards: int, max_parked: int): self._timer: asyncio.TimerHandle | None = None self._timer_at = float("inf") # Waiting tokens per key in arrival order. A request is granted only when all its keys are - # free and it heads every one of their queues: no newcomer can barge, and single-key - # requests cannot starve an older multi-key one. All queues share one arrival order, so the - # oldest waiter heads all of its queues and the wait can never form a cycle. + # free and it heads every one of their queues: no newcomer can barge (bar a waiter skipped + # as abandoned), and single-key requests cannot starve an older multi-key one. All queues + # share one arrival order, so the oldest waiter heads all of its queues and the wait can + # never form a cycle. self._queues: dict[tuple[str, str], deque[str]] = {} # token -> names, lease_s, info, done (set on grant or withdrawal), calls in flight, seen (last call end) self._waiting: dict[str, dict] = {} @@ -126,10 +128,11 @@ async def acquire(self, names, token, timeout, lease_s, holder_info, num_shards) Returns None if still queued after ``timeout``, or at once if ``max_parked`` calls are already waiting. Calling again with the same token keeps its place in the queues; only ``release`` withdraws it. A token that is neither queued nor held (its grant expired - before the caller saw it) queues afresh. + before the caller saw it, or it was skipped as abandoned) queues afresh. """ - # A caller that hashed over another shard count may send a key to the wrong shard, - # where nothing excludes the callers that sent it to the right one. + # A process whose cached shard count predates a restart with another lock.num_shards, + # yet finds these actors by name (e.g. a child forked after its parent's first lock), + # would send keys to the wrong shard, where nothing excludes the right shard's callers. if num_shards != self._num_shards: raise ValueError( f"kv_global_lock hashed keys over {num_shards} lock shards, but lock.num_shards is " From bea098b8b35e0ccedeaa075f38b0b52f51467349 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 12 Oct 2026 10:06:17 +0800 Subject: [PATCH 32/42] [fix] Stop a kv_global_lock wait that close() gave up An attached process's tq.close() releases its leases, which withdraws a thread still waiting at a later shard. That thread treated the early reply as a full actor, sent the same token again and queued afresh, then waited for a lease that was already lost. The wait loops now stop with LockLostError once the lease is lost, which also ends a wait whose earlier shard's lease ran out. release tombstones a withdrawn waiter too, so a call Ray runs after the release cannot queue it again. The lease keeps the handle each acquire went to, so its release still goes out after close() dropped this process's handles. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 40 ++++++++++++++++++++++++++++ transfer_queue/kv_lock.py | 31 +++++++++++---------- transfer_queue/lock_manager.py | 21 ++++++++------- 3 files changed, 68 insertions(+), 24 deletions(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index 4fa80bb6..97300a7a 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -16,6 +16,7 @@ """kv_global_lock across Ray actors: exclusion, timeouts, queue order, leases, and close().""" import importlib +import threading import time import pytest @@ -61,6 +62,25 @@ def hold(self, key, lease_s=30): def release(self): self.held.__exit__(None, None, None) + def wait_in_thread(self, keys): + """Wait for ``keys`` on another thread, so ``close()`` can run meanwhile; ``outcome()`` reports the end.""" + + def wait(): + try: + with tq.kv_global_lock(keys, P, timeout=TIMEOUT_S): + self.ended = "entered" + except Exception as e: + self.ended = type(e).__name__ + + self.ended = None + threading.Thread(target=wait, daemon=True).start() + + def outcome(self): + return self.ended + + def close(self): + tq.close() + @pytest.fixture(scope="module", autouse=True) def tq_session(): @@ -173,6 +193,9 @@ def test_waiters_are_granted_in_arrival_order(own_manager): # b is free, but granting it past the older [a, b] request would let single-key traffic starve it. assert ray.get(acquire(own_manager, ["b"], "newcomer", timeout=0), timeout=TIMEOUT_S) is None ray.get(own_manager.release.remote("newcomer"), timeout=TIMEOUT_S) # as kv_global_lock does on timeout + # A call Ray runs after the withdrawal must not queue it again. + assert ray.get(acquire(own_manager, ["b"], "newcomer", timeout=0), timeout=TIMEOUT_S) is None + assert n_waiters(own_manager) == 1 only_b = acquire(own_manager, ["b"], "only_b") wait_until(lambda: n_waiters(own_manager) == 2) ray.get(own_manager.release.remote("holder"), timeout=TIMEOUT_S) @@ -260,6 +283,23 @@ def actors_named(prefix): return [a["name"] for a in ray.util.list_named_actors(all_namespaces=True) if a["name"].startswith(prefix)] +def test_attached_close_gives_up_its_leases_and_leaves_tq_running(): + early, late = keys_on_two_shards("ac") + holder, attached = Worker.remote(), Worker.remote() + ray.get(holder.hold.remote(late), timeout=TIMEOUT_S) + ray.get(attached.wait_in_thread.remote([early, late]), timeout=TIMEOUT_S) + wait_until(lambda: holders() == {early, late} and tq.kv_lock_list(P)["waiters"] == 1) + + ray.get(attached.close.remote(), timeout=TIMEOUT_S) + # The waiter must stop, not queue its lost lease again behind the holder of `late`. + wait_until(lambda: ray.get(attached.outcome.remote(), timeout=TIMEOUT_S) == "LockLostError") + wait_until(lambda: tq.kv_lock_list(P)["waiters"] == 0 and holders() == {late}) + assert len(actors_named("TransferQueueLockManager_")) == kvl._shard_count() + assert actors_named("TransferQueueController") == ["TransferQueueController"] + ray.get(holder.release.remote(), timeout=TIMEOUT_S) + wait_until(lambda: tq.kv_lock_list(P) == {"holders": [], "waiters": 0}) + + def test_owner_close_kills_the_managers_and_fails_waiters(): """Tears TransferQueue down for this module, so it runs last.""" held = tq.kv_global_lock("z", P) diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 085589ea..98e58cf9 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -242,7 +242,9 @@ def __init__(self, names: list[tuple[str, str]], shards: dict[int, list[tuple[st self.token = f"{os.getpid()}:{uuid.uuid4().hex}" self.deadline = float("inf") # local monotonic time the lease is known to last until, once granted self.lost = False - self.pending: list[int] = [] # shards that may hold or await this token, so exit must release them + # shard -> the actor handle each acquire went to: those may hold or await this token, so + # exit must release it there, even after close() dropped this process's handles. + self.pending: dict = {} self.granted: list[int] = [] # shards that granted it, so the renewer keeps their leases alive def check(self) -> None: @@ -305,12 +307,9 @@ def _poll_window(deadline: float | None, lease_s: float) -> float: def _send_acquire(lease: GlobalLease, shard: int, window: float): - if shard not in lease.pending: - lease.pending.append(shard) + manager = lease.pending[shard] = _lock_manager(shard) holder = {"node_ip": ray.util.get_node_ip_address(), "pid": os.getpid(), "thread": threading.current_thread().name} - return _lock_manager(shard).acquire.remote( - lease.shards[shard], lease.token, window, lease.lease_s, holder, _shard_count() - ) + return manager.acquire.remote(lease.shards[shard], lease.token, window, lease.lease_s, holder, _shard_count()) def _take_grant(lease: GlobalLease, shard: int, remaining: float | None, sent: float, deadline: float | None) -> bool: @@ -337,13 +336,13 @@ def _take_grant(lease: GlobalLease, shard: int, remaining: float | None, sent: f def _release_lease(lease: GlobalLease) -> None: with _registry: _leases.pop(lease.token, None) - shards, lease.pending = lease.pending, [] - for shard in shards: + managers, lease.pending = lease.pending, {} + for manager in managers.values(): # Releasing by token also withdraws an acquire that is still waiting or in flight, so a # caller that left is never left holding a grant; ray.cancel would miss one already made. try: - _managers[shard].release.remote(lease.token) - except Exception: # close() dropped the handle, or the actor is gone + manager.release.remote(lease.token) + except Exception: # Ray is shut down pass @@ -416,9 +415,9 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No (``lock.enabled: false``) or are gone, when nested, taken while holding a ``kv_local_lock``, or called with a running event loop; ``ValueError`` for no keys, ``lease_s <= 0``, or a ``lock.num_shards`` that changed since this process read it. Raises - ``LockLostError`` on entry if an earlier key's lease was lost while a later one was - awaited. Exit always releases the lock, then raises ``LockLostError`` if the lease was - lost, unless the body already raised. + ``LockLostError`` while waiting once the lease is lost: an earlier key's lease ran out while + a later one was awaited, or this process's ``tq.close()`` gave it up. Exit always releases + the lock, then raises ``LockLostError`` if the lease was lost, unless the body already raised. """ _reject_running_loop("kv_global_lock") lease = _new_lease(keys, partition_id, lease_s) @@ -442,7 +441,10 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No except RayActorError as e: raise RuntimeError(_MANAGER_GONE) from e granted = _take_grant(lease, shard, remaining, sent, deadline) - if not granted: # a reply before the window ends means the actor was full + if not granted: + # Stop once an earlier shard's lease ran out, or close() gave the lease up and + # withdrew this wait. Otherwise a reply before the window ends means a full actor. + lease.check() time.sleep(max(0.0, sent + window - time.monotonic())) lease.check() yield lease @@ -478,6 +480,7 @@ async def async_kv_global_lock( raise RuntimeError(_MANAGER_GONE) from e granted = _take_grant(lease, shard, remaining, sent, deadline) if not granted: + lease.check() await asyncio.sleep(max(0.0, sent + window - time.monotonic())) lease.check() yield lease diff --git a/transfer_queue/lock_manager.py b/transfer_queue/lock_manager.py index b91e4061..3dd2775c 100644 --- a/transfer_queue/lock_manager.py +++ b/transfer_queue/lock_manager.py @@ -183,22 +183,23 @@ def renew_many(self, tokens: list[str]) -> dict[str, bool]: return alive async def release(self, token: str) -> None: - """Free ``token``'s locks, or withdraw its acquire if it is still waiting or has not arrived.""" + """Free ``token``'s locks, or withdraw its acquire for good if it is still waiting or has not arrived.""" if self._free(token): return if token in self._waiting: waiter = self._leave_queues(token) self._hand_off(waiter["names"]) waiter["done"].set() # its pending call finds no lease and returns None - else: - now = time.monotonic() - # One TTL for all, so insertion order is expiry order: expired tombstones lead. - while self._withdrawn: - oldest, until = next(iter(self._withdrawn.items())) - if until > now: - break - del self._withdrawn[oldest] - self._withdrawn[token] = now + _WITHDRAWN_TTL_S + # Even a withdrawn waiter may have a call still on its way (Ray may run it after this + # release), which must not queue the token afresh. + now = time.monotonic() + # One TTL for all, so insertion order is expiry order: expired tombstones lead. + while self._withdrawn: + oldest, until = next(iter(self._withdrawn.items())) + if until > now: + break + del self._withdrawn[oldest] + self._withdrawn[token] = now + _WITHDRAWN_TTL_S def list_locks(self, partition_id: str | None = None) -> dict: """Return current holders and the number of waiters, optionally for one partition.""" From 6d4101a86990f54c7c16a2a581a10fe0155e941b Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 12 Oct 2026 10:16:15 +0800 Subject: [PATCH 33/42] [fix] Keep one unresponsive lock shard from losing every kv_global_lock lease The renewer waited on each shard's renewal in turn, with a budget of 3 * period, which equals the shortest lease_s. One shard that did not answer used the whole budget, so every lease, even one living only on healthy shards, ran out. It now sends all renewals and collects them with one ray.wait bounded by the period. A shard that does not answer in time is not judged: its leases keep their deadline and run out only when it passes. Only an explicit refusal, or a dead lock actor, marks a lease lost. The deadline and lost updates now happen under _registry, as _take_grant's do, so the renewer cannot overwrite a shorter deadline that _take_grant just set for a newly granted shard. That race has no deterministic test. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 29 ++++++++++++++++++++++++ transfer_queue/kv_lock.py | 33 +++++++++++++++++----------- 2 files changed, 49 insertions(+), 13 deletions(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index 97300a7a..cfe5f64d 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -15,9 +15,11 @@ """kv_global_lock across Ray actors: exclusion, timeouts, queue order, leases, and close().""" +import asyncio import importlib import threading import time +from concurrent.futures import ThreadPoolExecutor import pytest import ray @@ -279,6 +281,33 @@ def test_lease_renews_and_lost_lease_raises(): wait_until(lambda: holders() == set()) +@ray.remote(num_cpus=0) +class Unresponsive: + """Stands in for a lock actor on a partitioned node: renewals never answer.""" + + async def renew_many(self, tokens): + await asyncio.sleep(3600) + + +def test_an_unresponsive_shard_costs_no_lease_on_another(monkeypatch): + healthy, hung = keys_on_two_shards("ur") + other = tq.kv_global_lock(hung, P, timeout=TIMEOUT_S, lease_s=2) + with ThreadPoolExecutor(1) as pool: # a lock of its own needs a thread of its own + other_lease = pool.submit(other.__enter__).result(TIMEOUT_S) + unresponsive = Unresponsive.remote() + with tq.kv_global_lock(healthy, P, timeout=TIMEOUT_S, lease_s=2) as lease: + monkeypatch.setitem(kvl._managers, shard_of(hung), unresponsive) + time.sleep(3) # past both leases + lease.check() + with pytest.raises(tq.LockLostError): # unanswered, it runs out once its deadline passes + other_lease.check() + with pytest.raises(tq.LockLostError): + other.__exit__(None, None, None) + monkeypatch.undo() + ray.kill(unresponsive) + wait_until(lambda: holders() == set()) + + def actors_named(prefix): return [a["name"] for a in ray.util.list_named_actors(all_namespaces=True) if a["name"].startswith(prefix)] diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 98e58cf9..b0f6ebe1 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -369,25 +369,32 @@ def _renew_loop() -> None: # which is later, so the local deadline never outlives the server's. last = time.monotonic() calls = {} - renewed: set[tuple[int, str]] = set() for shard, tokens in by_shard.items(): try: - calls[shard] = _managers[shard].renew_many.remote(tokens) + calls[_managers[shard].renew_many.remote(tokens)] = shard except Exception: # close() dropped the handle pass - for shard, call in calls.items(): + # One wait for all shards, so a shard that does not answer cannot delay the others' + # renewals. Its silence proves nothing: its leases just keep their deadline. + ready, _ = ray.wait(list(calls), num_returns=len(calls), timeout=period) + replies: dict[tuple[int, str], bool] = {} + for call in ready: + shard = calls[call] try: - alive = ray.get(call, timeout=max(0.0, last + 3 * period - time.monotonic())) - except Exception: # a dead lock manager or a timeout - continue - renewed.update((shard, token) for token, ok in alive.items() if ok) - # A lease is lost as soon as any one of its shards fails to renew it. A shard granted + alive = ray.get(call) + except Exception: # the lock manager died, and its leases with it + alive = {} + for token in by_shard[shard]: + replies[shard, token] = alive.get(token, False) + # A lease is lost as soon as any one of its shards refuses to renew it. A shard granted # since the snapshot was not renewed here, so it keeps the deadline _take_grant set. - for token, (lease, shards) in batch.items(): - if not all((shard, token) in renewed for shard in shards): - lease.lost = True - elif len(lease.granted) == len(shards): - lease.deadline = last + lease.lease_s + with _registry: # _take_grant also moves the deadline + for token, (lease, shards) in batch.items(): + answers = [replies.get((shard, token)) for shard in shards] + if False in answers: + lease.lost = True + elif None not in answers and len(lease.granted) == len(shards): + lease.deadline = last + lease.lease_s @contextmanager From 3dd543ee0e8e0426f6614b38772b96d8a4444fee Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 12 Oct 2026 10:19:23 +0800 Subject: [PATCH 34/42] [fix] Send the shard count a kv_global_lock lease's keys were hashed with Each acquire carried the process's current cached shard count rather than the one the lease's keys were hashed with. A waiter whose process ran tq.close() and tq.init() against a TransferQueue with another lock.num_shards re-read the count before its next call, so the actor's check passed while the keys went to the wrong shard. GlobalLease now records the count its keys were hashed over, and every acquire sends that. Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 11 +++++++++++ transfer_queue/kv_lock.py | 13 ++++++++----- transfer_queue/lock_manager.py | 7 ++++--- 3 files changed, 23 insertions(+), 8 deletions(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index cfe5f64d..5d3b34bc 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -169,6 +169,17 @@ def send_acquire(lease, shard, window): wait_until(lambda: tq.kv_lock_list(P) == {"holders": [], "waiters": 0}) +def test_acquire_carries_the_shard_count_its_keys_were_hashed_with(monkeypatch): + n = kvl._shard_count() + key = next(k for k in map(str, range(50)) if kvl._shard((P, k), n + 1) < n) + monkeypatch.setattr(kvl, "_num_shards", n + 1) # hashed before a restart with another lock.num_shards + lease = kvl._new_lease(key, P, 30) + monkeypatch.undo() + with pytest.raises(ValueError, match="restarted"): + ray.get(kvl._send_acquire(lease, next(iter(lease.shards)), 0), timeout=TIMEOUT_S) + kvl._release_lease(lease) + + @pytest.fixture def own_manager(): """A lock actor of its own, so these tests control every request it sees. diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index b0f6ebe1..bd1df167 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -237,8 +237,11 @@ class GlobalLease: write that outlasts the remaining lease can still land after another holder took the keys. """ - def __init__(self, names: list[tuple[str, str]], shards: dict[int, list[tuple[str, str]]], lease_s: float): - self.names, self.shards, self.lease_s = names, shards, lease_s + def __init__( + self, names: list[tuple[str, str]], shards: dict[int, list[tuple[str, str]]], num_shards: int, lease_s: float + ): + # num_shards is the count the keys were hashed over, which every acquire must carry. + self.names, self.shards, self.num_shards, self.lease_s = names, shards, num_shards, lease_s self.token = f"{os.getpid()}:{uuid.uuid4().hex}" self.deadline = float("inf") # local monotonic time the lease is known to last until, once granted self.lost = False @@ -296,7 +299,7 @@ def _new_lease(keys: str | list[str], partition_id: str, lease_s: float) -> Glob shards: dict[int, list[tuple[str, str]]] = {} for name in names: shards.setdefault(_shard(name, num_shards), []).append(name) - return GlobalLease(names, dict(sorted(shards.items())), lease_s) + return GlobalLease(names, dict(sorted(shards.items())), num_shards, lease_s) def _poll_window(deadline: float | None, lease_s: float) -> float: @@ -309,7 +312,7 @@ def _poll_window(deadline: float | None, lease_s: float) -> float: def _send_acquire(lease: GlobalLease, shard: int, window: float): manager = lease.pending[shard] = _lock_manager(shard) holder = {"node_ip": ray.util.get_node_ip_address(), "pid": os.getpid(), "thread": threading.current_thread().name} - return manager.acquire.remote(lease.shards[shard], lease.token, window, lease.lease_s, holder, _shard_count()) + return manager.acquire.remote(lease.shards[shard], lease.token, window, lease.lease_s, holder, lease.num_shards) def _take_grant(lease: GlobalLease, shard: int, remaining: float | None, sent: float, deadline: float | None) -> bool: @@ -421,7 +424,7 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No ``timeout`` does not bound. Raises ``RuntimeError`` if the lock actors were never started (``lock.enabled: false``) or are gone, when nested, taken while holding a ``kv_local_lock``, or called with a running event loop; ``ValueError`` for no keys, - ``lease_s <= 0``, or a ``lock.num_shards`` that changed since this process read it. Raises + ``lease_s <= 0``, or a ``lock.num_shards`` that changed since this call hashed its keys. Raises ``LockLostError`` while waiting once the lease is lost: an earlier key's lease ran out while a later one was awaited, or this process's ``tq.close()`` gave it up. Exit always releases the lock, then raises ``LockLostError`` if the lease was lost, unless the body already raised. diff --git a/transfer_queue/lock_manager.py b/transfer_queue/lock_manager.py index 3dd2775c..3bc0ebb1 100644 --- a/transfer_queue/lock_manager.py +++ b/transfer_queue/lock_manager.py @@ -130,9 +130,10 @@ async def acquire(self, names, token, timeout, lease_s, holder_info, num_shards) ``release`` withdraws it. A token that is neither queued nor held (its grant expired before the caller saw it, or it was skipped as abandoned) queues afresh. """ - # A process whose cached shard count predates a restart with another lock.num_shards, - # yet finds these actors by name (e.g. a child forked after its parent's first lock), - # would send keys to the wrong shard, where nothing excludes the right shard's callers. + # A caller whose keys were hashed before a restart with another lock.num_shards, yet + # finds these actors by name (e.g. a child forked after its parent's first lock, or a + # waiter whose process re-ran tq.init()), would send keys to the wrong shard, where + # nothing excludes the right shard's callers. if num_shards != self._num_shards: raise ValueError( f"kv_global_lock hashed keys over {num_shards} lock shards, but lock.num_shards is " From 788123e8989f510a0fc4e4359e09e102f4bb859c Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 12 Oct 2026 12:11:57 +0800 Subject: [PATCH 35/42] [fix] Keep GlobalLease's internal state private; expose only check() and names Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 8 ++-- transfer_queue/kv_lock.py | 69 +++++++++++++++------------- 2 files changed, 42 insertions(+), 35 deletions(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index 5d3b34bc..e3e7cd4f 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -176,7 +176,7 @@ def test_acquire_carries_the_shard_count_its_keys_were_hashed_with(monkeypatch): lease = kvl._new_lease(key, P, 30) monkeypatch.undo() with pytest.raises(ValueError, match="restarted"): - ray.get(kvl._send_acquire(lease, next(iter(lease.shards)), 0), timeout=TIMEOUT_S) + ray.get(kvl._send_acquire(lease, next(iter(lease._shards)), 0), timeout=TIMEOUT_S) kvl._release_lease(lease) @@ -273,7 +273,7 @@ def test_killed_holder_frees_the_key_after_its_lease(): def test_lease_renews_and_lost_lease_raises(): with tq.kv_global_lock("r", P, lease_s=2) as lease: - assert isinstance(lease, tq.GlobalLease) + assert isinstance(lease, tq.GlobalLease) and lease.names == [(P, "r")] time.sleep(3) # outlives one lease only through renewal lease.check() @@ -286,8 +286,8 @@ def test_lease_renews_and_lost_lease_raises(): with pytest.raises(ValueError, match="body"): # never masked by LockLostError with tq.kv_global_lock("r", P, lease_s=1) as lease: - ray.get(manager("r").release.remote(lease.token), timeout=TIMEOUT_S) # the server drops it - wait_until(lambda: lease.lost) + ray.get(manager("r").release.remote(lease._token), timeout=TIMEOUT_S) # the server drops it + wait_until(lambda: lease._lost) raise ValueError("body") wait_until(lambda: holders() == set()) diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index bd1df167..1d1dfab6 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -230,30 +230,34 @@ class LockLostError(RuntimeError): class GlobalLease: - """Yielded by ``kv_global_lock``; ``check()`` raises ``LockLostError`` once the lease is lost. + """Yielded by ``kv_global_lock``: call ``check()`` before writes; ``names`` lists the locked keys. - Call ``check()`` rather than reading ``lost``, which stays False when the lease merely ran - out locally. ``check()`` only narrows the window: KV calls take no fencing token, so a - write that outlasts the remaining lease can still land after another holder took the keys. + ``check()`` only narrows the window: KV calls take no fencing token, so a write that + outlasts the remaining lease can still land after another holder took the keys. """ def __init__( self, names: list[tuple[str, str]], shards: dict[int, list[tuple[str, str]]], num_shards: int, lease_s: float ): # num_shards is the count the keys were hashed over, which every acquire must carry. - self.names, self.shards, self.num_shards, self.lease_s = names, shards, num_shards, lease_s - self.token = f"{os.getpid()}:{uuid.uuid4().hex}" - self.deadline = float("inf") # local monotonic time the lease is known to last until, once granted - self.lost = False + self._names, self._shards, self._num_shards, self._lease_s = names, shards, num_shards, lease_s + self._token = f"{os.getpid()}:{uuid.uuid4().hex}" + self._deadline = float("inf") # local monotonic time the lease is known to last until, once granted + self._lost = False # shard -> the actor handle each acquire went to: those may hold or await this token, so # exit must release it there, even after close() dropped this process's handles. - self.pending: dict = {} - self.granted: list[int] = [] # shards that granted it, so the renewer keeps their leases alive + self._pending: dict = {} + self._granted: list[int] = [] # shards that granted it, so the renewer keeps their leases alive + + @property + def names(self) -> list[tuple[str, str]]: + """The locked ``(partition_id, key)`` pairs, deduplicated and sorted by key.""" + return list(self._names) def check(self) -> None: """Raise ``LockLostError`` if the lease was lost; call it right before writes in long sections.""" - if self.lost or time.monotonic() >= self.deadline: - raise LockLostError(f"kv_global_lock lease on {self.names} was lost") + if self._lost or time.monotonic() >= self._deadline: + raise LockLostError(f"kv_global_lock lease on {self._names} was lost") def _shard_count() -> int: @@ -310,9 +314,9 @@ def _poll_window(deadline: float | None, lease_s: float) -> float: def _send_acquire(lease: GlobalLease, shard: int, window: float): - manager = lease.pending[shard] = _lock_manager(shard) + manager = lease._pending[shard] = _lock_manager(shard) holder = {"node_ip": ray.util.get_node_ip_address(), "pid": os.getpid(), "thread": threading.current_thread().name} - return manager.acquire.remote(lease.shards[shard], lease.token, window, lease.lease_s, holder, lease.num_shards) + return manager.acquire.remote(lease._shards[shard], lease._token, window, lease._lease_s, holder, lease._num_shards) def _take_grant(lease: GlobalLease, shard: int, remaining: float | None, sent: float, deadline: float | None) -> bool: @@ -321,14 +325,14 @@ def _take_grant(lease: GlobalLease, shard: int, remaining: float | None, sent: f if remaining is None: if deadline is not None and time.monotonic() >= deadline: # Still queued there, so exit's release must withdraw it: the shard stays pending. - raise TimeoutError(f"Timed out waiting for kv_global_lock on {lease.names}") + raise TimeoutError(f"Timed out waiting for kv_global_lock on {lease._names}") return False # Renew from the first grant on: waiting on a later shard may outlast an earlier one's lease. with _registry: # the renewer also moves the deadline # The reply left the actor after `sent`, so this never outlives the shard's lease. - lease.deadline = min(lease.deadline, sent + remaining) - lease.granted.append(shard) - _leases[lease.token] = lease + lease._deadline = min(lease._deadline, sent + remaining) + lease._granted.append(shard) + _leases[lease._token] = lease _registry.notify() # a shorter lease_s shortens the renewal period if not _renewer_running: _renewer_running = True @@ -338,13 +342,13 @@ def _take_grant(lease: GlobalLease, shard: int, remaining: float | None, sent: f def _release_lease(lease: GlobalLease) -> None: with _registry: - _leases.pop(lease.token, None) - managers, lease.pending = lease.pending, {} + _leases.pop(lease._token, None) + managers, lease._pending = lease._pending, {} for manager in managers.values(): # Releasing by token also withdraws an acquire that is still waiting or in flight, so a # caller that left is never left holding a grant; ray.cancel would miss one already made. try: - manager.release.remote(lease.token) + manager.release.remote(lease._token) except Exception: # Ray is shut down pass @@ -358,12 +362,12 @@ def _renew_loop() -> None: if not _leases: _renewer_running = False return - period = min(lease.lease_s for lease in _leases.values()) / 3 + period = min(lease._lease_s for lease in _leases.values()) / 3 remaining = last + period - time.monotonic() if remaining <= 0: break _registry.wait(remaining) - batch = {token: (lease, list(lease.granted)) for token, lease in _leases.items()} + batch = {token: (lease, list(lease._granted)) for token, lease in _leases.items()} by_shard: dict[int, list[str]] = {} for token, (_, shards) in batch.items(): for shard in shards: @@ -395,9 +399,9 @@ def _renew_loop() -> None: for token, (lease, shards) in batch.items(): answers = [replies.get((shard, token)) for shard in shards] if False in answers: - lease.lost = True - elif None not in answers and len(lease.granted) == len(shards): - lease.deadline = last + lease.lease_s + lease._lost = True + elif None not in answers and len(lease._granted) == len(shards): + lease._deadline = last + lease._lease_s @contextmanager @@ -418,6 +422,9 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No others even if it then times out. Lock the fewest keys you can and keep the block short. The ``lock`` section of ``config.yaml`` describes capacity limits. + Call the yielded lease's ``check()`` right before each write; its ``names`` lists the + locked ``(partition_id, key)`` pairs. + Raises ``TimeoutError`` if the keys are not all granted within ``timeout`` seconds (``None`` waits forever). The wait can overrun ``timeout`` by up to 1 s, the time allowed for the last reply, and the first call in a process also looks up the lock actors, which @@ -436,10 +443,10 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No try: # Every caller visits shards in ascending order, so no two callers can each hold a # shard the other awaits: multi-shard requests cannot deadlock. - for shard in lease.shards: + for shard in lease._shards: call, granted = None, False while not granted: - window = _poll_window(deadline, lease.lease_s) + window = _poll_window(deadline, lease._lease_s) if call is None: # a late reply is awaited again, not re-sent, so calls never pile up call, sent = _send_acquire(lease, shard, window), time.monotonic() try: @@ -475,10 +482,10 @@ async def async_kv_global_lock( deadline = None if timeout is None else time.monotonic() + timeout held = _mark(_holding_global) try: - for shard in lease.shards: + for shard in lease._shards: call, granted = None, False while not granted: - window = _poll_window(deadline, lease.lease_s) + window = _poll_window(deadline, lease._lease_s) if call is None: call, sent = asyncio.ensure_future(_send_acquire(lease, shard, window)), time.monotonic() try: @@ -543,7 +550,7 @@ def _close_lock_managers() -> None: _leases.clear() _registry.notify() for lease in leases: - lease.lost = True + lease._lost = True if not _owned: # the owner kills the actors, so releasing would only hand keys to doomed waiters _release_lease(lease) for actor in _owned: From 735ec8c4735dbef45f63157e1a650c3361681804 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 12 Oct 2026 12:12:25 +0800 Subject: [PATCH 36/42] [fix] Rename lock.max_concurrency to lock.max_waiting_acquires Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 2 +- transfer_queue/config.yaml | 11 ++++++----- transfer_queue/kv_lock.py | 6 +++--- transfer_queue/lock_manager.py | 10 +++++----- 4 files changed, 15 insertions(+), 14 deletions(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index e3e7cd4f..6695f4ea 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -184,7 +184,7 @@ def test_acquire_carries_the_shard_count_its_keys_were_hashed_with(monkeypatch): def own_manager(): """A lock actor of its own, so these tests control every request it sees. - Sized as tq.init() sizes one for ``lock.max_concurrency: 2``. + Sized as tq.init() sizes one for ``lock.max_waiting_acquires: 2``. """ lock_manager = TransferQueueLockManager.options(max_concurrency=2 + kvl._RESERVED_SLOTS).remote(1, 2) yield lock_manager diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index db3cb78c..468908de 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -20,14 +20,15 @@ lock: # Whether tq.init() starts the lock actors; kv_global_lock raises while this is false. enabled: false # Number of lock actors that kv_global_lock spreads keys over by hash (>= 1). With keys - # hashing evenly, about num_shards * max_concurrency requests can wait without polling delay. + # hashing evenly, about num_shards * max_waiting_acquires requests can wait without polling delay. num_shards: 8 # Acquire calls each lock actor lets wait at once (>= 1). Past this, a request keeps its queue # place but polls once per window (5 s, or lease_s / 3 if shorter), so it may see its grant up - # to a window late; all waiters on one key share one actor. 16 extra slots keep parked calls - # from taking the slots renewal and release need, but every waiter still costs one call per - # window: past the actor's call throughput, renewals are delayed and leases can be lost. - max_concurrency: 10000 + # to a window late; all waiters on one key share one actor. The actor's Ray max_concurrency is + # this plus 16 slots that keep waiting calls from taking the ones renewal and release need, but + # every waiter still costs one call per window: past the actor's call throughput, renewals are + # delayed and leases can be lost. + max_waiting_acquires: 10000 backend: diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 1d1dfab6..3031ee7b 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -520,7 +520,7 @@ def kv_lock_list(partition_id: str | None = None) -> dict: def _check_lock_conf(lock_conf) -> None: """Called by ``tq.init()`` before it creates anything, so a bad config leaves nothing behind.""" - for field in ("num_shards", "max_concurrency"): + for field in ("num_shards", "max_waiting_acquires"): value = lock_conf[field] if lock_conf.enabled and (not isinstance(value, int) or value < 1): raise ValueError(f"lock.{field} must be an integer >= 1, got {value!r}") @@ -536,8 +536,8 @@ def _start_lock_managers(lock_conf) -> None: name=_manager_name(shard), namespace="transfer_queue", get_if_exists=True, - max_concurrency=lock_conf.max_concurrency + _RESERVED_SLOTS, - ).remote(lock_conf.num_shards, lock_conf.max_concurrency) + max_concurrency=lock_conf.max_waiting_acquires + _RESERVED_SLOTS, + ).remote(lock_conf.num_shards, lock_conf.max_waiting_acquires) for shard in range(lock_conf.num_shards) ] diff --git a/transfer_queue/lock_manager.py b/transfer_queue/lock_manager.py index 3bc0ebb1..ecc2d332 100644 --- a/transfer_queue/lock_manager.py +++ b/transfer_queue/lock_manager.py @@ -29,21 +29,21 @@ # lease. A live caller that stalled, or whose next call sat queued in Ray, is skipped too and # queues again at the back on its next call. _ABANDONED_S = 10.0 -# Slots beyond lock.max_concurrency, shared with no priority by calls that never wait: renew, +# Slots beyond lock.max_waiting_acquires, shared with no priority by calls that never wait: renew, # release, list_locks and acquires turned away. They only keep parked acquires from taking # every slot; past the actor's call throughput, renewals still run late and leases can be lost. _RESERVED_SLOTS = 16 -# tq.init() runs each actor with lock.max_concurrency + _RESERVED_SLOTS slots. +# tq.init() runs each actor with lock.max_waiting_acquires + _RESERVED_SLOTS slots. @ray.remote(num_cpus=0) class TransferQueueLockManager: """Exclusive leased locks on ``(partition_id, key)``. Runs on one event loop, so no thread locks.""" - def __init__(self, num_shards: int, max_parked: int): + def __init__(self, num_shards: int, max_waiting_acquires: int): # Kept so that a process that never ran tq.init() can learn how many shards to hash over. self._num_shards = num_shards - self._max_parked, self._parked = max_parked, 0 # acquire calls waiting inside the actor + self._max_parked, self._parked = max_waiting_acquires, 0 # acquire calls waiting inside the actor self._holder: dict[tuple[str, str], str] = {} # name -> token self._leases: dict[str, dict] = {} # token -> names, lease_s, expires_at, granted_at, info # Min-heap of (expires_at, token). Release and renewal leave stale entries behind, which @@ -125,7 +125,7 @@ def _expire(self) -> None: async def acquire(self, names, token, timeout, lease_s, holder_info, num_shards) -> float | None: """Grant all ``names`` to ``token`` at once; return the seconds left on its lease. - Returns None if still queued after ``timeout``, or at once if ``max_parked`` calls are + Returns None if still queued after ``timeout``, or at once if ``max_waiting_acquires`` calls are already waiting. Calling again with the same token keeps its place in the queues; only ``release`` withdraws it. A token that is neither queued nor held (its grant expired before the caller saw it, or it was skipped as abandoned) queues afresh. From c9f2a198807f289b9b0e4a10a7ef88f0588e6afc Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 12 Oct 2026 12:51:27 +0800 Subject: [PATCH 37/42] [feat] Enable kv_global_lock with the TQ_ENABLE_KV_GLOBAL_LOCK environment variable Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 8 +++++--- transfer_queue/config.yaml | 4 ++-- transfer_queue/kv_lock.py | 19 ++++++++++++------- 3 files changed, 19 insertions(+), 12 deletions(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index 6695f4ea..f97c751b 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -36,7 +36,6 @@ CONF = { "controller": {"polling_mode": True}, "backend": {"storage_backend": "SimpleStorage", "SimpleStorage": {"total_storage_size": 100}}, - "lock": {"enabled": True}, } @@ -88,7 +87,10 @@ def close(self): def tq_session(): if not ray.is_initialized(): ray.init(namespace="TestKVGlobalLockE2E") - tq.init(OmegaConf.create(CONF)) + # Only the owner's tq.init() reads it; the Worker actors attach without it. + with pytest.MonkeyPatch.context() as mp: + mp.setenv("TQ_ENABLE_KV_GLOBAL_LOCK", "1") + tq.init(OmegaConf.create(CONF)) yield tq.close() ray.shutdown() @@ -356,5 +358,5 @@ def test_owner_close_kills_the_managers_and_fails_waiters(): held.__exit__(None, None, None) wait_until(lambda: actors_named("TransferQueue") == []) # the controller and storage go too - with pytest.raises(RuntimeError, match="lock.enabled"): + with pytest.raises(RuntimeError, match="TQ_ENABLE_KV_GLOBAL_LOCK"): tq.kv_lock_list() diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index 468908de..81e11edf 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -16,9 +16,9 @@ controller: zmq_info: null # Cluster-wide KV key locks (kv_global_lock). The process-local kv_local_lock needs none of this. +# tq.init() starts the lock actors only if the environment variable TQ_ENABLE_KV_GLOBAL_LOCK=1 is +# set in the process whose call starts TransferQueue; kv_global_lock raises otherwise. lock: - # Whether tq.init() starts the lock actors; kv_global_lock raises while this is false. - enabled: false # Number of lock actors that kv_global_lock spreads keys over by hash (>= 1). With keys # hashing evenly, about num_shards * max_waiting_acquires requests can wait without polling delay. num_shards: 8 diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 3031ee7b..53659a4d 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -23,8 +23,9 @@ ``kv_global_lock`` and ``async_kv_global_lock`` do the same across the Ray cluster through ``TransferQueueLockManager`` actors that split keys by hash, under a lease that renews -automatically. ``tq.init()`` starts those actors only with ``lock.enabled: true``. Take a -global lock before a local one, never inside it. +automatically. ``tq.init()`` starts those actors only if ``TQ_ENABLE_KV_GLOBAL_LOCK`` is set +to a true value (e.g. ``1``) in the process whose call starts TransferQueue. Take a global +lock before a local one, never inside it. """ import asyncio @@ -41,6 +42,7 @@ from ray.exceptions import GetTimeoutError, RayActorError from transfer_queue.lock_manager import _RESERVED_SLOTS, TransferQueueLockManager +from transfer_queue.utils.common import get_env_bool # Not asyncio.Lock: it is bound to one event loop and is not thread-safe, while a process # may run several loops plus plain threads. A key is held exactly while it is in _table, @@ -219,8 +221,8 @@ async def async_kv_local_lock(keys: str | list[str], partition_id: str, timeout: _POLL_S = 5.0 _REPLY_SLACK_S = 1.0 # how late a reply may arrive past its window before the caller checks its timeout _NOT_INITIALIZED = ( - "TransferQueueLockManager not found; kv_global_lock needs `lock.enabled: true` " - "in the config of the tq.init() that starts TransferQueue" + "TransferQueueLockManager not found; kv_global_lock needs TQ_ENABLE_KV_GLOBAL_LOCK=1 " + "set before the tq.init() that starts TransferQueue, in that process" ) _MANAGER_GONE = "The TransferQueueLockManager actor is gone; did the process that ran tq.init() call tq.close()?" @@ -429,7 +431,7 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No (``None`` waits forever). The wait can overrun ``timeout`` by up to 1 s, the time allowed for the last reply, and the first call in a process also looks up the lock actors, which ``timeout`` does not bound. Raises ``RuntimeError`` if the lock actors were never started - (``lock.enabled: false``) or are gone, when nested, taken while holding a + (``TQ_ENABLE_KV_GLOBAL_LOCK`` unset) or are gone, when nested, taken while holding a ``kv_local_lock``, or called with a running event loop; ``ValueError`` for no keys, ``lease_s <= 0``, or a ``lock.num_shards`` that changed since this call hashed its keys. Raises ``LockLostError`` while waiting once the lease is lost: an earlier key's lease ran out while @@ -520,16 +522,19 @@ def kv_lock_list(partition_id: str | None = None) -> dict: def _check_lock_conf(lock_conf) -> None: """Called by ``tq.init()`` before it creates anything, so a bad config leaves nothing behind.""" + if not get_env_bool("TQ_ENABLE_KV_GLOBAL_LOCK"): + return for field in ("num_shards", "max_waiting_acquires"): value = lock_conf[field] - if lock_conf.enabled and (not isinstance(value, int) or value < 1): + if not isinstance(value, int) or value < 1: raise ValueError(f"lock.{field} must be an integer >= 1, got {value!r}") def _start_lock_managers(lock_conf) -> None: """Called by the ``tq.init()`` that created the controller: start the lock actors if enabled.""" global _owned - if lock_conf.enabled: + # Read only here, by the creator: other processes find the actors by name, env or not. + if get_env_bool("TQ_ENABLE_KV_GLOBAL_LOCK"): # Created here rather than on first lock: a non-detached actor dies with its creator. _owned = [ TransferQueueLockManager.options( # type: ignore[attr-defined] From 7e155820fa175ecc95d5574a80f0f7f2bf4ab3c7 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 12 Oct 2026 13:00:22 +0800 Subject: [PATCH 38/42] [feat] Start one kv_global_lock actor per storage unit instead of lock.num_shards Signed-off-by: OutstanderWang --- tests/e2e/test_kv_global_lock_e2e.py | 4 ++-- transfer_queue/config.yaml | 10 ++++----- transfer_queue/interface.py | 2 +- transfer_queue/kv_lock.py | 31 ++++++++++++++-------------- transfer_queue/lock_manager.py | 4 ++-- 5 files changed, 26 insertions(+), 25 deletions(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index f97c751b..213cdada 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -154,7 +154,7 @@ def test_multi_shard_timeout_releases_the_earlier_shard(): def test_a_late_reply_is_awaited_instead_of_sent_again(monkeypatch): lock_manager = manager("late") - ray.get(lock_manager.acquire.remote([(P, "late")], "holder", 0, 30, {}, 8), timeout=TIMEOUT_S) + ray.get(lock_manager.acquire.remote([(P, "late")], "holder", 0, 30, {}, kvl._shard_count()), timeout=TIMEOUT_S) real_send, sends = kvl._send_acquire, [] def send_acquire(lease, shard, window): @@ -174,7 +174,7 @@ def send_acquire(lease, shard, window): def test_acquire_carries_the_shard_count_its_keys_were_hashed_with(monkeypatch): n = kvl._shard_count() key = next(k for k in map(str, range(50)) if kvl._shard((P, k), n + 1) < n) - monkeypatch.setattr(kvl, "_num_shards", n + 1) # hashed before a restart with another lock.num_shards + monkeypatch.setattr(kvl, "_num_shards", n + 1) # hashed before a restart with another shard count lease = kvl._new_lease(key, P, 30) monkeypatch.undo() with pytest.raises(ValueError, match="restarted"): diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index 81e11edf..2dbc5877 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -17,14 +17,14 @@ controller: # Cluster-wide KV key locks (kv_global_lock). The process-local kv_local_lock needs none of this. # tq.init() starts the lock actors only if the environment variable TQ_ENABLE_KV_GLOBAL_LOCK=1 is -# set in the process whose call starts TransferQueue; kv_global_lock raises otherwise. +# set in the process whose call starts TransferQueue; kv_global_lock raises otherwise. It starts one +# lock actor per storage unit (SimpleStorage num_data_storage_units; 1 for other backends) and +# spreads keys over them by hash. lock: - # Number of lock actors that kv_global_lock spreads keys over by hash (>= 1). With keys - # hashing evenly, about num_shards * max_waiting_acquires requests can wait without polling delay. - num_shards: 8 # Acquire calls each lock actor lets wait at once (>= 1). Past this, a request keeps its queue # place but polls once per window (5 s, or lease_s / 3 if shorter), so it may see its grant up - # to a window late; all waiters on one key share one actor. The actor's Ray max_concurrency is + # to a window late. With keys hashing evenly, about (lock actors) * max_waiting_acquires requests + # can wait without that delay, but all waiters on one key share one actor. The actor's Ray max_concurrency is # this plus 16 slots that keep waiting calls from taking the ones renewal and release need, but # every waiter still costs one call per window: past the actor's call throughput, renewals are # delayed and leases can be lost. diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 86ba94ee..2d8c5a87 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -208,7 +208,7 @@ def init(conf: DictConfig | None = None) -> DictConfig | None: ).remote(sampler=sampler, polling_mode=final_conf.controller.polling_mode) _TQ_IS_OWNER = True logger.info("TransferQueueController has been created.") - _start_lock_managers(final_conf.lock) + _start_lock_managers(final_conf) except ValueError: logger.info("Some other rank has initialized TransferQueueController. Try to connect to existing controller.") _init_from_existing() diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 53659a4d..b037d94f 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -212,7 +212,7 @@ async def async_kv_local_lock(keys: str | list[str], partition_id: str, timeout: _registry = threading.Condition() _leases: dict[str, "GlobalLease"] = {} _renewer_running = False -_num_shards: int | None = None # lock.num_shards, asked of shard 0 on first use +_num_shards: int | None = None # the lock actor count, asked of shard 0 on first use _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() # Longest one call waits inside a lock actor; waiting longer takes another call that keeps @@ -263,7 +263,7 @@ def check(self) -> None: def _shard_count() -> int: - """Return lock.num_shards, looking every lock actor up on first use.""" + """Return the lock actor count, looking every lock actor up on first use.""" global _num_shards if _num_shards is None: try: @@ -410,7 +410,7 @@ def _renew_loop() -> None: def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | None = None, lease_s: float = 30): """Hold an exclusive cluster-wide lock on ``keys`` in ``partition_id``; yield a ``GlobalLease``. - The locks live in the ``lock.num_shards`` actors that ``tq.init()`` creates, + The locks live in the actors that ``tq.init()`` creates, one per storage unit, each key in the one its hash picks, and are held under a ``lease_s`` lease that a background thread renews. Each actor grants in arrival order: a request waits behind every earlier one that shares a key with it, except one with no call in flight for 10 s: @@ -419,7 +419,7 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No Keys on several actors are taken one actor at a time, each all at once, so a request holds the keys already granted while it waits for a later actor's. The hash decides which keys share an actor, so whether a request is all-or-nothing is not up to the - caller: with 8 shards, two keys land on different actors 7 times in 8. A request queues + caller: with N shards, two keys land on different actors N-1 times in N. A request queues afresh at each later actor, so its waits add up, and the keys it holds meanwhile block others even if it then times out. Lock the fewest keys you can and keep the block short. The ``lock`` section of ``config.yaml`` describes capacity limits. @@ -433,7 +433,7 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No ``timeout`` does not bound. Raises ``RuntimeError`` if the lock actors were never started (``TQ_ENABLE_KV_GLOBAL_LOCK`` unset) or are gone, when nested, taken while holding a ``kv_local_lock``, or called with a running event loop; ``ValueError`` for no keys, - ``lease_s <= 0``, or a ``lock.num_shards`` that changed since this call hashed its keys. Raises + ``lease_s <= 0``, or a shard count that changed since this call hashed its keys. Raises ``LockLostError`` while waiting once the lease is lost: an earlier key's lease ran out while a later one was awaited, or this process's ``tq.close()`` gave it up. Exit always releases the lock, then raises ``LockLostError`` if the lease was lost, unless the body already raised. @@ -522,28 +522,29 @@ def kv_lock_list(partition_id: str | None = None) -> dict: def _check_lock_conf(lock_conf) -> None: """Called by ``tq.init()`` before it creates anything, so a bad config leaves nothing behind.""" - if not get_env_bool("TQ_ENABLE_KV_GLOBAL_LOCK"): - return - for field in ("num_shards", "max_waiting_acquires"): - value = lock_conf[field] - if not isinstance(value, int) or value < 1: - raise ValueError(f"lock.{field} must be an integer >= 1, got {value!r}") + value = lock_conf.max_waiting_acquires + if get_env_bool("TQ_ENABLE_KV_GLOBAL_LOCK") and (not isinstance(value, int) or value < 1): + raise ValueError(f"lock.max_waiting_acquires must be an integer >= 1, got {value!r}") -def _start_lock_managers(lock_conf) -> None: +def _start_lock_managers(conf) -> None: """Called by the ``tq.init()`` that created the controller: start the lock actors if enabled.""" global _owned # Read only here, by the creator: other processes find the actors by name, env or not. if get_env_bool("TQ_ENABLE_KV_GLOBAL_LOCK"): + # One lock actor per storage unit; backends without storage units get a single one. + backend = conf.backend + num_shards = backend.SimpleStorage.num_data_storage_units if backend.storage_backend == "SimpleStorage" else 1 + max_waiting_acquires = conf.lock.max_waiting_acquires # Created here rather than on first lock: a non-detached actor dies with its creator. _owned = [ TransferQueueLockManager.options( # type: ignore[attr-defined] name=_manager_name(shard), namespace="transfer_queue", get_if_exists=True, - max_concurrency=lock_conf.max_waiting_acquires + _RESERVED_SLOTS, - ).remote(lock_conf.num_shards, lock_conf.max_waiting_acquires) - for shard in range(lock_conf.num_shards) + max_concurrency=max_waiting_acquires + _RESERVED_SLOTS, + ).remote(num_shards, max_waiting_acquires) + for shard in range(num_shards) ] diff --git a/transfer_queue/lock_manager.py b/transfer_queue/lock_manager.py index ecc2d332..10bac9f6 100644 --- a/transfer_queue/lock_manager.py +++ b/transfer_queue/lock_manager.py @@ -130,13 +130,13 @@ async def acquire(self, names, token, timeout, lease_s, holder_info, num_shards) ``release`` withdraws it. A token that is neither queued nor held (its grant expired before the caller saw it, or it was skipped as abandoned) queues afresh. """ - # A caller whose keys were hashed before a restart with another lock.num_shards, yet + # A caller whose keys were hashed before a restart with another shard count, yet # finds these actors by name (e.g. a child forked after its parent's first lock, or a # waiter whose process re-ran tq.init()), would send keys to the wrong shard, where # nothing excludes the right shard's callers. if num_shards != self._num_shards: raise ValueError( - f"kv_global_lock hashed keys over {num_shards} lock shards, but lock.num_shards is " + f"kv_global_lock hashed keys over {num_shards} lock shards, but there are now " f"{self._num_shards}; TransferQueue restarted, so call tq.close() and tq.init() again" ) if token in self._withdrawn: # kept until it expires: a retried call may still arrive From ece5c46e0c4cab23041a52b1862d1f94a63b2885 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 12 Oct 2026 13:01:26 +0800 Subject: [PATCH 39/42] [feat] Set kv_global_lock's wait limit with TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES and document the lock Signed-off-by: OutstanderWang --- README.md | 20 ++++++++++++++++++++ tests/e2e/test_kv_global_lock_e2e.py | 2 +- transfer_queue/config.yaml | 16 ---------------- transfer_queue/interface.py | 6 +++--- transfer_queue/kv_lock.py | 26 +++++++++++++++----------- transfer_queue/lock_manager.py | 8 ++++---- 6 files changed, 43 insertions(+), 35 deletions(-) diff --git a/README.md b/README.md index 3c08a6ad..acde20de 100644 --- a/README.md +++ b/README.md @@ -129,6 +129,26 @@ To simplify the usage of TransferQueue, we provide a Redis-style high-level API Refer to [tutorials/basic.ipynb](https://github.com/Ascend/TransferQueue/blob/main/tutorial/basic.ipynb) and [tutorials/02_kv_interface.py](https://github.com/Ascend/TransferQueue/blob/main/tutorial/02_kv_interface.py) for detailed usage examples. +**KV Key Locks** + +Locks are **advisory**: KV calls never check them, so every writer that needs mutual exclusion must take the lock. `(async_)kv_local_lock` excludes threads and asyncio tasks of one process and needs no setup. + +- **(async_)kv_global_lock(keys, partition_id, timeout=None, lease_s=30)**: Hold an exclusive cluster-wide lock on `(partition_id, key)` pairs under a lease that a background thread renews. Yields a **GlobalLease**: call `check()` right before writes (it raises **LockLostError** once the lease is lost); `names` lists the locked pairs. Exiting the block also raises `LockLostError` if the lease was lost. +- **kv_lock_list(partition_id=None)**: Return `{"holders": [...], "waiters": n}` for live global locks. + +```python +import os +os.environ["TQ_ENABLE_KV_GLOBAL_LOCK"] = "1" # only in the process that creates TransferQueue, before tq.init() +tq.init() +with tq.kv_global_lock(["a", "b"], "train", timeout=30) as lease: + lease.check() + tq.kv_put("a", "train", fields={"v": torch.tensor([1])}) +``` + +- **Enabling**: set `TQ_ENABLE_KV_GLOBAL_LOCK` (`1`/`true`/`yes`/`y`/`on`; default off) before `tq.init()` in the process that creates TransferQueue. Processes and Ray workers that attach find the lock actors by name and need nothing. Without it, `kv_global_lock` raises `RuntimeError`. +- **Shards**: keys are spread by hash over one `TransferQueueLockManager` actor per storage unit (`backend.SimpleStorage.num_data_storage_units`). Backends without storage units (MooncakeStore, Yuanrong) get a single lock actor. Keys on several actors are taken one actor at a time, so a multi-key request may hold some keys while it waits for others. +- **Capacity**: `TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES` (default `10000`, integer >= 1, read by the creating process) caps the acquire calls each actor lets wait at once. A request past the cap keeps its queue place but polls once per window (5 s, or `lease_s / 3` if shorter), so it may see its grant up to a window late. About shards × this value requests can wait without that delay, but all waiters on one key share one actor, so a hot key is bounded by this value alone. Each actor's Ray `max_concurrency` is this value plus 16 reserved slots shared, without priority, by renew, release, `kv_lock_list` and turned-away acquires. Every waiter still costs one call per window: past an actor's call throughput, renewals are delayed and leases can be lost. + #### StreamingDataLoader API Designed as a drop-in replacement for the standard PyTorch `DataLoader`, this API allows each rank to automatically consume data without single-controller intervention. diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index 213cdada..74c08ace 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -186,7 +186,7 @@ def test_acquire_carries_the_shard_count_its_keys_were_hashed_with(monkeypatch): def own_manager(): """A lock actor of its own, so these tests control every request it sees. - Sized as tq.init() sizes one for ``lock.max_waiting_acquires: 2``. + Sized as tq.init() sizes one for ``TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES=2``. """ lock_manager = TransferQueueLockManager.options(max_concurrency=2 + kvl._RESERVED_SLOTS).remote(1, 2) yield lock_manager diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index 2dbc5877..e0ac6599 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -15,22 +15,6 @@ controller: # ZMQ Server IP & Ports (automatically generated during init) zmq_info: null -# Cluster-wide KV key locks (kv_global_lock). The process-local kv_local_lock needs none of this. -# tq.init() starts the lock actors only if the environment variable TQ_ENABLE_KV_GLOBAL_LOCK=1 is -# set in the process whose call starts TransferQueue; kv_global_lock raises otherwise. It starts one -# lock actor per storage unit (SimpleStorage num_data_storage_units; 1 for other backends) and -# spreads keys over them by hash. -lock: - # Acquire calls each lock actor lets wait at once (>= 1). Past this, a request keeps its queue - # place but polls once per window (5 s, or lease_s / 3 if shorter), so it may see its grant up - # to a window late. With keys hashing evenly, about (lock actors) * max_waiting_acquires requests - # can wait without that delay, but all waiters on one key share one actor. The actor's Ray max_concurrency is - # this plus 16 slots that keep waiting calls from taking the ones renewal and release need, but - # every waiter still costs one call per window: past the actor's call throughput, renewals are - # delayed and leases can be lost. - max_waiting_acquires: 10000 - - backend: # Pluggable storage/transport backend of TransferQueue. Choose from: # SimpleStorage, Yuanrong, MooncakeStore, ... diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 2d8c5a87..6c6a03c7 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -45,7 +45,7 @@ from transfer_queue.client import TransferQueueClient from transfer_queue.controller import TransferQueueController -from transfer_queue.kv_lock import _check_lock_conf, _close_lock_managers, _start_lock_managers +from transfer_queue.kv_lock import _close_lock_managers, _lock_max_waiting_acquires, _start_lock_managers from transfer_queue.metadata import KVBatchMeta from transfer_queue.sampler import * # noqa: F401 from transfer_queue.sampler import BaseSampler @@ -199,7 +199,7 @@ def init(conf: DictConfig | None = None) -> DictConfig | None: sampler = globals()[final_conf.controller.sampler] except KeyError: raise ValueError(f"Could not find sampler {final_conf.controller.sampler}") from None - _check_lock_conf(final_conf.lock) + max_waiting_acquires = _lock_max_waiting_acquires() try: global _TQ_CONTROLLER, _TQ_IS_OWNER @@ -208,7 +208,7 @@ def init(conf: DictConfig | None = None) -> DictConfig | None: ).remote(sampler=sampler, polling_mode=final_conf.controller.polling_mode) _TQ_IS_OWNER = True logger.info("TransferQueueController has been created.") - _start_lock_managers(final_conf) + _start_lock_managers(max_waiting_acquires, final_conf.backend) except ValueError: logger.info("Some other rank has initialized TransferQueueController. Try to connect to existing controller.") _init_from_existing() diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index b037d94f..63f67cd7 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -422,7 +422,7 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No caller: with N shards, two keys land on different actors N-1 times in N. A request queues afresh at each later actor, so its waits add up, and the keys it holds meanwhile block others even if it then times out. Lock the fewest keys you can and keep the block short. - The ``lock`` section of ``config.yaml`` describes capacity limits. + The kv_global_lock section of README.md describes capacity limits. Call the yielded lease's ``check()`` right before each write; its ``names`` lists the locked ``(partition_id, key)`` pairs. @@ -520,22 +520,26 @@ def kv_lock_list(partition_id: str | None = None) -> dict: return {"holders": [h for r in replies for h in r["holders"]], "waiters": sum(r["waiters"] for r in replies)} -def _check_lock_conf(lock_conf) -> None: - """Called by ``tq.init()`` before it creates anything, so a bad config leaves nothing behind.""" - value = lock_conf.max_waiting_acquires - if get_env_bool("TQ_ENABLE_KV_GLOBAL_LOCK") and (not isinstance(value, int) or value < 1): - raise ValueError(f"lock.max_waiting_acquires must be an integer >= 1, got {value!r}") +def _lock_max_waiting_acquires() -> int | None: + """Return the acquire calls each lock actor lets wait at once, or None if kv_global_lock is off. + + Called by ``tq.init()`` before it creates anything, so a bad value leaves nothing behind. + Only the process that creates TransferQueue reads these variables: others find the actors by name. + """ + if not get_env_bool("TQ_ENABLE_KV_GLOBAL_LOCK"): + return None + value = int(os.environ.get("TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES", 10000)) + if value < 1: + raise ValueError(f"TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES must be >= 1, got {value}") + return value -def _start_lock_managers(conf) -> None: +def _start_lock_managers(max_waiting_acquires: int | None, backend) -> None: """Called by the ``tq.init()`` that created the controller: start the lock actors if enabled.""" global _owned - # Read only here, by the creator: other processes find the actors by name, env or not. - if get_env_bool("TQ_ENABLE_KV_GLOBAL_LOCK"): + if max_waiting_acquires is not None: # One lock actor per storage unit; backends without storage units get a single one. - backend = conf.backend num_shards = backend.SimpleStorage.num_data_storage_units if backend.storage_backend == "SimpleStorage" else 1 - max_waiting_acquires = conf.lock.max_waiting_acquires # Created here rather than on first lock: a non-detached actor dies with its creator. _owned = [ TransferQueueLockManager.options( # type: ignore[attr-defined] diff --git a/transfer_queue/lock_manager.py b/transfer_queue/lock_manager.py index 10bac9f6..af0ff2b7 100644 --- a/transfer_queue/lock_manager.py +++ b/transfer_queue/lock_manager.py @@ -29,13 +29,13 @@ # lease. A live caller that stalled, or whose next call sat queued in Ray, is skipped too and # queues again at the back on its next call. _ABANDONED_S = 10.0 -# Slots beyond lock.max_waiting_acquires, shared with no priority by calls that never wait: renew, -# release, list_locks and acquires turned away. They only keep parked acquires from taking -# every slot; past the actor's call throughput, renewals still run late and leases can be lost. +# Slots beyond TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES, shared with no priority by calls that never +# wait: renew, release, list_locks and acquires turned away. They only keep parked acquires from +# taking every slot; past the actor's call throughput, renewals still run late and leases can be lost. _RESERVED_SLOTS = 16 -# tq.init() runs each actor with lock.max_waiting_acquires + _RESERVED_SLOTS slots. +# tq.init() runs each actor with TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES + _RESERVED_SLOTS slots. @ray.remote(num_cpus=0) class TransferQueueLockManager: """Exclusive leased locks on ``(partition_id, key)``. Runs on one event loop, so no thread locks.""" From 5d44725ec92115260c809abfc76bedb9043752bb Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 12 Oct 2026 13:21:39 +0800 Subject: [PATCH 40/42] [feat] Run 8 kv_global_lock actors spread evenly across the cluster's nodes Signed-off-by: OutstanderWang --- README.md | 4 ++-- tests/test_kv_lock.py | 19 +++++++++++++++++++ transfer_queue/interface.py | 2 +- transfer_queue/kv_lock.py | 20 +++++++++++--------- 4 files changed, 33 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index acde20de..8bc8dcbf 100644 --- a/README.md +++ b/README.md @@ -146,8 +146,8 @@ with tq.kv_global_lock(["a", "b"], "train", timeout=30) as lease: ``` - **Enabling**: set `TQ_ENABLE_KV_GLOBAL_LOCK` (`1`/`true`/`yes`/`y`/`on`; default off) before `tq.init()` in the process that creates TransferQueue. Processes and Ray workers that attach find the lock actors by name and need nothing. Without it, `kv_global_lock` raises `RuntimeError`. -- **Shards**: keys are spread by hash over one `TransferQueueLockManager` actor per storage unit (`backend.SimpleStorage.num_data_storage_units`). Backends without storage units (MooncakeStore, Yuanrong) get a single lock actor. Keys on several actors are taken one actor at a time, so a multi-key request may hold some keys while it waits for others. -- **Capacity**: `TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES` (default `10000`, integer >= 1, read by the creating process) caps the acquire calls each actor lets wait at once. A request past the cap keeps its queue place but polls once per window (5 s, or `lease_s / 3` if shorter), so it may see its grant up to a window late. About shards × this value requests can wait without that delay, but all waiters on one key share one actor, so a hot key is bounded by this value alone. Each actor's Ray `max_concurrency` is this value plus 16 reserved slots shared, without priority, by renew, release, `kv_lock_list` and turned-away acquires. Every waiter still costs one call per window: past an actor's call throughput, renewals are delayed and leases can be lost. +- **Shards**: keys are spread by hash over 8 `TransferQueueLockManager` actors, which `tq.init()` pins round-robin to the alive Ray nodes (with fewer than 8 nodes, per-node counts differ by at most 1). Keys on several actors are taken one actor at a time, so a multi-key request may hold some keys while it waits for others. +- **Capacity**: `TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES` (default `10000`, integer >= 1, read by the creating process) caps the acquire calls each actor lets wait at once. A request past the cap keeps its queue place but polls once per window (5 s, or `lease_s / 3` if shorter), so it may see its grant up to a window late. About 8 × this value (80,000 by default) requests can wait without that delay, but all waiters on one key share one actor, so a hot key is bounded by this value alone. Each actor's Ray `max_concurrency` is this value plus 16 reserved slots shared, without priority, by renew, release, `kv_lock_list` and turned-away acquires. Every waiter still costs one call per window: past an actor's call throughput, renewals are delayed and leases can be lost. #### StreamingDataLoader API diff --git a/tests/test_kv_lock.py b/tests/test_kv_lock.py index d9ed2ae9..9f4a8aa2 100644 --- a/tests/test_kv_lock.py +++ b/tests/test_kv_lock.py @@ -17,6 +17,8 @@ import importlib import threading import time +from collections import Counter +from unittest.mock import MagicMock import pytest @@ -126,3 +128,20 @@ def test_timeout_is_a_total_deadline_and_partial_acquisition_is_released(): for holder in holders: holder.join(TIMEOUT) assert not holder.is_alive() + + +def test_lock_actors_are_spread_evenly_over_the_alive_nodes(monkeypatch): + a, b, c, dead = ("01" * 28, "02" * 28, "03" * 28, "04" * 28) # Ray checks node IDs are hex + nodes = [{"NodeID": node, "Alive": True} for node in (c, a, b)] + [{"NodeID": dead, "Alive": False}] + monkeypatch.setattr(kvl.ray, "nodes", lambda: nodes) + manager_cls = MagicMock() + monkeypatch.setattr(kvl, "TransferQueueLockManager", manager_cls) + monkeypatch.setattr(kvl, "_owned", []) + + kvl._start_lock_managers(max_waiting_acquires=2) + + strategies = [call.kwargs["scheduling_strategy"] for call in manager_cls.options.call_args_list] + assert [s.node_id for s in strategies] == [a, b, c, a, b, c, a, b] + assert Counter(s.node_id for s in strategies) == {a: 3, b: 3, c: 2} + assert not any(s.soft for s in strategies) + manager_cls.options.return_value.remote.assert_called_with(kvl._NUM_SHARDS, 2) diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 6c6a03c7..04dddb4a 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -208,7 +208,7 @@ def init(conf: DictConfig | None = None) -> DictConfig | None: ).remote(sampler=sampler, polling_mode=final_conf.controller.polling_mode) _TQ_IS_OWNER = True logger.info("TransferQueueController has been created.") - _start_lock_managers(max_waiting_acquires, final_conf.backend) + _start_lock_managers(max_waiting_acquires) except ValueError: logger.info("Some other rank has initialized TransferQueueController. Try to connect to existing controller.") _init_from_existing() diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 63f67cd7..ee062eef 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -42,7 +42,7 @@ from ray.exceptions import GetTimeoutError, RayActorError from transfer_queue.lock_manager import _RESERVED_SLOTS, TransferQueueLockManager -from transfer_queue.utils.common import get_env_bool +from transfer_queue.utils.common import get_env_bool, get_node_round_robin_scheduling_strategies # Not asyncio.Lock: it is bound to one event loop and is not thread-safe, while a process # may run several loops plus plain threads. A key is held exactly while it is in _table, @@ -215,6 +215,7 @@ async def async_kv_local_lock(keys: str | list[str], partition_id: str, timeout: _num_shards: int | None = None # the lock actor count, asked of shard 0 on first use _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() +_NUM_SHARDS = 8 # lock actors tq.init() starts; other processes ask shard 0, so only the creator reads this # 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. @@ -410,7 +411,7 @@ def _renew_loop() -> None: def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | None = None, lease_s: float = 30): """Hold an exclusive cluster-wide lock on ``keys`` in ``partition_id``; yield a ``GlobalLease``. - The locks live in the actors that ``tq.init()`` creates, one per storage unit, + The locks live in the 8 actors that ``tq.init()`` spreads evenly over the nodes, each key in the one its hash picks, and are held under a ``lease_s`` lease that a background thread renews. Each actor grants in arrival order: a request waits behind every earlier one that shares a key with it, except one with no call in flight for 10 s: @@ -419,7 +420,7 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No Keys on several actors are taken one actor at a time, each all at once, so a request holds the keys already granted while it waits for a later actor's. The hash decides which keys share an actor, so whether a request is all-or-nothing is not up to the - caller: with N shards, two keys land on different actors N-1 times in N. A request queues + caller: two keys land on different actors 7 times in 8. A request queues afresh at each later actor, so its waits add up, and the keys it holds meanwhile block others even if it then times out. Lock the fewest keys you can and keep the block short. The kv_global_lock section of README.md describes capacity limits. @@ -534,21 +535,22 @@ def _lock_max_waiting_acquires() -> int | None: return value -def _start_lock_managers(max_waiting_acquires: int | None, backend) -> None: +def _start_lock_managers(max_waiting_acquires: int | None) -> None: """Called by the ``tq.init()`` that created the controller: start the lock actors if enabled.""" global _owned if max_waiting_acquires is not None: - # One lock actor per storage unit; backends without storage units get a single one. - num_shards = backend.SimpleStorage.num_data_storage_units if backend.storage_backend == "SimpleStorage" else 1 - # Created here rather than on first lock: a non-detached actor dies with its creator. + # Round-robin over the alive nodes, as for storage units, so one node's loss or load + # hits the fewest shards. Created here rather than on first lock: a non-detached actor + # dies with its creator. _owned = [ TransferQueueLockManager.options( # type: ignore[attr-defined] name=_manager_name(shard), namespace="transfer_queue", get_if_exists=True, max_concurrency=max_waiting_acquires + _RESERVED_SLOTS, - ).remote(num_shards, max_waiting_acquires) - for shard in range(num_shards) + scheduling_strategy=strategy, + ).remote(_NUM_SHARDS, max_waiting_acquires) + for shard, strategy in enumerate(get_node_round_robin_scheduling_strategies(_NUM_SHARDS)) ] From b3aa1b3383c446763cc2209bd593331b6b6fd625 Mon Sep 17 00:00:00 2001 From: OutstanderWang Date: Mon, 12 Oct 2026 13:47:42 +0800 Subject: [PATCH 41/42] [feat] Set the kv_global_lock shard count with TQ_KV_GLOBAL_LOCK_NUM_SHARDS Signed-off-by: OutstanderWang --- README.md | 13 +++++++--- tests/test_kv_lock.py | 49 ++++++++++++++++++++++++++++++------- transfer_queue/interface.py | 6 ++--- transfer_queue/kv_lock.py | 42 ++++++++++++++++++------------- 4 files changed, 78 insertions(+), 32 deletions(-) diff --git a/README.md b/README.md index 8bc8dcbf..671b0c9c 100644 --- a/README.md +++ b/README.md @@ -145,9 +145,16 @@ with tq.kv_global_lock(["a", "b"], "train", timeout=30) as lease: tq.kv_put("a", "train", fields={"v": torch.tensor([1])}) ``` -- **Enabling**: set `TQ_ENABLE_KV_GLOBAL_LOCK` (`1`/`true`/`yes`/`y`/`on`; default off) before `tq.init()` in the process that creates TransferQueue. Processes and Ray workers that attach find the lock actors by name and need nothing. Without it, `kv_global_lock` raises `RuntimeError`. -- **Shards**: keys are spread by hash over 8 `TransferQueueLockManager` actors, which `tq.init()` pins round-robin to the alive Ray nodes (with fewer than 8 nodes, per-node counts differ by at most 1). Keys on several actors are taken one actor at a time, so a multi-key request may hold some keys while it waits for others. -- **Capacity**: `TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES` (default `10000`, integer >= 1, read by the creating process) caps the acquire calls each actor lets wait at once. A request past the cap keeps its queue place but polls once per window (5 s, or `lease_s / 3` if shorter), so it may see its grant up to a window late. About 8 × this value (80,000 by default) requests can wait without that delay, but all waiters on one key share one actor, so a hot key is bounded by this value alone. Each actor's Ray `max_concurrency` is this value plus 16 reserved slots shared, without priority, by renew, release, `kv_lock_list` and turned-away acquires. Every waiter still costs one call per window: past an actor's call throughput, renewals are delayed and leases can be lost. +- **Enabling**: set these before `tq.init()` in the process that creates TransferQueue; only that process reads them. Processes and Ray workers that attach find the lock actors by name and need nothing. Without the switch, `kv_global_lock` raises `RuntimeError`. + +| Variable | Default | Meaning | +| --- | --- | --- | +| `TQ_ENABLE_KV_GLOBAL_LOCK` | off | `1`/`true`/`yes`/`y`/`on` starts the lock actors | +| `TQ_KV_GLOBAL_LOCK_NUM_SHARDS` | `8` | Number of lock actors (integer >= 1) | +| `TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES` | `10000` | Acquire calls each lock actor lets wait at once (integer >= 1) | + +- **Shards**: keys are spread by hash over `TQ_KV_GLOBAL_LOCK_NUM_SHARDS` `TransferQueueLockManager` actors, which `tq.init()` pins round-robin to the alive Ray nodes (per-node counts differ by at most 1). Keys on several actors are taken one actor at a time, so a multi-key request may hold some keys while it waits for others. +- **Capacity**: `TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES` caps the acquire calls each actor lets wait at once. A request past the cap keeps its queue place but polls once per window (5 s, or `lease_s / 3` if shorter), so it may see its grant up to a window late. About `TQ_KV_GLOBAL_LOCK_NUM_SHARDS` × this value (80,000 by default) requests can wait without that delay. More shards spread the load, but all waiters on one key share one actor, so a hot key is bounded by this value alone. Each actor's Ray `max_concurrency` is this value plus 16 reserved slots shared, without priority, by renew, release, `kv_lock_list` and turned-away acquires. Every waiter still costs one call per window: past an actor's call throughput, renewals are delayed and leases can be lost. #### StreamingDataLoader API diff --git a/tests/test_kv_lock.py b/tests/test_kv_lock.py index 9f4a8aa2..1a239ed9 100644 --- a/tests/test_kv_lock.py +++ b/tests/test_kv_lock.py @@ -17,7 +17,6 @@ import importlib import threading import time -from collections import Counter from unittest.mock import MagicMock import pytest @@ -130,18 +129,50 @@ def test_timeout_is_a_total_deadline_and_partial_acquisition_is_released(): assert not holder.is_alive() -def test_lock_actors_are_spread_evenly_over_the_alive_nodes(monkeypatch): - a, b, c, dead = ("01" * 28, "02" * 28, "03" * 28, "04" * 28) # Ray checks node IDs are hex - nodes = [{"NodeID": node, "Alive": True} for node in (c, a, b)] + [{"NodeID": dead, "Alive": False}] +A, B, C = "01" * 28, "02" * 28, "03" * 28 # Ray checks node IDs are hex + + +@pytest.fixture +def fake_lock_actors(monkeypatch): + """Three alive nodes and one dead one; yields the mocked lock actor class instead of starting actors.""" + nodes = [{"NodeID": node, "Alive": True} for node in (C, A, B)] + [{"NodeID": "04" * 28, "Alive": False}] monkeypatch.setattr(kvl.ray, "nodes", lambda: nodes) manager_cls = MagicMock() monkeypatch.setattr(kvl, "TransferQueueLockManager", manager_cls) monkeypatch.setattr(kvl, "_owned", []) + monkeypatch.setenv("TQ_ENABLE_KV_GLOBAL_LOCK", "1") + monkeypatch.setenv("TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES", "2") + yield manager_cls + + +@pytest.mark.parametrize( + "num_shards, expected", + [(None, [A, B, C, A, B, C, A, B]), ("3", [A, B, C])], # None: unset, so the default 8 +) +def test_lock_actors_are_spread_evenly_over_the_alive_nodes(fake_lock_actors, monkeypatch, num_shards, expected): + if num_shards is not None: + monkeypatch.setenv("TQ_KV_GLOBAL_LOCK_NUM_SHARDS", num_shards) + else: + monkeypatch.delenv("TQ_KV_GLOBAL_LOCK_NUM_SHARDS", raising=False) - kvl._start_lock_managers(max_waiting_acquires=2) + kvl._start_lock_managers(kvl._lock_settings()) - strategies = [call.kwargs["scheduling_strategy"] for call in manager_cls.options.call_args_list] - assert [s.node_id for s in strategies] == [a, b, c, a, b, c, a, b] - assert Counter(s.node_id for s in strategies) == {a: 3, b: 3, c: 2} + strategies = [call.kwargs["scheduling_strategy"] for call in fake_lock_actors.options.call_args_list] + assert [s.node_id for s in strategies] == expected # per-node counts differ by at most 1 assert not any(s.soft for s in strategies) - manager_cls.options.return_value.remote.assert_called_with(kvl._NUM_SHARDS, 2) + fake_lock_actors.options.return_value.remote.assert_called_with(len(expected), 2) + + +@pytest.mark.parametrize("name", ["TQ_KV_GLOBAL_LOCK_NUM_SHARDS", "TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES"]) +@pytest.mark.parametrize("value", ["0", "-1", "abc", "true", "2.5"]) +def test_invalid_lock_settings_raise_before_any_actor_starts(fake_lock_actors, monkeypatch, name, value): + monkeypatch.setenv(name, value) + with pytest.raises(ValueError, match=f"{name} must be an integer >= 1"): + kvl._lock_settings() + fake_lock_actors.options.assert_not_called() + + +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") + assert kvl._lock_settings() is None diff --git a/transfer_queue/interface.py b/transfer_queue/interface.py index 04dddb4a..cb399ca0 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -45,7 +45,7 @@ from transfer_queue.client import TransferQueueClient from transfer_queue.controller import TransferQueueController -from transfer_queue.kv_lock import _close_lock_managers, _lock_max_waiting_acquires, _start_lock_managers +from transfer_queue.kv_lock import _close_lock_managers, _lock_settings, _start_lock_managers from transfer_queue.metadata import KVBatchMeta from transfer_queue.sampler import * # noqa: F401 from transfer_queue.sampler import BaseSampler @@ -199,7 +199,7 @@ def init(conf: DictConfig | None = None) -> DictConfig | None: sampler = globals()[final_conf.controller.sampler] except KeyError: raise ValueError(f"Could not find sampler {final_conf.controller.sampler}") from None - max_waiting_acquires = _lock_max_waiting_acquires() + lock_settings = _lock_settings() try: global _TQ_CONTROLLER, _TQ_IS_OWNER @@ -208,7 +208,7 @@ def init(conf: DictConfig | None = None) -> DictConfig | None: ).remote(sampler=sampler, polling_mode=final_conf.controller.polling_mode) _TQ_IS_OWNER = True logger.info("TransferQueueController has been created.") - _start_lock_managers(max_waiting_acquires) + _start_lock_managers(lock_settings) except ValueError: logger.info("Some other rank has initialized TransferQueueController. Try to connect to existing controller.") _init_from_existing() diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index ee062eef..68f891a1 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -215,7 +215,6 @@ async def async_kv_local_lock(keys: str | list[str], partition_id: str, timeout: _num_shards: int | None = None # the lock actor count, asked of shard 0 on first use _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() -_NUM_SHARDS = 8 # lock actors tq.init() starts; other processes ask shard 0, so only the creator reads this # 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. @@ -411,16 +410,17 @@ def _renew_loop() -> None: def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | None = None, lease_s: float = 30): """Hold an exclusive cluster-wide lock on ``keys`` in ``partition_id``; yield a ``GlobalLease``. - The locks live in the 8 actors that ``tq.init()`` spreads evenly over the nodes, - each key in the one its hash picks, and are held under a ``lease_s`` lease that a - background thread renews. Each actor grants in arrival order: a request waits behind - every earlier one that shares a key with it, except one with no call in flight for 10 s: - that waiter is skipped as dead, and if it was alive, its next call queues it at the back. + The locks live in the ``TQ_KV_GLOBAL_LOCK_NUM_SHARDS`` actors (8 by default) that ``tq.init()`` + spreads evenly over the nodes, each key in the one its hash picks, and are held under a + ``lease_s`` lease that a background thread renews. Each actor grants in arrival order: a + request waits behind every earlier one that shares a key with it, except one with no call in + flight for 10 s: that waiter is skipped as dead, and if it was alive, its next call queues it + at the back. Keys on several actors are taken one actor at a time, each all at once, so a request holds the keys already granted while it waits for a later actor's. The hash decides which keys share an actor, so whether a request is all-or-nothing is not up to the - caller: two keys land on different actors 7 times in 8. A request queues + caller: with N shards, two keys land on different actors N-1 times in N. A request queues afresh at each later actor, so its waits add up, and the keys it holds meanwhile block others even if it then times out. Lock the fewest keys you can and keep the block short. The kv_global_lock section of README.md describes capacity limits. @@ -521,24 +521,32 @@ def kv_lock_list(partition_id: str | None = None) -> dict: return {"holders": [h for r in replies for h in r["holders"]], "waiters": sum(r["waiters"] for r in replies)} -def _lock_max_waiting_acquires() -> int | None: - """Return the acquire calls each lock actor lets wait at once, or None if kv_global_lock is off. +def _positive_int_env(name: str, default: int) -> int: + value = os.environ.get(name, str(default)) + if not value.strip().isdecimal() or int(value) < 1: + raise ValueError(f"{name} must be an integer >= 1, got {value!r}") + return int(value) + + +def _lock_settings() -> tuple[int, int] | None: + """Return ``(num_shards, max_waiting_acquires)`` for the lock actors, or None if kv_global_lock is off. Called by ``tq.init()`` before it creates anything, so a bad value leaves nothing behind. Only the process that creates TransferQueue reads these variables: others find the actors by name. """ if not get_env_bool("TQ_ENABLE_KV_GLOBAL_LOCK"): return None - value = int(os.environ.get("TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES", 10000)) - if value < 1: - raise ValueError(f"TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES must be >= 1, got {value}") - return value + return ( + _positive_int_env("TQ_KV_GLOBAL_LOCK_NUM_SHARDS", 8), + _positive_int_env("TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES", 10000), + ) -def _start_lock_managers(max_waiting_acquires: int | None) -> None: +def _start_lock_managers(settings: tuple[int, int] | None) -> None: """Called by the ``tq.init()`` that created the controller: start the lock actors if enabled.""" global _owned - if max_waiting_acquires is not None: + if settings is not None: + num_shards, max_waiting_acquires = settings # Round-robin over the alive nodes, as for storage units, so one node's loss or load # hits the fewest shards. Created here rather than on first lock: a non-detached actor # dies with its creator. @@ -549,8 +557,8 @@ def _start_lock_managers(max_waiting_acquires: int | None) -> None: get_if_exists=True, max_concurrency=max_waiting_acquires + _RESERVED_SLOTS, scheduling_strategy=strategy, - ).remote(_NUM_SHARDS, max_waiting_acquires) - for shard, strategy in enumerate(get_node_round_robin_scheduling_strategies(_NUM_SHARDS)) + ).remote(num_shards, max_waiting_acquires) + for shard, strategy in enumerate(get_node_round_robin_scheduling_strategies(num_shards)) ] From ba63698b3f5cd4358ab8f113e9c06f455c4eb028 Mon Sep 17 00:00:00 2001 From: 0oshowero0 Date: Mon, 12 Oct 2026 15:14:35 +0800 Subject: [PATCH 42/42] [fix] Stop a kv_global_lock wait that tq.close() gave up before any shard grant tq.close() only marked leases already in the renewer registry as lost, so a waiter still queued for its first shard kept going and could take the keys afterwards. Count close() calls and have lease.check() fail once that count moves; the existing attached-close e2e now holds the waiter's first shard. Signed-off-by: 0oshowero0 Co-authored-by: Cursor --- tests/e2e/test_kv_global_lock_e2e.py | 8 ++++---- transfer_queue/kv_lock.py | 12 +++++++----- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/tests/e2e/test_kv_global_lock_e2e.py b/tests/e2e/test_kv_global_lock_e2e.py index 74c08ace..661bf7fa 100644 --- a/tests/e2e/test_kv_global_lock_e2e.py +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -328,14 +328,14 @@ def actors_named(prefix): def test_attached_close_gives_up_its_leases_and_leaves_tq_running(): early, late = keys_on_two_shards("ac") holder, attached = Worker.remote(), Worker.remote() - ray.get(holder.hold.remote(late), timeout=TIMEOUT_S) + ray.get(holder.hold.remote(early), timeout=TIMEOUT_S) ray.get(attached.wait_in_thread.remote([early, late]), timeout=TIMEOUT_S) - wait_until(lambda: holders() == {early, late} and tq.kv_lock_list(P)["waiters"] == 1) + wait_until(lambda: holders() == {early} and tq.kv_lock_list(P)["waiters"] == 1) ray.get(attached.close.remote(), timeout=TIMEOUT_S) - # The waiter must stop, not queue its lost lease again behind the holder of `late`. + # Still queued for its first shard, the waiter must stop rather than take the keys later. wait_until(lambda: ray.get(attached.outcome.remote(), timeout=TIMEOUT_S) == "LockLostError") - wait_until(lambda: tq.kv_lock_list(P)["waiters"] == 0 and holders() == {late}) + wait_until(lambda: tq.kv_lock_list(P)["waiters"] == 0 and holders() == {early}) assert len(actors_named("TransferQueueLockManager_")) == kvl._shard_count() assert actors_named("TransferQueueController") == ["TransferQueueController"] ray.get(holder.release.remote(), timeout=TIMEOUT_S) diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py index 68f891a1..4ddff4f1 100644 --- a/transfer_queue/kv_lock.py +++ b/transfer_queue/kv_lock.py @@ -215,6 +215,7 @@ async def async_kv_local_lock(keys: str | list[str], partition_id: str, timeout: _num_shards: int | None = None # the lock actor count, asked of shard 0 on first use _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 # 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. @@ -246,6 +247,7 @@ def __init__( self._token = f"{os.getpid()}:{uuid.uuid4().hex}" self._deadline = float("inf") # local monotonic time the lease is known to last until, once granted self._lost = False + self._closes = _closes # shard -> the actor handle each acquire went to: those may hold or await this token, so # exit must release it there, even after close() dropped this process's handles. self._pending: dict = {} @@ -258,7 +260,7 @@ def names(self) -> list[tuple[str, str]]: def check(self) -> None: """Raise ``LockLostError`` if the lease was lost; call it right before writes in long sections.""" - if self._lost or time.monotonic() >= self._deadline: + if self._lost or self._closes != _closes or time.monotonic() >= self._deadline: raise LockLostError(f"kv_global_lock lease on {self._names} was lost") @@ -462,8 +464,8 @@ def kv_global_lock(keys: str | list[str], partition_id: str, timeout: float | No raise RuntimeError(_MANAGER_GONE) from e granted = _take_grant(lease, shard, remaining, sent, deadline) if not granted: - # Stop once an earlier shard's lease ran out, or close() gave the lease up and - # withdrew this wait. Otherwise a reply before the window ends means a full actor. + # Stop once an earlier shard's lease ran out, or this process's tq.close() gave + # the lease up. Otherwise a reply before the window ends means a full actor. lease.check() time.sleep(max(0.0, sent + window - time.monotonic())) lease.check() @@ -564,13 +566,13 @@ def _start_lock_managers(settings: tuple[int, int] | None) -> None: def _close_lock_managers() -> None: """Called by ``tq.close()``: give up this process's leases, then kill the actors it created.""" - global _num_shards, _managers, _owned + global _num_shards, _managers, _owned, _closes with _registry: # emptied so the renewer stops now, not when the blocks holding them exit leases = list(_leases.values()) _leases.clear() _registry.notify() + _closes += 1 for lease in leases: - lease._lost = True if not _owned: # the owner kills the actors, so releasing would only hand keys to doomed waiters _release_lease(lease) for actor in _owned: