diff --git a/modelq/app/base.py b/modelq/app/base.py index e9b4a7d..bc74e8f 100644 --- a/modelq/app/base.py +++ b/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}", diff --git a/modelq/app/tasks/base.py b/modelq/app/tasks/base.py index 6c9543d..589efd5 100644 --- a/modelq/app/tasks/base.py +++ b/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" diff --git a/pyproject.toml b/pyproject.toml index 3507535..c2d5fea 100644 --- a/pyproject.toml +++ b/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 "] readme = "README.md" diff --git a/tests/test_stage_timings.py b/tests/test_stage_timings.py new file mode 100644 index 0000000..ce21e9c --- /dev/null +++ b/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() == {}