[feat] support the kv local and global locks

[OutstanderWang]

Summary

Adds advisory locks for TransferQueue KV keys, so read-modify-write sequences such as kv_batch_get → compute → kv_put don't lose updates when several threads or Ray workers touch the same key.

• kv_local_lock: a process-local mutex per (partition_id, key), shared by all threads and asyncio tasks of one process. It needs no configuration and makes no Ray calls.
• kv_global_lock: a cluster-wide mutex per (partition_id, key), held under an auto-renewed lease by TransferQueueLockManager Ray actors. It is opt-in: setting the environment variable TQ_ENABLE_KV_GLOBAL_LOCK=1 (parsed like other TQ boolean flags: 1/true/yes/y/on) before the owning tq.init() starts the actors, and keys are sharded across TQ_KV_GLOBAL_LOCK_NUM_SHARDS of them (default 8), spread evenly over the cluster's nodes.
• close() now tears TransferQueue down only in the process whose init() created it. Before this change, close() in any attached worker killed the controller for the whole cluster. The global lock relies on the same ownership rule for its actors.

Both locks are advisory: KV calls never check them, so every writer that needs mutual exclusion has to take the lock.

Usage

import os

import torch
import transfer_queue as tq

os.environ["TQ_ENABLE_KV_GLOBAL_LOCK"] = "1"  # only kv_global_lock needs this; set before the owning tq.init()
tq.init()

# Within one process (threads / asyncio tasks):
with tq.kv_local_lock("counter", "train", timeout=5):
    v = int(tq.kv_batch_get("counter", "train")["v"][0])
    tq.kv_put("counter", "train", fields={"v": torch.tensor([v + 1])})

# Across the Ray cluster (any actor or task, any node):
with tq.kv_global_lock(["a", "b"], "train", timeout=30) as lease:
    ...                 # long section: call lease.check() right before writing
    lease.check()       # raises tq.LockLostError if the lease may have passed to someone else
    tq.kv_batch_put(["a", "b"], "train", fields=td)

# Inspect live global locks:
tq.kv_lock_list("train")  # {"holders": [{"key", "holder", "held_s", "lease_remaining_s", ...}], "waiters": n}

Async context managers: async with tq.async_kv_local_lock(...) and async with tq.async_kv_global_lock(...). kv_global_lock yields a tq.GlobalLease, whose public surface is check() (call it right before writes) and the read-only names (the locked (partition_id, key) pairs); its other state is private.

Configuration

There is no config.yaml section; two environment variables, read only by the tq.init() that creates TransferQueue (set them before that call), configure the lock. Attaching processes and Ray workers find the lock actors by name and need not set them.

Variable	Default	Meaning
TQ_ENABLE_KV_GLOBAL_LOCK	off	1/true/yes/y/on starts the lock actors; while off, kv_global_lock raises a RuntimeError telling the user to set it
TQ_KV_GLOBAL_LOCK_NUM_SHARDS	8	number of lock actors (integer >= 1, validated only when enabled)
TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES	10000	acquire calls each lock actor lets wait at once (integer >= 1, validated only when enabled); see Capacity limits

The number of lock actors (shards) is TQ_KV_GLOBAL_LOCK_NUM_SHARDS, 8 by default, for every backend; an earlier revision of this PR tied the count to the storage unit count, which is no longer the case. Both integer variables must be decimal integers >= 1; anything else (0, -1, abc, true, 2.5) raises a ValueError naming the variable before tq.init() creates any actor. tq.init() pins them round-robin to the alive Ray nodes, sorted by node ID, with hard NodeAffinitySchedulingStrategy through the helper that already places SimpleStorage units: with fewer nodes than shards, per-node counts differ by at most 1; with more, each shard gets its own node. README.md documents the lock under the KV API.

Capacity limits:

• A waiting request keeps one acquire call waiting, on the actor it is currently waiting at. A request that spans several shards takes them one at a time, so it uses one waiting call, not one per shard. A reply that comes late is awaited again rather than re-sent, so a request never has more than one call in flight.
• Each actor lets at most TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES acquire calls wait at once. tq.init() sets the actor's Ray max_concurrency to that plus 16 slots (_RESERVED_SLOTS) for calls that never wait: renew, release, list_locks and acquires turned away. Parked acquires therefore cannot occupy the slots renewal and release need. The reserved slots give no priority, though: they are shared with turned-away acquires and list_locks, and calls beyond all slots queue in Ray's FIFO, behind turned-away acquires.
• Past the limit, an acquire keeps its place in the queue but returns at once, and the caller polls again after its window (5 s, or lease_s / 3 if shorter). Such a request may see its grant up to one window late.
• With keys hashing evenly, about shards × TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES requests can wait without that delay (80,000 by default). All waiters on one key land on one actor, so a hot key is bounded by TQ_KV_GLOBAL_LOCK_MAX_WAITING_ACQUIRES alone.
• Every waiter, parked or polling, still costs its actor one call per window, so the limit does not reduce call volume. Past an actor's call throughput, renewals are delayed and leases can be lost. With the earlier-revision measurement of a few thousand acquire calls/s per actor (see Performance) and a 5 s window, that is on the order of tens of thousands of waiters on one shard, which a single hot key can reach.

Rules

• keys is a key or a list of keys in one partition; duplicates are ignored and an empty list raises ValueError.
• timeout=None waits forever; otherwise TimeoutError is raised and nothing stays held. For the global lock, 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.
• Nesting the same kind of lock raises RuntimeError; pass all keys to one call instead. Tasks created inside an async block cannot take the same kind of lock while the block lasts.
• Take a global lock before a local one: taking a global lock while holding a local one raises RuntimeError.
• The sync versions raise RuntimeError on a thread with a running event loop; use the async versions there.

Design

Local lock

• One table of (partition_id, key) → FIFO waiter queue, guarded by a threading.Lock. No Ray calls.
• A release hands the key directly to the oldest waiter, whether that waiter is a thread or an asyncio task (woken through its loop).
• A multi-key request takes keys in sorted order under one deadline, so overlapping requests cannot deadlock.
• If a timeout races a grant, the grant wins. A cancelled waiter passes the key on to the next one.
• Forked children start with an empty table.

Global lock

• Actors
  • With TQ_ENABLE_KV_GLOBAL_LOCK set, the owning init() creates TransferQueueLockManager_0..N-1, N being TQ_KV_GLOBAL_LOCK_NUM_SHARDS (async, num_cpus=0), placed round-robin over the alive nodes. Their naming, creation, lookup and teardown all live in kv_lock.py.
  • A key goes to actor crc32(partition_id + key) % N. Python's hash() is salted per process, so it can't be used here.
  • Other processes find the actors by name on first use and read N from actor 0, so a worker that only takes locks does not need tq.init(). The async path does this lookup in a worker thread, so it never blocks the event loop.
  • Every acquire carries the N its keys were hashed with, and an actor rejects a mismatch with ValueError. A process that attached before TransferQueue restarted gets the "actor is gone" error first, because its cached handles point to the dead actors. The check covers the narrower case of a process whose cached N predates a restart with another TQ_KV_GLOBAL_LOCK_NUM_SHARDS but which finds the new actors by name, such as a child forked after its parent's first lock, or a waiter whose own process runs close() and init() mid-wait.
