From cb7e0f5c4844a20bfe57c8e2b77b96044005b52b Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Mon, 26 Oct 2026 23:19:12 -0700 Subject: [PATCH] fix(autotag): keep each download route's error shape and harden the fake download --- py/api/autotag_routes.py | 45 ++++++++++++++++------ tests/browser/workspace_tagging_fixture.py | 10 ++++- tests/test_api_autotag.py | 33 +++++++++++----- tests/test_job_admission.py | 2 +- 4 files changed, 66 insertions(+), 24 deletions(-) diff --git a/py/api/autotag_routes.py b/py/api/autotag_routes.py index ffa2264..e5521dd 100644 --- a/py/api/autotag_routes.py +++ b/py/api/autotag_routes.py @@ -257,13 +257,20 @@ class AutotagRoutesMixin: if model_type not in MODELS: invalid = {"type": "error", "message": f"Invalid model type: {model_type}"} return await self._stream_sse(request, _single(invalid)) - job = self._jobs.active(MODEL_DOWNLOAD_JOB_KIND) - if job is not None and job.scope.get("model_type") != model_type: - return model_busy_response() - if job is None: - job, refusal = self._start_model_download(model_type, get_autotag_service()) - if refusal is not None: - return refusal + try: + job = self._jobs.active(MODEL_DOWNLOAD_JOB_KIND) + if job is not None and job.scope.get("model_type") != model_type: + return model_busy_response() + if job is None: + job, refusal = self._start_model_download( + model_type, get_autotag_service() + ) + if refusal is not None: + return refusal + except Exception: + self.logger.exception("Download model error") + failed = {"type": "error", "message": _INTERNAL_ERROR} + return await self._stream_sse(request, _single(failed)) return await self._stream_sse(request, self._download_frames(job)) @staticmethod @@ -300,7 +307,13 @@ class AutotagRoutesMixin: return self._bad_request( "model_type must be one of " + ", ".join(sorted(MODELS)) ) - job, refusal = self._start_model_download(model_type, get_autotag_service()) + try: + job, refusal = self._start_model_download(model_type, get_autotag_service()) + except Exception as e: + self.logger.exception("Model download start error") + return web.json_response( + {"success": False, "error": self._public_error(e)}, status=500 + ) if refusal is not None: return refusal return web.json_response({"success": True, "job": job.summary()}, status=202) @@ -340,10 +353,20 @@ class AutotagRoutesMixin: self._run_model_download, service, model_type, owner ), ) - except JobBusy as refused: + except BaseException as refused: + # JobBusy, or anything else the registry raised: give the gate back. model_gate.release(owner) - return None, busy_response(refused) - job.task.add_done_callback(lambda _task: model_gate.release(owner)) + _download_titles.pop(owner, None) + if isinstance(refused, JobBusy): + return None, busy_response(refused) + raise + + def ended(_task): + model_gate.release(owner) + if model_gate.holder() != owner: # not while a thread still runs + _download_titles.pop(owner, None) + + job.task.add_done_callback(ended) return job, None async def _run_model_download(self, service, model_type, owner, job): diff --git a/tests/browser/workspace_tagging_fixture.py b/tests/browser/workspace_tagging_fixture.py index 52d2dce..387fa82 100644 --- a/tests/browser/workspace_tagging_fixture.py +++ b/tests/browser/workspace_tagging_fixture.py @@ -185,14 +185,20 @@ def register_tagging_commands( # The fake download (``/__test__/tagging/download``): waits on ``event``, # then answers ``result``. Every reset sets the event and puts back the # model's raising default, so no test can leave a worker hanging. - download = {"event": threading.Event(), "result": True} + # ``generation`` grows on every reset: a download a reset released (or one + # that timed out) fails without touching the next test's model state. + download = {"event": threading.Event(), "result": True, "generation": 0} def restore_download() -> None: + download["generation"] += 1 download["event"].set() model.__dict__.pop("download_model", None) def fake_download(model_type: str, progress_callback=None) -> bool: - download["event"].wait(DOWNLOAD_HOLD_LIMIT_S) + event, generation = download["event"], download["generation"] + released = event.wait(DOWNLOAD_HOLD_LIMIT_S) + if not released or download["generation"] != generation: + return False if download["result"]: model.unavailable = model.unavailable - {model_type} return download["result"] diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py index 11c91a5..3d90786 100644 --- a/tests/test_api_autotag.py +++ b/tests/test_api_autotag.py @@ -1476,6 +1476,13 @@ class DownloadJobTestCase(AutotagAPITestCase): "/prompt_manager/jobs/model-download", {"model_type": model_type} ) + async def _until(self, condition, limit=10): + """Yield to the loop until ``condition()`` holds (fails after ``limit`` s).""" + deadline = asyncio.get_running_loop().time() + limit + while not condition(): + self.assertLess(asyncio.get_running_loop().time(), deadline) + await asyncio.sleep(0) + async def _ended(self, job_id): job = self.api._jobs.get(job_id) self.proceed.set() @@ -1633,6 +1640,8 @@ class TestLegacyDownloadAttaches(DownloadJobTestCase): resp = await self.client.post("/prompt_manager/autotag/download/wd14-vit") self.assertEqual(resp.status, 200) self.assertEqual(self._sse_events(await resp.text()), self.GOLDEN_OK) + # The gate is free by the time the stream says the download ended. + self.assertIsNone(model_gate.holder()) jobs = self.api._jobs.list() self.assertEqual( [(j.kind, j.title, j.state) for j in jobs], @@ -1651,13 +1660,8 @@ class TestLegacyDownloadAttaches(DownloadJobTestCase): self.client.post("/prompt_manager/autotag/download/wd14-vit") ) try: - # Wait until the legacy stream is subscribed to the running job. job = self.api._jobs.get(summary["id"]) - for _ in range(300): - if job.subscriber_count: - break - await asyncio.sleep(0.01) - self.assertEqual(job.subscriber_count, 1) + await self._until(lambda: job.subscriber_count == 1) finally: await self._ended(summary["id"]) resp = await legacy @@ -1687,11 +1691,20 @@ class TestLegacyDownloadAttaches(DownloadJobTestCase): await self._started() resp.close() # the page closed the modal (job,) = self.api._jobs.list() - for _ in range(300): - if not job.subscriber_count: - break - await asyncio.sleep(0.01) + # The closed stream lets go of the job; the download goes on. + await self._until(lambda: job.subscriber_count == 0) self.assertTrue(job.running) + self.assertTrue(model_gate.holder().startswith("model-download:")) job = await self._ended(job.id) self.assertEqual(job.state, "finished") self.assertIsNone(model_gate.holder()) + + async def test_unexpected_start_error_keeps_each_route_shape(self): + with patch("py.autotag.get_autotag_service", side_effect=RuntimeError("x")): + legacy = await self.client.post("/prompt_manager/autotag/download/wd14-vit") + events = self._sse_events(await legacy.text()) + native = await self._start() + self.assertEqual(events[-1]["type"], "error") + self.assertEqual(native.status, 500) + self.assertFalse((await native.json())["success"]) + self.assertIsNone(model_gate.holder()) diff --git a/tests/test_job_admission.py b/tests/test_job_admission.py index 0bfb11c..44d47ff 100644 --- a/tests/test_job_admission.py +++ b/tests/test_job_admission.py @@ -460,7 +460,7 @@ class TestLegacyRoutesHoldTheRegistry(_LegacyHoldCase): async def test_second_legacy_autotag_batch_keeps_its_json_409(self): from py.api.autotag_routes import _BUSY_ERROR - self._run_job("autotag", "Auto-tag images (legacy page)") + self._run_job("autotag", "Tag output images (legacy page)") resp = await self.client.request( "POST", "/prompt_manager/autotag/start", json={} )