diff --git a/README.md b/README.md index 3c08a6ad..671b0c9c 100644 --- a/README.md +++ b/README.md @@ -129,6 +129,33 @@ 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 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 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 new file mode 100644 index 00000000..661bf7fa --- /dev/null +++ b/tests/e2e/test_kv_global_lock_e2e.py @@ -0,0 +1,362 @@ +# 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, timeouts, queue order, leases, and close().""" + +import asyncio +import importlib +import threading +import time +from concurrent.futures import ThreadPoolExecutor + +import pytest +import ray +import torch +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" +# 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): + 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) + self.held.__enter__() + + 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(): + if not ray.is_initialized(): + ray.init(namespace="TestKVGlobalLockE2E") + # 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() + + +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 and waiters to drain. + 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_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}) + 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) + wait_until(lambda: tq.kv_lock_list(P) == {"holders": [], "waiters": 0}) + + +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, {}, kvl._shard_count()), 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) + + monkeypatch.setattr(kvl, "_send_acquire", send_acquire) + with pytest.raises(TimeoutError): # 0.5 s polls, each awaited for up to 1.5 s + 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_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 shard count + 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. + + 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 + ray.kill(lock_manager) + + +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_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 + 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) + 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_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_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) + # 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 + + +def test_renewal_and_release_run_while_waiters_fill_the_actor(own_manager): + ray.get(acquire(own_manager, ["a"], "holder"), timeout=TIMEOUT_S) + 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 + 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) + 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: + assert isinstance(lease, tq.GlobalLease) and lease.names == [(P, "r")] + 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()) + + +@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)] + + +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(early), timeout=TIMEOUT_S) + ray.get(attached.wait_in_thread.remote([early, late]), timeout=TIMEOUT_S) + wait_until(lambda: holders() == {early} and tq.kv_lock_list(P)["waiters"] == 1) + + ray.get(attached.close.remote(), timeout=TIMEOUT_S) + # 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() == {early}) + 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) + held.__enter__() + worker = Worker.remote() + waiting = worker.cycle.remote(["z"], 1) + 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): + held.__exit__(None, None, None) + + wait_until(lambda: actors_named("TransferQueue") == []) # the controller and storage go too + with pytest.raises(RuntimeError, match="TQ_ENABLE_KV_GLOBAL_LOCK"): + tq.kv_lock_list() diff --git a/tests/test_kv_lock.py b/tests/test_kv_lock.py new file mode 100644 index 00000000..1a239ed9 --- /dev/null +++ b/tests/test_kv_lock.py @@ -0,0 +1,178 @@ +# 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 importlib +import threading +import time +from unittest.mock import MagicMock + +import pytest + +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 + + +@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) + + +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_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_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() + + +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(kvl._lock_settings()) + + 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) + 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/__init__.py b/transfer_queue/__init__.py index 754bb4d8..1d349a27 100644 --- a/transfer_queue/__init__.py +++ b/transfer_queue/__init__.py @@ -37,6 +37,15 @@ load_checkpoint, save_checkpoint, ) +from .kv_lock import ( + GlobalLease, + LockLostError, + async_kv_global_lock, + async_kv_local_lock, + kv_global_lock, + kv_local_lock, + kv_lock_list, +) from .metadata import BatchMeta, KVBatchMeta from .sampler import BaseSampler from .sampler.grpo_group_n_sampler import GRPOGroupNSampler @@ -65,6 +74,16 @@ "async_kv_clear", "KVBatchMeta", ] + + [ + # KV key locks: process-local and cluster-wide + "kv_local_lock", + "async_kv_local_lock", + "kv_global_lock", + "async_kv_global_lock", + "kv_lock_list", + "GlobalLease", + "LockLostError", + ] + [ # Checkpoint Interface "save_checkpoint", diff --git a/transfer_queue/config.yaml b/transfer_queue/config.yaml index c4da5fba..e0ac6599 100644 --- a/transfer_queue/config.yaml +++ b/transfer_queue/config.yaml @@ -15,7 +15,6 @@ controller: # ZMQ Server IP & Ports (automatically generated during init) zmq_info: null - 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 eca8762e..cb399ca0 100644 --- a/transfer_queue/interface.py +++ b/transfer_queue/interface.py @@ -45,6 +45,7 @@ from transfer_queue.client import TransferQueueClient from transfer_queue.controller import TransferQueueController +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 @@ -59,6 +60,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 +151,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`. @@ -194,13 +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 + lock_settings = _lock_settings() 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.") + _start_lock_managers(lock_settings) except ValueError: logger.info("Some other rank has initialized TransferQueueController. Try to connect to existing controller.") _init_from_existing() @@ -241,10 +249,10 @@ 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 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. @@ -252,6 +260,7 @@ def close(): global _TQ_CLIENT global _TQ_STORAGE global _TQ_CONTROLLER + global _TQ_IS_OWNER try: if _TQ_STORAGE: @@ -311,12 +320,14 @@ def close(): _TQ_CLIENT.close() _TQ_CLIENT = None - if _TQ_CONTROLLER: + _close_lock_managers() + if _TQ_IS_OWNER: try: ray.kill(_TQ_CONTROLLER) except Exception: pass - _TQ_CONTROLLER = None + _TQ_CONTROLLER = None + _TQ_IS_OWNER = False # ==================== Metrics API ==================== diff --git a/transfer_queue/kv_lock.py b/transfer_queue/kv_lock.py new file mode 100644 index 00000000..4ddff4f1 --- /dev/null +++ b/transfer_queue/kv_lock.py @@ -0,0 +1,593 @@ +# 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 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_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 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 +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 GetTimeoutError, RayActorError + +from transfer_queue.lock_manager import _RESERVED_SLOTS, TransferQueueLockManager +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, +# 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; 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: + 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, + 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() == [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))] + + +def _release_all(names: list[tuple[str, str]]) -> None: + with _state_lock: + for name in reversed(names): + _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``. + + 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 + held = _mark(_holding) + try: + for name in names: + _acquire(name, deadline) + acquired.append(name) + yield + finally: + _release_all(acquired) + 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 while it lasts.""" + names, acquired = _lock_names(keys, partition_id), [] + deadline = None if timeout is None else time.monotonic() + timeout + held = _mark(_holding) + try: + for name in names: + await _acquire_async(name, deadline) + acquired.append(name) + yield + finally: + _release_all(acquired) + held[0] = 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"] = {} +_renewer_running = False +_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. +_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 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()?" + + +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``: call ``check()`` before writes; ``names`` lists the locked 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._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 = {} + 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 self._closes != _closes or time.monotonic() >= self._deadline: + raise LockLostError(f"kv_global_lock lease on {self._names} was lost") + + +def _shard_count() -> int: + """Return the lock actor count, 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()) + 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 + + +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 _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(_manager_name(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() == [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") + 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())), num_shards, lease_s) + + +def _poll_window(deadline: float | None, lease_s: float) -> float: + # 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) + + +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, lease._num_shards) + + +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 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 + # 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 + _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() + return True + + +def _release_lease(lease: GlobalLease) -> None: + with _registry: + _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) + except Exception: # Ray is shut down + 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 = {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: + 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 = {} + for shard, tokens in by_shard.items(): + try: + calls[_managers[shard].renew_many.remote(tokens)] = shard + except Exception: # close() dropped the handle + pass + # 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) + 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. + 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 +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 ``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: 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. + + 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 + ``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 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. + """ + _reject_running_loop("kv_global_lock") + lease = _new_lease(keys, partition_id, lease_s) + deadline = None if timeout is None else time.monotonic() + timeout + 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. + for shard in lease._shards: + call, granted = None, False + while not granted: + 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(call, timeout=window + _REPLY_SLACK_S) + call = None + except GetTimeoutError: + remaining = None + except RayActorError as e: + 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 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() + yield lease + lease.check() + finally: + _release_lease(lease) + 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 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 + held = _mark(_holding_global) + try: + for shard in lease._shards: + call, granted = None, False + while not granted: + 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(asyncio.shield(call), window + _REPLY_SLACK_S) + call = None + except asyncio.TimeoutError: + remaining = None + except RayActorError as e: + 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 + lease.check() + finally: + _release_lease(lease) + held[0] = 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``. 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)} + + +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 + 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(settings: tuple[int, int] | None) -> None: + """Called by the ``tq.init()`` that created the controller: start the lock actors if enabled.""" + global _owned + 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. + _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, + scheduling_strategy=strategy, + ).remote(num_shards, max_waiting_acquires) + for shard, strategy in enumerate(get_node_round_robin_scheduling_strategies(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, _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: + if not _owned: # the owner kills the actors, so releasing would only hand keys to doomed waiters + _release_lease(lease) + 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, _owned + _state_lock, _table = threading.Lock(), {} + _registry, _leases, _renewer_running, _managers, _owned = 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..af0ff2b7 --- /dev/null +++ b/transfer_queue/lock_manager.py @@ -0,0 +1,222 @@ +# 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 +from collections import deque + +import ray + +# A withdrawal whose acquire never arrives (e.g. the caller died first) is dropped after this. +_WITHDRAWN_TTL_S = 300 +# 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 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 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.""" + + 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_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 + # _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") + # 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 (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] = {} + 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._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.""" + 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] + 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: + lease = self._leases.pop(token, None) + if lease is None: + return False + for name in lease["names"]: + del self._holder[name] + self._hand_off(lease["names"]) + return True + + 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, 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_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. + """ + # 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 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 + return None + 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(), 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: + 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() + + 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]: + """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 lease is not None and 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 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 + # 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.""" + 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, waiter["names"][0][0]) for waiter in self._waiting.values()) + return {"holders": holders, "waiters": waiters}