fix(data-process): 发布精确三路数据切分
This commit is contained in:
caoxiaozhu
@@ -15,7 +15,7 @@ import psycopg
|
||||
from psycopg.rows import dict_row
|
||||
|
||||
from app.core.config import get_settings
|
||||
from app.modules.data_process.algorithms import estimate_token_count, stable_split
|
||||
from app.modules.data_process.algorithms import estimate_token_count, stable_split_assignments
|
||||
|
||||
TASK_STATUSES = {"pending", "running", "completed", "failed", "stopped"}
|
||||
EDITABLE_STATUSES = {"pending", "failed", "stopped", "completed"}
|
||||
@@ -1190,20 +1190,20 @@ class DataProcessStore:
|
||||
return _decode_row(row) or {}
|
||||
|
||||
def publish(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
"""发布有效结果；任务行锁保证重复请求返回同一数据集。"""
|
||||
"""按精确配额发布三个切分文件；重复发布会同步修复既有发布物。"""
|
||||
with self.connect() as conn:
|
||||
task = self._task_in_connection(conn, task_id, for_update=True)
|
||||
existing_dataset = None
|
||||
if task.get("output_dataset_id"):
|
||||
dataset = conn.execute(
|
||||
existing_dataset = conn.execute(
|
||||
"SELECT * FROM datasets WHERE id=%s", (task["output_dataset_id"],)
|
||||
).fetchone()
|
||||
if dataset:
|
||||
return {"dataset": _decode_row(dataset), "created": False}
|
||||
# 数据集被外部流程清理后，解除断链并重新发布。
|
||||
conn.execute(
|
||||
"UPDATE data_process_tasks SET output_dataset_id=NULL WHERE id=%s",
|
||||
(task_id,),
|
||||
)
|
||||
if not existing_dataset:
|
||||
# 数据集被外部流程清理后，解除断链并重新发布。
|
||||
conn.execute(
|
||||
"UPDATE data_process_tasks SET output_dataset_id=NULL WHERE id=%s",
|
||||
(task_id,),
|
||||
)
|
||||
if task["status"] != "completed":
|
||||
raise InvalidStateError("only a completed task can be published")
|
||||
rows = conn.execute(
|
||||
@@ -1225,183 +1225,250 @@ class DataProcessStore:
|
||||
if invalid_count:
|
||||
raise InvalidStateError(f"task contains {invalid_count} invalid results")
|
||||
|
||||
dataset_id = new_id("dataset")
|
||||
file_id = new_id("dfile")
|
||||
version_id = new_id("dfv")
|
||||
dataset_id = (
|
||||
str(existing_dataset["id"]) if existing_dataset else new_id("dataset")
|
||||
)
|
||||
now = utcnow()
|
||||
requested_split = payload.get("split") or {
|
||||
"train": 80,
|
||||
"validation": 10,
|
||||
"test": 10,
|
||||
}
|
||||
assignments = stable_split_assignments(
|
||||
[str(row["id"]) for row in rows],
|
||||
requested_split,
|
||||
seed=task_id,
|
||||
)
|
||||
records = [
|
||||
{
|
||||
"instruction": row["instruction"],
|
||||
"input": row["input"],
|
||||
"output": row["output"],
|
||||
"split": stable_split(
|
||||
str(row["id"]),
|
||||
requested_split,
|
||||
seed=task_id,
|
||||
),
|
||||
"split": assignment,
|
||||
}
|
||||
for row in rows
|
||||
for row, assignment in zip(rows, assignments, strict=True)
|
||||
]
|
||||
content = "".join(json_dumps(record) + "\n" for record in records)
|
||||
raw = content.encode("utf-8")
|
||||
checksum = hashlib.sha256(raw).hexdigest()
|
||||
storage_object_id = f"db://data-process/{task_id}/{file_id}/v1"
|
||||
split_order = ("train", "validation", "test")
|
||||
split_counts = {
|
||||
split_name: assignments.count(split_name) for split_name in split_order
|
||||
}
|
||||
file_specs: list[dict[str, Any]] = []
|
||||
for split_name in split_order:
|
||||
split_records = [
|
||||
(source_row, record)
|
||||
for source_row, record in zip(rows, records, strict=True)
|
||||
if record["split"] == split_name
|
||||
]
|
||||
if not split_records:
|
||||
continue
|
||||
file_id = new_id("dfile")
|
||||
version_id = new_id("dfv")
|
||||
content = "".join(
|
||||
json_dumps(record) + "\n" for _, record in split_records
|
||||
)
|
||||
raw = content.encode("utf-8")
|
||||
file_specs.append(
|
||||
{
|
||||
"split": split_name,
|
||||
"records": split_records,
|
||||
"file_id": file_id,
|
||||
"version_id": version_id,
|
||||
"content": content,
|
||||
"raw": raw,
|
||||
"checksum": hashlib.sha256(raw).hexdigest(),
|
||||
"storage_object_id": (
|
||||
f"db://data-process/{task_id}/{file_id}/v1"
|
||||
),
|
||||
}
|
||||
)
|
||||
total_size = sum(len(spec["raw"]) for spec in file_specs)
|
||||
source_result_ids = [row["id"] for row in rows]
|
||||
metadata = {
|
||||
"source": "data_process",
|
||||
"storage_backend": "database",
|
||||
"storage_object_id": storage_object_id,
|
||||
"source_task_id": task_id,
|
||||
"source_file_ids": [item["id"] for item in self._source_ids(conn, task_id)],
|
||||
"source_result_ids": source_result_ids,
|
||||
"format": payload.get("format") or "alpaca_jsonl",
|
||||
"split": payload.get("split") or {},
|
||||
"split": requested_split,
|
||||
"split_counts": split_counts,
|
||||
}
|
||||
try:
|
||||
dataset = conn.execute(
|
||||
"""
|
||||
INSERT INTO datasets
|
||||
(id, name, type, storage_type, source, task_id, source_task_id,
|
||||
size, size_bytes, count, record_count, description, metadata,
|
||||
tenant_id, project_id, owner_id, created_by, create_time,
|
||||
created_at, updated_at)
|
||||
VALUES (%s, %s, %s, %s, 'task', %s, %s, %s, %s, %s, %s, %s, %s,
|
||||
%s, %s, %s, %s, %s, %s, %s)
|
||||
RETURNING *
|
||||
""",
|
||||
(
|
||||
dataset_id,
|
||||
payload["dataset_name"],
|
||||
payload.get("dataset_type") or "train",
|
||||
payload.get("storage_type") or "local",
|
||||
task_id,
|
||||
task_id,
|
||||
f"{len(raw)} B",
|
||||
len(raw),
|
||||
len(records),
|
||||
len(records),
|
||||
payload.get("description") or task.get("description") or "",
|
||||
json_dumps(metadata),
|
||||
task.get("tenant_id"),
|
||||
task.get("project_id"),
|
||||
task.get("owner_id"),
|
||||
payload.get("created_by") or task.get("created_by"),
|
||||
now,
|
||||
now,
|
||||
now,
|
||||
),
|
||||
).fetchone()
|
||||
version = {
|
||||
"id": version_id,
|
||||
"version_no": 1,
|
||||
"description": "data process publish",
|
||||
"checksum_sha256": checksum,
|
||||
"size_bytes": len(raw),
|
||||
"record_count": len(records),
|
||||
"created_at": now,
|
||||
"source_task_id": task_id,
|
||||
"storage_object_id": storage_object_id,
|
||||
}
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO dataset_files
|
||||
(id, dataset_id, name, storage_object_id, size, content,
|
||||
active_version_id, versions, create_time, current_version_id,
|
||||
size_bytes, record_count, file_format, checksum_sha256, version_no,
|
||||
source_task_id, tenant_id, project_id, created_by, metadata,
|
||||
created_at, updated_at)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
|
||||
%s, %s, 1, %s, %s, %s, %s, %s, %s, %s)
|
||||
""",
|
||||
(
|
||||
file_id,
|
||||
dataset_id,
|
||||
f"{payload['dataset_name']}.jsonl",
|
||||
storage_object_id,
|
||||
f"{len(raw)} B",
|
||||
content,
|
||||
version_id,
|
||||
json_dumps([version]),
|
||||
now,
|
||||
version_id,
|
||||
len(raw),
|
||||
len(records),
|
||||
"jsonl",
|
||||
checksum,
|
||||
task_id,
|
||||
task.get("tenant_id"),
|
||||
task.get("project_id"),
|
||||
payload.get("created_by") or task.get("created_by"),
|
||||
json_dumps(metadata),
|
||||
now,
|
||||
now,
|
||||
),
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO dataset_file_versions
|
||||
(id, dataset_file_id, version_no, storage_object_id, content_preview,
|
||||
description, size_bytes, record_count, checksum_sha256,
|
||||
source_task_id, metadata, created_by, created_at)
|
||||
VALUES (%s, %s, 1, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||||
""",
|
||||
(
|
||||
version_id,
|
||||
file_id,
|
||||
storage_object_id,
|
||||
content[:2000],
|
||||
"data process publish",
|
||||
len(raw),
|
||||
len(records),
|
||||
checksum,
|
||||
task_id,
|
||||
json_dumps(metadata),
|
||||
payload.get("created_by") or task.get("created_by"),
|
||||
now,
|
||||
),
|
||||
)
|
||||
for line_number, (source_row, record) in enumerate(
|
||||
zip(rows, records, strict=True), start=1
|
||||
):
|
||||
if existing_dataset:
|
||||
# 用户显式重新发布时，在同一数据集 ID 下修复旧的混合文件。
|
||||
conn.execute(
|
||||
"DELETE FROM dataset_records WHERE dataset_id=%s", (dataset_id,)
|
||||
)
|
||||
conn.execute(
|
||||
"""DELETE FROM dataset_file_versions
|
||||
WHERE dataset_file_id IN
|
||||
(SELECT id FROM dataset_files WHERE dataset_id=%s)""",
|
||||
(dataset_id,),
|
||||
)
|
||||
conn.execute("DELETE FROM dataset_files WHERE dataset_id=%s", (dataset_id,))
|
||||
dataset = conn.execute(
|
||||
"""
|
||||
INSERT INTO dataset_records
|
||||
(id, dataset_id, dataset_file_id, version_id, line_no, split,
|
||||
instruction, input, output, raw, status, source_task_id,
|
||||
source_result_id, preview_item_id, created_at)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
|
||||
%s, %s, %s, %s)
|
||||
UPDATE datasets
|
||||
SET size=%s, size_bytes=%s, count=%s, record_count=%s,
|
||||
description=%s, metadata=%s, updated_at=%s
|
||||
WHERE id=%s RETURNING *
|
||||
""",
|
||||
(
|
||||
new_id("drec"),
|
||||
f"{total_size} B",
|
||||
total_size,
|
||||
len(records),
|
||||
len(records),
|
||||
payload.get("description") or task.get("description") or "",
|
||||
json_dumps(metadata),
|
||||
now,
|
||||
dataset_id,
|
||||
file_id,
|
||||
version_id,
|
||||
line_number,
|
||||
record["split"],
|
||||
record["instruction"],
|
||||
record["input"],
|
||||
record["output"],
|
||||
json_dumps(
|
||||
{
|
||||
**record,
|
||||
"source_task_id": task_id,
|
||||
"source_result_id": source_row["id"],
|
||||
"preview_item_id": source_row.get("preview_item_id"),
|
||||
}
|
||||
),
|
||||
source_row["status"],
|
||||
),
|
||||
).fetchone()
|
||||
else:
|
||||
dataset = conn.execute(
|
||||
"""
|
||||
INSERT INTO datasets
|
||||
(id, name, type, storage_type, source, task_id, source_task_id,
|
||||
size, size_bytes, count, record_count, description, metadata,
|
||||
tenant_id, project_id, owner_id, created_by, create_time,
|
||||
created_at, updated_at)
|
||||
VALUES (%s, %s, %s, %s, 'task', %s, %s, %s, %s, %s, %s, %s, %s,
|
||||
%s, %s, %s, %s, %s, %s, %s)
|
||||
RETURNING *
|
||||
""",
|
||||
(
|
||||
dataset_id,
|
||||
payload["dataset_name"],
|
||||
payload.get("dataset_type") or "train",
|
||||
payload.get("storage_type") or "local",
|
||||
task_id,
|
||||
source_row["id"],
|
||||
source_row.get("preview_item_id"),
|
||||
task_id,
|
||||
f"{total_size} B",
|
||||
total_size,
|
||||
len(records),
|
||||
len(records),
|
||||
payload.get("description") or task.get("description") or "",
|
||||
json_dumps(metadata),
|
||||
task.get("tenant_id"),
|
||||
task.get("project_id"),
|
||||
task.get("owner_id"),
|
||||
payload.get("created_by") or task.get("created_by"),
|
||||
now,
|
||||
now,
|
||||
now,
|
||||
),
|
||||
).fetchone()
|
||||
|
||||
for spec in file_specs:
|
||||
file_metadata = {**metadata, "file_split": spec["split"]}
|
||||
version = {
|
||||
"id": spec["version_id"],
|
||||
"version_no": 1,
|
||||
"version": 1,
|
||||
"description": f"data process {spec['split']} publish",
|
||||
"checksum_sha256": spec["checksum"],
|
||||
"size_bytes": len(spec["raw"]),
|
||||
"record_count": len(spec["records"]),
|
||||
"created_at": now,
|
||||
"create_time": now,
|
||||
"source_task_id": task_id,
|
||||
"storage_object_id": spec["storage_object_id"],
|
||||
}
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO dataset_files
|
||||
(id, dataset_id, name, storage_object_id, size, content,
|
||||
active_version_id, versions, create_time, current_version_id,
|
||||
size_bytes, record_count, file_format, checksum_sha256, version_no,
|
||||
source_task_id, tenant_id, project_id, created_by, metadata,
|
||||
created_at, updated_at)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
|
||||
%s, %s, 1, %s, %s, %s, %s, %s, %s, %s)
|
||||
""",
|
||||
(
|
||||
spec["file_id"],
|
||||
dataset_id,
|
||||
f"{payload['dataset_name']}.{spec['split']}.jsonl",
|
||||
spec["storage_object_id"],
|
||||
f"{len(spec['raw'])} B",
|
||||
spec["content"],
|
||||
spec["version_id"],
|
||||
json_dumps([version]),
|
||||
now,
|
||||
spec["version_id"],
|
||||
len(spec["raw"]),
|
||||
len(spec["records"]),
|
||||
"jsonl",
|
||||
spec["checksum"],
|
||||
task_id,
|
||||
task.get("tenant_id"),
|
||||
task.get("project_id"),
|
||||
payload.get("created_by") or task.get("created_by"),
|
||||
json_dumps(file_metadata),
|
||||
now,
|
||||
now,
|
||||
),
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO dataset_file_versions
|
||||
(id, dataset_file_id, version_no, storage_object_id, content_preview,
|
||||
description, size_bytes, record_count, checksum_sha256,
|
||||
source_task_id, metadata, created_by, created_at)
|
||||
VALUES (%s, %s, 1, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
||||
""",
|
||||
(
|
||||
spec["version_id"],
|
||||
spec["file_id"],
|
||||
spec["storage_object_id"],
|
||||
spec["content"][:2000],
|
||||
f"data process {spec['split']} publish",
|
||||
len(spec["raw"]),
|
||||
len(spec["records"]),
|
||||
spec["checksum"],
|
||||
task_id,
|
||||
json_dumps(file_metadata),
|
||||
payload.get("created_by") or task.get("created_by"),
|
||||
now,
|
||||
),
|
||||
)
|
||||
for line_number, (source_row, record) in enumerate(
|
||||
spec["records"], start=1
|
||||
):
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO dataset_records
|
||||
(id, dataset_id, dataset_file_id, version_id, line_no, split,
|
||||
instruction, input, output, raw, status, source_task_id,
|
||||
source_result_id, preview_item_id, created_at)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
|
||||
%s, %s, %s, %s)
|
||||
""",
|
||||
(
|
||||
new_id("drec"),
|
||||
dataset_id,
|
||||
spec["file_id"],
|
||||
spec["version_id"],
|
||||
line_number,
|
||||
record["split"],
|
||||
record["instruction"],
|
||||
record["input"],
|
||||
record["output"],
|
||||
json_dumps(
|
||||
{
|
||||
**record,
|
||||
"source_task_id": task_id,
|
||||
"source_result_id": source_row["id"],
|
||||
"preview_item_id": source_row.get("preview_item_id"),
|
||||
}
|
||||
),
|
||||
source_row["status"],
|
||||
task_id,
|
||||
source_row["id"],
|
||||
source_row.get("preview_item_id"),
|
||||
now,
|
||||
),
|
||||
)
|
||||
except psycopg.errors.UniqueViolation as exc:
|
||||
raise ConflictError("dataset name already exists") from exc
|
||||
conn.execute(
|
||||
@@ -1412,7 +1479,11 @@ class DataProcessStore:
|
||||
""",
|
||||
(dataset_id, now, payload.get("created_by"), task_id),
|
||||
)
|
||||
return {"dataset": _decode_row(dataset), "created": True}
|
||||
return {
|
||||
"dataset": _decode_row(dataset),
|
||||
"created": existing_dataset is None,
|
||||
"split_counts": split_counts,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _source_ids(
|
||||
|
||||
Reference in New Issue
Block a user