• Granting
  • Within one actor, all keys of a request are granted at once or not at all.
  • A request that spans actors visits them in ascending shard order, which keeps it deadlock-free. The cost is that keys already granted on earlier shards stay held while later shards are awaited. The hash decides which keys share an actor (with N shards, two keys land on different actors N-1 times in N), so the docstring spells out what this means for callers.
  • Each actor serves requests in arrival order. It keeps a FIFO queue per key and grants a request only when all its keys are free and it is first in each of their queues. Newcomers cannot barge ahead of older waiters (except past a waiter skipped as abandoned, see below), and single-key traffic cannot starve an older multi-key request. The cost is head-of-line blocking: a request waits behind every earlier one that shares a key with it, even while its own keys are free.
  • Release, expiry and withdrawal hand the keys straight to the next eligible waiter and wake only that waiter. Expiries are tracked in a min-heap, and a timer frees each lease when it runs out.
  • Acquires are long-polled. One call waits inside the actor for at most 5 s, capped at lease_s / 3; the client then repeats acquire with the same token, which keeps the queue position. Each reply is awaited with a bound, so an unresponsive actor cannot outlast the caller's timeout by more than 1 s; a reply later than that is awaited again, not re-sent, so calls never pile up in the actor. Parked acquires never take the reserved slots, but renew and release share them with no priority (see Capacity limits).
  • If a grant expires before its caller sees it (for example, the caller stalls between two polls), the repeated acquire queues the token again instead of reporting "not granted" forever.
  • A waiter with no call in flight for 10 s is taken for dead, because a live caller calls again once per poll window (5 s at most), even when a full actor turned it away. Hand-off skips it rather than granting it a lease that would only sit idle until it expired. 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; this is the one way a later request can be granted before an earlier one.
• Leases
  • The default lease is 30 s. One daemon thread per process renews all of that process's leases every lease_s / 3, with one batched call per shard. It sends all shards' calls at once and waits at most one period for their replies, so a shard that does not answer delays no other shard's renewals; its own leases keep their deadline and run out once it passes. Renewal starts as soon as the first shard grants, so waiting on a later shard cannot let an earlier shard's lease run out.
  • The client's local deadline counts from the send time of the call that granted or last renewed the lease, so it never outlives the actor's lease.
  • A refused renewal, a dead lock actor, or a passed local deadline loses the lease: check() raises, waiting for later shards stops with LockLostError, and exiting the block raises unless the body already raised. The lease's state is private, so check() is the only way to ask.
  • 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.
• Withdrawal: a caller that gives up (timeout, KeyboardInterrupt, task cancellation) sends release(token). This withdraws a request that is still queued or in flight, so a caller that has left is never left holding the lock; a grant made just before the release arrives is freed by it, and a withdrawn token cannot queue again even if one of its calls runs after the release. Exit sends the release without waiting, so kv_lock_list may briefly still list a holder that just left.
• Lifecycle
  • The owner's close() kills all lock actors.
  • An attached process's close() releases its own leases and leaves the actors running. A thread of that process still waiting for keys on a later shard stops with LockLostError rather than queuing again.
  • If the actors die, waiters get RuntimeError and holders get LockLostError. close() stops this process's renewer right away, even while blocks are still open.

Close ownership

init() records ownership only when it actually creates the controller. In other processes, close() closes the client and drops the handles, so they can attach again later. Storage cleanup was already owner-only.

Performance

All numbers below were measured on earlier revisions of this PR and have not been re-measured on the current one. None includes the last two changes (one acquire in flight per shard and the reserved slots), and the 2-node cluster numbers also predate the FIFO queue, long-polling and the abandoned-waiter skip.

Measured on 2 nodes (about 375 CPUs each) with Ray 2.51.1. Client actors were spread across both nodes; each lock is an acquire plus a release round trip.

Uncontended single-key locks, distinct keys per client (locks/s, with p50 latency):

clients	1 lock actor	8 lock actors (default)
1	961 (1.0 ms)	884 (1.1 ms)
4	1,936 (2.1 ms)	3,820 (1.0 ms)
16	1,778 (8.9 ms)	11,272 (1.3 ms)
64	1,785 (35.8 ms)	13,642 (2.0 ms)

• A single actor plateaus at about 1.8k locks/s. With 8 shards, throughput keeps scaling with the number of clients.
• With 1000 held leases and 1000 blocked waiters on other keys, an unrelated client still gets about 1,000 locks/s, the same as when the actors are idle. Without per-key wakeups, every release woke every waiter; on a single machine this cut the same scenario from about 1,600 to 40 locks/s. These cluster numbers predate the FIFO change.
• A single hot key is served by one actor, so sharding doesn't help it. On the 2-node cluster, before the FIFO change, 10, 100 and 1000 waiters on the same key got 1,162, 677 and 136 locks/s: every release woke all of them. Now a release wakes only the waiter it grants. On one 18-core machine (4 client actors, 5 s per run, two runs each), 12, 100 and 1000 waiters went from 1,277–1,395, 943–944 and 253–256 locks/s to 1,400–1,408, 1,411–1,441 and 1,429–1,450 locks/s. Long-polling did not change this: in one back-to-back session, 12 and 100 waiters got 1,208–1,285 and 1,124–1,223 locks/s before it and 1,258–1,281 and 1,244–1,311 locks/s after. I have not re-measured this on the cluster yet.
• Idle polling has a cost: one actor serves about 4,700 acquire calls/s. On one 18-core machine, waiters polled a key held for the whole run while another client locked and unlocked a free key on the same actor. With no waiters it got 1,629 locks/s. With a 1 s poll window, 1,000, 4,000 and 8,000 waiters cut that to 1,282, 175 and 1 locks/s, which would starve renewals too. With a 5 s window they cut it only to 1,573, 1,223 and 1,087 locks/s, which is why the window is 5 s.

Tests

The suite keeps only the core cases: 16 tests in two files.

• tests/test_kv_lock.py covers the local lock: exclusion across threads and asyncio tasks on several event loops (an exact read-modify-write counter), FIFO order, and a timeout that is one deadline for all keys and releases the keys already taken. It also checks, against 3 fake nodes, that the default 8 lock actors are pinned 3/3/2 and TQ_KV_GLOBAL_LOCK_NUM_SHARDS=3 gives one per node, all with hard node affinity, and that invalid shard or waiting-limit values raise.
• tests/e2e/test_kv_global_lock_e2e.py covers the actor-backed global lock:
  • an exact counter across Ray actors that lock [a, b] and [b, a] on two shards, without deadlock;
  • a multi-shard timeout releases the earlier shard and leaves no waiter;
  • waiters are granted in arrival order: a single-key request cannot pass an older multi-key one, and a withdrawn request cannot queue again;
  • a grant that expires before its caller sees it is queued again, and a waiter that stopped polling is skipped;
  • renewal and release still run while waiting acquires fill an actor, whose extra waiters are answered at once and keep their places, and a late reply is awaited again instead of re-sent;
  • an acquire carries the shard count its keys were hashed with;
  • a killed holder's keys free up once its lease expires, leases renew, and a lost lease raises without masking an exception from the body;
  • a shard that stops answering renewals costs no lease that lives only on another shard;
  • an attached process's close() stops its thread that waits at a later shard with LockLostError, and leaves the lock actors and controller running;
  • the owner's close() removes every TransferQueue actor, fails waiters, and stops the renewer while a block is still open.
• The focused local runs pass these 16 tests in about 60 s; CI runs the full suite. On an earlier revision, before the FIFO queue, long-polling, the abandoned-waiter skip and the reserved slots, a 2-node global check also stayed exact: 16 actors each did 50 counter increments and 50 random-order two-key updates, with no locks left afterwards. I have not rerun it on the cluster since.

[0oshowero0]

[AI Review] Summary

Reviewed for simplicity, maintainability, interface clarity and reliability. All 40 tests added by this PR pass locally (Ray 2.56, Python 3.11, SimpleStorage). Items 1–8 are inline comments; repro snippets are attached where a finding was reproduced.

