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-24 20:43:47 +08:00
parent e6a5a36bc0
commit 9cb77c251a
8 changed files with 405 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
backend/app/api/v1/endpoints/platform.py
View File

@@ -534,6 +534,12 @@ async def _sync_training_dataset_to_compute_node(
files = store.training_dataset_files(dataset_id)
if not files:
raise RuntimeError(f"dataset has no uploaded file: {dataset_id}")
split_aware = any(item.get("split") for item in files)
files = [
item
for item in files
if not split_aware or item.get("split") in {"train", "validation"}
]
client = ComputeNodeClient(node["api_base_url"])
results: list[dict[str, Any]] = []
for item in files:
@@ -1533,4 +1539,3 @@ async def log_content(file: str = Query(...)) -> dict[str, Any]:
@router.post("/web-log")
async def web_log(payload: dict[str, Any] = Body(...)) -> dict[str, Any]:
return ok({"received": True, **payload})

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

@@ -1099,11 +1099,28 @@ class PlatformStore:
def _dataset(self, conn: PgConnection, row: PgRow) -> dict[str, Any]:
files = conn.execute(
"SELECT id, name, size, active_version_id, create_time FROM dataset_files WHERE dataset_id=? ORDER BY create_time",
"""SELECT id, name, size, active_version_id, create_time,
record_count, metadata
FROM dataset_files WHERE dataset_id=? ORDER BY create_time, id""",
(row["id"],),
).fetchall()
dataset_metadata = json_loads(row.get("metadata"), {})
split_counts = dict(dataset_metadata.get("split_counts") or {})
if row.get("source") == "task" and not split_counts:
split_rows = conn.execute(
"""SELECT split, COUNT(*) AS count FROM dataset_records
WHERE dataset_id=? GROUP BY split""",
(row["id"],),
).fetchall()
split_counts = {str(item["split"]): int(item["count"]) for item in split_rows}
return {
**dict(row),
"metadata": dataset_metadata,
"split_counts": {
"train": int(split_counts.get("train", 0) or 0),
"validation": int(split_counts.get("validation", 0) or 0),
"test": int(split_counts.get("test", 0) or 0),
},
"files": [
{
"id": f["id"],
@@ -1111,6 +1128,8 @@ class PlatformStore:
"size": f["size"],
"active_version_id": f["active_version_id"],
"create_time": f["create_time"],
"record_count": int(f.get("record_count") or 0),
"split": json_loads(f.get("metadata"), {}).get("file_split"),
}
for f in files
],
@@ -1212,14 +1231,22 @@ class PlatformStore:
raise KeyError(dataset_id)
rows = conn.execute(
"""
SELECT id, dataset_id, name, size, content, active_version_id, create_time
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
""",
(dataset_id,),
).fetchall()
return [dict(row) for row in rows]
return [
{
**dict(row),
"metadata": json_loads(row.get("metadata"), {}),
"split": json_loads(row.get("metadata"), {}).get("file_split"),
}
for row in rows
]
def file_versions(self, file_id: str) -> dict[str, Any]:
row = self.dataset_file(file_id)
@@ -1504,15 +1531,36 @@ class PlatformStore:
model = conn.execute("SELECT * FROM models WHERE id=?", (base_model_id,)).fetchone()
dataset = conn.execute("SELECT * FROM datasets WHERE id=?", (dataset_id,)).fetchone()
files = conn.execute(
"SELECT id, name, size, active_version_id, create_time FROM dataset_files WHERE dataset_id=? ORDER BY create_time",
"""SELECT id, name, size, active_version_id, create_time, metadata
FROM dataset_files WHERE dataset_id=? ORDER BY create_time, id""",
(dataset_id,),
).fetchall()
model_path = (model and model.get("path")) or task.get("model_name_or_path") or base_model_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))
dataset_file_names = [Path(str(row["name"] or row["id"])).name for row in files]
dataset_keys = llama_dataset_keys(dataset_key, dataset_file_names)
file_entries = [
{
**dict(row),
"name": Path(str(row["name"] or row["id"])).name,
"split": json_loads(row.get("metadata"), {}).get("file_split"),
}
for row in files
]
split_aware = any(item["split"] for item in file_entries)
training_files = [
item for item in file_entries if not split_aware or item["split"] == "train"
]
validation_files = [
item for item in file_entries if split_aware and item["split"] == "validation"
]
if not training_files:
raise RuntimeError(f"dataset has no training split: {dataset_id}")
runtime_files = [*training_files, *validation_files]
runtime_file_names = [str(item["name"]) for item in runtime_files]
runtime_keys = llama_dataset_keys(dataset_key, runtime_file_names)
training_keys = runtime_keys[: len(training_files)]
validation_keys = runtime_keys[len(training_files) :]
dataset_format = str(task.get("dataset_format") or (dataset and dataset.get("formatting")) or "alpaca").lower()
health_detail = node.get("health_detail") or {}
dataset_root = str(health_detail.get("dataset_root") or f"{node['data_root'].rstrip('/')}/datasets")
@@ -1526,23 +1574,26 @@ class PlatformStore:
"name": task["name"],
"base_model": model_path,
"model_name_or_path": model_path,
"dataset": ",".join(dataset_keys),
"dataset": ",".join(training_keys),
"dataset_key": dataset_key,
"dataset_keys": dataset_keys,
"dataset_keys": training_keys,
"eval_dataset": ",".join(validation_keys) or None,
"eval_dataset_keys": validation_keys,
"dataset_display_name": (dataset and dataset.get("name")) or dataset_id,
"dataset_dir": dataset_dir,
"dataset_info": llama_dataset_info(dataset_key, dataset_file_names, dataset_format),
"dataset_info": llama_dataset_info(dataset_key, runtime_file_names, dataset_format),
"dataset_files": [
{
"id": row["id"],
"name": Path(str(row["name"] or row["id"])).name,
"relative_path": f"{dataset_id}/{Path(str(row['name'] or row['id'])).name}",
"local_path": f"{dataset_dir.rstrip('/')}/{Path(str(row['name'] or row['id'])).name}",
"active_version_id": row["active_version_id"],
"size": row["size"],
"create_time": row["create_time"],
"id": item["id"],
"name": item["name"],
"relative_path": f"{dataset_id}/{item['name']}",
"local_path": f"{dataset_dir.rstrip('/')}/{item['name']}",
"active_version_id": item["active_version_id"],
"size": item["size"],
"create_time": item["create_time"],
"split": item["split"],
}
for row in files
for item in runtime_files
],
"output_dir": output_dir,
"gpus": selected_gpus or task.get("gpus") or [0],
@@ -2536,4 +2587,3 @@ def get_platform_store() -> PlatformStore:
if _store is None:
_store = PlatformStore()
return _store

