diff --git a/backend/app/db/platform_store.py b/backend/app/db/platform_store.py index e0828f0..d640ae0 100644 --- a/backend/app/db/platform_store.py +++ b/backend/app/db/platform_store.py @@ -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)) diff --git a/backend/app/modules/data_process/store.py b/backend/app/modules/data_process/store.py index 8bcde57..6945597 100644 --- a/backend/app/modules/data_process/store.py +++ b/backend/app/modules/data_process/store.py @@ -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, } diff --git a/backend/tests/test_data_process_store.py b/backend/tests/test_data_process_store.py index 004504b..eec8e4d 100644 --- a/backend/tests/test_data_process_store.py +++ b/backend/tests/test_data_process_store.py @@ -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"