fix(data-process): 恢复提前中断的重新生成任务
This commit is contained in:
caoxiaozhu
@@ -408,6 +408,173 @@ class DataProcessStore:
|
||||
if task.get("output_dataset_id") and not _is_regeneration_prepared(task):
|
||||
raise InvalidStateError("published task cannot be edited")
|
||||
|
||||
def recover_legacy_aborted_regeneration(self, task_id: str) -> dict[str, Any]:
|
||||
"""恢复旧版在真正开始生成前误删的上一轮结果。
|
||||
|
||||
旧实现会在 ``POST /regenerate`` 时立即把已发布任务置为 pending、
|
||||
清空结果并解除输出指针。三个已发布数据集仍是独立完整产物，因此只在
|
||||
这个特征完全匹配时，使用其记录恢复结果和任务状态。该操作幂等，不会
|
||||
触碰正常的新建待生成任务或已经开始的新一轮生成。
|
||||
"""
|
||||
|
||||
with self.connect() as conn:
|
||||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||||
if (
|
||||
task.get("status") != "pending"
|
||||
or task.get("generation_run_id")
|
||||
or task.get("output_dataset_id")
|
||||
or int(task.get("output_count") or 0) != 0
|
||||
):
|
||||
return {"recovered": False, "result_count": 0}
|
||||
|
||||
result_count = int(
|
||||
(
|
||||
conn.execute(
|
||||
"SELECT COUNT(*) AS count FROM data_process_results WHERE task_id=%s",
|
||||
(task_id,),
|
||||
).fetchone()
|
||||
or {}
|
||||
).get("count")
|
||||
or 0
|
||||
)
|
||||
if result_count:
|
||||
return {"recovered": False, "result_count": result_count}
|
||||
|
||||
datasets = conn.execute(
|
||||
"""
|
||||
SELECT id, type, count, created_at
|
||||
FROM datasets
|
||||
WHERE source='task' AND deleted_at IS NULL
|
||||
AND (
|
||||
source_task_id=%s
|
||||
OR (source_task_id IS NULL AND task_id=%s)
|
||||
)
|
||||
ORDER BY CASE type
|
||||
WHEN 'train' THEN 1 WHEN 'val' THEN 2 WHEN 'test' THEN 3 ELSE 4
|
||||
END, created_at, id
|
||||
""",
|
||||
(task_id, task_id),
|
||||
).fetchall()
|
||||
train_dataset = next(
|
||||
(dataset for dataset in datasets if dataset.get("type") == "train"),
|
||||
None,
|
||||
)
|
||||
if not train_dataset:
|
||||
return {"recovered": False, "result_count": 0}
|
||||
|
||||
dataset_ids = [str(dataset["id"]) for dataset in datasets]
|
||||
records = conn.execute(
|
||||
"""
|
||||
SELECT id, dataset_id, line_no, split, instruction, input, output,
|
||||
raw, status, source_result_id, preview_item_id, created_at
|
||||
FROM dataset_records
|
||||
WHERE dataset_id = ANY(%s)
|
||||
ORDER BY created_at, dataset_id, line_no NULLS LAST, id
|
||||
""",
|
||||
(dataset_ids,),
|
||||
).fetchall()
|
||||
if not records:
|
||||
return {"recovered": False, "result_count": 0}
|
||||
|
||||
preview_rows = conn.execute(
|
||||
"SELECT id FROM data_process_preview_items WHERE task_id=%s",
|
||||
(task_id,),
|
||||
).fetchall()
|
||||
preview_ids = {str(row["id"]) for row in preview_rows}
|
||||
used_result_ids: set[str] = set()
|
||||
recovered_count = 0
|
||||
for record in records:
|
||||
raw = _json_value(record.get("raw"), {})
|
||||
raw = raw if isinstance(raw, dict) else {}
|
||||
candidate_id = str(
|
||||
record.get("source_result_id")
|
||||
or raw.get("source_result_id")
|
||||
or ""
|
||||
)
|
||||
result_id = (
|
||||
candidate_id
|
||||
if candidate_id and candidate_id not in used_result_ids
|
||||
else new_id("dpr")
|
||||
)
|
||||
used_result_ids.add(result_id)
|
||||
candidate_preview_id = str(
|
||||
record.get("preview_item_id")
|
||||
or raw.get("preview_item_id")
|
||||
or ""
|
||||
)
|
||||
preview_item_id = (
|
||||
candidate_preview_id if candidate_preview_id in preview_ids else None
|
||||
)
|
||||
instruction = str(record.get("instruction") or raw.get("instruction") or "")
|
||||
input_text = str(record.get("input") or raw.get("input") or "")
|
||||
output = str(record.get("output") or raw.get("output") or "")
|
||||
split = str(record.get("split") or raw.get("split") or "") or None
|
||||
status = str(record.get("status") or "valid")
|
||||
if status not in {"valid", "modified", "invalid"}:
|
||||
status = "valid"
|
||||
created_at = record.get("created_at") or utcnow()
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO data_process_results
|
||||
(id, task_id, preview_item_id, instruction, input, output,
|
||||
original_instruction, original_input, original_output, status,
|
||||
error, split, quality_score, created_at, updated_at)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
|
||||
NULL, %s, '{}', %s, %s)
|
||||
""",
|
||||
(
|
||||
result_id,
|
||||
task_id,
|
||||
preview_item_id,
|
||||
instruction,
|
||||
input_text,
|
||||
output,
|
||||
instruction,
|
||||
input_text,
|
||||
output,
|
||||
status,
|
||||
split,
|
||||
created_at,
|
||||
created_at,
|
||||
),
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
UPDATE dataset_records
|
||||
SET source_result_id=%s, preview_item_id=%s
|
||||
WHERE id=%s
|
||||
""",
|
||||
(result_id, preview_item_id, record["id"]),
|
||||
)
|
||||
recovered_count += 1
|
||||
|
||||
now = utcnow()
|
||||
published_at = max(
|
||||
(dataset.get("created_at") for dataset in datasets if dataset.get("created_at")),
|
||||
default=None,
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
UPDATE data_process_tasks
|
||||
SET status='completed', progress=100, output_dataset_id=%s,
|
||||
output_count=%s, filtered_count=0, duplicate_count=0,
|
||||
error_count=(SELECT COUNT(*) FROM data_process_results
|
||||
WHERE task_id=%s AND status='invalid'),
|
||||
failure_reason=NULL, completed_at=COALESCE(completed_at, %s),
|
||||
updated_at=%s
|
||||
WHERE id=%s
|
||||
""",
|
||||
(
|
||||
train_dataset["id"],
|
||||
recovered_count,
|
||||
task_id,
|
||||
published_at,
|
||||
now,
|
||||
task_id,
|
||||
),
|
||||
)
|
||||
return {"recovered": True, "result_count": recovered_count}
|
||||
|
||||
def update_task(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
allowed = {
|
||||
"name",
|
||||
@@ -504,6 +671,8 @@ class DataProcessStore:
|
||||
生成时，才在同一事务内切换运行状态并清理上一轮结果。
|
||||
"""
|
||||
|
||||
# 先修复曾被旧版 prepare 提前清空的任务，再建立新的非破坏性草稿标记。
|
||||
self.recover_legacy_aborted_regeneration(task_id)
|
||||
try:
|
||||
with self.connect() as conn:
|
||||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||||
|
||||
Reference in New Issue
Block a user