fix(data-process): 恢复提前中断的重新生成任务

This commit is contained in:
caoxiaozhu
2026-07-27 11:07:18 +08:00
parent 3f5fedb9ed
commit 8caaaa5bbc
4 changed files with 323 additions and 0 deletions

View File

@@ -95,6 +95,10 @@ class FakeDataProcessStore:
)
return task
def recover_legacy_aborted_regeneration(self, task_id: str) -> dict[str, Any]:
self.get_task(task_id)
return {"recovered": False, "result_count": 0}
def update_task(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]:
self.get_task(task_id)
self.tasks[task_id].update(deepcopy(payload))

View File

@@ -357,6 +357,120 @@ class _TaskListStore(DataProcessStore):
yield self._conn
class _LegacyRecoveryConnection:
def __init__(self) -> None:
self.task = {
**_regeneration_task(
status="pending",
progress=20,
output_dataset_id=None,
output_count=0,
generation_run_id=None,
started_at=None,
completed_at=None,
),
}
self.datasets = [
{"id": "dataset_train", "type": "train", "count": 1, "created_at": "2026-07-25T18:12:00Z"},
{"id": "dataset_test", "type": "test", "count": 1, "created_at": "2026-07-25T18:12:00Z"},
]
self.records = [
{
"id": "record_train",
"dataset_id": "dataset_train",
"line_no": 1,
"split": "train",
"instruction": "训练问题",
"input": "",
"output": "训练答案",
"raw": json.dumps(
{
"source_result_id": "result_train",
"preview_item_id": "preview_1",
}
),
"status": "valid",
"source_result_id": None,
"preview_item_id": None,
"created_at": "2026-07-25T18:12:00Z",
},
{
"id": "record_test",
"dataset_id": "dataset_test",
"line_no": 1,
"split": "test",
"instruction": "测试问题",
"input": "输入",
"output": "测试答案",
"raw": json.dumps({"source_result_id": "result_test"}),
"status": "modified",
"source_result_id": None,
"preview_item_id": None,
"created_at": "2026-07-25T18:12:00Z",
},
]
self.previews = [{"id": "preview_1"}]
self.results: list[dict[str, Any]] = []
def execute(self, sql: str, params: Any = None) -> _Result:
normalized = " ".join(sql.split())
if normalized.startswith("SELECT COUNT(*) AS count FROM data_process_results"):
return _Result(row={"count": len(self.results)})
if normalized.startswith("SELECT id, type, count, created_at FROM datasets"):
return _Result(rows=list(self.datasets))
if normalized.startswith("SELECT id, dataset_id, line_no, split"):
return _Result(rows=list(self.records))
if normalized.startswith("SELECT id FROM data_process_preview_items"):
return _Result(rows=list(self.previews))
if normalized.startswith("INSERT INTO data_process_results"):
self.results.append(
{
"id": params[0],
"task_id": params[1],
"preview_item_id": params[2],
"instruction": params[3],
"input": params[4],
"output": params[5],
"status": params[9],
"split": params[10],
}
)
return _Result()
if normalized.startswith("UPDATE dataset_records SET source_result_id="):
record = next(item for item in self.records if item["id"] == params[2])
record["source_result_id"] = params[0]
record["preview_item_id"] = params[1]
return _Result()
if normalized.startswith("UPDATE data_process_tasks SET status='completed'"):
self.task.update(
{
"status": "completed",
"progress": 100,
"output_dataset_id": params[0],
"output_count": params[1],
"completed_at": params[3],
}
)
return _Result()
raise AssertionError(f"unexpected SQL: {normalized}")
class _LegacyRecoveryStore(DataProcessStore):
def __init__(self, conn: _LegacyRecoveryConnection) -> None:
self._conn = conn
@contextmanager
def connect(self) -> Iterator[_LegacyRecoveryConnection]:
yield self._conn
def _task_in_connection(
self, conn: Any, task_id: str, *, for_update: bool = False
) -> dict[str, Any]:
assert task_id == "task-1"
assert for_update is True
return dict(self._conn.task)
class _StartGenerationConnection:
def __init__(self, *, published_prepared: bool = False, preview_count: int = 1) -> None:
config = {"chunk_method": "fixed", "temperature": 0.7}
@@ -616,6 +730,40 @@ def test_prepared_published_task_survives_generation_preflight_failure() -> None
assert conn.results == [{"id": "old-result"}]
def test_legacy_aborted_regeneration_recovers_results_and_published_state() -> None:
conn = _LegacyRecoveryConnection()
store = _LegacyRecoveryStore(conn)
recovered = store.recover_legacy_aborted_regeneration("task-1")
assert recovered == {"recovered": True, "result_count": 2}
assert conn.task["status"] == "completed"
assert conn.task["progress"] == 100
assert conn.task["output_dataset_id"] == "dataset_train"
assert conn.task["output_count"] == 2
assert [item["id"] for item in conn.results] == ["result_train", "result_test"]
assert conn.results[0]["preview_item_id"] == "preview_1"
assert conn.results[1]["preview_item_id"] is None
assert conn.records[0]["source_result_id"] == "result_train"
assert conn.records[0]["preview_item_id"] == "preview_1"
repeated = store.recover_legacy_aborted_regeneration("task-1")
assert repeated == {"recovered": False, "result_count": 0}
assert len(conn.results) == 2
def test_normal_pending_task_is_not_mistaken_for_legacy_regeneration() -> None:
conn = _LegacyRecoveryConnection()
conn.datasets = []
conn.records = []
result = _LegacyRecoveryStore(conn).recover_legacy_aborted_regeneration("task-1")
assert result == {"recovered": False, "result_count": 0}
assert conn.task["status"] == "pending"
assert conn.results == []
def _legacy_published_datasets(task_id: str = "task-1") -> list[dict[str, Any]]:
specs = (
("dataset_train", "制度问答-训练集", "train", "train", 22),