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

45
backend/app/db/platform_store.py
View File

@@ -1226,18 +1226,26 @@ class PlatformStore:
def training_dataset_files(self, dataset_id: str) -> list[dict[str, Any]]:
with self.connect() as conn:
dataset = conn.execute("SELECT id FROM datasets WHERE id=?", (dataset_id,)).fetchone()
dataset = conn.execute(
"SELECT id, metadata FROM datasets WHERE id=?", (dataset_id,)
).fetchone()
if not dataset:
raise KeyError(dataset_id)
dataset_metadata = json_loads(dataset.get("metadata"), {})
related_ids = dataset_metadata.get("split_dataset_ids") or {}
runtime_dataset_ids = [dataset_id]
validation_dataset_id = related_ids.get("validation")
if dataset_metadata.get("dataset_split") == "train" and validation_dataset_id:
runtime_dataset_ids.append(str(validation_dataset_id))
rows = conn.execute(
"""
SELECT id, dataset_id, name, size, content, active_version_id,
create_time, record_count, metadata
FROM dataset_files
WHERE dataset_id=?
ORDER BY create_time
WHERE dataset_id = ANY(%s)
ORDER BY CASE WHEN dataset_id=%s THEN 0 ELSE 1 END, create_time, id
""",
(dataset_id,),
(runtime_dataset_ids, dataset_id),
).fetchall()
return [
{
@@ -1350,6 +1358,18 @@ class PlatformStore:
raise ValueError("base_model or base_model_id is required")
if not train_dataset_id:
raise ValueError("train_dataset_id is required")
with self.connect() as conn:
train_dataset = conn.execute(
"SELECT id, type, metadata FROM datasets WHERE id=?",
(train_dataset_id,),
).fetchone()
if not train_dataset:
raise ValueError("training dataset not found")
train_metadata = json_loads(train_dataset.get("metadata"), {})
if train_dataset.get("type") != "train" or train_metadata.get(
"dataset_split"
) in {"validation", "test"}:
raise ValueError("train_dataset_id must reference a training dataset")
now = utcnow()
task = {
"id": task_id,
@@ -1535,7 +1555,24 @@ class PlatformStore:
FROM dataset_files WHERE dataset_id=? ORDER BY create_time, id""",
(dataset_id,),
).fetchall()
dataset_metadata = json_loads(dataset.get("metadata"), {}) if dataset else {}
related_ids = dataset_metadata.get("split_dataset_ids") or {}
validation_dataset_id = related_ids.get("validation")
if dataset_metadata.get("dataset_split") == "train" and validation_dataset_id:
files = [
*files,
*conn.execute(
"""SELECT id, name, size, active_version_id, create_time, metadata
FROM dataset_files WHERE dataset_id=? ORDER BY create_time, id""",
(str(validation_dataset_id),),
).fetchall(),
]
model_path = (model and model.get("path")) or task.get("model_name_or_path") or base_model_id
dataset_metadata = json_loads(dataset.get("metadata"), {}) if dataset else {}
if not dataset or dataset.get("type") != "train" or dataset_metadata.get(
"dataset_split"
) in {"validation", "test"}:
raise RuntimeError(f"training dataset is invalid: {dataset_id}")
if not files:
raise RuntimeError(f"dataset has no uploaded file: {dataset_id}")
dataset_key = str(task.get("dataset_key") or llama_dataset_key(dataset_id))

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,
}

183
backend/tests/test_data_process_store.py
View File

@@ -1,16 +1,112 @@
from __future__ import annotations
from contextlib import contextmanager
from decimal import Decimal
from typing import Any, Iterator
import pytest
from app.modules.data_process.store import (
DataProcessStore,
DataProcessStoreError,
_decode_row,
_source_storage_descriptor,
)
class _Result:
def __init__(self, *, row: dict[str, Any] | None = None, rows: list[dict[str, Any]] | None = None):
self.row = row
self.rows = rows or []
def fetchone(self) -> dict[str, Any] | None:
return self.row
def fetchall(self) -> list[dict[str, Any]]:
return self.rows
class _PublishConnection:
def __init__(self, results: list[dict[str, Any]]):
self.results = results
self.datasets: list[dict[str, Any]] = []
self.files: list[dict[str, Any]] = []
self.records: list[dict[str, Any]] = []
def execute(self, sql: str, params: Any = None) -> _Result:
normalized = " ".join(sql.split())
if normalized.startswith("SELECT * FROM data_process_results"):
return _Result(rows=self.results)
if normalized.startswith("SELECT * FROM datasets WHERE source_task_id"):
return _Result(rows=self.datasets)
if normalized.startswith("INSERT INTO datasets"):
dataset = {
"id": params[0],
"name": params[1],
"type": params[2],
"count": params[8],
"record_count": params[9],
"metadata": params[11],
}
self.datasets.append(dataset)
return _Result(row=dataset)
if normalized.startswith("UPDATE datasets SET name="):
dataset = next(item for item in self.datasets if item["id"] == params[10])
dataset.update(
{
"name": params[0],
"type": params[1],
"count": params[5],
"record_count": params[6],
"metadata": params[8],
}
)
return _Result(row=dataset)
if normalized.startswith("DELETE FROM dataset_records WHERE dataset_id"):
self.records = [item for item in self.records if item["dataset_id"] != params[0]]
if normalized.startswith("DELETE FROM dataset_files WHERE dataset_id"):
self.files = [item for item in self.files if item["dataset_id"] != params[0]]
if normalized.startswith("INSERT INTO dataset_files"):
self.files.append(
{
"id": params[0],
"dataset_id": params[1],
"name": params[2],
"record_count": params[11],
}
)
if normalized.startswith("INSERT INTO dataset_records"):
self.records.append(
{"dataset_id": params[1], "line_no": params[4], "split": params[5]}
)
return _Result()
class _PublishStore(DataProcessStore):
def __init__(self, conn: _PublishConnection):
self._conn = conn
@contextmanager
def connect(self) -> Iterator[_PublishConnection]:
yield self._conn
def _task_in_connection(self, conn: Any, task_id: str, *, for_update: bool = False) -> dict[str, Any]:
train_dataset = next(
(item for item in self._conn.datasets if item["type"] == "train"), None
)
return {
"id": task_id,
"status": "completed",
"description": "",
"config": {},
"output_dataset_id": train_dataset and train_dataset["id"],
}
@staticmethod
def _source_ids(conn: Any, task_id: str) -> list[dict[str, Any]]:
return []
def test_decode_row_serializes_postgres_numeric_values_as_json_numbers() -> None:
decoded = _decode_row(
{
@@ -22,6 +118,93 @@ def test_decode_row_serializes_postgres_numeric_values_as_json_numbers() -> None
assert decoded == {"progress": 100.0, "duration_seconds": 389.0}
def test_publish_creates_three_independent_datasets_with_exact_counts() -> None:
results = [
{
"id": f"result-{index}",
"status": "valid",
"instruction": f"问题 {index}",
"input": "",
"output": f"答案 {index}",
"preview_item_id": f"preview-{index}",
}
for index in range(28)
]
conn = _PublishConnection(results)
published = _PublishStore(conn).publish(
"task-1",
{
"dataset_name": "制度问答",
"storage_type": "local",
"format": "alpaca_jsonl",
"split": {"train": 80, "validation": 10, "test": 10},
},
)
assert [(item["name"], item["type"], item["count"]) for item in conn.datasets] == [
("制度问答-训练集", "train", 22),
("制度问答-验证集", "val", 3),
("制度问答-测试集", "test", 3),
]
assert len(conn.files) == 3
assert {item["dataset_id"] for item in conn.files} == {
item["id"] for item in conn.datasets
}
assert len(conn.records) == 28
assert published["dataset"]["type"] == "train"
assert len(published["datasets"]) == 3
assert published["split_counts"] == {"train": 22, "validation": 3, "test": 3}
original_ids = [item["id"] for item in conn.datasets]
republished = _PublishStore(conn).publish(
"task-1",
{
"dataset_name": "制度问答-训练集",
"storage_type": "local",
"format": "alpaca_jsonl",
"split": {"train": 80, "validation": 10, "test": 10},
},
)
assert [item["id"] for item in conn.datasets] == original_ids
assert len(conn.datasets) == 3
assert len(conn.files) == 3
assert len(conn.records) == 28
assert republished["created"] is False
def test_publish_keeps_all_three_datasets_when_a_small_split_is_empty() -> None:
conn = _PublishConnection(
[
{
"id": "result-only",
"status": "valid",
"instruction": "唯一问题",
"input": "",
"output": "唯一答案",
"preview_item_id": "preview-only",
}
]
)
published = _PublishStore(conn).publish(
"task-small",
{
"dataset_name": "小样本",
"storage_type": "local",
"format": "alpaca_jsonl",
"split": {"train": 80, "validation": 10, "test": 10},
},
)
assert [(item["type"], item["count"]) for item in conn.datasets] == [
("train", 1),
("val", 0),
("test", 0),
]
assert len(published["datasets"]) == 3
assert len(conn.files) == 3
def test_source_storage_descriptor_accepts_owned_local_and_legacy_db_references() -> None:
task_id = "dpt_task"
source_file_id = "dpsf_source"