• Must fix: 1) a cross-shard wait silently loses the earlier shards' locks; 2) kv_local_lock fails in any process that attached via tq.init().
• Suggested before merge: 3) fencing tokens; 4) cross-shard hold-and-wait; 5) fairness; 6) the *_locked wrappers; 7) make the lock actors opt-in; 8) module boundaries.

What works well: the process-local queue handles the hard edges carefully (a grant racing a timeout, waiters on closed loops, cancelled waiters passing the key on); withdraw-by-token, crc32 routing and per-key wakeups are solid; ascending shard order does make cross-shard requests deadlock-free; and the close-ownership fix is correct.

Minor

9. _holding is copied into asyncio tasks created inside an async_kv_local_lock block and stays True after the block exits, so such a task later gets "Nested kv_local_lock" even for an unrelated key (reproduced). (kv_lock.py L54, L200)
10. On first use the async paths block the event loop: _shard_count() calls ray.get and _lock_manager() calls ray.get_actor (kv_lock.py L294–319, reached from async_kv_global_lock via _new_lease).
11. The sync acquire calls ray.get without a client-side timeout (kv_lock.py L428), so if a lock actor is alive but unresponsive the caller's timeout is not honored.
12. With max_concurrency=10_000 (lock_manager.py L30), once that many acquires wait on one actor, release and renew_many queue behind them and holders lose their leases.
13. _num_shards is cached per process (kv_lock.py L246). If the owner re-inits with a different num_lock_shards, processes route the same key to different actors and exclusion breaks silently.
14. _expire only runs inside acquire (lock_manager.py L78) while every renewal pushes a heap entry (L125), so the heap keeps growing while locks are renewed but no new acquires arrive. _withdrawn is also rebuilt on every release that frees nothing (L132).
15. _reset_after_fork (kv_lock.py L539–546) targets forked children, which cannot use Ray safely anyway (see StreamingDataset._create_client), so the lock path it resets is not usable there.
16. GlobalLease is what kv_global_lock yields, but it is not exported from transfer_queue.

Related, outside this PR

17. The description says each KV call is safe on its own, but concurrent tag updates on one key already lose updates on main: kv_put(tag=...) reads the current tags, merges on the client, and the controller's set_custom_meta replaces the per-sample dict. Two threads writing different tag names to the same 300 keys left every key with only one of the two tags (sequential calls keep both). This matters if retry logic keeps an attempt id in a tag; merging per key in the controller, in a separate PR, would fix it.

Repro for 17

keys = [f"k{i}" for i in range(300)]
for key in keys:
    tq.kv_put(key, P, fields={"x": torch.tensor([0])})
barriers = [threading.Barrier(2) for _ in keys]

def writer(tag_name):
    for key, barrier in zip(keys, barriers):
        barrier.wait()
        tq.kv_put(key, P, tag={tag_name: 1})

threads = [threading.Thread(target=writer, args=(name,)) for name in ("a", "b")]
for t in threads:
    t.start()
for t in threads:
    t.join()
tags = tq.kv_list(P)[P]
print(sum(not {"a", "b"} <= set(tags[k]) for k in keys))  # 300

[0oshowero0]

[AI Review] 1 · Must fix: a cross-shard wait silently loses the earlier shards' locks

_start_lease runs only after every shard has been granted, so while this loop waits on a later shard, the leases already granted on earlier shards are not renewed. If that wait exceeds lease_s (30 s by default: a hot key, a long critical section, or a crashed holder whose lease has to run out), the earlier keys expire on their actor and can be granted to someone else. kv_global_lock then yields without checking and kv_global_locked calls fn without checking, so the body writes without exclusion and LockLostError only surfaces after the write. With 8 shards, two keys land on different shards 7/8 of the time.

Reproduced with lease_s=1: inside fn, kv_lock_list showed the earlier key held by another process, and LockLostError was raised only after fn returned.

Suggestion: renew each shard as soon as it is granted (the renewer then has to renew only granted shards; registering the whole lease up front would make ungranted shards report it lost), call lease.check() before yielding, and add a test that waits on a later shard for longer than lease_s.

Repro

P = "repro"
shard_of = lambda k: kvl._shard((P, k), kvl._shard_count())  # kvl = transfer_queue.kv_lock
keys = sorted((f"k{i}" for i in range(50)), key=shard_of)
early, late = keys[0], keys[-1]  # shard 0 and shard 7

@ray.remote
class Holder:
    def __init__(self):
        tq.init()

    def hold(self, key, secs):
        with tq.kv_global_lock(key, P, timeout=30):
            time.sleep(secs)

    def grab_after(self, key, delay, hold_s):
        time.sleep(delay)
        with tq.kv_global_lock(key, P, timeout=30):
            time.sleep(hold_s)

a, b = Holder.remote(), Holder.remote()
a.hold.remote(late, 4)              # the later shard stays busy for 4 s
time.sleep(0.5)
b.grab_after.remote(early, 2.0, 5)  # another client asks for the earlier key at ~2.5 s

def fn(keys, partition_id):
    holders = {h["key"]: h["holder"]["pid"] for h in tq.kv_lock_list(P)["holders"]}
    print("early key held by", holders.get(early), "| my pid", os.getpid())

tq.kv_global_locked(fn, [early, late], P, lease_s=1, lock_timeout=30)
# early key held by 67541 | my pid 67283
# then: LockLostError: kv_global_lock lease on [...] was lost

[OutstanderWang]

Fixed in 7bc0ee0. Each shard's lease is now registered with the renewer as soon as that shard grants, and the renewer renews only the shards that have granted. If a shard is granted after a renewal batch was snapshotted, that batch does not extend the local deadline. kv_global_lock calls lease.check() before yielding, so a lease lost while later shards were awaited raises LockLostError on entry instead of after the write.

test_waiting_on_a_later_shard_keeps_the_earlier_keys_held holds the later shard for 3 s against a 2 s lease and checks that the earlier key stays ours and the block is entered. It fails without the fix (the earlier key drops out of kv_lock_list).

[0oshowero0]

[AI Review] 2 · Must fix: kv_local_lock fails in every process that attached via tq.init()

The process-family actor is created only in the owner's init() (interface.py L222–225). In a Ray actor or task that calls tq.init() and attaches to the existing controller, _process_family_id() mints a new id, the actor lookup fails, and kv_local_lock raises RuntimeError: TransferQueue lock manager not found; call tq.init() first even though tq.init() was called. It also writes that orphan id into os.environ, so any child process spawned later inherits it. Forked children such as DataLoader workers cannot use Ray safely either (see StreamingDataset._create_client), so in practice the process family covers only the owner plus children that run ray.init themselves.

Suggestion: go back to the purely in-process lock from the first commit, scoped to the threads and tasks of one process, and point cross-process users to kv_global_lock. That also removes the per-acquire actor round trip (0.56 ms p50 uncontended, measured locally), one resident actor, the renewer dependency, and the LockLostError the local lock can currently raise on exit without giving the caller a lease to check.

Repro

@ray.remote
def local_lock_in_attached_worker():
    tq.init()  # attaches to the controller created by the driver
    with tq.kv_local_lock("k", "p", timeout=5):
        return "acquired"

ray.get(local_lock_in_attached_worker.remote())
# RuntimeError: TransferQueue lock manager not found; call tq.init() first

[OutstanderWang]

Fixed in 0dcbcb0: kv_local_lock is back to the purely in-process lock (a FIFO queue per key under a threading.Lock, with direct handoff), with no actor, no round trip and no LockLostError. Cross-process exclusion is kv_global_lock's job, as the module docstring now says. The SimpleStorage read-modify-write test now runs under the default config, where no lock actors exist at all.

