Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 20 additions & 7 deletions app/services/image_generation_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -193,28 +193,41 @@ def restore_recovery(self, workspaces):
"created_at": task["created_at"], **deepcopy(runtime)})

def _recovery_task(self, record):
"""Link one leftover to its admission, or withhold it.

Non-image leftovers return None so the native queue can recover them.
A linked image leftover returns ``(registry, task)``. An image row that
cannot be matched — deleted workspace, missing admission, corrupt
snapshot, or job-id drift — returns False so this one row is skipped.
Raising here would take down list/resume/discard for every other job.
"""
provenance = record.get("provenance") or {}
if provenance.get("capability") != "generation.image":
return None
registry = self._registry(record["workspace"])
intent_id = provenance.get("command", {}).get("command_id")
entry = registry.command_admission(intent_id)
if entry is None or entry["receipt"]["result"]["job_id"] != record["id"]:
raise command_error(503, "recovery_mismatch", "Recovery does not match a durable image admission")
return registry, registry.get(entry["task_id"])
try:
registry = self._registry(record.get("workspace"))
intent_id = (provenance.get("command") or {}).get("command_id")
entry = registry.command_admission(intent_id)
if entry is None or entry["receipt"]["result"]["job_id"] != record.get("id"):
return False
return registry, registry.get(entry["task_id"])
except (HTTPException, OSError, sqlite3.Error, TypeError, KeyError):
return False

def filter_recovery(self, records):
retained = []
for record in records:
linked = self._recovery_task(record)
if linked is False:
continue
if linked is None or (linked[1] and linked[1]["status"] == "interrupted"):
retained.append(record)
return retained

def discard_recovery(self, records):
for record in records:
linked = self._recovery_task(record)
if linked is not None and linked[1] and linked[1]["status"] == "interrupted":
if linked and linked[1] and linked[1]["status"] == "interrupted":
registry, task = linked
registry.update(task["id"], status="cancelled", phase="recovery_discarded",
message="Recovery discarded", completed_at=time.time(), recoverable=False)
43 changes: 43 additions & 0 deletions tests/test_image_generation_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -441,6 +441,49 @@ def test_restart_marks_queued_task_interrupted_and_projects_recovery_without_sta
assert receipt["task"]["status"] == "interrupted"


def _queue_record(job_id, workspace, *, capability, command_id=None, params=None):
provenance = {"capability": capability}
if command_id is not None:
provenance["command"] = {"command_id": command_id}
return {
"id": job_id,
"status": "interrupted",
"workspace": workspace,
"params": params or {"model_type": "pi_flux2", "prompt": job_id},
"provenance": provenance,
}


def test_orphaned_image_leftover_does_not_block_other_recovery(tmp_path):
native = FakeNative(tmp_path)
command = _command("linked-recovery")
first = _run(native.service().submit(command))
restarted = FakeNative(tmp_path, interrupt_stale=True)
service = restarted.service()
service.restore_recovery(["workspace-a"])
linked = restarted.persisted[first["receipt"]["result"]["job_id"]]
video = _queue_record("video-leftover", "workspace-b", capability="generation.video")
orphan = _queue_record(
"orphan-image", "deleted-workspace",
capability="generation.image", command_id="missing-admission",
)
drifted = _queue_record(
"drifted-job", "workspace-a",
capability="generation.image", command_id=command["intent_id"],
)
invalid_workspace = _queue_record(
"bad-workspace", "../outside",
capability="generation.image", command_id="any-intent",
)

retained = service.filter_recovery([video, orphan, drifted, invalid_workspace, linked])

assert [record["id"] for record in retained] == ["video-leftover", linked["id"]]
service.discard_recovery([video, orphan, drifted, invalid_workspace, linked])
assert restarted.registry("workspace-a").get(first["receipt"]["taskIds"][0])["status"] == "cancelled"
assert service.filter_recovery([video, orphan, linked]) == [video]


def test_queued_admission_is_not_a_recovery_candidate_while_dispatch_is_pending(tmp_path):
native = FakeNative(tmp_path)
service, registry, _entry = _seed_undispatched(native, _command("queued-not-recovery"))
Expand Down