276 lines
11 KiB
Python
276 lines
11 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,
|
||
|
|
DataProcessStoreError,
|
||
|
|
NotFoundError,
|
||
|
|
ConflictError,
|
||
|
|
InvalidStateError,
|
||
|
|
EDITABLE_STATUSES,
|
||
|
|
ACTIVE_PREVIEW_STATUSES,
|
||
|
|
WORKFLOW_STEPS,
|
||
|
|
_REGENERATION_MARKER_KEY,
|
||
|
|
_REPEAT_SOURCE_TASK_KEY,
|
||
|
|
_REPEAT_REQUEST_KEY,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class GenerationMixin:
|
||
|
|
"""生成管理 Mixin。"""
|
||
|
|
|
||
|
|
def _invalidate_results(
|
||
|
|
self,
|
||
|
|
conn: psycopg.Connection[dict[str, Any]],
|
||
|
|
task: dict[str, Any],
|
||
|
|
task_id: str,
|
||
|
|
now: str,
|
||
|
|
) -> None:
|
||
|
|
if _is_regeneration_prepared(task):
|
||
|
|
conn.execute(
|
||
|
|
"UPDATE data_process_tasks SET updated_at=%s WHERE id=%s",
|
||
|
|
(now, task_id),
|
||
|
|
)
|
||
|
|
return
|
||
|
|
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
|
||
|
|
conn.execute(
|
||
|
|
"""
|
||
|
|
UPDATE data_process_tasks
|
||
|
|
SET status='pending', progress=20, output_count=0, filtered_count=0,
|
||
|
|
duplicate_count=0, error_count=0, failure_reason=NULL,
|
||
|
|
generation_run_id=NULL, results_confirmed=FALSE, updated_at=%s
|
||
|
|
WHERE id=%s
|
||
|
|
""",
|
||
|
|
(now, task_id),
|
||
|
|
)
|
||
|
|
|
||
|
|
def start_generation(self, task_id: str, *, replace_existing: bool = True) -> dict[str, Any]:
|
||
|
|
if not replace_existing:
|
||
|
|
raise DataProcessStoreError("incremental generation is not supported")
|
||
|
|
with self.connect() as conn:
|
||
|
|
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
|
|
regeneration_prepared = _is_regeneration_prepared(task)
|
||
|
|
if task.get("output_dataset_id") and not regeneration_prepared:
|
||
|
|
raise InvalidStateError("published task cannot be regenerated")
|
||
|
|
if task["status"] == "running":
|
||
|
|
raise ConflictError("data process task is already running")
|
||
|
|
if task.get("preview_status") in ACTIVE_PREVIEW_STATUSES:
|
||
|
|
raise ConflictError("preview is still running")
|
||
|
|
preview_count = conn.execute(
|
||
|
|
"SELECT COUNT(*) AS count FROM data_process_preview_items WHERE task_id=%s",
|
||
|
|
(task_id,),
|
||
|
|
).fetchone()["count"]
|
||
|
|
if not preview_count:
|
||
|
|
raise InvalidStateError("preview must be built before generation")
|
||
|
|
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
|
||
|
|
now = utcnow()
|
||
|
|
generation_run_id = new_id("dprun")
|
||
|
|
next_config = dict(task.get("config") or {})
|
||
|
|
next_config.pop(_REGENERATION_MARKER_KEY, None)
|
||
|
|
row = conn.execute(
|
||
|
|
"""
|
||
|
|
UPDATE data_process_tasks
|
||
|
|
SET config=%s, status='running', progress=30, failure_reason=NULL,
|
||
|
|
started_at=%s, completed_at=NULL, output_dataset_id=NULL,
|
||
|
|
output_count=0, filtered_count=0, duplicate_count=0, error_count=0,
|
||
|
|
generation_run_id=%s, results_confirmed=FALSE,
|
||
|
|
workflow_step='generate', updated_at=%s
|
||
|
|
WHERE id=%s
|
||
|
|
RETURNING *
|
||
|
|
""",
|
||
|
|
(json_dumps(next_config), now, generation_run_id, now, task_id),
|
||
|
|
).fetchone()
|
||
|
|
return _decode_row(row) or {}
|
||
|
|
|
||
|
|
def stop_task(self, task_id: str) -> dict[str, Any]:
|
||
|
|
with self.connect() as conn:
|
||
|
|
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
|
|
if task["status"] != "running":
|
||
|
|
raise InvalidStateError("only a running task can be stopped")
|
||
|
|
now = utcnow()
|
||
|
|
row = conn.execute(
|
||
|
|
"""
|
||
|
|
UPDATE data_process_tasks
|
||
|
|
SET status='stopped', failure_reason=NULL, generation_run_id=NULL,
|
||
|
|
updated_at=%s
|
||
|
|
WHERE id=%s RETURNING *
|
||
|
|
""",
|
||
|
|
(now, task_id),
|
||
|
|
).fetchone()
|
||
|
|
return _decode_row(row) or {}
|
||
|
|
|
||
|
|
def generation_is_running(self, task_id: str, generation_run_id: str) -> bool:
|
||
|
|
with self.connect() as conn:
|
||
|
|
row = conn.execute(
|
||
|
|
"""
|
||
|
|
SELECT status, generation_run_id
|
||
|
|
FROM data_process_tasks
|
||
|
|
WHERE id=%s AND deleted_at IS NULL
|
||
|
|
""",
|
||
|
|
(task_id,),
|
||
|
|
).fetchone()
|
||
|
|
return bool(
|
||
|
|
row
|
||
|
|
and row.get("status") == "running"
|
||
|
|
and row.get("generation_run_id") == generation_run_id
|
||
|
|
)
|
||
|
|
|
||
|
|
def update_generation_progress(
|
||
|
|
self,
|
||
|
|
task_id: str,
|
||
|
|
generation_run_id: str,
|
||
|
|
processed_count: int,
|
||
|
|
total_count: int,
|
||
|
|
) -> bool:
|
||
|
|
ratio = processed_count / max(1, total_count)
|
||
|
|
progress = min(95.0, 30.0 + ratio * 65.0)
|
||
|
|
with self.connect() as conn:
|
||
|
|
row = conn.execute(
|
||
|
|
"""
|
||
|
|
UPDATE data_process_tasks
|
||
|
|
SET progress=%s, updated_at=%s
|
||
|
|
WHERE id=%s AND status='running' AND generation_run_id=%s
|
||
|
|
RETURNING id
|
||
|
|
""",
|
||
|
|
(progress, utcnow(), task_id, generation_run_id),
|
||
|
|
).fetchone()
|
||
|
|
return row is not None
|
||
|
|
|
||
|
|
def complete_generation(
|
||
|
|
self,
|
||
|
|
task_id: str,
|
||
|
|
results: Sequence[dict[str, Any]],
|
||
|
|
*,
|
||
|
|
generation_run_id: str,
|
||
|
|
filtered_count: int = 0,
|
||
|
|
duplicate_count: int = 0,
|
||
|
|
error_count: int = 0,
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
now = utcnow()
|
||
|
|
with self.connect() as conn:
|
||
|
|
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
|
|
if (
|
||
|
|
task["status"] != "running"
|
||
|
|
or task.get("generation_run_id") != generation_run_id
|
||
|
|
):
|
||
|
|
return task
|
||
|
|
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
|
||
|
|
for result in results:
|
||
|
|
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, %s, %s, %s, %s, %s)
|
||
|
|
""",
|
||
|
|
(
|
||
|
|
result.get("id") or new_id("dpr"),
|
||
|
|
task_id,
|
||
|
|
result.get("preview_item_id"),
|
||
|
|
result.get("instruction") or "",
|
||
|
|
result.get("input") or "",
|
||
|
|
result.get("output") or "",
|
||
|
|
result.get("chosen") or "",
|
||
|
|
result.get("rejected") or "",
|
||
|
|
result.get("original_instruction", result.get("instruction") or ""),
|
||
|
|
result.get("original_input", result.get("input") or ""),
|
||
|
|
result.get("original_output", result.get("output") or ""),
|
||
|
|
result.get("original_chosen", result.get("chosen") or ""),
|
||
|
|
result.get("original_rejected", result.get("rejected") or ""),
|
||
|
|
result.get("status") or "valid",
|
||
|
|
result.get("error"),
|
||
|
|
result.get("split"),
|
||
|
|
json_dumps(result.get("quality_score") or {}),
|
||
|
|
now,
|
||
|
|
now,
|
||
|
|
),
|
||
|
|
)
|
||
|
|
row = conn.execute(
|
||
|
|
"""
|
||
|
|
UPDATE data_process_tasks
|
||
|
|
SET status='completed', progress=100, output_count=%s, filtered_count=%s,
|
||
|
|
duplicate_count=%s, error_count=%s, failure_reason=NULL,
|
||
|
|
completed_at=%s, generation_run_id=NULL, results_confirmed=FALSE,
|
||
|
|
updated_at=%s
|
||
|
|
WHERE id=%s AND generation_run_id=%s RETURNING *
|
||
|
|
""",
|
||
|
|
(
|
||
|
|
len(results),
|
||
|
|
filtered_count,
|
||
|
|
duplicate_count,
|
||
|
|
error_count,
|
||
|
|
now,
|
||
|
|
now,
|
||
|
|
task_id,
|
||
|
|
generation_run_id,
|
||
|
|
),
|
||
|
|
).fetchone()
|
||
|
|
return _decode_row(row) or {}
|
||
|
|
|
||
|
|
def mark_failed(
|
||
|
|
self, task_id: str, reason: str, *, generation_run_id: str
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
with self.connect() as conn:
|
||
|
|
task = self._task_in_connection(conn, task_id, for_update=True)
|
||
|
|
if (
|
||
|
|
task["status"] != "running"
|
||
|
|
or task.get("generation_run_id") != generation_run_id
|
||
|
|
):
|
||
|
|
return task
|
||
|
|
now = utcnow()
|
||
|
|
row = conn.execute(
|
||
|
|
"""
|
||
|
|
UPDATE data_process_tasks
|
||
|
|
SET status='failed', failure_reason=%s, completed_at=%s,
|
||
|
|
generation_run_id=NULL, updated_at=%s
|
||
|
|
WHERE id=%s AND generation_run_id=%s RETURNING *
|
||
|
|
""",
|
||
|
|
(reason[:4000], now, now, task_id, generation_run_id),
|
||
|
|
).fetchone()
|
||
|
|
return _decode_row(row) or {}
|
||
|
|
|
||
|
|
def progress(self, task_id: str) -> dict[str, Any]:
|
||
|
|
task = self.get_task(task_id)
|
||
|
|
return {
|
||
|
|
"task_id": task["id"],
|
||
|
|
"status": task["status"],
|
||
|
|
"progress": float(task.get("progress") or 0),
|
||
|
|
"input_count": int(task.get("input_count") or 0),
|
||
|
|
"output_count": int(task.get("output_count") or 0),
|
||
|
|
"filtered_count": int(task.get("filtered_count") or 0),
|
||
|
|
"duplicate_count": int(task.get("duplicate_count") or 0),
|
||
|
|
"error_count": int(task.get("error_count") or 0),
|
||
|
|
"failure_reason": task.get("failure_reason"),
|
||
|
|
"results_confirmed": bool(task.get("results_confirmed")),
|
||
|
|
"started_at": task.get("started_at"),
|
||
|
|
"completed_at": task.get("completed_at"),
|
||
|
|
}
|