6f0e82f351
- 平台治理: 租户用户权限层次、资源ACL、审批中心与审批模板、访问申请 - 存储: MinIO 存储进度迁移、对象存储安全加固与测试 - 计算: GPU 资源预留、compute 轮询与同步增强 - 权限: permission v2 迁移、权限安全验收测试 - 日志: 后端运行日志中文说明、操作日志整合 - 数据处理/评测: 数据转换与模型评测优化 Co-Authored-By: Claude <noreply@anthropic.com>
880 lines
38 KiB
Python
880 lines
38 KiB
Python
"""数据处理存储层 - 任务管理。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from typing import Any
|
||
from collections.abc import Sequence
|
||
|
||
import psycopg
|
||
|
||
from .base import (
|
||
StoreBase,
|
||
utcnow,
|
||
new_id,
|
||
repeat_task_id,
|
||
json_dumps,
|
||
_json_value,
|
||
_decode_row,
|
||
_public_task,
|
||
_business_config,
|
||
_preview_config_value,
|
||
_preview_config_changed,
|
||
_preview_config_projection,
|
||
_normalized_preprocess_options,
|
||
_regeneration_marker,
|
||
_is_regeneration_prepared,
|
||
_task_output_type,
|
||
_task_reasoning_detail,
|
||
_reasoning_output_is_valid,
|
||
_dpo_fields_are_valid,
|
||
_source_storage_descriptor,
|
||
_serialize_value,
|
||
NotFoundError,
|
||
ConflictError,
|
||
InvalidStateError,
|
||
EDITABLE_STATUSES,
|
||
ACTIVE_PREVIEW_STATUSES,
|
||
WORKFLOW_STEPS,
|
||
_REGENERATION_MARKER_KEY,
|
||
_REPEAT_SOURCE_TASK_KEY,
|
||
_REPEAT_REQUEST_KEY,
|
||
_INTERNAL_CONFIG_KEYS,
|
||
)
|
||
|
||
from ..algorithms import estimate_token_count
|
||
|
||
class TasksMixin:
|
||
"""任务管理 Mixin。"""
|
||
|
||
def list_tasks(
|
||
self,
|
||
*,
|
||
page: int = 1,
|
||
page_size: int = 20,
|
||
keyword: str | None = None,
|
||
status: str | None = None,
|
||
process_type: str | None = None,
|
||
tenant_id: str | None = None,
|
||
project_id: str | None = None,
|
||
) -> dict[str, Any]:
|
||
clauses = ["task.deleted_at IS NULL"]
|
||
params: list[Any] = []
|
||
if keyword:
|
||
clauses.append("(task.name ILIKE %s OR COALESCE(task.description, '') ILIKE %s)")
|
||
pattern = f"%{keyword.strip()}%"
|
||
params.extend([pattern, pattern])
|
||
if status:
|
||
clauses.append("task.status = %s")
|
||
params.append(status)
|
||
if process_type:
|
||
clauses.append("task.process_type = %s")
|
||
params.append(process_type)
|
||
if tenant_id:
|
||
clauses.append("task.tenant_id = %s")
|
||
params.append(tenant_id)
|
||
if project_id:
|
||
clauses.append("task.project_id = %s")
|
||
params.append(project_id)
|
||
where = " AND ".join(clauses)
|
||
with self.connect() as conn:
|
||
total = conn.execute(
|
||
f"SELECT COUNT(*) AS count FROM data_process_tasks task WHERE {where}",
|
||
params,
|
||
).fetchone()["count"]
|
||
rows = conn.execute(
|
||
f"""
|
||
SELECT task.*,
|
||
creator.display_name AS creator_name,
|
||
processor.display_name AS processor_name,
|
||
(SELECT COUNT(*) FROM data_process_source_files source_file
|
||
WHERE source_file.task_id=task.id
|
||
AND source_file.deleted_at IS NULL) AS source_file_count
|
||
FROM data_process_tasks task
|
||
LEFT JOIN users creator ON creator.id=task.created_by
|
||
LEFT JOIN users processor ON processor.id=task.updated_by
|
||
WHERE {where}
|
||
ORDER BY task.created_at DESC, task.id DESC
|
||
LIMIT %s OFFSET %s
|
||
""",
|
||
[*params, page_size, (page - 1) * page_size],
|
||
).fetchall()
|
||
return {
|
||
"items": [_public_task(_decode_row(row)) for row in rows],
|
||
"total": int(total),
|
||
"page": page,
|
||
"page_size": page_size,
|
||
}
|
||
|
||
def create_task(self, payload: dict[str, Any]) -> dict[str, Any]:
|
||
task_id = new_id("dpt")
|
||
now = utcnow()
|
||
try:
|
||
with self.connect() as conn:
|
||
row = conn.execute(
|
||
"""
|
||
INSERT INTO data_process_tasks
|
||
(id, name, description, status, process_type, source_dataset_id, config,
|
||
progress, results_confirmed, tenant_id, project_id, owner_id, created_by, updated_by,
|
||
created_at, updated_at)
|
||
VALUES (%s, %s, %s, 'pending', %s, %s, %s, 0, FALSE,
|
||
%s, %s, %s, %s, %s, %s, %s)
|
||
RETURNING *
|
||
""",
|
||
(
|
||
task_id,
|
||
payload["name"],
|
||
payload.get("description") or "",
|
||
payload["process_type"],
|
||
payload.get("source_dataset_id"),
|
||
json_dumps(_business_config(payload.get("config"))),
|
||
payload.get("tenant_id"),
|
||
payload.get("project_id"),
|
||
payload.get("owner_id"),
|
||
payload.get("created_by"),
|
||
payload.get("created_by"),
|
||
now,
|
||
now,
|
||
),
|
||
).fetchone()
|
||
except psycopg.errors.UniqueViolation as exc:
|
||
raise ConflictError("data process task name already exists") from exc
|
||
return _public_task(_decode_row(row)) or {}
|
||
|
||
@staticmethod
|
||
def _repeat_response(
|
||
conn: psycopg.Connection[dict[str, Any]],
|
||
row: dict[str, Any],
|
||
*,
|
||
source_task_id: str,
|
||
created: bool,
|
||
) -> dict[str, Any]:
|
||
task_id = str(row["id"])
|
||
counts = conn.execute(
|
||
"""
|
||
SELECT
|
||
(SELECT COUNT(*) FROM data_process_source_files
|
||
WHERE task_id=%s AND deleted_at IS NULL) AS source_file_count,
|
||
(SELECT COUNT(*) FROM data_process_preview_items
|
||
WHERE task_id=%s) AS preview_count
|
||
""",
|
||
(task_id, task_id),
|
||
).fetchone() or {}
|
||
task = _public_task(_decode_row(row)) or {}
|
||
task["source_file_count"] = int(counts.get("source_file_count") or 0)
|
||
task["preview_count"] = int(counts.get("preview_count") or 0)
|
||
return {
|
||
"task": task,
|
||
"source_task_id": source_task_id,
|
||
"created": created,
|
||
"copied_source_file_count": task["source_file_count"],
|
||
"copied_preview_count": task["preview_count"],
|
||
}
|
||
|
||
def find_repeated_task(
|
||
self,
|
||
source_task_id: str,
|
||
request_id: str,
|
||
) -> dict[str, Any] | None:
|
||
"""查找同一幂等请求已创建的新任务。"""
|
||
|
||
task_id = repeat_task_id(source_task_id, request_id)
|
||
with self.connect() as conn:
|
||
row = conn.execute(
|
||
"SELECT * FROM data_process_tasks WHERE id=%s",
|
||
(task_id,),
|
||
).fetchone()
|
||
if row is None:
|
||
return None
|
||
decoded = _decode_row(row) or {}
|
||
config = decoded.get("config") or {}
|
||
if (
|
||
config.get(_REPEAT_SOURCE_TASK_KEY) != source_task_id
|
||
or config.get(_REPEAT_REQUEST_KEY) != request_id
|
||
):
|
||
raise ConflictError("再次生成请求与现有任务冲突")
|
||
if decoded.get("deleted_at"):
|
||
raise ConflictError("此次再次生成创建的任务已被删除，请重新发起")
|
||
return self._repeat_response(
|
||
conn,
|
||
row,
|
||
source_task_id=source_task_id,
|
||
created=False,
|
||
)
|
||
|
||
def repeat_task(
|
||
self,
|
||
source_task_id: str,
|
||
*,
|
||
expected_updated_at: str,
|
||
request_id: str,
|
||
file_copies: dict[str, dict[str, str]],
|
||
) -> dict[str, Any]:
|
||
"""复制已确认任务的配置、源文件和预览，结果与发布数据保持独立。"""
|
||
|
||
task_id = repeat_task_id(source_task_id, request_id)
|
||
now = utcnow()
|
||
try:
|
||
with self.connect() as conn:
|
||
existing = conn.execute(
|
||
"SELECT * FROM data_process_tasks WHERE id=%s FOR UPDATE",
|
||
(task_id,),
|
||
).fetchone()
|
||
if existing is not None:
|
||
decoded = _decode_row(existing) or {}
|
||
config = decoded.get("config") or {}
|
||
if (
|
||
config.get(_REPEAT_SOURCE_TASK_KEY) != source_task_id
|
||
or config.get(_REPEAT_REQUEST_KEY) != request_id
|
||
):
|
||
raise ConflictError("再次生成请求与现有任务冲突")
|
||
if decoded.get("deleted_at"):
|
||
raise ConflictError("此次再次生成创建的任务已被删除，请重新发起")
|
||
return self._repeat_response(
|
||
conn,
|
||
existing,
|
||
source_task_id=source_task_id,
|
||
created=False,
|
||
)
|
||
|
||
source_task = self._task_in_connection(
|
||
conn,
|
||
source_task_id,
|
||
for_update=True,
|
||
)
|
||
if (
|
||
source_task.get("status") != "completed"
|
||
or source_task.get("results_confirmed") is False
|
||
):
|
||
raise InvalidStateError("只有已完成并确认结果的任务可以再次生成")
|
||
if source_task.get("preview_status") in ACTIVE_PREVIEW_STATUSES:
|
||
raise ConflictError("源任务仍在处理切分，暂时不能再次生成")
|
||
if expected_updated_at != _serialize_value(source_task.get("updated_at")):
|
||
raise ConflictError("源任务已被其他操作修改，请刷新后重试")
|
||
|
||
source_files = conn.execute(
|
||
"""
|
||
SELECT * FROM data_process_source_files
|
||
WHERE task_id=%s AND deleted_at IS NULL
|
||
ORDER BY created_at, id
|
||
""",
|
||
(source_task_id,),
|
||
).fetchall()
|
||
source_file_ids = {str(row["id"]) for row in source_files}
|
||
if source_file_ids != set(file_copies):
|
||
raise ConflictError("源文件快照已变化，请刷新后重试")
|
||
previews = conn.execute(
|
||
"""
|
||
SELECT * FROM data_process_preview_items
|
||
WHERE task_id=%s
|
||
ORDER BY source_file_id NULLS LAST, source_start NULLS LAST,
|
||
created_at, id
|
||
""",
|
||
(source_task_id,),
|
||
).fetchall()
|
||
if not previews:
|
||
raise InvalidStateError("源任务没有可用于再次生成的切分结果")
|
||
|
||
suffix = f"（再次生成-{task_id[-6:]}）"
|
||
base_name = str(source_task.get("name") or "数据处理任务")
|
||
repeated_name = f"{base_name[: max(1, 150 - len(suffix))]}{suffix}"
|
||
repeated_config = _business_config(source_task.get("config") or {})
|
||
repeated_config[_REPEAT_SOURCE_TASK_KEY] = source_task_id
|
||
repeated_config[_REPEAT_REQUEST_KEY] = request_id
|
||
input_count = sum(int(row.get("record_count") or 0) for row in source_files)
|
||
task_row = conn.execute(
|
||
"""
|
||
INSERT INTO data_process_tasks
|
||
(id, name, description, status, process_type, source_dataset_id,
|
||
output_dataset_id, config, progress, input_count, output_count,
|
||
filtered_count, duplicate_count, error_count, failure_reason,
|
||
generation_run_id, results_confirmed, workflow_step,
|
||
preview_status, preview_progress, preview_run_id,
|
||
preview_failure_reason, preview_total_files,
|
||
preview_completed_files, tenant_id, project_id, owner_id,
|
||
approval_status, created_by, updated_by, created_at, updated_at)
|
||
VALUES
|
||
(%s, %s, %s, 'pending', %s, %s, NULL, %s, 20, %s, 0,
|
||
0, 0, 0, NULL, NULL, FALSE, 'preview', 'completed', 100,
|
||
NULL, NULL, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||
RETURNING *
|
||
""",
|
||
(
|
||
task_id,
|
||
repeated_name,
|
||
source_task.get("description") or "",
|
||
source_task["process_type"],
|
||
source_task.get("source_dataset_id"),
|
||
json_dumps(repeated_config),
|
||
input_count,
|
||
len(source_files),
|
||
len(source_files),
|
||
source_task.get("tenant_id"),
|
||
source_task.get("project_id"),
|
||
source_task.get("owner_id"),
|
||
source_task.get("approval_status") or "not_required",
|
||
source_task.get("created_by"),
|
||
source_task.get("created_by"),
|
||
now,
|
||
now,
|
||
),
|
||
).fetchone()
|
||
|
||
file_id_map: dict[str, str] = {}
|
||
for source in source_files:
|
||
old_file_id = str(source["id"])
|
||
copy = file_copies[old_file_id]
|
||
new_file_id = str(copy["id"])
|
||
storage_object_id, metadata = _source_storage_descriptor(
|
||
{
|
||
"storage_object_id": copy["storage_object_id"],
|
||
"metadata": _json_value(source.get("metadata"), {}),
|
||
},
|
||
task_id,
|
||
new_file_id,
|
||
)
|
||
file_id_map[old_file_id] = new_file_id
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO data_process_source_files
|
||
(id, task_id, storage_object_id, name, size_bytes, record_count,
|
||
file_format, checksum_sha256, version_no, content,
|
||
content_preview, metadata, tenant_id, project_id, created_by,
|
||
created_at, updated_at)
|
||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, 1, %s, %s, %s,
|
||
%s, %s, %s, %s, %s)
|
||
""",
|
||
(
|
||
new_file_id,
|
||
task_id,
|
||
storage_object_id,
|
||
source["name"],
|
||
source.get("size_bytes") or 0,
|
||
source.get("record_count") or 0,
|
||
source.get("file_format"),
|
||
source["checksum_sha256"],
|
||
source.get("content") or "",
|
||
source.get("content_preview"),
|
||
json_dumps(metadata),
|
||
source_task.get("tenant_id"),
|
||
source_task.get("project_id"),
|
||
source.get("created_by") or source_task.get("created_by"),
|
||
now,
|
||
now,
|
||
),
|
||
)
|
||
|
||
for preview in previews:
|
||
old_source_file_id = preview.get("source_file_id")
|
||
conn.execute(
|
||
"""
|
||
INSERT INTO data_process_preview_items
|
||
(id, task_id, source_file_id, original_content, edited_content,
|
||
source_start, source_end, source_start_line, source_end_line,
|
||
token_count, status, quality_score, created_at, updated_at)
|
||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
|
||
%s, %s)
|
||
""",
|
||
(
|
||
new_id("dpp"),
|
||
task_id,
|
||
file_id_map.get(str(old_source_file_id))
|
||
if old_source_file_id
|
||
else None,
|
||
preview.get("original_content") or "",
|
||
preview.get("edited_content") or "",
|
||
preview.get("source_start"),
|
||
preview.get("source_end"),
|
||
preview.get("source_start_line"),
|
||
preview.get("source_end_line"),
|
||
max(0, int(preview.get("token_count") or 0)),
|
||
preview.get("status") or "original",
|
||
json_dumps(_json_value(preview.get("quality_score"), {})),
|
||
now,
|
||
now,
|
||
),
|
||
)
|
||
return self._repeat_response(
|
||
conn,
|
||
task_row or {},
|
||
source_task_id=source_task_id,
|
||
created=True,
|
||
)
|
||
except psycopg.errors.UniqueViolation as exc:
|
||
raise ConflictError("再次生成任务名称或请求发生冲突，请重试") from exc
|
||
|
||
def get_task(self, task_id: str, *, for_update: bool = False) -> dict[str, Any]:
|
||
lock = " FOR UPDATE" if for_update else ""
|
||
with self.connect() as conn:
|
||
if for_update:
|
||
row = conn.execute(
|
||
f"SELECT * FROM data_process_tasks WHERE id=%s AND deleted_at IS NULL{lock}",
|
||
(task_id,),
|
||
).fetchone()
|
||
else:
|
||
row = conn.execute(
|
||
"""
|
||
SELECT task.*,
|
||
creator.display_name AS creator_name,
|
||
processor.display_name AS processor_name,
|
||
(SELECT COUNT(*) FROM data_process_source_files source
|
||
WHERE source.task_id=task.id AND source.deleted_at IS NULL)
|
||
AS source_file_count,
|
||
(SELECT COUNT(*) FROM data_process_preview_items preview
|
||
WHERE preview.task_id=task.id) AS preview_count,
|
||
(SELECT COALESCE(json_agg(json_build_object(
|
||
'id', dataset.id,
|
||
'name', dataset.name,
|
||
'type', dataset.type,
|
||
'count', dataset.count,
|
||
'dataset_split', CASE dataset.type
|
||
WHEN 'train' THEN 'train'
|
||
WHEN 'val' THEN 'validation'
|
||
WHEN 'test' THEN 'test'
|
||
ELSE NULL
|
||
END
|
||
) ORDER BY CASE dataset.type
|
||
WHEN 'train' THEN 1 WHEN 'val' THEN 2 WHEN 'test' THEN 3 ELSE 4 END), '[]'::json)
|
||
FROM datasets dataset
|
||
WHERE dataset.source='task'
|
||
AND dataset.deleted_at IS NULL
|
||
AND (
|
||
dataset.source_task_id=task.id
|
||
OR (dataset.source_task_id IS NULL AND dataset.task_id=task.id)
|
||
))
|
||
AS output_datasets,
|
||
CASE
|
||
WHEN task.started_at IS NOT NULL AND task.completed_at IS NOT NULL
|
||
THEN EXTRACT(EPOCH FROM (task.completed_at - task.started_at))
|
||
ELSE NULL
|
||
END AS duration_seconds
|
||
FROM data_process_tasks task
|
||
LEFT JOIN users creator ON creator.id=task.created_by
|
||
LEFT JOIN users processor ON processor.id=task.updated_by
|
||
WHERE task.id=%s AND task.deleted_at IS NULL
|
||
""",
|
||
(task_id,),
|
||
).fetchone()
|
||
if not row:
|
||
raise NotFoundError("data process task not found")
|
||
return _public_task(_decode_row(row)) or {}
|
||
|
||
def _task_in_connection(
|
||
self,
|
||
conn: psycopg.Connection[dict[str, Any]],
|
||
task_id: str,
|
||
*,
|
||
for_update: bool = False,
|
||
) -> dict[str, Any]:
|
||
lock = " FOR UPDATE" if for_update else ""
|
||
row = conn.execute(
|
||
f"SELECT * FROM data_process_tasks WHERE id=%s AND deleted_at IS NULL{lock}",
|
||
(task_id,),
|
||
).fetchone()
|
||
if not row:
|
||
raise NotFoundError("data process task not found")
|
||
return _decode_row(row) or {}
|
||
|
||
@staticmethod
|
||
def _ensure_editable(task: dict[str, Any]) -> None:
|
||
if task["status"] not in EDITABLE_STATUSES:
|
||
raise InvalidStateError(f"task cannot be edited while status is {task['status']}")
|
||
if task.get("preview_status") in ACTIVE_PREVIEW_STATUSES:
|
||
raise InvalidStateError("task cannot be edited while preview is running")
|
||
if task.get("output_dataset_id") and not _is_regeneration_prepared(task):
|
||
raise InvalidStateError("published task cannot be edited")
|
||
|
||
def update_workflow_step(self, task_id: str, workflow_step: str) -> dict[str, Any]:
|
||
"""独立保存向导位置，不触发配置或结果失效逻辑。"""
|
||
|
||
if workflow_step not in WORKFLOW_STEPS:
|
||
raise ValueError("invalid data process workflow step")
|
||
with self.connect() as conn:
|
||
row = conn.execute(
|
||
"""
|
||
UPDATE data_process_tasks
|
||
SET workflow_step=%s, updated_at=%s
|
||
WHERE id=%s AND deleted_at IS NULL
|
||
RETURNING *
|
||
""",
|
||
(workflow_step, utcnow(), task_id),
|
||
).fetchone()
|
||
if not row:
|
||
raise NotFoundError("data process task not found")
|
||
return _public_task(_decode_row(row)) or {}
|
||
|
||
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 "")
|
||
chosen = str(raw.get("chosen") or "")
|
||
rejected = str(raw.get("rejected") or "")
|
||
output = str(
|
||
record.get("output") or raw.get("output") or chosen 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,
|
||
chosen, rejected, original_instruction, original_input,
|
||
original_output, original_chosen, original_rejected, status,
|
||
error, split, quality_score, created_at, updated_at)
|
||
VALUES (%s, %s, %s, %s, %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,
|
||
chosen,
|
||
rejected,
|
||
instruction,
|
||
input_text,
|
||
output,
|
||
chosen,
|
||
rejected,
|
||
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()
|
||
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, results_confirmed=TRUE,
|
||
workflow_step='results', updated_at=%s
|
||
WHERE id=%s
|
||
""",
|
||
(
|
||
train_dataset["id"],
|
||
recovered_count,
|
||
task_id,
|
||
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",
|
||
"description",
|
||
"process_type",
|
||
"source_dataset_id",
|
||
}
|
||
values: dict[str, Any] = {key: value for key, value in payload.items() if key in allowed}
|
||
if not values and payload.get("config") is None:
|
||
return self.get_task(task_id)
|
||
try:
|
||
with self.connect() as conn:
|
||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
self._ensure_editable(task)
|
||
regeneration_prepared = _is_regeneration_prepared(task)
|
||
if regeneration_prepared and any(
|
||
key in payload and payload.get(key) != task.get(key)
|
||
for key in ("process_type", "source_dataset_id")
|
||
):
|
||
raise InvalidStateError(
|
||
"process type and source dataset cannot change during regeneration"
|
||
)
|
||
if payload.get("config") is not None:
|
||
next_config = _business_config(payload["config"])
|
||
current_config = dict(task.get("config") or {})
|
||
for key in _INTERNAL_CONFIG_KEYS:
|
||
if key in current_config:
|
||
next_config[key] = current_config[key]
|
||
values["config"] = json_dumps(next_config)
|
||
invalidates_results = (
|
||
("config" in payload and payload.get("config") != task.get("config"))
|
||
or (
|
||
"process_type" in payload
|
||
and payload.get("process_type") != task.get("process_type")
|
||
)
|
||
or (
|
||
"source_dataset_id" in payload
|
||
and payload.get("source_dataset_id") != task.get("source_dataset_id")
|
||
)
|
||
)
|
||
if invalidates_results and not regeneration_prepared:
|
||
values.update(
|
||
{
|
||
"status": "pending",
|
||
"progress": 0,
|
||
"output_count": 0,
|
||
"filtered_count": 0,
|
||
"duplicate_count": 0,
|
||
"error_count": 0,
|
||
"failure_reason": None,
|
||
"generation_run_id": None,
|
||
"results_confirmed": False,
|
||
"preview_status": "idle",
|
||
"preview_progress": 0,
|
||
"preview_run_id": None,
|
||
"preview_failure_reason": None,
|
||
"preview_total_files": 0,
|
||
"preview_completed_files": 0,
|
||
"started_at": None,
|
||
"completed_at": None,
|
||
}
|
||
)
|
||
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
|
||
conn.execute(
|
||
"DELETE FROM data_process_preview_items WHERE task_id=%s", (task_id,)
|
||
)
|
||
if (
|
||
"process_type" in payload
|
||
and payload.get("process_type") != task.get("process_type")
|
||
):
|
||
conn.execute(
|
||
"""
|
||
UPDATE data_process_source_files
|
||
SET deleted_at=%s, updated_at=%s
|
||
WHERE task_id=%s AND deleted_at IS NULL
|
||
""",
|
||
(utcnow(), utcnow(), task_id),
|
||
)
|
||
values["input_count"] = 0
|
||
values["updated_at"] = utcnow()
|
||
assignments = ", ".join(f"{key}=%s" for key in values)
|
||
row = conn.execute(
|
||
f"UPDATE data_process_tasks SET {assignments} WHERE id=%s RETURNING *",
|
||
[*values.values(), task_id],
|
||
).fetchone()
|
||
except psycopg.errors.UniqueViolation as exc:
|
||
raise ConflictError("data process task name already exists") from exc
|
||
return _public_task(_decode_row(row)) or {}
|
||
|
||
def prepare_regeneration(
|
||
self,
|
||
task_id: str,
|
||
payload: dict[str, Any],
|
||
) -> dict[str, Any]:
|
||
"""非破坏性地保存重新生成配置。
|
||
|
||
准备阶段保留任务当前状态、结果、切片及已发布数据集。真正开始
|
||
生成时，才在同一事务内切换运行状态并清理上一轮结果。
|
||
"""
|
||
|
||
# 先修复曾被旧版 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)
|
||
if task["status"] == "running":
|
||
raise ConflictError("running task cannot be prepared for regeneration")
|
||
if task.get("preview_status") in ACTIVE_PREVIEW_STATUSES:
|
||
raise ConflictError("running preview cannot be prepared for regeneration")
|
||
|
||
current_updated_at = _serialize_value(task.get("updated_at"))
|
||
if payload["expected_updated_at"] != current_updated_at:
|
||
raise ConflictError("data process task was modified by another request")
|
||
|
||
process_type = str(payload["process_type"])
|
||
if process_type != str(task["process_type"]):
|
||
raise InvalidStateError("process_type cannot be changed during regeneration")
|
||
|
||
current_config = dict(task.get("config") or {})
|
||
next_config = _business_config(payload.get("config"))
|
||
for key in (_REPEAT_SOURCE_TASK_KEY, _REPEAT_REQUEST_KEY):
|
||
if key in current_config:
|
||
next_config[key] = current_config[key]
|
||
preview_invalidated = _preview_config_changed(
|
||
process_type,
|
||
current_config,
|
||
next_config,
|
||
)
|
||
now = utcnow()
|
||
next_config[_REGENERATION_MARKER_KEY] = {
|
||
"prepared": True,
|
||
"preview_invalidated": preview_invalidated,
|
||
"prepared_at": now,
|
||
}
|
||
# 002 迁移前发布的数据集只有 task_id。先补齐新关联字段，保证
|
||
# 解除任务输出指针后，详情和后续重新发布仍能定位原来的三份数据集。
|
||
conn.execute(
|
||
"""
|
||
UPDATE datasets
|
||
SET source_task_id=%s, updated_at=%s
|
||
WHERE source='task' AND source_task_id IS NULL AND task_id=%s
|
||
AND deleted_at IS NULL
|
||
""",
|
||
(task_id, now, task_id),
|
||
)
|
||
published_row = conn.execute(
|
||
"""
|
||
SELECT EXISTS(
|
||
SELECT 1 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)
|
||
)
|
||
) AS exists
|
||
""",
|
||
(task_id, task_id),
|
||
).fetchone()
|
||
published_outputs_preserved = bool(task.get("output_dataset_id")) or bool(
|
||
published_row and published_row.get("exists")
|
||
)
|
||
|
||
row = conn.execute(
|
||
"""
|
||
UPDATE data_process_tasks
|
||
SET name=%s, description=%s, config=%s, updated_at=%s
|
||
WHERE id=%s
|
||
RETURNING *
|
||
""",
|
||
(
|
||
payload["name"],
|
||
payload.get("description") or "",
|
||
json_dumps(next_config),
|
||
now,
|
||
task_id,
|
||
),
|
||
).fetchone()
|
||
except psycopg.errors.UniqueViolation as exc:
|
||
raise ConflictError("data process task name already exists") from exc
|
||
return {
|
||
"task": _public_task(_decode_row(row)) or {},
|
||
"preview_invalidated": preview_invalidated,
|
||
"published_outputs_preserved": published_outputs_preserved,
|
||
}
|
||
|
||
def delete_task(self, task_id: str, *, deleted_by: str | None = None) -> None:
|
||
with self.connect() as conn:
|
||
self._task_in_connection(conn, task_id, for_update=True)
|
||
now = utcnow()
|
||
conn.execute(
|
||
"""
|
||
UPDATE data_process_tasks
|
||
SET status=CASE WHEN status='running' THEN 'stopped' ELSE status END,
|
||
generation_run_id=NULL,
|
||
preview_status=CASE
|
||
WHEN preview_status IN ('queued', 'running') THEN 'cancelled'
|
||
ELSE preview_status
|
||
END,
|
||
preview_run_id=NULL,
|
||
deleted_at=%s, deleted_by=%s, updated_at=%s
|
||
WHERE id=%s
|
||
""",
|
||
(now, deleted_by, now, task_id),
|
||
)
|