diff --git a/py/api/image_tag_jobs.py b/py/api/image_tag_jobs.py index c2596af6..4b670f32 100644 --- a/py/api/image_tag_jobs.py +++ b/py/api/image_tag_jobs.py @@ -375,13 +375,11 @@ def _busy_refusal(busy: JobBusy) -> Refusal: return Refusal(409, "busy", f"{title} already running", busy_response(busy)) -def _batch_cancelled(job): - """End the batch as cancelled; None (no service to hand back).""" +def _batch_cancelled(job) -> str: + """Say the batch was cancelled; its end state (the runner ends the job).""" job.update(status_line=BATCH_CANCELLED) job.log_line(BATCH_CANCELLED) - job.publish({"type": "cancelled", "progress": job.progress}) - job.finish("cancelled") - return None + return "cancelled" class ImageTagJobsMixin: @@ -513,16 +511,17 @@ class ImageTagJobsMixin: ``job.cancelled`` is checked between phases (and per pre-pass chunk and per image). Each phase publishes ``{"type": "phase", "phase"}``. - With ``keep_in_memory`` false the model is unloaded under the gate - before the gate is released, also after a cancel or a failure. + With ``keep_in_memory`` false the model is unloaded under the gate. + The job ends only after that cleanup and the gate release, so a job + that reads as ended never still holds the gate. """ + state = "failed" try: - await self._autotag_batch_phases(run, job) + state = await self._autotag_batch_phases(run, job) except Exception: self.logger.exception("Tag output images failed") job.update(status_line=_INTERNAL_ERROR) job.log_line(f"Tagging failed · {_INTERNAL_ERROR}") - job.finish("failed") finally: try: if not run.keep_in_memory: @@ -531,6 +530,8 @@ class ImageTagJobsMixin: await self._prune_backup_pool(job, self.db) finally: model_gate.release(run.owner) + job.publish({"type": state, "progress": job.progress}) + job.finish(state) async def _unload_batch_model(self, owner) -> None: from ..autotag import get_autotag_service @@ -542,11 +543,12 @@ class ImageTagJobsMixin: except Exception: self.logger.exception("The tagging model could not be unloaded") - async def _autotag_batch_phases(self, run: AutotagRun, job): + async def _autotag_batch_phases(self, run: AutotagRun, job) -> str: + """Run the phases; the job's end state (the caller ends the job).""" job.publish({"type": "phase", "phase": "backup"}) ok, _backup = await self._safety_backup_step(job, self.db, run.backup_on) if not ok: - return None + return "failed" if job.cancelled: return _batch_cancelled(job) job.publish({"type": "phase", "phase": "model"}) @@ -554,8 +556,7 @@ class ImageTagJobsMixin: spec = _Spec("auto", (), model_type=run.model_type, options=run.options) service = await self._prepare_tagging_model(job, spec, run.owner) if service is None: - job.finish("failed") - return None + return "failed" if job.cancelled: return _batch_cancelled(job) job.publish({"type": "phase", "phase": "listing"}) @@ -567,8 +568,7 @@ class ImageTagJobsMixin: self.logger.exception("Tag output images could not list the outputs") job.update(status_line=LISTING_FAILED) job.log_line(f"! {LISTING_FAILED}") - job.finish("failed") - return None + return "failed" total = len(targets) job.scope["images"] = total job.update(total=total) @@ -592,16 +592,13 @@ class ImageTagJobsMixin: if not await self._tag_each_file(loop_run, done_before=done): job.update(status_line=STORAGE_RESTORED) job.log_line(f"Tagging stopped · {STORAGE_RESTORED}") - job.finish("failed") - return service + return "failed" if job.cancelled: return _batch_cancelled(job) finished = f"Tagged {_plural(job.counters.get('tagged', 0), 'image')}" job.update(status_line=finished, progress=100) job.log_line(finished) - job.publish({"type": "finished", "progress": 100}) - job.finish("finished") - return service + return "finished" async def _autotag_prepass(self, job, targets, skip_marked: bool): """The targets left for the model, or None when cancelled meanwhile. @@ -637,6 +634,7 @@ class ImageTagJobsMixin: ) skipped += len(skips) job.update(processed=skipped, progress=skipped * 100 // total) + job.publish({"type": "progress", "progress": job.progress}) if skipped: job.log_line(f"Skipped {_plural(skipped, 'image')} before tagging") return remaining diff --git a/tests/test_autotag_batch.py b/tests/test_autotag_batch.py index ba557d39..c6703ecb 100644 --- a/tests/test_autotag_batch.py +++ b/tests/test_autotag_batch.py @@ -372,6 +372,68 @@ class TestBatchLifecycle(BatchTestCase): self.assertEqual(self.tagger.calls, []) self.assertIsNone(model_gate.holder()) + async def test_cancel_between_prepass_chunks_keeps_counts(self): + for n in range(3): + self.image(f"{n}.png", (n, n, n)) + await self.run_batch() + seen = [] + real = jobs_module._prepass_chunk + + def chunk(db, items, skip_marked, max_file): + result = real(db, items, skip_marked, max_file) + job = self.api._jobs.active("autotag") + seen.append(job.processed) + job.cancelled = True # cancel after the first chunk + return result + + with ( + patch.object(jobs_module, "PREPASS_CHUNK", 1), + patch.object(jobs_module, "_prepass_chunk", side_effect=chunk), + ): + job = await self.run_batch() + self.assertEqual(job.state, "cancelled") + self.assertEqual(job.counters["skipped"], 1) + self.assertEqual((job.processed, job.total), (1, 3)) + self.assertEqual(seen, [0]) + + async def test_the_job_ends_only_after_its_cleanup(self): + self.image("a.png", (1, 1, 1)) + PromptManagerConfig.SAFETY_BACKUP = True + states = [] + + async def prune(job, _db): + states.append((job.state, model_gate.holder() is not None)) + + with ( + patch("py.api.maintenance.take_safety_backup", return_value="db.backup_x"), + patch.object(self.api, "_prune_backup_pool", side_effect=prune), + ): + job = await self.run_batch() + self.assertEqual(job.state, "finished") + self.assertEqual(states, [("running", True)]) + self.assertIsNone(model_gate.holder()) + + async def test_prepass_publishes_its_progress(self): + for n in range(2): + self.image(f"{n}.png", (n, n, n)) + await self.run_batch() + published = [] + real = type(self.api._jobs).start + + def start(registry, *args, **kwargs): + job = real(registry, *args, **kwargs) + original = job.publish + job.publish = lambda payload: ( + published.append(dict(payload)), + original(payload), + ) + return job + + with patch.object(type(self.api._jobs), "start", start): + await self.run_batch() + self.assertIn({"type": "progress", "progress": 100}, published) + self.assertEqual(published[-1], {"type": "finished", "progress": 100}) + async def test_restore_mid_run_fails_the_job(self): self.image("a.png", (1, 1, 1)) self.image("b.png", (2, 2, 2))