55
backend/app/modules/data_process/algorithms.py
View File

@@ -2871,6 +2871,51 @@ def stable_split(
return "test"
def stable_split_assignments(
values: Sequence[str | int],
split: Mapping[str, int] | None = None,
*,
seed: str = "",
) -> list[DatasetSplit]:
"""按稳定顺序和精确配额批量划分数据集。
单条哈希分桶只能在大样本下近似比例。这里先按哈希稳定排序，再用
最大余数法计算各切分配额，确保小数据集也严格遵循配置比例。
"""
ratios = dict(split or {"train": 80, "validation": 10, "test": 10})
# 复用单条划分的参数校验，避免两套规则逐渐漂移。
stable_split("validation", ratios, seed=seed)
if not values:
return []
split_order: tuple[DatasetSplit, ...] = ("train", "validation", "test")
exact = {name: len(values) * ratios[name] / 100 for name in split_order}
quotas = {name: math.floor(exact[name]) for name in split_order}
remaining = len(values) - sum(quotas.values())
remainder_order = sorted(
split_order,
key=lambda name: (-(exact[name] - quotas[name]), split_order.index(name)),
)
for name in remainder_order[:remaining]:
quotas[name] += 1
ranked_indices = sorted(
range(len(values)),
key=lambda index: (
hashlib.sha256(f"{seed}:{values[index]}".encode("utf-8")).digest(),
index,
),
)
assignments: list[DatasetSplit] = ["train"] * len(values)
cursor = 0
for name in split_order:
for index in ranked_indices[cursor : cursor + quotas[name]]:
assignments[index] = name
cursor += quotas[name]
return assignments
def _preview_content(item: Mapping[str, Any]) -> str:
for field in ("edited_content", "editedContent", "original_content", "originalContent", "content"):
value = item.get(field)
@@ -2970,9 +3015,16 @@ def generate_standard_records(
"original_input": input_text,
"original_output": output,
"status": status,
"split": stable_split(result_id, split, seed=split_seed),
"split": "train",
}
)
assignments = stable_split_assignments(
[str(result["id"]) for result in results],
split,
seed=split_seed,
)
for result, assignment in zip(results, assignments, strict=True):
result["split"] = assignment
return results
@@ -3021,4 +3073,5 @@ __all__ = [
"remove_document_noise",
"score_quality",
"stable_split",
"stable_split_assignments",
]

13
backend/app/modules/data_process/generation.py
View File

@@ -11,7 +11,7 @@ from urllib.parse import urlsplit, urlunsplit
import httpx
from app.modules.data_process.algorithms import normalize_text, stable_split
from app.modules.data_process.algorithms import normalize_text, stable_split_assignments
class ModelGenerationError(ValueError):
@@ -203,7 +203,7 @@ def generate_model_records(
"original_output": "",
"status": "invalid",
"error": error_message,
"split": stable_split(result_id, split, seed=task_id),
"split": "train",
}
)
if on_progress:
@@ -236,7 +236,7 @@ def generate_model_records(
"original_output": output,
"status": "valid" if valid else "invalid",
"error": None if valid else "model result is missing instruction or output",
"split": stable_split(result_id, split, seed=task_id),
"split": "train",
}
)
if on_progress:
@@ -244,6 +244,13 @@ def generate_model_records(
finally:
if owns_client:
http_client.close()
assignments = stable_split_assignments(
[str(result["id"]) for result in results],
split,
seed=task_id,
)
for result, assignment in zip(results, assignments, strict=True):
result["split"] = assignment
return results

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

@@ -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(