[0oshowero0]

[AI Review] 3 · Fencing tokens: enforce them on the write path, or drop them

Fencing is the only defense against a holder that stalls past its lease and keeps writing (for example, a slow rollout attempt that has already been retried), but nothing checks lease.fence today, so it provides no protection yet. Meanwhile _fences keeps an entry for every (partition_id, key) ever locked, even after kv_clear, which grows without bound in long jobs keyed by sample uid. Either have the KV write path reject stale fences, or remove the field until it does.

[OutstanderWang]

Dropped in a8b52ac. Nothing enforced the fence, and _fences grew with every key ever locked. Callers who need protection against a stalled holder can call lease.check() right before writing; real fencing would need the KV write path to check it, which is a separate change.

[0oshowero0]

[AI Review] 4 · Spell out what cross-shard hold-and-wait means for callers

This docstring notes that granted keys stay held, but the impact is larger than it reads:

• Whether a request is all-or-nothing depends on hashing the caller cannot see: [a, b] on one shard never holds a while waiting for b; on two shards it does. Changing num_lock_shards changes this.
• Cross-shard is the common case: with 8 shards, 2 keys span two shards 7/8 of the time, and a 16-key batch touches about 7 shards, taking 7 sequential round trips while holding the earlier ones.
• Holding while waiting chains: H held z (shard 7) for 3 s, X asked for [a, z] (a on shard 0), and Y, which only needed a, waited 2.41 s. With a and z on the same shard, Y waited 0 s.

Suggestion: state these consequences in the kv_global_lock / async_kv_global_lock docstrings, and advise locking the smallest key set and keeping critical sections short.

[0oshowero0]

[AI Review] 4 (follow-up) Since each shard only sees its own keys, three more effects:

• A request re-queues at every later shard: time already waited on earlier shards gives it no priority there, so its total wait adds up shard by shard. Granting in arrival order (5, follow-up) does not change this.
• If it times out on a later shard, the keys it held on earlier shards blocked others for nothing.
• kv_lock_list shows it as the holder of the earlier keys with a growing held_s, which reads like a long critical section; consider marking holders that are still acquiring.

[OutstanderWang]

Done in d36d05e. The kv_global_lock docstring now states that:

• 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 span two actors 7 times in 8);
• a request queues afresh at each later actor, so its waits add up;
• the keys it holds meanwhile block others even if it then times out.

It advises locking the fewest keys possible and keeping the block short. kv_lock_list also notes that a holder may still be waiting for its keys on other actors.

Related: 7bc0ee0 (renew from the first grant) fixes the case where such a wait outlasted the earlier shards' leases.

[0oshowero0]

[AI Review] 5 · No fairness guarantee on the actor side

A release sets every waiter of the freed keys and whoever re-checks first wins; nothing reserves the key for the oldest waiter, and a newly arriving acquire can take it too. So "wake oldest first" orders wakeups, not grants. Locally, with two threads cycling on each of k1 and k2 (holds of 10–30 ms), a caller asking for [k1, k2] waited 0.3–1.5 s while each key was granted to others 20–80 times; FIFO would bound that at one or two holds (about 20–40 ms). Multi-key waiters suffer most, since all their keys must be free at the moment they re-check. Please document that grants are not fair; if hot keys are an expected workload, granting in arrival order would fix it (sketch in the reply below).

[0oshowero0]

[AI Review] 5 (follow-up) · Cost of waking every waiter, and granting in arrival order

Measured by running the lock-manager class directly under asyncio (Python 3.11), one hot key:

waiters on the key	actor CPU per release	release → winner's reply handed to Ray
1	0.22 ms	0.30 ms
1,000	7.1 ms	8.5 ms
10,000	101 ms	108 ms

Every woken waiter re-checks and suspends again (a new task and timer, about 3–10 µs each), and the winner's result reaches Ray only after all losers have done so, because its task-done callback queues behind them on the loop. This matches the 136 locks/s with 1,000 waiters in the description. It never takes more than one core per lock actor, but while it runs every other key on that shard, renewals included, waits. With 2–3 contenders per key (OPD, a retried rollout) the cost is negligible.

If hot keys are expected, a sketch that keeps all-or-nothing per shard:

• Keep waiters in arrival order. On every release, expiry or withdrawal, walk them oldest first and grant each one whose keys are neither held nor wanted by an older waiter that is still waiting; wake only the ones granted. The walk is set checks only, so a release no longer resumes every waiter.
• A new acquire follows the same rule: granted at once only if it conflicts with no holder and no waiter, otherwise it joins the tail.

This stays deadlock-free: a waiter waits only on a holder that is running or waiting on a later shard, or on an older waiter on the same shard, so wait chains cannot cycle (given no await while enqueueing and one shard count everywhere). The parts that need care are:

• a timeout racing a grant: record the grant before waking, and have the woken side check it, as _settle does in the in-process queue;
• a grant to a waiter that has left: release(token) must also dequeue it and pass its keys on; a crashed client blocks the queue until its lease expires;
• an older waiter leaving without a release (timeout, withdrawal) can unblock younger ones, so the walk must run then too;
• a holder that dies without releasing: something must wake at its lease expiry to run the walk.

The trade-off: a free key can sit idle while an older waiter that wants it waits for its other keys, and cross-shard order is unchanged, since each later shard sees the request as a newcomer.

Suggestion: confirm the expected per-key contention for OPD and rollout retries first. If hot keys are not expected, documenting "not fair" is enough and 1 remains the priority.

[OutstanderWang]

Fixed in 191327a with your sketch. Each actor keeps a FIFO queue per key and grants a request only when all its keys are free and it heads every one of their queues. A newcomer joins the tail unless it conflicts with no holder and no waiter. Release, expiry and withdrawal hand keys straight to the next eligible waiter and wake only that one. Expiry runs on a timer instead of waiter polling.

On one 18-core machine (4 client actors, one hot key), 12, 100 and 1000 waiters went from 1,277–1,395, 943–944 and 253–256 locks/s to 1,400–1,408, 1,411–1,441 and 1,429–1,450 locks/s. The e2e tests cover a single-key request not passing an older multi-key one, a withdrawing waiter letting the next one through, and expired leases passing the key down the queue.

Since 420b822, a timed-out call keeps its queue position until the client withdraws it, so long waits no longer pin actor slots (item 12 in the summary).

[0oshowero0]

[AI Review] 6 · Consider dropping the four *_locked wrappers

They save repeating keys and partition_id, but add an implicit fn(keys, partition_id, **kwargs) convention, let lock_timeout / lease_s collide with fn's own keyword arguments, and leave the caller no point to call lease.check() before writing (which makes 1 worse). with tq.kv_global_lock(...) as lease: is already one line, and dropping them shrinks the new public surface from ten names to six. The same applies to kv_local_locked / async_kv_local_locked (L216–238).

[OutstanderWang]

Dropped in d8909fe. The public lock API is now kv_local_lock, async_kv_local_lock, kv_global_lock, async_kv_global_lock, kv_lock_list, plus the LockLostError and GlobalLease types.

[0oshowero0]

[AI Review] 7 · Make the lock actors opt-in

Every owner init() now starts num_lock_shards + 1 resident actor processes whether or not locks are used; locally that was 9 processes and about 877 MiB RSS. Suggest a config switch so users enable the lock feature explicitly (for example disabled by default: no lock actors are created and the global lock APIs raise a clear error, while an in-process kv_local_lock as proposed in 2 keeps working). Since the lock actors are not part of the controller, a top-level lock: section in config.yaml holding the switch and the shard count would read better than controller.num_lock_shards.

