From 5f5e9936d64e984b7616dbfe22a8529e5832551c Mon Sep 17 00:00:00 2001 From: Adhik Joshi Date: Thu, 20 Aug 2026 12:13:40 +0530 Subject: [PATCH] fix: keep started_at through completion and stop re-queuing live tasks Two coupled defects in in-flight task tracking. 1. started_at was erased on every completion. The worker writes it into task_dict at pickup but never onto the Task object, and _store_final_task_state() serialises from the object -- so the finished record came back with started_at: null. Across 6,953 completed tasks sampled from production, it survived on 127 (1.8%), and on zero tasks in 18 of 24 queues. Nothing downstream could separate queue wait from run time, which is exactly the number you need to decide whether a slow endpoint needs more replicas or a faster model. 2. requeue_stuck_processing_tasks() decided "stuck" from age alone, so a healthy long job was re-queued and rendered a second time on another GPU. On the wan22 text2video queue the median job takes 499s and p90 takes 949s against a 180s threshold; 7.8% of its tasks carry a rewritten queued_at, the signature of a re-queue. ltx_2.3_server is at 1.0%, flux_klein at 1.5%. They have to ship together. Fixing (1) alone would make finished strays in processing_tasks look like old stuck ones to the age rule, and re-running a completed generation is worse than leaking a set member. Liveness now comes from a heartbeat the owning worker republishes on the existing heartbeat thread, so: - a job that runs for 20 minutes is left alone as long as its worker is alive, and - a task whose worker actually died is picked up once its heartbeat goes stale, which is sooner than the old age rule managed. Tasks in a terminal state are drained from processing_tasks rather than re-queued. That also clears the backlog the erased started_at had been hiding: flux_klein's processing_tasks currently holds 61 ids of which 10 are real -- 47 have already completed, the oldest is 23.7h old -- which inflates every processing count that reads the set, including the one the Verda autoscaler scales on. A task with no heartbeat at all falls back to the age rule, so a fleet running mixed versions behaves exactly as it does today and converges as workers roll. Side effect worth having: the sweep no longer GETs every in-flight payload once a minute per worker just to decide to do nothing. A fresh heartbeat short-circuits before the fetch. --- modelq/app/base.py | 360 +++++++++++++++++++++++++++---- pyproject.toml | 2 +- tests/test_inflight_custody.py | 19 ++ tests/test_inflight_tracking.py | 370 ++++++++++++++++++++++++++++++++ 4 files changed, 711 insertions(+), 40 deletions(-) create mode 100644 tests/test_inflight_tracking.py diff --git a/modelq/app/base.py b/modelq/app/base.py index e04f8e5..e9b4a7d 100644 --- a/modelq/app/base.py +++ b/modelq/app/base.py @@ -136,6 +136,14 @@ class ModelQ: SOCKET_TIMEOUT = 60 # seconds: hard ceiling on any single read SOCKET_CONNECT_TIMEOUT = 10 # seconds: fail fast when the host is unreachable HEALTH_CHECK_INTERVAL = 30 # seconds: PING a pooled connection idle this long + + # --- in-flight task liveness -------------------------------------------- + # A worker republishes a heartbeat for every task it is actually running, so + # "still working" can be told apart from "the worker died holding this". + # Refreshed on the heartbeat thread's cadence; a task counts as abandoned + # once its heartbeat is older than INFLIGHT_STALE_AFTER. + INFLIGHT_HEARTBEAT_KEY = "task_heartbeats" + INFLIGHT_STALE_AFTER = 180 # seconds without a heartbeat before we re-queue BACKGROUND_LOOP_BACKOFF = 5 # seconds: pause before retrying a crashed loop body # --- in-flight task custody ------------------------------------------------ @@ -197,6 +205,10 @@ def __init__( ) self.worker_threads = [] + # task_id -> started_at for every task this process is currently running. + # Guarded because worker threads mutate it while the heartbeat thread reads it. + self._inflight_tasks = {} + self._inflight_lock = threading.Lock() if server_id is None: # Attempt to load the server_id from a local file: server_id = self._get_or_create_server_id_file() @@ -356,48 +368,251 @@ def _update_task_history(self, task_id: str, task_dict: dict) -> None: ex=self.TASK_HISTORY_RETENTION ) - def requeue_stuck_processing_tasks(self, threshold: float = 180.0): + # ------------------------------------------------------------------ # + # In-flight task liveness # + # ------------------------------------------------------------------ # + + 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: + self._inflight_tasks[task_id] = started_at + try: + self.redis_client.hset( + self.INFLIGHT_HEARTBEAT_KEY, task_id, str(started_at) + ) + except Exception as e: # never let bookkeeping kill a task + logger.warning(f"Could not publish in-flight heartbeat for {task_id}: {e}") + + def _clear_task_inflight(self, task_id: str) -> None: + """Forget `task_id`; it is no longer running here.""" + with self._inflight_lock: + self._inflight_tasks.pop(task_id, None) + try: + self.redis_client.hdel(self.INFLIGHT_HEARTBEAT_KEY, task_id) + except Exception as e: + logger.warning(f"Could not clear in-flight heartbeat for {task_id}: {e}") + + def _publish_inflight_heartbeats(self) -> None: + """ + Refresh the heartbeat of every task running in this process. + + Without this, `requeue_stuck_processing_tasks` can only ask "how long ago + did this start?", which re-queues healthy long jobs: a 500s video + generation crosses a 180s threshold every single time and gets run again + on another GPU while the first worker is still producing the same output. + A heartbeat separates "running long" from "nobody is holding this". + """ + with self._inflight_lock: + snapshot = dict(self._inflight_tasks) + + if not snapshot: + return + + now = time.time() + try: + self.redis_client.hset( + self.INFLIGHT_HEARTBEAT_KEY, + mapping={task_id: str(now) for task_id in snapshot}, + ) + except Exception as e: + logger.warning(f"Could not refresh in-flight heartbeats: {e}") + + def _inflight_heartbeats(self) -> dict: + """task_id -> last heartbeat timestamp, across every worker on this queue.""" + try: + raw = self.redis_client.hgetall(self.INFLIGHT_HEARTBEAT_KEY) or {} + except Exception as e: + logger.warning(f"Could not read in-flight heartbeats: {e}") + return {} + + beats = {} + for key, value in raw.items(): + task_id = key.decode("utf-8") if isinstance(key, bytes) else key + try: + beats[task_id] = float(value) + except (TypeError, ValueError): + continue + return beats + + @staticmethod + def _is_terminal_status(status: Optional[str]) -> bool: + return status in ("completed", "failed", "cancelled") + + def _requeue_custodied_task( + self, task_id: str, task_dict: dict, queued_at: float + ) -> bool: + """Atomically move a stale task out of worker custody and back to queue. + + Returns ``False`` for mixed-version workers that use no custody list, so + the caller can retain the original processing-set recovery path. + """ + recovered_task = dict(task_dict) + recovered_task["status"] = "queued" + recovered_task["queued_at"] = queued_at + recovered_blob = json.dumps(recovered_task) + + for raw_key in self.redis_client.smembers(self.INFLIGHT_REGISTRY) or []: + inflight_key = ( + raw_key.decode("utf-8") if isinstance(raw_key, bytes) else raw_key + ) + while True: + pipe = self.redis_client.pipeline() + try: + pipe.watch(inflight_key) + matching_item = None + for item in pipe.lrange(inflight_key, 0, -1) or []: + try: + candidate = json.loads(item) + except Exception: + continue + if candidate.get("task_id") == task_id: + matching_item = item + break + + if matching_item is None: + pipe.unwatch() + break + + # Keep the list registered: its owning worker/server may + # still be alive and will reuse this custody list. + pipe.multi() + pipe.lrem(inflight_key, 1, matching_item) + pipe.rpush("ml_tasks", recovered_blob) + pipe.set( + f"task:{task_id}", recovered_blob, ex=self.task_ttl + ) + pipe.zadd("queued_requests", {task_id: queued_at}) + pipe.srem("processing_tasks", task_id) + pipe.hdel(self.INFLIGHT_HEARTBEAT_KEY, task_id) + results = pipe.execute() + if results[0] == 1: + return True + break + except redis.WatchError: + continue + finally: + pipe.reset() + return False + + def requeue_stuck_processing_tasks(self, threshold: float = INFLIGHT_STALE_AFTER): """ - Re-queues any tasks that have been in 'processing' for more than 'threshold' seconds. + Re-queue tasks that no worker is holding any more. + + "Stuck" used to mean "started more than `threshold` seconds ago", which + is not the same question. A 500s video generation crosses a 180s + threshold while it is running perfectly well, gets pushed back onto + `ml_tasks`, and a second GPU renders the same output again. Liveness is + now decided by the heartbeat the owning worker republishes + (`_publish_inflight_heartbeats`), so long jobs are left alone and a task + whose worker actually died is picked up as soon as its heartbeat goes + stale -- which is sooner than the old rule managed. + + Tasks that already reached a terminal state are drained rather than + re-queued: `process_task` removes them on the way out, but a worker + killed between finishing and cleaning up leaves the id behind, and those + strays otherwise sit in `processing_tasks` until their task key expires + a day later, inflating every processing count that reads the set. """ - if self.requeue_threshold : + if self.requeue_threshold: threshold = self.requeue_threshold processing_task_ids = self.redis_client.smembers("processing_tasks") + if not processing_task_ids: + return + now = time.time() + heartbeats = self._inflight_heartbeats() for pid in processing_task_ids: - task_id = pid.decode("utf-8") + task_id = pid.decode("utf-8") if isinstance(pid, bytes) else pid + + # A fresh heartbeat means a worker is holding this right now. Skip it + # before fetching the task, so a busy queue does not re-read every + # in-flight payload once a minute just to decide to do nothing. + last_beat = heartbeats.get(task_id) + if last_beat is not None and now - last_beat <= threshold: + continue + task_data = self.redis_client.get(f"task:{task_id}") if not task_data: # If there's no data in Redis for that task, remove it from processing set. self.redis_client.srem("processing_tasks", task_id) + self.redis_client.hdel(self.INFLIGHT_HEARTBEAT_KEY, task_id) logger.warning( f"No record found for in-progress task {task_id}. Removing from 'processing_tasks'." ) continue task_dict = json.loads(task_data) + + if self._is_terminal_status(task_dict.get("status")): + self.redis_client.srem("processing_tasks", task_id) + self.redis_client.hdel(self.INFLIGHT_HEARTBEAT_KEY, task_id) + logger.info( + f"Draining finished task {task_id} " + f"(status={task_dict.get('status')}) from 'processing_tasks'." + ) + continue + started_at = task_dict.get("started_at", 0) - if started_at: - if now - started_at > threshold: - logger.info( - f"Re-queuing stuck task {task_id} which has been 'processing' for {now - started_at:.2f} seconds." - ) - # Update status, queued_at, etc. - task_dict["status"] = "queued" - task_dict["queued_at"] = now - - # Store the updated dict back in Redis - self.redis_client.set(f"task:{task_id}", json.dumps(task_dict),ex=86400) - - # Push it back into ml_tasks - self.redis_client.rpush("ml_tasks", json.dumps(task_dict)) - self.redis_client.zadd("queued_requests", {task_id: now}) + if not started_at: + continue + + age = now - started_at + if last_beat is None and age <= threshold: + # No heartbeat at all: either a worker on an older build, or one + # that has not reached its first heartbeat tick yet. Fall back to + # the age rule so mixed-version fleets behave as they did before. + continue + + if last_beat is not None: + logger.info( + f"Re-queuing abandoned task {task_id}: no heartbeat for " + f"{now - last_beat:.2f}s." + ) + else: + logger.info( + f"Re-queuing stuck task {task_id} which has been " + f"'processing' for {age:.2f} seconds." + ) + + # New workers keep a recoverable custody copy in Redis. Move that + # exact copy instead of creating a second queue entry and leaving + # the original to be recovered again later. False means an older + # worker with no custody list, so the established fallback below is + # intentionally preserved for rolling deployments. + if self._requeue_custodied_task(task_id, task_dict, now): + continue - # Remove from processing set - self.redis_client.srem("processing_tasks", task_id) + # Custody can disappear because the worker completed while we were + # deciding. Re-check the terminal state before using the legacy + # fallback so a just-finished task is never resurrected. + refreshed_data = self.redis_client.get(f"task:{task_id}") + if not refreshed_data: + self.redis_client.srem("processing_tasks", task_id) + self.redis_client.hdel(self.INFLIGHT_HEARTBEAT_KEY, task_id) + continue + task_dict = json.loads(refreshed_data) + if self._is_terminal_status(task_dict.get("status")): + self.redis_client.srem("processing_tasks", task_id) + self.redis_client.hdel(self.INFLIGHT_HEARTBEAT_KEY, task_id) + continue + + # Update status, queued_at, etc. + task_dict["status"] = "queued" + task_dict["queued_at"] = now + + # Store the updated dict back in Redis + self.redis_client.set(f"task:{task_id}", json.dumps(task_dict), ex=self.task_ttl) + + # Push it back into ml_tasks + self.redis_client.rpush("ml_tasks", json.dumps(task_dict)) + self.redis_client.zadd("queued_requests", {task_id: now}) + + # Remove from processing set + self.redis_client.srem("processing_tasks", task_id) + self.redis_client.hdel(self.INFLIGHT_HEARTBEAT_KEY, task_id) SCAN_BATCH = 500 @@ -865,10 +1080,12 @@ def worker_loop(worker_id): # Release custody only when we are finished with it, whatever # the outcome. If this process dies mid-task the entry stays # put on purpose — that is what makes it recoverable. + claimed_task_id = None try: self.update_server_status(f"worker_{worker_id}: busy") task_dict = json.loads(task_json) task = Task.from_dict(task_dict) + claimed_task_id = task.task_id # Mark task as 'processing' added = self.redis_client.sadd("processing_tasks", task.task_id) @@ -885,13 +1102,27 @@ def worker_loop(worker_id): # inflates queue_num / queue_time. self.redis_client.zrem("queued_requests", task.task_id) - # Set started_at - task_dict["started_at"] = time.time() + # Set started_at on BOTH the dict we persist and the Task + # object. _store_final_task_state() serialises from the Task, + # so leaving it off the object silently overwrote the real + # pickup time with null on completion. + started_at = time.time() + task_dict["status"] = "processing" + task_dict["started_at"] = started_at + task.started_at = started_at # Update in Redis - self.redis_client.set(f"task:{task.task_id}", json.dumps(task_dict),ex=86400) + self.redis_client.set( + f"task:{task.task_id}", + json.dumps(task_dict), + ex=self.task_ttl, + ) if task.task_name in self.allowed_tasks: + # Only a worker that can actually execute this task + # advertises it as live. Routing it elsewhere must not + # leave a heartbeat that suppresses future recovery. + self._mark_task_inflight(task.task_id, started_at) try: logger.info(f"Worker {worker_id} started processing: {task.task_name}") @@ -947,6 +1178,10 @@ def worker_loop(worker_id): self.redis_client.zadd("queued_requests", {task.task_id: task_dict.get("queued_at", time.time())}) self.redis_client.srem("processing_tasks", task.task_id) finally: + if claimed_task_id is not None: + # Idempotent with process_task's cleanup and covers + # routing/worker-loop errors before process_task runs. + self._clear_task_inflight(claimed_task_id) self.redis_client.lrem(inflight_key, 1, task_json) except Exception as e: @@ -1008,17 +1243,35 @@ def drain_inflight(self, inflight_key: str) -> int: recovered_item = item task_id = None + terminal = False queued_at = time.time() try: task_dict = json.loads(item) task_id = task_dict.get("task_id") if task_id: - # Match the established stuck-task recovery contract: - # recovered work is queued again with a fresh queue - # timestamp and a bounded task record. - task_dict["status"] = "queued" - task_dict["queued_at"] = queued_at - recovered_item = json.dumps(task_dict) + task_key = f"task:{task_id}" + # A worker can finish and publish its terminal state, + # then die before the custody-finally removes the list + # entry. Watch the record with the list so recovery + # never turns that completed task back into queued work. + pipe.watch(task_key) + current_blob = pipe.get(task_key) + try: + current_task = ( + json.loads(current_blob) if current_blob else {} + ) + except Exception: + current_task = {} + terminal = self._is_terminal_status( + current_task.get("status") + ) + if not terminal: + # Match the established stuck-task recovery contract: + # recovered work is queued again with a fresh queue + # timestamp and a bounded task record. + task_dict["status"] = "queued" + task_dict["queued_at"] = queued_at + recovered_item = json.dumps(task_dict) except Exception as exc: logger.warning( f"Could not decode in-flight entry from '{inflight_key}' " @@ -1027,18 +1280,28 @@ def drain_inflight(self, inflight_key: str) -> int: pipe.multi() pipe.rpop(inflight_key) - pipe.lpush("ml_tasks", recovered_item) if task_id: - pipe.set( - f"task:{task_id}", - recovered_item, - ex=self.task_ttl, - ) - pipe.zadd("queued_requests", {task_id: queued_at}) pipe.srem("processing_tasks", task_id) + pipe.hdel(self.INFLIGHT_HEARTBEAT_KEY, task_id) + if terminal: + pipe.zrem("queued_requests", task_id) + else: + pipe.lpush("ml_tasks", recovered_item) + if task_id: + pipe.set( + f"task:{task_id}", + recovered_item, + ex=self.task_ttl, + ) + pipe.zadd("queued_requests", {task_id: queued_at}) results = pipe.execute() - if results[0] is not None: + if results[0] is not None and not terminal: moved += 1 + elif results[0] is not None: + logger.info( + f"Discarded terminal task {task_id} from abandoned " + f"in-flight list '{inflight_key}'." + ) break except redis.WatchError: # Another recovery worker won the race. Re-read the new tail @@ -1114,6 +1377,8 @@ def _heartbeat_loop(self): while True: with self._guarded_iteration("heartbeat"): self.heartbeat() + with self._guarded_iteration("inflight_heartbeat"): + self._publish_inflight_heartbeats() time.sleep(self.HEARTBEAT_INTERVAL) def _pruning_loop(self): @@ -1126,7 +1391,9 @@ def _pruning_loop(self): # Runs after the prune so dead servers are already deregistered # and their in-flight lists read as abandoned on this same pass. self.recover_abandoned_inflight_tasks() - self.requeue_stuck_processing_tasks(threshold=180) + self.requeue_stuck_processing_tasks( + threshold=self.INFLIGHT_STALE_AFTER + ) # prune_old_task_results() is deliberately NOT called here. Every # task_result key is written with a TTL, so Redis expires it on # its own; running the scan every PRUNE_CHECK_INTERVAL only @@ -1338,6 +1605,7 @@ def process_task(self, task: Task) -> None: finally: self.redis_client.srem("processing_tasks", task.task_id) + self._clear_task_inflight(task.task_id) def _store_final_task_state(self, task: Task, success: bool, error: Optional[Exception] = None): @@ -1350,6 +1618,20 @@ def _store_final_task_state(self, task: Task, success: bool, error: Optional[Exc # Mark finished_at task_dict["finished_at"] = time.time() + # Stop advertising this task as in-flight before the terminal blob is + # visible. Ordering matters: the blob now carries a real started_at, and + # a sweeper that saw it while the id was still in `processing_tasks` + # could read "started long ago" and re-queue work that is already done. + try: + self.redis_client.srem("processing_tasks", task.task_id) + self.redis_client.hdel(self.INFLIGHT_HEARTBEAT_KEY, task.task_id) + except Exception as e: + logger.warning( + f"Could not clear in-flight state for {task.task_id}: {e}" + ) + with self._inflight_lock: + self._inflight_tasks.pop(task.task_id, None) + # Add error details if failed if not success and error: task_dict["error"] = { diff --git a/pyproject.toml b/pyproject.toml index 041d8b9..3507535 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "modelq" -version = "1.0.16" +version = "1.0.17" description = "Celery-like task queue for ML inference." authors = ["Tanmaypatil123 "] readme = "README.md" diff --git a/tests/test_inflight_custody.py b/tests/test_inflight_custody.py index 22765cc..eafb599 100644 --- a/tests/test_inflight_custody.py +++ b/tests/test_inflight_custody.py @@ -90,6 +90,25 @@ def test_drain_keeps_malformed_entries_recoverable(mq): assert mq.redis_client.lpop("ml_tasks") == b"not-json" +def test_drain_does_not_rerun_task_that_finished_before_custody_release(mq): + """A crash after terminal persistence must not resurrect completed work.""" + inflight = mq._inflight_key(0) + task_id = "finished-before-release" + mq.redis_client.rpush(inflight, _task(task_id)) + mq.redis_client.sadd(mq.INFLIGHT_REGISTRY, inflight) + mq.redis_client.sadd("processing_tasks", task_id) + mq.redis_client.set( + f"task:{task_id}", + json.dumps({"task_id": task_id, "status": "completed", "result": "ok"}), + ) + + assert mq.drain_inflight(inflight) == 0 + assert mq.redis_client.llen(inflight) == 0 + assert mq.redis_client.llen("ml_tasks") == 0 + assert not mq.redis_client.sismember("processing_tasks", task_id) + assert json.loads(mq.redis_client.get(f"task:{task_id}"))["status"] == "completed" + + def test_drain_is_bounded_and_terminates_on_empty(mq): """An unbounded drain loop spins forever the day the list refills.""" assert mq.drain_inflight(mq._inflight_key(9)) == 0 diff --git a/tests/test_inflight_tracking.py b/tests/test_inflight_tracking.py new file mode 100644 index 0000000..e3abd38 --- /dev/null +++ b/tests/test_inflight_tracking.py @@ -0,0 +1,370 @@ +""" +Tests for in-flight task liveness. + +Two production failures motivate this file: + + 1. `started_at` was written into the task dict at pickup but never onto the + Task object, so `_store_final_task_state()` -- which serialises from the + object -- overwrote the real pickup time with null on every completion. + Nothing downstream could separate queue wait from run time, and the + requeue sweep (which keys off `started_at`) skipped finished strays + forever, leaving `processing_tasks` to grow until the task keys expired. + + 2. The sweep decided "stuck" from age alone. A 500s video generation crosses + a 180s threshold while running perfectly well, so it was pushed back onto + `ml_tasks` and rendered a second time on another GPU. +""" + +import json +import time + +import fakeredis +import pytest + +from modelq import ModelQ + + +@pytest.fixture +def mock_redis(): + return fakeredis.FakeStrictRedis() + + +@pytest.fixture +def mq(mock_redis): + return ModelQ(redis_client=mock_redis) + + +def _blob(raw): + return json.loads(raw.decode() if isinstance(raw, (bytes, bytearray)) else raw) + + +def _place_in_processing(mq, task_id, *, status, started_at, heartbeat=None): + """Put a task into `processing_tasks` exactly as a worker pickup would.""" + mq.redis_client.sadd("processing_tasks", task_id) + mq.redis_client.set( + f"task:{task_id}", + json.dumps( + { + "task_id": task_id, + "task_name": "render", + "payload": {"data": {"args": []}}, + "status": status, + "result": None, + "created_at": started_at, + "queued_at": started_at, + "started_at": started_at, + "finished_at": None, + "stream": False, + } + ), + ) + if heartbeat is not None: + mq.redis_client.hset(ModelQ.INFLIGHT_HEARTBEAT_KEY, task_id, str(heartbeat)) + + +# --------------------------------------------------------------------------- +# the sweep +# --------------------------------------------------------------------------- + + +def test_long_running_task_with_fresh_heartbeat_is_not_requeued(mq): + """A 10-minute job whose worker is alive must be left alone.""" + now = time.time() + _place_in_processing( + mq, "video-1", status="processing", started_at=now - 600, heartbeat=now - 5 + ) + + mq.requeue_stuck_processing_tasks(threshold=180) + + assert mq.redis_client.llen("ml_tasks") == 0, "healthy long job was re-queued" + assert mq.redis_client.sismember("processing_tasks", "video-1") + assert _blob(mq.redis_client.get("task:video-1"))["status"] == "processing" + + +def test_task_is_requeued_once_its_heartbeat_goes_stale(mq): + """The worker died holding it; nobody is producing this output.""" + now = time.time() + _place_in_processing( + mq, "video-2", status="processing", started_at=now - 600, heartbeat=now - 600 + ) + + mq.requeue_stuck_processing_tasks(threshold=180) + + assert mq.redis_client.llen("ml_tasks") == 1 + assert _blob(mq.redis_client.lindex("ml_tasks", 0))["task_id"] == "video-2" + assert not mq.redis_client.sismember("processing_tasks", "video-2") + assert mq.redis_client.hget(ModelQ.INFLIGHT_HEARTBEAT_KEY, "video-2") is None + + +def test_stale_task_moves_out_of_custody_instead_of_leaving_a_duplicate(mq): + """Heartbeat recovery and BLMOVE custody must produce one queue copy.""" + now = time.time() + task_id = "custodied-stale" + _place_in_processing( + mq, + task_id, + status="processing", + started_at=now - 600, + heartbeat=now - 600, + ) + inflight = mq._inflight_key(0) + mq.redis_client.rpush(inflight, mq.redis_client.get(f"task:{task_id}")) + mq.redis_client.sadd(mq.INFLIGHT_REGISTRY, inflight) + + mq.requeue_stuck_processing_tasks(threshold=180) + + assert mq.redis_client.llen("ml_tasks") == 1 + assert mq.redis_client.llen(inflight) == 0 + assert not mq.redis_client.sismember("processing_tasks", task_id) + assert mq.redis_client.sismember(mq.INFLIGHT_REGISTRY, inflight) + assert _blob(mq.redis_client.lindex("ml_tasks", 0))["task_id"] == task_id + + +def test_heartbeat_freshness_is_measured_against_the_threshold(mq): + """A job younger than the threshold is safe even with no heartbeat yet.""" + now = time.time() + _place_in_processing(mq, "fresh", status="processing", started_at=now - 30) + + mq.requeue_stuck_processing_tasks(threshold=180) + + assert mq.redis_client.llen("ml_tasks") == 0 + assert mq.redis_client.sismember("processing_tasks", "fresh") + + +def test_task_without_heartbeat_falls_back_to_the_age_rule(mq): + """ + Mixed-version fleets: a worker on an older build publishes no heartbeat, so + the sweep must behave exactly as it did before rather than never acting. + """ + now = time.time() + _place_in_processing(mq, "legacy", status="processing", started_at=now - 600) + + mq.requeue_stuck_processing_tasks(threshold=180) + + assert mq.redis_client.llen("ml_tasks") == 1 + assert not mq.redis_client.sismember("processing_tasks", "legacy") + + +def test_finished_stray_is_drained_not_rerun(mq): + """ + The regression guard for fix (1). Once `started_at` survives completion, a + finished task left behind in `processing_tasks` looks exactly like an old + stuck one. Re-queueing it would re-run a generation that already produced + its output. + """ + now = time.time() + _place_in_processing(mq, "done-1", status="completed", started_at=now - 600) + + mq.requeue_stuck_processing_tasks(threshold=180) + + assert mq.redis_client.llen("ml_tasks") == 0, "completed task was re-run" + assert not mq.redis_client.sismember("processing_tasks", "done-1") + + +@pytest.mark.parametrize("status", ["completed", "failed", "cancelled"]) +def test_every_terminal_status_is_drained(mq, status): + now = time.time() + _place_in_processing(mq, f"t-{status}", status=status, started_at=now - 600) + + mq.requeue_stuck_processing_tasks(threshold=180) + + assert mq.redis_client.llen("ml_tasks") == 0 + assert not mq.redis_client.sismember("processing_tasks", f"t-{status}") + + +def test_missing_task_key_is_dropped_from_the_set(mq): + mq.redis_client.sadd("processing_tasks", "ghost") + + mq.requeue_stuck_processing_tasks(threshold=180) + + assert not mq.redis_client.sismember("processing_tasks", "ghost") + assert mq.redis_client.llen("ml_tasks") == 0 + + +def test_live_task_is_not_read_back_from_redis(mq): + """ + The sweep runs once a minute on every worker and used to GET every in-flight + payload -- hundreds of KB each on image queues -- only to decide to do + nothing. A fresh heartbeat should short-circuit before the fetch. + """ + now = time.time() + _place_in_processing( + mq, "busy", status="processing", started_at=now - 600, heartbeat=now - 5 + ) + + reads = [] + real_get = mq.redis_client.get + + def counting_get(key, *args, **kwargs): + reads.append(key) + return real_get(key, *args, **kwargs) + + mq.redis_client.get = counting_get + try: + mq.requeue_stuck_processing_tasks(threshold=180) + finally: + mq.redis_client.get = real_get + + assert reads == [], f"payload was fetched anyway: {reads}" + + +# --------------------------------------------------------------------------- +# heartbeat bookkeeping +# --------------------------------------------------------------------------- + + +def test_marking_a_task_inflight_publishes_a_heartbeat(mq): + mq._mark_task_inflight("abc", 1234.5) + + assert mq._inflight_tasks == {"abc": 1234.5} + assert float(mq.redis_client.hget(ModelQ.INFLIGHT_HEARTBEAT_KEY, "abc")) == 1234.5 + + +def test_clearing_a_task_removes_the_heartbeat(mq): + mq._mark_task_inflight("abc", 1234.5) + mq._clear_task_inflight("abc") + + assert mq._inflight_tasks == {} + assert mq.redis_client.hget(ModelQ.INFLIGHT_HEARTBEAT_KEY, "abc") is None + + +def test_publishing_advances_every_inflight_heartbeat(mq): + mq._mark_task_inflight("a", time.time() - 500) + mq._mark_task_inflight("b", time.time() - 500) + + mq._publish_inflight_heartbeats() + + beats = mq._inflight_heartbeats() + now = time.time() + assert set(beats) == {"a", "b"} + assert all(now - v < 5 for v in beats.values()) + + +def test_publishing_with_nothing_inflight_touches_nothing(mq): + mq._publish_inflight_heartbeats() + assert mq.redis_client.hgetall(ModelQ.INFLIGHT_HEARTBEAT_KEY) == {} + + +def test_unreadable_heartbeat_values_are_ignored(mq): + mq.redis_client.hset(ModelQ.INFLIGHT_HEARTBEAT_KEY, "junk", "not-a-float") + mq.redis_client.hset(ModelQ.INFLIGHT_HEARTBEAT_KEY, "good", "100.0") + + assert mq._inflight_heartbeats() == {"good": 100.0} + + +# --------------------------------------------------------------------------- +# end-to-end: the started_at fix, through a real worker +# --------------------------------------------------------------------------- + + +def _wait_for(predicate, timeout=10.0, interval=0.05): + deadline = time.time() + timeout + while time.time() < deadline: + if predicate(): + return True + time.sleep(interval) + return False + + +def test_started_at_survives_a_real_completion(mock_redis): + """ + The headline regression. Run one task through an actual worker and assert + the finished record still carries the pickup time. + + Before the fix this came back as null on every completed task, because + `_store_final_task_state()` re-serialised a Task object whose `started_at` + had only ever been written to a separate dict. + """ + mq = ModelQ(redis_client=mock_redis) + + @mq.task(timeout=30) + def render(): + time.sleep(0.25) + return {"ok": True} + + render(_task_id="e2e-1") + mq.start_workers(no_of_workers=1) + + assert _wait_for( + lambda: _blob(mock_redis.get("task:e2e-1")).get("status") == "completed" + ), "task never completed" + + record = _blob(mock_redis.get("task:e2e-1")) + assert record["started_at"] is not None, "started_at was erased on completion" + assert record["finished_at"] is not None + assert record["finished_at"] >= record["started_at"] + assert record["started_at"] >= record["queued_at"] + + # run time is now derivable, which is the whole point + assert 0.2 <= record["finished_at"] - record["started_at"] < 5.0 + + # and the worker cleaned up after itself + assert _wait_for(lambda: not mock_redis.sismember("processing_tasks", "e2e-1")) + assert mock_redis.hget(ModelQ.INFLIGHT_HEARTBEAT_KEY, "e2e-1") is None + + +def test_a_task_is_marked_inflight_while_it_runs(mock_redis): + """The heartbeat must exist during the run, not only after it.""" + mq = ModelQ(redis_client=mock_redis) + seen = {} + + @mq.task(timeout=30) + def slow(): + seen["inflight"] = dict(mq._inflight_tasks) + seen["heartbeat"] = mock_redis.hget(ModelQ.INFLIGHT_HEARTBEAT_KEY, "e2e-2") + seen["processing"] = mock_redis.sismember("processing_tasks", "e2e-2") + return {"ok": True} + + slow(_task_id="e2e-2") + mq.start_workers(no_of_workers=1) + + assert _wait_for(lambda: "inflight" in seen), "task never ran" + + assert "e2e-2" in seen["inflight"], "task was not registered in-flight" + assert seen["heartbeat"] is not None, "no heartbeat published while running" + assert seen["processing"] + + +def test_task_routed_to_another_worker_does_not_leak_a_heartbeat(mock_redis): + """Requeued work is not live here and must never suppress recovery.""" + import threading + + mq = ModelQ(redis_client=mock_redis, server_id="routing-worker") + now = time.time() + task_id = "wrong-worker-task" + task = { + "task_id": task_id, + "task_name": "not-allowed-here", + "payload": {}, + "status": "queued", + "result": None, + "created_at": now, + "queued_at": now, + "started_at": None, + "finished_at": None, + "stream": False, + } + mock_redis.lpush("ml_tasks", json.dumps(task)) + + requeued = threading.Event() + real_rpush = mock_redis.rpush + + def stop_after_requeue(name, *values): + result = real_rpush(name, *values) + if name == "ml_tasks": + mq.worker_healthy = False + requeued.set() + return result + + mock_redis.rpush = stop_after_requeue + try: + mq.start_workers(no_of_workers=1) + assert requeued.wait(timeout=5), "task was not routed back to the queue" + finally: + mock_redis.rpush = real_rpush + + assert task_id not in mq._inflight_tasks + assert mock_redis.hget(ModelQ.INFLIGHT_HEARTBEAT_KEY, task_id) is None + assert not mock_redis.sismember("processing_tasks", task_id) + assert mock_redis.llen(mq._inflight_key(0)) == 0 + assert mock_redis.llen("ml_tasks") == 1