diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index aea34d3..7d167eb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -21,7 +21,7 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip - pip install pytest fakeredis click requests redis Pillow pydantic fastapi uvicorn typer + pip install pytest "fakeredis[lua]" click requests redis Pillow pydantic fastapi uvicorn typer - name: Run tests run: | diff --git a/modelq/app/base.py b/modelq/app/base.py index bc74e8f..8abdfd5 100644 --- a/modelq/app/base.py +++ b/modelq/app/base.py @@ -157,6 +157,25 @@ class ModelQ: # Registry of every in-flight list, so recovery never has to SCAN for them. INFLIGHT_REGISTRY = "inflight_lists" + # BLMOVE needs Redis >= 6.2. Plenty of deployed workers still run 6.0, where + # the command does not exist and every claim raises -- the worker crash-loops + # at ~30 restarts/second, registers nothing, and the queue silently never + # drains while the process still looks healthy. Fall back to an atomic Lua + # LPOP+RPUSH polled on a short interval: same FIFO order, same "taking the + # task and recording who took it is one step" guarantee, just not blocking. + # + # BRPOPLPUSH is NOT a valid substitute: it pops the tail, which turns the + # queue LIFO against an rpush producer. + _CLAIM_TASK_LUA = """ + local task = redis.call('LPOP', KEYS[1]) + if task then + redis.call('RPUSH', KEYS[2], task) + end + return task + """ + # How often the fallback re-checks the queue while waiting. + CLAIM_POLL_INTERVAL = 0.1 + def __init__( self, host: str = "localhost", @@ -209,6 +228,9 @@ def __init__( # Guarded because worker threads mutate it while the heartbeat thread reads it. self._inflight_tasks = {} self._inflight_lock = threading.Lock() + # Probed lazily on the first claim; see _blmove_supported(). + self._blmove_available = None + self._claim_script = None if server_id is None: # Attempt to load the server_id from a local file: server_id = self._get_or_create_server_id_file() @@ -372,6 +394,59 @@ def _update_task_history(self, task_id: str, task_dict: dict) -> None: # In-flight task liveness # # ------------------------------------------------------------------ # + def _blmove_supported(self) -> bool: + """ + Whether this Redis has BLMOVE (>= 6.2). Probed once, on a key that cannot + exist, and cached -- an unsupported server must not cost a round trip and + an exception on every single claim. + """ + if self._blmove_available is None: + try: + self.redis_client.blmove( + "modelq:blmove:probe", "modelq:blmove:probe", 0.01, "LEFT", "RIGHT" + ) + self._blmove_available = True + except redis.exceptions.ResponseError as e: + if "unknown command" in str(e).lower(): + logger.warning( + "Redis has no BLMOVE (needs >= 6.2); claiming tasks with the " + "polled Lua fallback instead. Queue order and in-flight " + "custody are unchanged." + ) + self._blmove_available = False + else: + # A different server-side error says nothing about support. + self._blmove_available = True + except (AttributeError, TypeError): + # redis-py older than 3.5 has no blmove() binding at all. + self._blmove_available = False + return self._blmove_available + + def _claim_task(self, inflight_key: str): + """ + Atomically move one task from `ml_tasks` into this worker's in-flight + list, blocking up to BLPOP_TIMEOUT. Returns the raw task JSON, or None if + nothing arrived in time. + """ + if self._blmove_supported(): + return self.redis_client.blmove( + "ml_tasks", inflight_key, self.BLPOP_TIMEOUT, "LEFT", "RIGHT" + ) + + # Polled equivalent for Redis < 6.2. The Lua body is atomic, so the task + # is never in neither list; only the waiting is emulated. + if self._claim_script is None: + self._claim_script = self.redis_client.register_script(self._CLAIM_TASK_LUA) + + deadline = time.time() + self.BLPOP_TIMEOUT + while True: + task_json = self._claim_script(keys=["ml_tasks", inflight_key]) + if task_json: + return task_json + if time.time() >= deadline: + return None + time.sleep(self.CLAIM_POLL_INTERVAL) + def _mark_task_inflight(self, task_id: str, started_at: float) -> None: """Record that this process is actively running `task_id`.""" with self._inflight_lock: @@ -1076,9 +1151,7 @@ def worker_loop(worker_id): # taking the task and recording who took it one atomic step, so # a task in transit to a worker that never receives it stays in # `inflight_key` until a sweep returns it to the queue. - task_json = self.redis_client.blmove( - "ml_tasks", inflight_key, self.BLPOP_TIMEOUT, "LEFT", "RIGHT" - ) + task_json = self._claim_task(inflight_key) if not task_json: continue diff --git a/pyproject.toml b/pyproject.toml index c2d5fea..e612457 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "modelq" -version = "1.0.18" +version = "1.0.19" description = "Celery-like task queue for ML inference." authors = ["Tanmaypatil123 "] readme = "README.md" diff --git a/requirements-test.txt b/requirements-test.txt index 48f7782..b8a6139 100644 --- a/requirements-test.txt +++ b/requirements-test.txt @@ -1,3 +1,3 @@ -fakeredis +fakeredis[lua] pytest click \ No newline at end of file diff --git a/tests/test_blmove_fallback.py b/tests/test_blmove_fallback.py new file mode 100644 index 0000000..dbd1873 --- /dev/null +++ b/tests/test_blmove_fallback.py @@ -0,0 +1,172 @@ +""" +Tests for claiming a task on Redis versions without BLMOVE. + +BLMOVE needs Redis >= 6.2. Every GPU worker container in production was still on +Redis 6.0.16, where the command does not exist. Upgrading such a worker to a +ModelQ that calls BLMOVE unconditionally produced: + + ERROR - Worker 0 crashed with error: unknown command `BLMOVE`, + with args beginning with: `ml_tasks`, `inflight:...`, `LEFT`, ... + +roughly thirty times a second, forever. The worker never registered, the queue +never drained, and nothing about the process looked unhealthy from the outside -- +`ps` showed it running and burning CPU. Requests simply queued and were answered +with an ever-growing ETA. + +The fallback must keep both properties the BLMOVE claim was introduced for: +atomicity (the task is never in neither list) and FIFO order. BRPOPLPUSH, the +obvious Redis 6.0 substitute, satisfies the first and breaks the second -- it +pops the tail, which turns an rpush-fed queue LIFO. Hence the Lua LPOP+RPUSH. +""" + +import json + +import fakeredis +import pytest +import redis as redis_lib + +from modelq import ModelQ + + +class RedisWithoutBlmove(fakeredis.FakeStrictRedis): + """A Redis 6.0-era server: everything works except BLMOVE.""" + + def blmove(self, *args, **kwargs): + raise redis_lib.exceptions.ResponseError( + "unknown command `BLMOVE`, with args beginning with: `ml_tasks`, " + "`inflight:abc:0`, `LEFT`, `RIGHT`, `15`, " + ) + + +@pytest.fixture +def old_redis(): + return RedisWithoutBlmove() + + +@pytest.fixture +def new_redis(): + return fakeredis.FakeStrictRedis() + + +def _queue(mq, *task_ids): + for task_id in task_ids: + mq.redis_client.rpush( + "ml_tasks", + json.dumps({"task_id": task_id, "task_name": "generate", "payload": {}}), + ) + + +def _ids(raw_items): + out = [] + for raw in raw_items: + blob = raw.decode() if isinstance(raw, (bytes, bytearray)) else raw + out.append(json.loads(blob)["task_id"]) + return out + + +def test_claims_a_task_when_redis_has_no_blmove(old_redis): + mq = ModelQ(redis_client=old_redis) + mq.BLPOP_TIMEOUT = 1 + _queue(mq, "task-1") + + claimed = mq._claim_task("inflight:worker:0") + + assert claimed is not None, "worker cannot pick up any work at all" + assert _ids([claimed]) == ["task-1"] + + +def test_fallback_keeps_the_task_in_exactly_one_list(old_redis): + """ + Atomicity is the whole reason the claim stopped being a plain BLPOP: the task + must never be absent from both the queue and the in-flight list. + """ + mq = ModelQ(redis_client=old_redis) + mq.BLPOP_TIMEOUT = 1 + _queue(mq, "task-1") + + mq._claim_task("inflight:worker:0") + + assert mq.redis_client.llen("ml_tasks") == 0 + assert _ids(mq.redis_client.lrange("inflight:worker:0", 0, -1)) == ["task-1"] + + +def test_fallback_preserves_fifo_order(old_redis): + """ + The regression BRPOPLPUSH would have introduced. Producers rpush, so claims + must come off the head -- otherwise the newest request jumps the queue and + the oldest starves. + """ + mq = ModelQ(redis_client=old_redis) + mq.BLPOP_TIMEOUT = 1 + _queue(mq, "first", "second", "third") + + claimed = [mq._claim_task("inflight:worker:0") for _ in range(3)] + + assert _ids(claimed) == ["first", "second", "third"] + + +def test_fallback_returns_none_on_an_empty_queue_without_hanging(old_redis): + mq = ModelQ(redis_client=old_redis) + mq.BLPOP_TIMEOUT = 0.3 + mq.CLAIM_POLL_INTERVAL = 0.05 + + assert mq._claim_task("inflight:worker:0") is None + + +def test_modern_redis_still_uses_blmove(new_redis): + """The control case: a capable server must not be downgraded to polling.""" + mq = ModelQ(redis_client=new_redis) + mq.BLPOP_TIMEOUT = 1 + _queue(mq, "task-1") + + calls = [] + real_blmove = new_redis.blmove + + def spy(*args, **kwargs): + calls.append(args) + return real_blmove(*args, **kwargs) + + new_redis.blmove = spy + claimed = mq._claim_task("inflight:worker:0") + + assert calls, "BLMOVE should still be used where it exists" + assert _ids([claimed]) == ["task-1"] + + +def test_support_is_probed_once_not_per_claim(old_redis): + """ + An unsupported server must not cost a failed round trip on every claim -- the + crash loop was thirty exceptions a second. + """ + probes = [] + original = old_redis.blmove + + def counting_blmove(*args, **kwargs): + probes.append(args) + return original(*args, **kwargs) + + old_redis.blmove = counting_blmove + + mq = ModelQ(redis_client=old_redis) + mq.BLPOP_TIMEOUT = 1 + _queue(mq, "a", "b", "c") + + for _ in range(3): + mq._claim_task("inflight:worker:0") + + assert len(probes) == 1 + + +def test_an_unrelated_redis_error_does_not_disable_blmove(new_redis): + """ + Only "unknown command" means the server lacks BLMOVE. A transient server-side + error must not permanently drop the worker into polling mode. + """ + mq = ModelQ(redis_client=new_redis) + + def transient(*args, **kwargs): + raise redis_lib.exceptions.ResponseError("LOADING Redis is loading the dataset") + + new_redis.blmove = transient + + assert mq._blmove_supported() is True