423 lines
15 KiB
Python
423 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any, Optional
|
|
|
|
from sqlalchemy import and_, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from ...models import (
|
|
QualityAsset,
|
|
QualityAssetPair,
|
|
TrainedModelArtifact,
|
|
TrainedModelDeployment,
|
|
TrainingDataset,
|
|
TrainingDatasetItem,
|
|
TrainingEvaluation,
|
|
TrainingJob,
|
|
TrainingJobRun,
|
|
)
|
|
from ...utils.timezone import now_cn_naive
|
|
|
|
|
|
EVALUATION_TYPES = ("offline_judge", "structure_compliance", "shadow_replay")
|
|
|
|
|
|
class TrainingPipelineService:
|
|
def __init__(self, db: AsyncSession):
|
|
self.db = db
|
|
|
|
@staticmethod
|
|
def _dataset_snapshot(row: TrainingDataset) -> dict[str, Any]:
|
|
return {
|
|
"id": int(row.id),
|
|
"tenant_id": int(row.tenant_id),
|
|
"project_id": int(row.project_id) if row.project_id is not None else None,
|
|
"dataset_type": row.dataset_type,
|
|
"version": row.version,
|
|
"status": row.status,
|
|
"sample_count": int(row.sample_count or 0),
|
|
"selection_rules": row.selection_rules or {},
|
|
"channel_scope": row.channel_scope or {},
|
|
"notes": row.notes,
|
|
}
|
|
|
|
def _job_snapshot(self, row: TrainingJob, *, artifact: Optional[TrainedModelArtifact] = None) -> dict[str, Any]:
|
|
payload = {
|
|
"id": int(row.id),
|
|
"tenant_id": int(row.tenant_id),
|
|
"project_id": int(row.project_id) if row.project_id is not None else None,
|
|
"dataset_id": int(row.dataset_id),
|
|
"job_type": row.job_type,
|
|
"base_model_family": row.base_model_family,
|
|
"target_slot": row.target_slot,
|
|
"status": row.status,
|
|
"provider_backend": row.provider_backend,
|
|
"artifact_ref": row.artifact_ref,
|
|
"metrics_json": row.metrics_json or {},
|
|
}
|
|
if artifact is not None:
|
|
payload["artifact"] = self._artifact_snapshot(artifact)
|
|
return payload
|
|
|
|
@staticmethod
|
|
def _artifact_snapshot(row: TrainedModelArtifact) -> dict[str, Any]:
|
|
return {
|
|
"id": int(row.id),
|
|
"tenant_id": int(row.tenant_id),
|
|
"training_job_id": int(row.training_job_id),
|
|
"artifact_ref": row.artifact_ref,
|
|
"base_model_family": row.base_model_family,
|
|
"target_slot": row.target_slot,
|
|
"status": row.status,
|
|
"metrics_json": row.metrics_json or {},
|
|
}
|
|
|
|
@staticmethod
|
|
def _deployment_snapshot(row: TrainedModelDeployment) -> dict[str, Any]:
|
|
return {
|
|
"id": int(row.id),
|
|
"tenant_id": int(row.tenant_id),
|
|
"artifact_id": int(row.artifact_id),
|
|
"deployment_status": row.deployment_status,
|
|
"bound_tier": row.bound_tier,
|
|
"effective_scope": row.effective_scope or {},
|
|
"decision": row.decision,
|
|
"notes": row.notes,
|
|
"decided_by": int(row.decided_by) if row.decided_by is not None else None,
|
|
"decided_at": row.decided_at,
|
|
}
|
|
|
|
@staticmethod
|
|
def _evaluation_snapshot(row: TrainingEvaluation) -> dict[str, Any]:
|
|
return {
|
|
"id": int(row.id),
|
|
"tenant_id": int(row.tenant_id),
|
|
"training_job_id": int(row.training_job_id),
|
|
"evaluation_type": row.evaluation_type,
|
|
"status": row.status,
|
|
"score": float(row.score or 0.0) if row.score is not None else None,
|
|
"evaluation_summary": row.evaluation_summary or {},
|
|
}
|
|
|
|
async def create_dataset(
|
|
self,
|
|
*,
|
|
tenant_id: int,
|
|
actor_id: Optional[int],
|
|
dataset_type: str,
|
|
version: str,
|
|
asset_ids: list[int],
|
|
pair_ids: list[int],
|
|
selection_rules: Optional[dict[str, Any]],
|
|
notes: Optional[str],
|
|
) -> dict[str, Any]:
|
|
dataset_type = str(dataset_type or "").strip().lower()
|
|
if dataset_type not in {"sft", "preference", "eval"}:
|
|
raise ValueError("unsupported dataset_type")
|
|
asset_ids = [int(x) for x in asset_ids or []]
|
|
pair_ids = [int(x) for x in pair_ids or []]
|
|
|
|
assets = []
|
|
if asset_ids:
|
|
assets = (
|
|
await self.db.execute(
|
|
select(QualityAsset).where(
|
|
QualityAsset.tenant_id == int(tenant_id),
|
|
QualityAsset.id.in_(asset_ids),
|
|
)
|
|
)
|
|
).scalars().all()
|
|
pairs = []
|
|
if pair_ids:
|
|
pairs = (
|
|
await self.db.execute(
|
|
select(QualityAssetPair).where(
|
|
QualityAssetPair.tenant_id == int(tenant_id),
|
|
QualityAssetPair.id.in_(pair_ids),
|
|
)
|
|
)
|
|
).scalars().all()
|
|
|
|
if dataset_type == "sft":
|
|
if not assets or any(asset.asset_type != "golden" or asset.status != "approved" for asset in assets):
|
|
raise ValueError("SFT dataset only accepts approved golden assets")
|
|
if dataset_type == "preference":
|
|
if not pairs or any(pair.status != "approved" for pair in pairs):
|
|
raise ValueError("Preference dataset only accepts approved asset pairs")
|
|
if dataset_type == "eval" and not (assets or pairs):
|
|
raise ValueError("Eval dataset requires assets or pairs")
|
|
|
|
project_id = None
|
|
if assets:
|
|
project_id = assets[0].project_id
|
|
elif pairs:
|
|
project_id = pairs[0].project_id
|
|
dataset = TrainingDataset(
|
|
tenant_id=int(tenant_id),
|
|
project_id=int(project_id) if project_id is not None else None,
|
|
dataset_type=dataset_type,
|
|
version=str(version or "").strip() or f"{dataset_type}-v1",
|
|
status="ready",
|
|
sample_count=len(assets) + len(pairs),
|
|
selection_rules=selection_rules or {},
|
|
channel_scope={},
|
|
notes=str(notes or "").strip() or None,
|
|
created_by=actor_id,
|
|
)
|
|
self.db.add(dataset)
|
|
await self.db.flush()
|
|
|
|
for asset in assets:
|
|
self.db.add(
|
|
TrainingDatasetItem(
|
|
tenant_id=int(tenant_id),
|
|
dataset_id=int(dataset.id),
|
|
quality_asset_id=int(asset.id),
|
|
redacted_snapshot={
|
|
"question_text": asset.question_text,
|
|
"answer_text": asset.answer_text,
|
|
"corrected_answer": asset.corrected_answer,
|
|
"prompt_key": asset.prompt_key,
|
|
},
|
|
)
|
|
)
|
|
for pair in pairs:
|
|
self.db.add(
|
|
TrainingDatasetItem(
|
|
tenant_id=int(tenant_id),
|
|
dataset_id=int(dataset.id),
|
|
quality_asset_pair_id=int(pair.id),
|
|
redacted_snapshot={
|
|
"pair_type": pair.pair_type,
|
|
"reason_text": pair.reason_text,
|
|
},
|
|
)
|
|
)
|
|
await self.db.commit()
|
|
await self.db.refresh(dataset)
|
|
return self._dataset_snapshot(dataset)
|
|
|
|
async def create_job(
|
|
self,
|
|
*,
|
|
tenant_id: int,
|
|
actor_id: Optional[int],
|
|
dataset_id: int,
|
|
job_type: str,
|
|
base_model_family: str,
|
|
target_slot: str,
|
|
provider_backend: str,
|
|
) -> dict[str, Any]:
|
|
normalized_job_type = str(job_type or "").strip().lower()
|
|
if normalized_job_type not in {"sft", "preference"}:
|
|
raise ValueError("unsupported training job_type")
|
|
normalized_target_slot = str(target_slot or "").strip().lower()
|
|
if normalized_target_slot not in {"low", "high"}:
|
|
raise ValueError("unsupported target_slot")
|
|
dataset = (
|
|
await self.db.execute(
|
|
select(TrainingDataset).where(
|
|
TrainingDataset.tenant_id == int(tenant_id),
|
|
TrainingDataset.id == int(dataset_id),
|
|
)
|
|
)
|
|
).scalar_one_or_none()
|
|
if dataset is None:
|
|
raise ValueError("training dataset not found")
|
|
|
|
job = TrainingJob(
|
|
tenant_id=int(tenant_id),
|
|
project_id=int(dataset.project_id) if dataset.project_id is not None else None,
|
|
dataset_id=int(dataset.id),
|
|
job_type=normalized_job_type,
|
|
base_model_family=str(base_model_family or "").strip(),
|
|
target_slot=normalized_target_slot,
|
|
status="queued",
|
|
provider_backend=str(provider_backend or "").strip(),
|
|
artifact_ref=None,
|
|
metrics_json={},
|
|
created_by=actor_id,
|
|
)
|
|
self.db.add(job)
|
|
await self.db.flush()
|
|
|
|
run = TrainingJobRun(
|
|
tenant_id=int(tenant_id),
|
|
training_job_id=int(job.id),
|
|
status="queued",
|
|
provider_run_ref=None,
|
|
metrics_json={},
|
|
)
|
|
self.db.add(run)
|
|
|
|
artifact = TrainedModelArtifact(
|
|
tenant_id=int(tenant_id),
|
|
training_job_id=int(job.id),
|
|
artifact_ref=f"artifact:{job.job_type}:{job.id}",
|
|
base_model_family=job.base_model_family,
|
|
target_slot=job.target_slot,
|
|
status="pending_eval",
|
|
metrics_json={},
|
|
)
|
|
self.db.add(artifact)
|
|
await self.db.flush()
|
|
job.artifact_ref = artifact.artifact_ref
|
|
|
|
for evaluation_type in EVALUATION_TYPES:
|
|
self.db.add(
|
|
TrainingEvaluation(
|
|
tenant_id=int(tenant_id),
|
|
training_job_id=int(job.id),
|
|
evaluation_type=evaluation_type,
|
|
status="pending",
|
|
score=None,
|
|
evaluation_summary={},
|
|
)
|
|
)
|
|
|
|
await self.db.commit()
|
|
await self.db.refresh(job)
|
|
await self.db.refresh(artifact)
|
|
return self._job_snapshot(job, artifact=artifact)
|
|
|
|
async def list_datasets(
|
|
self,
|
|
*,
|
|
tenant_id: int,
|
|
dataset_type: Optional[str] = None,
|
|
limit: int = 20,
|
|
offset: int = 0,
|
|
) -> dict[str, Any]:
|
|
stmt = select(TrainingDataset).where(TrainingDataset.tenant_id == int(tenant_id))
|
|
if dataset_type:
|
|
stmt = stmt.where(TrainingDataset.dataset_type == str(dataset_type).strip().lower())
|
|
rows = (
|
|
await self.db.execute(
|
|
stmt.order_by(TrainingDataset.id.desc()).offset(int(offset)).limit(int(limit))
|
|
)
|
|
).scalars().all()
|
|
return {"items": [self._dataset_snapshot(row) for row in rows], "total": len(rows), "limit": int(limit), "offset": int(offset)}
|
|
|
|
async def list_jobs(
|
|
self,
|
|
*,
|
|
tenant_id: int,
|
|
status: Optional[str] = None,
|
|
limit: int = 20,
|
|
offset: int = 0,
|
|
) -> dict[str, Any]:
|
|
stmt = select(TrainingJob).where(TrainingJob.tenant_id == int(tenant_id))
|
|
if status:
|
|
stmt = stmt.where(TrainingJob.status == str(status).strip().lower())
|
|
rows = (
|
|
await self.db.execute(
|
|
stmt.order_by(TrainingJob.id.desc()).offset(int(offset)).limit(int(limit))
|
|
)
|
|
).scalars().all()
|
|
return {"items": [self._job_snapshot(row) for row in rows], "total": len(rows), "limit": int(limit), "offset": int(offset)}
|
|
|
|
async def list_evaluations(self, *, tenant_id: int, job_id: int) -> dict[str, Any]:
|
|
rows = (
|
|
await self.db.execute(
|
|
select(TrainingEvaluation)
|
|
.where(
|
|
TrainingEvaluation.tenant_id == int(tenant_id),
|
|
TrainingEvaluation.training_job_id == int(job_id),
|
|
)
|
|
.order_by(TrainingEvaluation.id.asc())
|
|
)
|
|
).scalars().all()
|
|
return {"items": [self._evaluation_snapshot(row) for row in rows], "total": len(rows)}
|
|
|
|
async def promote_artifact(
|
|
self,
|
|
*,
|
|
tenant_id: int,
|
|
artifact_id: int,
|
|
actor_id: Optional[int],
|
|
) -> dict[str, Any]:
|
|
artifact = (
|
|
await self.db.execute(
|
|
select(TrainedModelArtifact).where(
|
|
TrainedModelArtifact.tenant_id == int(tenant_id),
|
|
TrainedModelArtifact.id == int(artifact_id),
|
|
)
|
|
)
|
|
).scalar_one_or_none()
|
|
if artifact is None:
|
|
raise ValueError("trained model artifact not found")
|
|
artifact.status = "candidate"
|
|
row = TrainedModelDeployment(
|
|
tenant_id=int(tenant_id),
|
|
artifact_id=int(artifact.id),
|
|
deployment_status="candidate",
|
|
bound_tier=str(artifact.target_slot or "").strip().lower() or "low",
|
|
effective_scope={},
|
|
decision=None,
|
|
notes=None,
|
|
decided_by=actor_id,
|
|
decided_at=now_cn_naive().isoformat(),
|
|
)
|
|
self.db.add(row)
|
|
await self.db.commit()
|
|
await self.db.refresh(row)
|
|
return self._deployment_snapshot(row)
|
|
|
|
async def decide_deployment(
|
|
self,
|
|
*,
|
|
tenant_id: int,
|
|
deployment_id: int,
|
|
actor_id: Optional[int],
|
|
decision: str,
|
|
notes: Optional[str],
|
|
) -> dict[str, Any]:
|
|
row = (
|
|
await self.db.execute(
|
|
select(TrainedModelDeployment).where(
|
|
and_(
|
|
TrainedModelDeployment.tenant_id == int(tenant_id),
|
|
TrainedModelDeployment.id == int(deployment_id),
|
|
)
|
|
)
|
|
)
|
|
).scalar_one_or_none()
|
|
if row is None:
|
|
raise ValueError("trained model deployment not found")
|
|
normalized = str(decision or "").strip().lower()
|
|
if normalized not in {"accepted", "rejected", "rolled_back"}:
|
|
raise ValueError("unsupported decision")
|
|
row.decision = normalized
|
|
row.notes = str(notes or "").strip() or None
|
|
row.decided_by = actor_id
|
|
row.decided_at = now_cn_naive().isoformat()
|
|
row.deployment_status = "active" if normalized == "accepted" else "rollback"
|
|
await self.db.commit()
|
|
await self.db.refresh(row)
|
|
return self._deployment_snapshot(row)
|
|
|
|
async def list_deployments(
|
|
self,
|
|
*,
|
|
tenant_id: int,
|
|
deployment_status: Optional[str] = None,
|
|
limit: int = 20,
|
|
offset: int = 0,
|
|
) -> dict[str, Any]:
|
|
stmt = select(TrainedModelDeployment).where(TrainedModelDeployment.tenant_id == int(tenant_id))
|
|
if deployment_status:
|
|
stmt = stmt.where(
|
|
TrainedModelDeployment.deployment_status == str(deployment_status).strip().lower()
|
|
)
|
|
rows = (
|
|
await self.db.execute(
|
|
stmt.order_by(TrainedModelDeployment.id.desc()).offset(int(offset)).limit(int(limit))
|
|
)
|
|
).scalars().all()
|
|
return {
|
|
"items": [self._deployment_snapshot(row) for row in rows],
|
|
"total": len(rows),
|
|
"limit": int(limit),
|
|
"offset": int(offset),
|
|
}
|