"""Tests for the cross-tenant job queue. Every case below corresponds to something that actually went wrong while building this, because the failure modes are not obvious from the code: * dispatching without room does not fail gracefully -- a CUDA OOM kills llama-server, and three queued jobs were destroyed in a row; * the VRAM requirement is a property of the job's model, not of the tenant, and a flat 4 GB let a 14.9 GB model be dispatched into 8 GB of free memory; * waiting forever is as wrong as failing immediately, when the job can never fit; * an exception during dispatch left the row RUNNING while the scheduler moved on. """ import asyncio import json import time import pytest import jobs as J import telemetry_store @pytest.fixture def queue(tmp_path, monkeypatch): monkeypatch.setattr(telemetry_store, "DB_PATH", str(tmp_path / "q.db")) J.init() return tmp_path class TestQueueBasics: def test_submit_and_read_back(self, queue): r = J.submit("ollama", {"model": "m"}, label="first") assert r["success"] job = J.get(r["id"]) assert job["state"] == J.PENDING and job["label"] == "first" assert job["payload"] == {"model": "m"} def test_unknown_tenant_is_rejected(self, queue): assert J.submit("nope", {})["success"] is False def test_priority_defaults_to_the_tenants_own(self, queue): import tenants as T r = J.submit("comfyui", {}) assert r["priority"] == T.get_tenant("comfyui").priority def test_pending_jobs_are_listed_in_execution_order(self, queue): J.submit("ollama", {}, priority=10, label="low") J.submit("ollama", {}, priority=90, label="high") J.submit("ollama", {}, priority=50, label="mid") order = [j["label"] for j in J.listing(J.PENDING)] assert order == ["high", "mid", "low"] def test_equal_priority_is_first_in_first_out(self, queue): a = J.submit("ollama", {}, priority=50, label="a")["id"] time.sleep(0.01) J.submit("ollama", {}, priority=50, label="b") assert J._next_job()["id"] == a def test_the_queue_has_no_depth_limit(self, queue): # "Unbounded" is the point; it lives on disk, not in memory. for i in range(500): J.submit("ollama", {}, label=f"j{i}") assert J.stats()["queue_depth"] == 500 def test_queue_survives_a_restart(self, queue): J.submit("ollama", {}, label="persisted") J._cache = None # nothing in-process is holding it assert [j["label"] for j in J.listing(J.PENDING)] == ["persisted"] class TestCancellation: def test_pending_jobs_can_be_cancelled(self, queue): jid = J.submit("ollama", {})["id"] assert J.cancel(jid)["success"] is True assert J.get(jid)["state"] == J.CANCELLED def test_running_work_is_never_cancelled(self, queue): # This service frees VRAM by asking, never by killing work in flight. jid = J.submit("ollama", {})["id"] J._mark(jid, J.RUNNING, started_at=time.time()) assert J.cancel(jid)["success"] is False assert J.get(jid)["state"] == J.RUNNING def test_clearing_the_queue_leaves_running_work_alone(self, queue): running = J.submit("ollama", {})["id"] J._mark(running, J.RUNNING, started_at=time.time()) J.submit("ollama", {}) J.submit("ollama", {}) assert J.clear_pending()["cancelled"] == 2 assert J.get(running)["state"] == J.RUNNING class TestOrphanRecovery: def test_jobs_left_running_by_a_dead_process_are_requeued(self, queue): """RUNNING means "this process is working on it". An exception during dispatch left the row RUNNING while the scheduler moved on, so the job never finished and never retried. """ jid = J.submit("ollama", {})["id"] J._mark(jid, J.RUNNING, started_at=time.time()) assert J.requeue_orphans() == 1 job = J.get(jid) assert job["state"] == J.PENDING and job["started_at"] is None def test_finished_jobs_are_untouched_by_recovery(self, queue): done = J.submit("ollama", {})["id"] J._mark(done, J.DONE, finished_at=time.time()) assert J.requeue_orphans() == 0 assert J.get(done)["state"] == J.DONE class TestPerJobVramRequirement: """A tenant-wide figure cannot be right for an LLM.""" def test_llm_requirement_comes_from_the_model_being_loaded(self, queue, monkeypatch): import vram_arbitrator monkeypatch.setattr(vram_arbitrator, "_model_size_bytes", lambda m: int(12.87 * 1024 ** 3)) s = J.Scheduler() # 12.87 GB on disk occupies ~14.9 GB once context and KV cache are allocated. assert 14.5 < s._job_vram_requirement("ollama", {"model": "big"}) < 15.5 def test_a_small_model_needs_correspondingly_less(self, queue, monkeypatch): import vram_arbitrator monkeypatch.setattr(vram_arbitrator, "_model_size_bytes", lambda m: int(1.96 * 1024 ** 3)) s = J.Scheduler() assert s._job_vram_requirement("ollama", {"model": "small"}) < 3.0 def test_falls_back_to_the_tenant_figure_when_the_model_is_unknown(self, queue, monkeypatch): import vram_arbitrator monkeypatch.setattr(vram_arbitrator, "_model_size_bytes", lambda m: 0) s = J.Scheduler() import tenants as T assert (s._job_vram_requirement("ollama", {"model": "?"}) == T.get_tenant("ollama").needs_vram_gb) def test_non_llm_tenants_use_their_declared_figure(self, queue): import tenants as T s = J.Scheduler() assert (s._job_vram_requirement("comfyui", {}) == T.get_tenant("comfyui").needs_vram_gb) class TestSchedulerStatus: def test_stats_report_depth_and_the_oldest_wait(self, queue): J.submit("ollama", {}) J.submit("comfyui", {}) st = J.stats() assert st["queue_depth"] == 2 assert st["pending_by_tenant"] == {"ollama": 1, "comfyui": 1} assert st["oldest_pending_s"] is not None def test_status_includes_queue_stats(self, queue): J.submit("ollama", {}) s = J.Scheduler().get_status() assert s["queue_depth"] == 1 and s["running"] is False def test_a_blocked_job_reports_why_it_is_waiting(self, queue): # Silence here is what left three jobs pending indefinitely with no explanation. s = J.Scheduler() s.blocked = {"id": "x", "tenant": "ollama", "reason": "needs 14.93 GB", "waited_s": 42.0, "since": time.time()} assert "14.93" in s.get_status()["blocked"]["reason"] class TestDispatchers: def test_every_tenant_kind_that_can_run_work_has_a_dispatcher(self): import tenants as T assert T.KIND_LLM in J.DISPATCHERS assert T.KIND_DIFFUSION in J.DISPATCHERS def test_ollama_dispatch_reports_failure_rather_than_raising(self, queue, monkeypatch): class _R: status_code = 500 text = '{"error":"llama-server process has terminated: cudaMalloc failed"}' class _C: async def __aenter__(self): return self async def __aexit__(self, *a): return False async def post(self, url, json=None): return _R() monkeypatch.setattr(J.httpx, "AsyncClient", lambda **k: _C()) res = asyncio.run(J._dispatch_ollama({"model": "m", "prompt": "hi"})) assert res["ok"] is False and "cudaMalloc" in res["error"]