"""数据处理存储层 - 数据集发布。""" from __future__ import annotations from typing import Any from collections.abc import Sequence import hashlib import psycopg from app.core.config import get_settings from app.modules.storage.minio_store import get_object_storage from app.modules.storage.policy import should_store_in_minio from .base import ( StoreBase, utcnow, new_id, repeat_task_id, json_dumps, _json_value, _decode_row, _public_task, _business_config, _preview_config_value, _preview_config_changed, _preview_config_projection, _normalized_preprocess_options, _regeneration_marker, _is_regeneration_prepared, _task_output_type, _task_reasoning_detail, _reasoning_output_is_valid, _dpo_fields_are_valid, _source_storage_descriptor, NotFoundError, ConflictError, InvalidStateError, EDITABLE_STATUSES, ACTIVE_PREVIEW_STATUSES, WORKFLOW_STEPS, _REGENERATION_MARKER_KEY, _REPEAT_SOURCE_TASK_KEY, _REPEAT_REQUEST_KEY, ) from ..algorithms import stable_split_assignments class DatasetsMixin: """数据集发布 Mixin。""" def get_generation_model(self, model_id: str) -> dict[str, Any]: with self.connect() as conn: row = conn.execute( """ SELECT id, name, type, purpose, model_source, description, path, api_url, api_key, online_model_name, create_time FROM models WHERE id=%s """, (model_id,), ).fetchone() if not row: raise NotFoundError("generation model not found") return _decode_row(row) or {} def save_generation_model_snapshot( self, task_id: str, model_snapshot: dict[str, Any], *, generation_run_id: str, ) -> dict[str, Any]: # API 密钥仅用于本次调用，绝不能进入任务配置、详情响应或审计快照。 safe_snapshot = { key: value for key, value in model_snapshot.items() if key != "api_key" } with self.connect() as conn: task = self._task_in_connection(conn, task_id, for_update=True) if ( task["status"] != "running" or task.get("generation_run_id") != generation_run_id ): raise InvalidStateError("generation run is no longer active") config = dict(task.get("config") or {}) config["generation_model_snapshot"] = safe_snapshot row = conn.execute( """ UPDATE data_process_tasks SET config=%s, updated_at=%s WHERE id=%s AND generation_run_id=%s RETURNING * """, (json_dumps(config), utcnow(), task_id, generation_run_id), ).fetchone() 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) if _is_regeneration_prepared(task): raise InvalidStateError( "regeneration must start and complete before publishing" ) if task["status"] != "completed": raise InvalidStateError("only a completed task can be published") if not task.get("results_confirmed"): raise InvalidStateError("results must be confirmed before publishing") rows = conn.execute( """ SELECT * FROM data_process_results WHERE task_id=%s ORDER BY created_at, id """, (task_id,), ).fetchall() if not rows: raise InvalidStateError("task has no results to publish") invalid_count = sum( 1 for row in rows if row["status"] == "invalid" or not str(row.get("instruction") or "").strip() or not str(row.get("output") or "").strip() or ( _task_output_type(task) == "reasoning" and not _reasoning_output_is_valid(row.get("output")) ) or ( _task_output_type(task) == "dpo" and not _dpo_fields_are_valid(row) ) ) if invalid_count: raise InvalidStateError(f"task contains {invalid_count} invalid results") 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, ) if _task_output_type(task) == "dpo": records = [ { "instruction": row["instruction"], "input": row["input"], "chosen": row["chosen"], "rejected": row["rejected"], "split": assignment, } for row, assignment in zip(rows, assignments, strict=True) ] else: records = [ { "instruction": row["instruction"], "input": row["input"], "output": row["output"], "split": assignment, } for row, assignment in zip(rows, assignments, strict=True) ] split_order = ("train", "validation", "test") split_counts = { split_name: assignments.count(split_name) for split_name in split_order } split_specs: list[dict[str, Any]] = [] # Explicit local is retained for old callers/tests that request the # legacy backend; all normal platform requests default to MinIO. allow_minio = bool(get_settings().minio_enabled) and str( payload.get("storage_type") or "minio" ).lower() != "local" 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 ] 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") split_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" ), "store_in_minio": allow_minio and should_store_in_minio( len(raw), content_type="application/jsonl", file_format="jsonl" ), } ) source_result_ids = [row["id"] for row in rows] common_metadata = { "source": "data_process", "storage_backend": "minio" if any(spec["store_in_minio"] for spec in split_specs) else "database", "source_task_id": task_id, "output_type": _task_output_type(task), "reasoning_detail": _task_reasoning_detail(task), "source_file_ids": [item["id"] for item in self._source_ids(conn, task_id)], "source_result_ids": source_result_ids, "format": ( "dpo" if _task_output_type(task) == "dpo" else payload.get("format") or "alpaca_jsonl" ), "split": requested_split, } existing_datasets = conn.execute( """ SELECT * FROM datasets WHERE source='task' AND deleted_at IS NULL AND ( source_task_id=%s OR (source_task_id IS NULL AND task_id=%s) ) ORDER BY created_at, id """, (task_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] storage_object_id = str(spec["storage_object_id"]) if spec["store_in_minio"]: file_name = f"{base_dataset_name}.{split_name}.jsonl" object_key = f"datasets/{dataset_id}/versions/{spec['version_id']}/{file_name}" uploaded = get_object_storage().put_bytes( object_key, spec["raw"], "application/jsonl", ) conn.execute( """ INSERT INTO storage_objects (id, resource_type, resource_id, version_id, bucket, object_key, file_name, content_type, checksum_sha256, byte_size, status, created_by, create_time) VALUES (%s, 'dataset', %s, %s, %s, %s, %s, %s, %s, %s, 'available', %s, %s) ON CONFLICT (resource_type, resource_id, version_id, object_key) DO UPDATE SET file_name=EXCLUDED.file_name, content_type=EXCLUDED.content_type, checksum_sha256=EXCLUDED.checksum_sha256, byte_size=EXCLUDED.byte_size, status='available', created_by=EXCLUDED.created_by """, ( new_id("object"), dataset_id, spec["version_id"], uploaded["bucket"], object_key, file_name, "application/jsonl", spec["checksum"], len(spec["raw"]), payload.get("created_by") or task.get("created_by"), now, ), ) storage_object_id = conn.execute( """ SELECT id FROM storage_objects WHERE resource_type='dataset' AND resource_id=%s AND version_id=%s AND object_key=%s """, (dataset_id, spec["version_id"], object_key), ).fetchone()["id"] 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], "minio" if spec["store_in_minio"] else "database", 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 ) RETURNING * """, ( dataset_id, dataset_name, dataset_types[split_name], "minio" if spec["store_in_minio"] else "database", 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, "storage_backend": "minio" if spec["store_in_minio"] else "database", } version = { "id": spec["version_id"], "version_no": 1, "version": 1, "description": f"data process {split_name} 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": 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"{base_dataset_name}.{split_name}.jsonl", 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"], storage_object_id, spec["content"][:2000], f"data process {split_name} 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.get("output") or record.get("chosen") or "", 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, ), ) 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 """, (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": train_dataset, "datasets": published_datasets, "output_datasets": published_datasets, "created": created_any, "split_counts": split_counts, } @staticmethod def _source_ids( conn: psycopg.Connection[dict[str, Any]], task_id: str ) -> list[dict[str, Any]]: return conn.execute( """ SELECT id FROM data_process_source_files WHERE task_id=%s AND deleted_at IS NULL ORDER BY created_at, id """, (task_id,), ).fetchall()