fix: fall back when Redis has no BLMOVE instead of crash-looping

.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: |

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
 

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 <tanmay@modelslab.com>"]
 readme = "README.md"

requirements-test.txt

@@ -1,3 +1,3 @@
-fakeredis
+fakeredis[lua]
 pytest
 click

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