feat(ai-studio): separate image pool from task states

app/image_studio.py

@@ -20,6 +20,14 @@ ASSET_STATUS_MISSING = "missing"
 ASSET_STATUSES = {ASSET_STATUS_AVAILABLE, ASSET_STATUS_MISSING}
 JOB_STATUSES = {"pending", "submitted", "running", "succeeded", "failed", "expired", "cancelled"}
 JOB_RESUMABLE_STATUSES = {"submitted", "running"}
+JOB_RECOVERY_NONE = "none"
+JOB_RECOVERY_RESUME = "resume"
+JOB_RECOVERY_REGENERATE = "regenerate"
+JOB_RECOVERY_ACTIONS = {
+    JOB_RECOVERY_NONE,
+    JOB_RECOVERY_RESUME,
+    JOB_RECOVERY_REGENERATE,
+}
 SELECTION_TYPES = {"main", "detail"}
 
 
@@ -70,6 +78,7 @@ class ImageStudioJob:
     status: str
     prompt: Optional[str]
     error: Optional[str]
+    recovery_action: str
     attempts: int
     points_cost: Optional[int]
     points_balance: Optional[int]
@@ -677,6 +686,7 @@ def set_job_submitted(job_id, task_id, *, call_id=None, points_cost=None, points
                 UPDATE image_studio_jobs
                 SET task_id = ?,
                     status = 'submitted',
+                    recovery_action = ?,
                     call_id = ?,
                     points_cost = ?,
                     points_balance = ?,
@@ -684,7 +694,16 @@ def set_job_submitted(job_id, task_id, *, call_id=None, points_cost=None, points
                     updated_at = ?
                 WHERE id = ?
                 """,
-                (str(task_id), call_id, points_cost, points_balance, now, now, int(job_id)),
+                (
+                    str(task_id),
+                    JOB_RECOVERY_RESUME,
+                    call_id,
+                    points_cost,
+                    points_balance,
+                    now,
+                    now,
+                    int(job_id),
+                ),
             )
         return get_job(job_id, conn=database)
 
@@ -696,12 +715,17 @@ def update_job_status(
     error=None,
     output_asset_id=None,
     points_balance=None,
+    recovery_action=None,
     increment_attempts=False,
     path=None,
     conn=None,
 ):
     if str(status) not in JOB_STATUSES:
         raise db.DbError("AI工场任务状态无效")
+    if recovery_action is not None and str(recovery_action) not in JOB_RECOVERY_ACTIONS:
+        raise db.DbError("AI工场任务恢复方式无效")
+    if recovery_action is None and str(status) == "succeeded":
+        recovery_action = JOB_RECOVERY_NONE
     now = _now()
     terminal = str(status) in {"succeeded", "failed", "expired", "cancelled"}
     with _connection(conn, path) as database:
@@ -716,6 +740,7 @@ def update_job_status(
                 UPDATE image_studio_jobs
                 SET status = ?,
                     error = ?,
+                    recovery_action = COALESCE(?, recovery_action),
                     output_asset_id = COALESCE(?, output_asset_id),
                     points_balance = COALESCE(?, points_balance),
                     attempts = attempts + ?,
@@ -726,6 +751,7 @@ def update_job_status(
                 (
                     str(status),
                     error,
+                    recovery_action,
                     output_asset_id,
                     points_balance,
                     1 if increment_attempts else 0,
@@ -742,12 +768,13 @@ def list_resumable_jobs(path=None, conn=None, project_id=None, include_failed_do
     if include_failed_downloads:
         clauses = [
             "task_id IS NOT NULL",
-            "(status IN (?, ?) OR (status = ? AND output_asset_id IS NULL))",
+            "recovery_action = ?",
+            "status IN (?, ?, ?, ?)",
         ]
-        params = ["submitted", "running", "failed"]
+        params = [JOB_RECOVERY_RESUME, "submitted", "running", "failed", "cancelled"]
     else:
-        clauses = ["status IN (?, ?)", "task_id IS NOT NULL"]
-        params = ["submitted", "running"]
+        clauses = ["status IN (?, ?)", "task_id IS NOT NULL", "recovery_action = ?"]
+        params = ["submitted", "running", JOB_RECOVERY_RESUME]
     if project_id is not None:
         clauses.append("project_id = ?")
         params.append(int(project_id))