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

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=?
+                WHERE dataset_id = ANY(%s)
-                ORDER BY create_time
+                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))

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:
+            existing_datasets = conn.execute(
-                    file_metadata = {**metadata, "file_split": spec["split"]}
+                """
+                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),
+            "dataset": train_dataset,
-            "created": existing_dataset is None,
+            "datasets": published_datasets,
+            "output_datasets": published_datasets,
+            "created": created_any,
             "split_counts": split_counts,
         }
 

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"