feat: expose per-task queue and run timings to the producer

modelq/app/base.py

@@ -986,6 +986,11 @@ def wrapper(*args, **kwargs):
                 now_ts = time.time()
                 task_dict["created_at"] = now_ts
                 task_dict["queued_at"]  = now_ts
+                # Mirror onto the Task the producer gets back, otherwise its
+                # queued_at stays None and the caller cannot tell queue wait apart
+                # from run time once the result lands.
+                task.created_at = now_ts
+                task.queued_at = now_ts
 
                 self.enqueue_task(task_dict, payload=payload)
                 self.redis_client.set(f"task:{task.task_id}",

modelq/app/tasks/base.py

@@ -77,6 +77,41 @@ def from_dict(data: dict) -> "Task":
         task.stream = data.get("stream", False)
         return task
 
+    def _absorb_timestamps(self, data: dict) -> None:
+        """
+        Copy queue/run timestamps from a task blob onto this Task, keeping whatever
+        is already set when the blob omits a field (an older worker, or a re-queue
+        that has not run yet).
+        """
+        for field in ("created_at", "queued_at", "started_at", "finished_at"):
+            value = data.get(field)
+            if value is not None:
+                setattr(self, field, value)
+
+    def stage_timings(self) -> Dict[str, float]:
+        """
+        Wall-clock split of this task's life, in seconds, for whichever stages have
+        both of their timestamps. Keys are omitted rather than zeroed when a stage
+        cannot be measured, so a caller never reports a fabricated 0.0.
+
+        - queue_time: enqueued until a worker picked it up
+        - run_time:   worker start until the result was persisted
+        - total_time: enqueued until the result was persisted
+        """
+        timings: Dict[str, float] = {}
+
+        def span(start: Optional[float], end: Optional[float], key: str) -> None:
+            if start is None or end is None:
+                return
+            delta = end - start
+            if delta >= 0:
+                timings[key] = round(delta, 3)
+
+        span(self.queued_at, self.started_at, "queue_time")
+        span(self.started_at, self.finished_at, "run_time")
+        span(self.queued_at, self.finished_at, "total_time")
+        return timings
+
     def _convert_to_string(self, data: Any) -> str:
         """
         Converts data to a string representation. If the data is a PIL image,
@@ -188,6 +223,12 @@ def get_result(
                 task_data = json.loads(task_json)
                 self.result = task_data.get("result")
                 self.status = task_data.get("status")
+                # The terminal blob carries the queue/run timestamps. Copy them back
+                # onto the Task so a caller that blocked on the result can report how
+                # much of the wait was queueing and how much was the run itself,
+                # without a second Redis read. Raising here (failed/cancelled) must
+                # still leave them populated, so this runs before the status checks.
+                self._absorb_timestamps(task_data)
 
                 if self.status == "failed":
                     error_message = self.result or "Task failed without an error message"

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "modelq"
-version = "1.0.17"
+version = "1.0.18"
 description = "Celery-like task queue for ML inference."
 authors = ["Tanmaypatil123 <tanmay@modelslab.com>"]
 readme = "README.md"

tests/test_stage_timings.py

@@ -0,0 +1,182 @@
+"""
+Tests for the queue/run split a producer sees after blocking on a result.
+
+Motivated by a production report where a generation API answered in ~37s while
+the only timer it exposed read 5.55s. That timer started after the task was
+enqueued, so everything before it -- and the queue wait itself -- was invisible,
+and there was no way to tell a slow queue apart from a slow run.
+
+Two gaps made the split unreportable:
+
+  1. The task decorator wrote `created_at`/`queued_at` into the dict it pushed to
+     Redis but never onto the Task object it handed back, so the producer's
+     `task.queued_at` stayed None.
+
+  2. `get_result()` copied only `result` and `status` off the terminal blob and
+     dropped `started_at`/`finished_at`, so even a worker that recorded them
+     could not surface them to the caller that was waiting.
+"""
+
+import json
+import time
+
+import fakeredis
+import pytest
+
+from modelq import ModelQ
+from modelq.app.tasks.base import Task
+
+
+@pytest.fixture
+def mock_redis():
+    return fakeredis.FakeStrictRedis()
+
+
+@pytest.fixture
+def mq(mock_redis):
+    return ModelQ(redis_client=mock_redis)
+
+
+def _complete(mq, task_id, *, queued_at, started_at, finished_at, result="done"):
+    """Write the terminal blob exactly as `_store_final_task_state()` would."""
+    mq.redis_client.set(
+        f"task_result:{task_id}",
+        json.dumps(
+            {
+                "task_id": task_id,
+                "task_name": "generate",
+                "payload": {},
+                "status": "completed",
+                "result": result,
+                "created_at": queued_at,
+                "queued_at": queued_at,
+                "started_at": started_at,
+                "finished_at": finished_at,
+            }
+        ),
+    )
+
+
+def test_enqueue_stamps_queued_at_on_the_returned_task(mq):
+    @mq.task()
+    def generate(payload=None):
+        return "ok"
+
+    before = time.time()
+    task = generate({"prompt": "a cat"})
+    after = time.time()
+
+    assert task.queued_at is not None, "producer cannot measure queue wait"
+    assert before <= task.queued_at <= after
+    assert task.created_at is not None
+
+
+def test_get_result_reports_queue_and_run_split(mq):
+    @mq.task()
+    def generate(payload=None):
+        return "ok"
+
+    task = generate({"prompt": "a cat"})
+
+    # 12s waiting behind other work, then a 5.5s run.
+    queued_at = task.queued_at
+    _complete(
+        mq,
+        task.task_id,
+        queued_at=queued_at,
+        started_at=queued_at + 12.0,
+        finished_at=queued_at + 17.5,
+    )
+
+    assert task.get_result(mq.redis_client, timeout=1) == "done"
+
+    timings = task.stage_timings()
+    assert timings["queue_time"] == pytest.approx(12.0, abs=0.01)
+    assert timings["run_time"] == pytest.approx(5.5, abs=0.01)
+    assert timings["total_time"] == pytest.approx(17.5, abs=0.01)
+
+
+def test_stage_timings_omits_stages_it_cannot_measure(mq):
+    """
+    A stage with a missing endpoint must be absent, not reported as 0.0 -- a
+    fabricated zero reads as "instant" and would hide the very stall this split
+    exists to expose.
+    """
+    task = Task(task_name="generate", payload={})
+    task.queued_at = 1000.0
+    task.started_at = None
+    task.finished_at = None
+
+    assert task.stage_timings() == {}
+
+    task.started_at = 1003.0
+    timings = task.stage_timings()
+    assert timings == {"queue_time": 3.0}
+
+
+def test_timestamps_survive_a_failed_task(mq):
+    """
+    A failure is exactly when the split matters most, so the timestamps must be
+    absorbed before get_result() raises.
+    """
+
+    @mq.task()
+    def generate(payload=None):
+        return "ok"
+
+    task = generate({"prompt": "a cat"})
+    queued_at = task.queued_at
+
+    mq.redis_client.set(
+        f"task_result:{task.task_id}",
+        json.dumps(
+            {
+                "task_id": task.task_id,
+                "task_name": "generate",
+                "payload": {},
+                "status": "failed",
+                "result": "boom",
+                "queued_at": queued_at,
+                "started_at": queued_at + 2.0,
+                "finished_at": queued_at + 9.0,
+            }
+        ),
+    )
+
+    with pytest.raises(Exception):
+        task.get_result(mq.redis_client, timeout=1)
+
+    timings = task.stage_timings()
+    assert timings["queue_time"] == pytest.approx(2.0, abs=0.01)
+    assert timings["run_time"] == pytest.approx(7.0, abs=0.01)
+
+
+def test_older_worker_blob_without_timestamps_does_not_clobber(mq):
+    """
+    A worker running an older build omits the run timestamps. The queued_at the
+    producer already holds must survive, rather than being reset to None.
+    """
+
+    @mq.task()
+    def generate(payload=None):
+        return "ok"
+
+    task = generate({"prompt": "a cat"})
+    queued_at = task.queued_at
+
+    mq.redis_client.set(
+        f"task_result:{task.task_id}",
+        json.dumps(
+            {
+                "task_id": task.task_id,
+                "task_name": "generate",
+                "payload": {},
+                "status": "completed",
+                "result": "done",
+            }
+        ),
+    )
+
+    assert task.get_result(mq.redis_client, timeout=1) == "done"
+    assert task.queued_at == queued_at
+    assert task.stage_timings() == {}