YG-Soft/YG_FT
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(data-process): 发布三个独立切分数据集

Browse Source
This commit is contained in:
caoxiaozhu
2026-07-25 17:04:14 +08:00
parent d4b9a76aa5
commit e9a121cfeb
3 changed files with 397 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files

268
backend/app/modules/data_process/store.py
View File

@@ -244,6 +244,17 @@ class DataProcessStore:
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', dataset.metadata::jsonb->>'dataset_split'
) 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_id=task.id AND dataset.source='task')
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))
@@ -1190,20 +1201,9 @@ 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"):
existing_dataset = conn.execute(
"SELECT * FROM datasets WHERE id=%s", (task["output_dataset_id"],)
).fetchone()
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,9 +1225,6 @@ class DataProcessStore:
if invalid_count:
raise InvalidStateError(f"task contains {invalid_count} invalid results")
dataset_id = (
str(existing_dataset["id"]) if existing_dataset else new_id("dataset")
)
now = utcnow()
requested_split = payload.get("split") or {
"train": 80,
@@ -1252,22 +1249,20 @@ class DataProcessStore:
split_counts = {
split_name: assignments.count(split_name) for split_name in split_order
}
file_specs: list[dict[str, Any]] = []
split_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_specs.append(
{
"split": split_name,
"records": split_records,
@@ -1281,9 +1276,8 @@ class DataProcessStore:
),
}
)
total_size = sum(len(spec["raw"]) for spec in file_specs)
source_result_ids = [row["id"] for row in rows]
metadata = {
common_metadata = {
"source": "data_process",
"storage_backend": "database",
"source_task_id": task_id,
@@ -1291,81 +1285,153 @@ class DataProcessStore:
"source_result_ids": source_result_ids,
"format": payload.get("format") or "alpaca_jsonl",
"split": requested_split,
"split_counts": split_counts,
}
try:
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(
"""
UPDATE datasets
SET size=%s, size_bytes=%s, count=%s, record_count=%s,
description=%s, metadata=%s, updated_at=%s
WHERE id=%s RETURNING *
""",
(
f"{total_size} B",
total_size,
len(records),
len(records),
payload.get("description") or task.get("description") or "",
json_dumps(metadata),
now,
dataset_id,
),
).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,
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"]}
existing_datasets = conn.execute(
"""
SELECT * FROM datasets
WHERE source_task_id=%s AND source='task'
ORDER BY created_at, id
""",
(task_id,),
).fetchall()
existing_by_split: dict[str, dict[str, Any]] = {}
primary_existing = None
for existing in existing_datasets:
existing_metadata = _json_value(existing.get("metadata"), {})
existing_split = str(existing_metadata.get("dataset_split") or "")
if existing_split in split_order:
existing_by_split[existing_split] = existing
if str(existing["id"]) == str(task.get("output_dataset_id") or ""):
primary_existing = existing
if primary_existing and "train" not in existing_by_split:
# 兼容旧版“一个数据集包含三个文件”的发布物，原数据集复用为训练集。
existing_by_split["train"] = primary_existing
existing_group_metadata = _json_value(
(primary_existing or {}).get("metadata"), {}
)
base_dataset_name = str(
existing_group_metadata.get("base_dataset_name")
or payload["dataset_name"]
).strip()
for suffix in ("-训练集", "-验证集", "-测试集"):
if base_dataset_name.endswith(suffix):
base_dataset_name = base_dataset_name[: -len(suffix)].rstrip()
break
split_group_id = str(
existing_group_metadata.get("split_group_id")
or f"dsg_{hashlib.sha256(task_id.encode()).hexdigest()[:20]}"
)
dataset_ids = {
spec["split"]: str(existing_by_split[spec["split"]]["id"])
if spec["split"] in existing_by_split
else new_id("dataset")
for spec in split_specs
}
created_any = any(
spec["split"] not in existing_by_split for spec in split_specs
)
split_labels = {
"train": "训练集",
"validation": "验证集",
"test": "测试集",
}
dataset_types = {"train": "train", "validation": "val", "test": "test"}
published_datasets: list[dict[str, Any]] = []
try:
for spec in split_specs:
split_name = str(spec["split"])
dataset_id = dataset_ids[split_name]
existing_dataset = existing_by_split.get(split_name)
dataset_metadata = {
**common_metadata,
"base_dataset_name": base_dataset_name,
"dataset_split": split_name,
"split_group_id": split_group_id,
"split_dataset_ids": dataset_ids,
"split_counts": {
name: split_counts[name] if name == split_name else 0
for name in split_order
},
}
dataset_name = f"{base_dataset_name}-{split_labels[split_name]}"
if existing_dataset:
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(
"""
UPDATE datasets
SET name=%s, type=%s, storage_type=%s, size=%s, size_bytes=%s,
count=%s, record_count=%s, description=%s, metadata=%s,
updated_at=%s
WHERE id=%s RETURNING *
""",
(
dataset_name,
dataset_types[split_name],
payload.get("storage_type") or "local",
f"{len(spec['raw'])} B",
len(spec["raw"]),
len(spec["records"]),
len(spec["records"]),
payload.get("description") or task.get("description") or "",
json_dumps(dataset_metadata),
now,
dataset_id,
),
).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, %s)
RETURNING *
""",
(
dataset_id,
dataset_name,
dataset_types[split_name],
payload.get("storage_type") or "local",
task_id,
task_id,
f"{len(spec['raw'])} B",
len(spec["raw"]),
len(spec["records"]),
len(spec["records"]),
payload.get("description") or task.get("description") or "",
json_dumps(dataset_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()
file_metadata = {**dataset_metadata, "file_split": split_name}
version = {
"id": spec["version_id"],
"version_no": 1,
"version": 1,
"description": f"data process {spec['split']} publish",
"description": f"data process {split_name} publish",
"checksum_sha256": spec["checksum"],
"size_bytes": len(spec["raw"]),
"record_count": len(spec["records"]),
@@ -1388,7 +1454,7 @@ class DataProcessStore:
(
spec["file_id"],
dataset_id,
f"{payload['dataset_name']}.{spec['split']}.jsonl",
f"{base_dataset_name}.{split_name}.jsonl",
spec["storage_object_id"],
f"{len(spec['raw'])} B",
spec["content"],
@@ -1422,7 +1488,7 @@ class DataProcessStore:
spec["file_id"],
spec["storage_object_id"],
spec["content"][:2000],
f"data process {spec['split']} publish",
f"data process {split_name} publish",
len(spec["raw"]),
len(spec["records"]),
spec["checksum"],
@@ -1469,19 +1535,31 @@ class DataProcessStore:
now,
),
)
published_datasets.append(_decode_row(dataset) or {})
except psycopg.errors.UniqueViolation as exc:
raise ConflictError("dataset name already exists") from exc
train_dataset_id = dataset_ids.get("train")
if not train_dataset_id:
raise InvalidStateError("published split does not contain training data")
conn.execute(
"""
UPDATE data_process_tasks
SET output_dataset_id=%s, updated_at=%s, updated_by=%s
WHERE id=%s
""",
(dataset_id, now, payload.get("created_by"), task_id),
(train_dataset_id, now, payload.get("created_by"), task_id),
)
train_dataset = next(
item
for item in published_datasets
if _json_value(item.get("metadata"), {}).get("dataset_split") == "train"
)
return {
"dataset": _decode_row(dataset),
"created": existing_dataset is None,
"dataset": train_dataset,
"datasets": published_datasets,
"output_datasets": published_datasets,
"created": created_any,
"split_counts": split_counts,
}