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

A generation API answered a request in 37.7s while the only timer in its
response read `generationTime: 5.55`. That timer is measured from enqueue to
result, so it hid everything the HTTP handler did beforehand and merged queue
wait into the same number as the run. Nothing in the response could account for
the missing 32s.

ModelQ already records created_at/queued_at/started_at/finished_at and persists
them in the terminal blob, but a producer could not read any of it:

- The task decorator wrote created_at/queued_at into the dict it pushed to Redis
  and never onto the Task object it handed back, so `task.queued_at` was None.
- `get_result()` copied only `result` and `status` off the terminal blob and
  dropped the run timestamps.

Mirror the enqueue stamps onto the returned Task, absorb the terminal blob's
timestamps in `get_result()` before it can raise, and add `Task.stage_timings()`
returning the queue/run/total split. Stages that cannot be measured are omitted
rather than reported as 0.0 -- a fabricated zero reads as "instant" and would
hide the very stall this exists to expose.

Backwards compatible: a result blob from an older worker leaves the timestamps
the producer already holds untouched, and stage_timings() simply returns {}.

Author: adhikjoshi

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() == {}