[OutstanderWang]

Done in 510a48b. There is a new top-level section:

lock:
  enabled: false   # true starts the lock actors
  num_shards: 8

It replaces controller.num_lock_shards. With the default, tq.init() creates no lock actors, kv_local_lock works as before, and kv_global_lock raises RuntimeError naming lock.enabled: true as the fix. lock.num_shards is validated before the controller is created, so a bad value leaves nothing behind.

[0oshowero0]

[AI Review] 8 · Keep the lock-actor lifecycle in one module

interface.py imports private helpers from kv_lock (_detach_lock_managers, _process_lock_manager_name), the actor name TransferQueueLockManager_{i} is spelled in both files, and the actors are created here but looked up in kv_lock.py. Moving creation into kv_lock (one function that init() calls, next to the existing detach) keeps naming, creation, lookup and teardown in one place.

[OutstanderWang]

Done in 510a48b. kv_lock.py now owns the actor name (_manager_name), creation (_start_lock_managers, called by the tq.init() that created the controller), lookup, and teardown (_close_lock_managers, which releases this process's leases and kills the actors it created). interface.py no longer imports TransferQueueLockManager or keeps lock-actor handles; it calls _check_lock_conf, _start_lock_managers and _close_lock_managers.

[OutstanderWang]

Thanks for the thorough review. Items 1–8 are answered in their threads; here are the minor items from the summary. Each fix is its own commit on the branch.

9. Fixed in 224d15b. The "held" context variable now stores a cell that exit clears, instead of a bool. A task created inside the block copies the context, so it now sees the lock as released once the block exits. While the block lasts it is still refused, so it cannot deadlock waiting on a key its creator holds. The nesting test now also runs such a task after the block exits.
10. Fixed in f4566fe. On first use, _shard_count() now looks up every lock actor at once, and async_kv_global_lock runs that lookup through asyncio.to_thread while the cache is empty. A new test spies on ray.get and ray.get_actor and fails if either runs on the event-loop thread.
11. and 12. Fixed together in 420b822. acquire now waits inside the actor for at most 1 s. The client then calls a new wait(token) method, which keeps the queue position, so a waiter never holds a concurrency slot for longer than that. With 16 pollers on an actor limited to 4 slots, release got through in about 1.1 s; before this change it never would. Every reply is awaited with a bound, so an unresponsive actor cannot outlast the caller's timeout. A grant now returns the lease time left rather than the time waited, because a grant can land between two polls. Hot-key throughput was unchanged in a back-to-back run.
12. Fixed in 795a451. Every acquire carries the shard count its caller hashed over, and the actor raises ValueError on a mismatch, telling the caller to tq.close() and tq.init() again.
13. Fixed. The heap part was already fixed in 191327a: a timer now fires _expire at the earliest expiry, so stale heap entries drain even when no acquires arrive. d1d7a85 fixes the tombstone part. All tombstones share one TTL, so expired ones are always the oldest and are now pruned from the front instead of rebuilding the dict on every release.
14. Keeping _reset_after_fork. Since 0dcbcb0, kv_local_lock is purely in-process, so a forked child (a DataLoader worker, for example) can use it without Ray. The child would otherwise inherit _state_lock possibly held by a parent thread, plus keys held by threads it does not have. The global-state reset is a cheap extra that keeps a child from renewing or releasing the parent's leases.
15. Done in cca8aae. GlobalLease is now exported from transfer_queue.
16. Opened Concurrent kv_put(tag=...) on one key loses tag updates #188 with the repro and a suggested controller-side per-key merge.

Also removed after review: the *_locked wrappers (6) and fencing tokens (3). The description is updated to match the current design.

[0oshowero0]

[AI Review] Round 2

Re-checked at d36d05e. Items 1–16 are fixed as described and hold up locally: the 41 tests pass, and the earlier repros now behave as intended. For example, the multi-key request in 5 is granted in 29–46 ms instead of 0.3–1.5 s, and a hot-key handoff takes under 0.1 ms whether 1 or 10,000 waiters queue (it was 0.3 ms and 108 ms). Thanks for the thorough follow-up.

Three new items, all from the long-poll in 420b822: 18 (must fix), 19 (suggested), 20 (minor).

[0oshowero0]

[AI Review] 18 · Must fix: a grant the actor has forgotten makes the client spin

wait returns None both while the token is still queued and once the actor no longer knows it, and _take_grant treats every None as "still queued" (kv_lock.py L311–315). If a grant lands between two polls and the client does not poll again within lease_s (a stalled process, or calls backed up on the actor), the unrenewed lease expires and the actor drops the token. From then on every wait returns None at once, so the client loops on RPCs until its deadline, or forever with timeout=None, while the key sits free.

Locally, a client paused for 3.5 s with lease_s=1 made 9,421 wait calls in about 4.5 s after resuming, then raised TimeoutError at 8.1 s; another client took the free key in 2 ms. Suggest returning a distinct result for an unknown token, and raising LockLostError (or re-queueing) on the client.

Repro

# Child process: ray.init(address="auto"), kv_lock._send_acquire wrapped to count calls.
with tq.kv_global_lock("spin", "spin_p", timeout=8, lease_s=1):
    ...
# -> TimeoutError after 8.1s, calls={'acquire': 1, 'wait': 9421}

# Driver: hold "spin", start the child, wait until kv_lock_list shows it queued, then
os.kill(child.pid, signal.SIGSTOP)  # the client stalls
time.sleep(1.5)                     # its in-flight wait call times out on the actor
held.__exit__(None, None, None)     # the actor hands "spin" to the stalled client...
time.sleep(2.0)                     # ...whose 1 s lease runs out unrenewed
os.kill(child.pid, signal.SIGCONT)

[OutstanderWang]

Fixed in 5d2b183, by re-queueing rather than raising. wait is gone: every re-poll repeats acquire with the same token. While the token is queued, that keeps its place; once granted, it returns the time left on the lease; if the actor no longer knows the token (its grant expired before the caller saw it), it queues it afresh, or grants it at once if the keys are free. Each call blocks on the actor for up to one poll window, so the client can no longer loop on instant replies. I chose re-queueing over LockLostError because the caller never entered the block, so it has not lost anything it wrote under.

test_a_grant_that_expires_while_the_caller_stalls_is_taken_again stalls the client past its lease between two polls (it times out on the old code), and test_a_grant_that_expires_before_the_caller_sees_it_queues_afresh checks that the re-poll queues behind the current holder instead of returning at once.

[0oshowero0]

[AI Review] 19 · A waiter that dies while queued stalls the queue for a full lease

Since only release withdraws a waiter, a client killed while queued (a retried rollout worker, for example) keeps its place. When it reaches the head it is granted, and the key moves on only when that lease expires. Locally, with lease_s=5, the next waiter got the key 5.0 s after it was released; with the default that is 30 s.

Waiters now poll every second, which gives the actor a free heartbeat: dropping a waiter that has not polled for a few windows (from the _expire timer, for example) would bound this to a few seconds. With 18 in place, a live but slow client whose entry was dropped would then re-queue or raise instead of spinning.

[OutstanderWang]

Fixed in 7110613, using the poll as a heartbeat as you suggested. For each waiter the actor tracks how many of its calls are in flight and when the last one ended. A queued client sends its next acquire as soon as one returns, so a waiter with no call in flight for 5 s is treated as dead: hand-off drops it and moves on to the next waiter instead of granting it a lease that would only expire. The check runs at hand-off rather than from the _expire timer, since that is the only point where a dead waiter costs anything, and it needs no extra timer. A client that was only stalled re-queues at the back on its next call, through the idempotent acquire from 18.

One case remains: a waiter that died only a few seconds before reaching the head can still be granted, and the key then idles for that caller's lease. test_a_waiter_that_stopped_polling_is_skipped covers the dead-waiter case.

[0oshowero0]

[AI Review] 20 · Minor: polling cost of long waits

Each waiter now makes one call per second while it waits. With 1,000 waiters on a held key, the lock actor used about 32% of a core (0% when idle), so roughly 3,000 long waiters would saturate one shard. This should not matter at the contention OPD or rollout retries are expected to see. If long waits are expected, a longer window would cut the cost: renewals only need to land within lease_s, so a few seconds is still safe, at the price of a slower heartbeat for 19.

[OutstanderWang]

Addressed in fce1c2b: the poll window is now 5 s, capped at lease_s / 3 and at the remaining timeout, so a long waiter costs about a fifth of what it did. One actor serves about 4,700 acquire calls/s. With 1 s polls, 8,000 idle waiters on one shard cut an unrelated client's locks there from 1,629/s to 1/s; with 5 s polls it keeps about 1,100/s. The lease_s / 3 cap keeps renewals safe: even if waiters take every actor slot, a renewal waits at most one window for a slot, the same margin as the renewer's period.

ef84850 makes lock.max_concurrency configurable (default 10000) and documents the capacity in the kv_global_lock docstring. About num_shards * max_concurrency requests can wait at once in total, while all waiters on one hot key share one actor and are bounded by max_concurrency alone. Past that, calls queue in Ray, so renewal and release are delayed by at most one poll window; they are never stuck and no lock is lost.

The full per-lock call count (one acquire + one release per shard, plus polls and batched renewals) and uncontended overhead are in the queue-design comparison: #186 (comment)

[OutstanderWang]

Follow-up after re-reading the long-poll change critically. Four commits fix what I found; the PR description is updated to match.

• Spin on an unseen expired grant (5d2b183). If a grant expired between two polls, the actor no longer knew the token, so wait(token) returned "not granted" instantly and the client spun until its timeout (about 4.3M calls in 1 s in a local repro). wait is gone: re-polls repeat acquire with the same token, which keeps its queue place or queues it afresh. A client-level test stalls the caller past its lease between polls; it times out on the old code and passes now.
• Dead waiters idled the key (7110613). A waiter whose caller died was still granted the key, which then sat idle for that caller's whole lease, once per dead waiter in line (1.5 s behind three dead waiters with lease_s=0.5). The actor now tracks each waiter's calls in flight and when its last call ended. Hand-off skips one that has had no call in flight for 5 s; a caller that was only stalled queues again on its next call.
• Poll window (fce1c2b). One actor serves about 4,700 acquire calls/s. With 1 s polls, 8,000 idle waiters on one shard cut an unrelated client's locks there from 1,629/s to 1/s, renewals included; with 5 s polls it keeps about 1,100/s. The window is now 5 s, capped at lease_s / 3, so a renewal waiting for a slot behind every waiter still gets one well within its lease.
• Minor (04d2230). close() empties the lease registry and wakes the renewer, so it stops at once instead of calling dead actors until each open block exits. A grant's deadline is now set under the same lock the renewer uses. The kv_global_lock docstring notes that an unresponsive actor can stretch timeout by about 1 s.

Local verification: pre-commit, license and docstring checks pass. The full suite passes (709 passed) apart from the Yuanrong tests that need datasystem. The lock suites also passed three repeated runs.

[OutstanderWang]

Per-key FIFO queues vs. one FIFO per shard + scan

I compared the current TransferQueueLockManager (ef84850, one FIFO deque per key) with the alternative of a single FIFO per shard that is scanned on every hand-off.

Conclusion: keep the per-key queues. FIFO+scan never wins: it ties (within run-to-run noise) when few waiters queue on a shard and loses once many do. With a blocked key set it grants exactly the same requests, so it buys no fairness. It is also no smaller (214 lines vs. 209). Its hand-off cost grows with the number of waiters on the shard; the per-key hand-off cost does not. Uncontended latency is dominated by Ray call dispatch and by the number of shards a request touches, not by this data structure.

Designs compared

• perkey (this PR): _queues: dict[key, deque[token]]. A request is granted only when all its keys are free and it heads every one of their queues. _hand_off looks only at the head of each freed key's queue.
• fifo: one deque[token] per shard plus a _wanted counter (key → number of waiters). On hand-off it scans from the head with a blocked set: it grants a waiter whose keys are neither held nor blocked, otherwise adds the waiter's keys to blocked. Frees and arrivals on keys nobody waits for skip the scan.
• fifo_ee: fifo, but the scan stops once every freed key is held or blocked again.
• first-fit: fifo without the blocked set (grants any waiter whose keys are free). Included only to show why blocked is needed.

Cost analysis

N = waiters on the shard, k = keys per request, q = length of one key's queue (q ≤ N; q = N for a single hot key), H = entries in the expiry heap (live leases plus stale entries left by renewals; about 4 per lease with renewal every lease_s / 3).

Operation	perkey (current)	fifo (full scan)	fifo_ee
New acquire (grant or enqueue)	O(k)	O(k)	O(k)
Re-poll acquire with the same token	O(1) dict lookups	O(1)	O(1)
Hand-off when a lease is released or expires (plus O(k + log H) per grant in every design)	O(k²): one head per freed key, independent of N and q	Θ(N·k) whenever a freed key is wanted: always a full pass that rebuilds the deque	O(N·k) worst case; stops early once freed keys are re-blocked; still O(N·k) for head-of-line
Withdrawal (release of a waiting token)	O(k·q): deque.remove per key, plus a hand-off that skips held keys	O(N) deque.remove, then a full O(N·k) scan if anyone else wants its keys, even if they are held	O(N) + scan
Abandoned waiter (no call for 5 s)	removed when it reaches the head of a freed key: O(k·q), since it need not head its other queues; its keys are then re-inspected	dropped during the next scan at no extra cost	dropped if the scan reaches it before exiting
renew_many of T tokens	O(T log H)	same	same
Expiry	O(log H) per heap entry + hand-off	same	same
Memory	O(N·k) waiter state + one deque per waited-on key	O(N·k) waiter state + one deque + one counter per waited-on key	same as fifo

Notes:

• The per-key hand-off is O(k²) per release regardless of queue length: a granted waiter heads all its queues, so removing it from each is O(1). The two O(q) paths are withdrawals (a waiter timing out or exiting while queued) and abandoned waiters, because _leave_queues uses deque.remove. On a single hot key that is O(N), the same order as fifo's _fifo.remove. Both happen once per waiter that gives up, not on every lock/unlock, and perkey does not follow them with a scan of the whole shard. An abandoned waiter is only removed once it reaches the head of a freed key, so in a long chain of dead waiters one hand-off removes several of them; each waiter is still removed only once.
• fifo's scan is not just a worst case: every free of a wanted key walks all N waiters. Early exit fixes the hot-key case (the next waiter re-blocks the freed key at once), but not head-of-line: if the queue head waits on a key that stays held and the freed key belongs to the last waiter, fifo_ee still walks all N. It also needs an extra invariant ("only a waiter on a pending key can settle one") that is easy to break.
• Fairness is identical. perkey's "heads every one of its queues" and fifo's "no earlier queued waiter wants any of its keys" are the same condition, and both re-evaluate exactly when a key is freed or a waiter leaves. Without blocked (first-fit), a multi-key request can be starved by a stream of single-key requests.
• Common to both: every lock costs 1 acquire + 1 release Ray call per shard it touches (the release is fire-and-forget). A waiting caller adds one acquire per poll window (min(5 s, lease_s / 3), capped by the remaining timeout), and a holder's process adds one batched renew_many per granted shard every lease_s / 3. Shards are acquired one after another in ascending order, so latency grows linearly with the number of shards a request touches; several keys on one shard cost the same as one.

Measurements

Pure-Python cost of one release, including hand-off and the resulting grant, with N = 50,000 waiters queued (no Ray):

Scenario	perkey	fifo	fifo_ee
One hot key	1.7 µs	12.5 ms	1.9 µs
Multi-key: 2–4 of 100 keys	3.5 µs	17.0 ms	27 µs
Head-of-line: head blocked, freed key belongs to the last waiter	1.6 µs	16.2 ms	17.4 ms
N distinct keys, one waiter each	4.2 µs	16.5 ms	8.4 ms
Free a key nobody waits for	0.44 µs	0.46 µs	0.39 µs

fifo costs about 0.25–0.35 µs per queued waiter per release; perkey stays at or below 4.2 µs in every scenario from N = 10 to N = 50,000.

End-to-end through Ray, one lock shard, client actors driving the real async_kv_global_lock (locks/s, median of 2 runs):

Scenario	perkey	fifo	fifo_ee
Uncontended, each task its own key, 1 client × 16 tasks / 8 clients × 32 tasks	2768 / 2436	2769 / 2407	–
One hot key, N = 1000 / 5000 waiters	1355 / 1331	923 (−32 %) / 294 (−78 %)	1361 / 1347
Random key of 100, N = 1000 / 5000 waiters	1648 / 1707	902 (−45 %) / 344 (−80 %)	1384 (−16 %) / 1565 (−8 %)
Wait p50, 100 keys, N = 5000	2.77 s	≥ 4.0 s (longer than the run)	2.76 s

At N = 100 on one hot key the three land at 1118 / 1102 / 1336 locks/s. There the in-actor hand-off costs differ by at most 24 µs per release (0.3 µs between perkey and fifo_ee), a few percent of the time per lock at those rates, so that spread is Ray run-to-run variation, not the data structure.

Uncontended overhead of the current lock (p50 / p99):

Operation	Latency
kv_global_lock, 1 key (1 shard)	0.61 / 0.99 ms
kv_global_lock, 4 keys on 1 shard	0.62 / 0.92 ms
kv_global_lock, 2 / 4 / 8 shards	1.2 / 2.5 / 5.8 ms (p50)
Empty Ray actor call	0.26 / 0.32 ms
tq.kv_put, one small tensor	1.62 / 2.80 ms
tq.kv_batch_get, 1 key	1.01 / 1.52 ms
kv_local_lock, 1 key	2.5 µs

Ray calls counted inside the actors match the analysis: exactly 1 acquire + 1 release per shard per lock (1,000 back-to-back locks, 1 or 4 shards, no renew_many with lease_s = 30). Renewals are batched per process and shard: a 3 s body on 4 shards with lease_s = 0.3 made 112 renew_many calls, about 30 rounds × 4 shards. One shard sustains about 2.4k–2.8k uncontended locks/s, while the same actor serves about 8.3k trivial calls/s. Since a lock is two calls, the lock logic is not the bottleneck; Ray dispatch is.

Fairness check: after each of 20,000 random acquire / release / withdraw operations, perkey, fifo and fifo_ee had granted the same set of tokens (2,010 grants compared). In a starvation test, a waiter for {a, b} competes with a stream of single-key requests on a and on b: perkey and fifo grant it after 1 release; first-fit had not granted it after 10,000 releases.

Code size (lock manager only; the client side is shared): perkey 209 lines, fifo 214 lines, fifo_ee fifo + a 30-line override. _hand_off is 15 vs. 16 lines. FIFO+scan saves no code and replaces one structure with two (_fifo + _wanted).

Environment and method

Apple M5 Pro (18 cores), 48 GiB RAM, macOS (Darwin 25.6.0), Python 3.13.13, Ray 2.57.0. All processes ran on one local Ray cluster, so there is no cross-machine network latency. A multi-node cluster adds network latency to every Ray call, which leaves the in-actor lock logic an even smaller share of the total.

The benchmark is a local harness, not part of this PR. fifo is a full copy of TransferQueueLockManager with the per-key queues swapped for one FIFO plus _wanted; fifo_ee and first-fit override only its _hand_off. The pure-Python benchmark drives the actor classes directly on one event loop and times a single release with N waiters queued. The Ray benchmark uses lease_s = 30 and one lock shard for the throughput and contention runs; contention runs spread the N looping tasks over 4 client actors. Each 100-key run lasts 7 s, with the first 2.5 s and the last 0.5 s discarded. The exact run length of the hot-key runs was not recorded. The overhead runs use 8 shards and compare against a no-op, kv_local_lock, an empty actor call, kv_put and kv_batch_get.

[OutstanderWang]

Thanks for round 2. 18, 19 and 20 are answered in their threads (5d2b183, 7110613, fce1c2b + ef84850).

Two updates to my earlier reply on items 11 and 12, which describe a 1 s window and a wait(token) call:

• 11. wait was removed in 5d2b183; re-polls repeat acquire with the same token. The window is now min(5 s, lease_s / 3) (fce1c2b). Every reply is still awaited with a bound, so an unresponsive actor can stretch the caller's timeout by about 1 s, as the kv_global_lock docstring now says (04d2230).
• 12. ef84850 replaces the hard-coded max_concurrency=10_000 with lock.max_concurrency (default 10000, validated like lock.num_shards) and documents the limit: about num_shards * max_concurrency waiting requests in total, and max_concurrency for any one hot key. Past it, renewal and release are delayed by at most one poll window, so they are never stuck and holders keep their leases.

[0oshowero0]

[AI Review] Round 3

Re-checked at ef84850: 18, 19 and 20 hold up locally (48 tests pass). After an unseen grant expires, a stalled client re-queues with 3 calls and takes the key at 4.0 s instead of spinning; a waiter dead for 8 s is skipped at hand-off (0.0 s instead of 5.0 s; the dead-for-1 s case remains, as you noted); 1,000 long waiters cost the actor 10% of a core instead of 32%. The per-key vs FIFO+scan comparison is convincing, and keeping the per-key queues makes sense.

One new item, 21, inline.

[0oshowero0]

[AI Review] 21 · Past the concurrency limit, holders can lose their leases

This docstring (and the reply on 12) says that past the limit renewal and release are delayed by at most one poll window. Locally they got stuck instead. With lock.num_shards: 1, lock.max_concurrency: 4, one holder (lease_s=15) holding for 40 s and 16 waiters, none of the holder's renewals got through: its lease was lost at 25 s and exit raised LockLostError. After it released, none of the 16 waiters got the key before their 100 s timeouts. The same run with max_concurrency: 10000 keeps the lease.

Two causes:

• Ray serves queued calls in arrival order, so a renewal or release waits behind every acquire queued ahead of it, roughly (calls queued ahead / slots) windows rather than one. With the client patched never to re-send (below), the holder still lost its lease at 25 s.
• When a reply is late (window + _REPLY_SLACK_S), the client sends a new acquire while the old one stays queued and later takes a slot, so the backlog keeps growing. With that patch, 7 of the 16 waiters got the key instead of 0.

From reading the code (not measured separately): a call still queued in Ray does not count in calls, so past the limit a live waiter can look abandoned at hand-off and lose its place.

By the same arithmetic, the defaults (30 s lease, 5 s window) need roughly 7× max_concurrency waiters on one actor before a lease is lost, so 10,000 is far away. But max_concurrency is now configurable, and a hot key puts every waiter on one actor. Suggestions:

• Don't stack calls: while the previous call is pending, keep waiting on it instead of sending another.
• Keep slots free for renewal and release: have the actor cap parked acquires a little below max_concurrency and reply at once past the cap, with the client backing off one window before it polls again. Ray concurrency groups would also isolate them, but in an async actor each group runs on its own thread (checked locally), which breaks the single-loop assumption the lock state relies on.
• If neither is done, the docstring should say that past the limit a holder may lose its lease.

Repro

tq.init(OmegaConf.create({..., "lock": {"enabled": True, "num_shards": 1, "max_concurrency": 4}}))

def holder():
    with tq.kv_global_lock("sat", P, timeout=10, lease_s=15) as lease:
        time.sleep(40)  # lease.lost turns True at ~25 s; exit raises LockLostError

def waiter():
    with tq.kv_global_lock("sat", P, timeout=100, lease_s=15):
        pass  # all 16 raise TimeoutError

# Start the holder, then 16 waiter threads.

# No-resend patch: keep waiting on the pending call while it is not ready.
def send_once(lease, shard, window):
    ref = pending.get((lease.token, shard))
    if ref is not None and not ray.wait([ref], timeout=0)[0]:
        return ref
    ref = pending[(lease.token, shard)] = original_send_acquire(lease, shard, window)
    return ref

[OutstanderWang]

Confirmed. Our reply on 12 and the docstring were wrong: past the limit, renewal and release were not "delayed by at most one poll window, never stuck". Ray serves an actor's queued calls in arrival order, so once waiting acquires took every slot, a renewal waited behind all of them.

Reproduced at ef84850 with your setup (num_shards: 1, max_concurrency: 4, one holder with lease_s=15 holding for 40 s, 16 waiters with 100 s timeouts):

Run	Holder	Waiters that got the key
ef84850	no renewal got through; its deadline passed at 15.1 s and exit raised LockLostError	0/16, all timed out
ef84850 + your no-resend patch	the same, lost at 15.1 s	6/16
ef84850, max_concurrency: 10000	kept its lease	16/16, at 40.1 s
b256abe, max_concurrency: 4	kept its lease, clean exit	16/16, between 40.1 s and 60.1 s

Two things are worse than the comment says:

• In the no-resend run, a waiter entered its block at 25.0 s while the original holder stayed in its block until 40 s, so two holders were in the critical section at once. The first holder's check() would have raised from 15 s on, but a body that does not call it keeps writing.
• The threshold under the defaults is lower than 7×. Each stuck waiter added a call every window + slack (6 s), while an actor finishes max_concurrency parked calls per 5 s window. So the backlog grows without bound once one actor has more than about 1.2 × max_concurrency waiters, roughly 12,000 at the default.

Your third point holds too. In an actor-level probe with max_concurrency: 2, a waiter whose next call sat in Ray's queue for more than 5 s was skipped at hand-off as abandoned, and two later waiters got the key before it.

Fixed in three commits, along the lines of your first two suggestions:

• efe17ea corrects the docstring and the comments that claimed "never stuck".
• 5810475 keeps one acquire in flight per shard: a reply later than its window plus slack is awaited again instead of re-sent, so calls no longer pile up in the actor. test_a_late_reply_is_awaited_instead_of_sent_again (sync and async) counts 2 sends on the old code and 1 now.
• b256abe reserves slots. lock.max_concurrency now caps the acquire calls an actor lets wait, and tq.init() gives the actor _RESERVED_SLOTS (16) more for calls that never wait: renew, release, list_locks and acquires turned away. Past the cap, acquire keeps the caller's queue position but returns at once, and the client sleeps out the rest of its window before polling again, so it does not spin. Since such a waiter has no call in flight for up to a window, _ABANDONED_S goes from 5 s to 10 s. Waiting calls now stay within their slots, so Ray's queue stays short and a live waiter no longer looks abandoned because its call is stuck there. test_renewal_and_release_run_while_waiters_fill_the_actor fills both waiting slots of an actor sized for max_concurrency: 2; on the old code the extra acquire and the renewal time out, and now both return within 1 s.

I left concurrency groups out, for the single-loop reason you gave. The cost of the cap is that a waiter turned away by a full actor sees its grant up to one window late (in the repro above, the last waiters entered in 5 s steps). The docstring, the config.yaml comment and the PR description now say this.

[0oshowero0]

Test size and run time

At ef84850 the tests for this PR are 1,064 lines (43 test functions, 48 cases with parameters), against 747 lines of lock code (kv_lock.py 538 + lock_manager.py 209):

File	Lines	Tests
tests/test_kv_lock.py	413	14
tests/e2e/test_kv_global_lock_e2e.py	468	24
tests/test_interface_close.py	114	4 (8 cases)
tests/e2e/test_close_ownership_e2e.py	69	1

The four files take 126 s locally (Ray 2.56, Python 3.11). Most of that is cluster setup and sleeps: the SimpleStorage read-modify-write test in test_kv_lock.py takes about 16 s because it starts Ray and TransferQueue for a process-local lock, test_close_ownership_e2e.py about 15 s, the shard-hash check about 6 s (two subprocesses), and several e2e cases sleep 3–6 s each.

Several cases check the same behavior. For example, the single-key read-modify-write and the two-shard opposite-order test both assert an exact counter, the single-shard timeout test is a subset of the multi-shard one, and the client-level and actor-level tests for an unseen expired grant cover the same fix. Please cut these tests down as far as you can. A small set that exercises the main lock paths is enough; there is no need to keep a dedicated test for every edge case or every fix made during review.

[OutstanderWang]

Thanks for the test size review. The tests are cut down to the core cases in 175b6af and 9261e6a.

	Before (ef84850)	After (9261e6a)
Test lines	1,064 in 4 files	406 in 2 files
Test functions / cases	43 / 48	13 / 13
Local run time, same machine (Ray 2.57, Python 3.13)	161 s	46–50 s over three runs

What remains:

• tests/test_kv_lock.py, 3 tests for kv_local_lock:
  • test_threads_and_tasks_on_several_loops_serialize_read_modify_write
  • test_waiters_are_served_fifo
  • test_timeout_is_a_total_deadline_and_partial_acquisition_is_released
• tests/e2e/test_kv_global_lock_e2e.py, 10 tests for kv_global_lock:
  • test_opposite_key_orders_on_two_shards_do_not_deadlock
  • test_multi_shard_timeout_releases_the_earlier_shard
  • test_a_late_reply_is_awaited_instead_of_sent_again
  • test_waiters_are_granted_in_arrival_order
  • test_a_grant_that_expires_before_the_caller_sees_it_queues_afresh
  • test_a_waiter_that_stopped_polling_is_skipped
  • test_renewal_and_release_run_while_waiters_fill_the_actor
  • test_killed_holder_frees_the_key_after_its_lease
  • test_lease_renews_and_lost_lease_raises
  • test_owner_close_kills_the_managers_and_fails_waiters

Removed: tests/e2e/test_close_ownership_e2e.py, tests/test_interface_close.py (the non-owner close and lock config validation tests), the shard-hash and poll-window unit tests, the SimpleStorage local-lock test, and the other local and global edge-case tests.

The dead-waiter test no longer sleeps past the 10 s abandonment threshold: it backdates the waiter's last call through __ray_call__. The reserved-slot test now queues its two waiters one at a time (f61aded), because Ray may run two calls to an async actor out of order. When the second waiter went first, the test still passed but waited out a 30 s lease.