128 lines
5.1 KiB
Python
128 lines
5.1 KiB
Python
"""
|
||
模型训练业务编排（移植自模型服务 projects/backend 的 training_service）
|
||
|
||
- preset 参数预设（quick / standard / high）
|
||
- train_type → stage 映射（sft/dpo/cpt/cot）
|
||
- 训练任务的启动 / 暂停 / 恢复 / 取消
|
||
- compute_gateway 状态轮询线程（当任务派发到算力时自动同步状态）
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import threading
|
||
import time
|
||
from typing import Any
|
||
|
||
from app.core.config import get_settings
|
||
from app.db.platform_store import get_platform_store
|
||
from app.modules.fine_tune import runner
|
||
|
||
PRESETS: dict[str, dict[str, Any]] = {
|
||
"quick": {"learning_rate": "5e-5", "n_epochs": 1, "batch_size": 4, "lora_rank": 8},
|
||
"standard": {"learning_rate": "2e-5", "n_epochs": 3, "batch_size": 8, "lora_rank": 16},
|
||
"high": {"learning_rate": "1e-5", "n_epochs": 5, "batch_size": 4, "lora_rank": 32},
|
||
}
|
||
|
||
# ── compute sync 轮询线程 ──────────────────────────────────────────────
|
||
_sync_thread: threading.Thread | None = None
|
||
_sync_thread_stop = threading.Event()
|
||
|
||
|
||
def _compute_sync_loop() -> None:
|
||
"""后台线程：周期性轮询算力节点，同步训练任务状态/日志/指标。"""
|
||
interval = get_settings().compute_poll_interval_seconds or 3
|
||
while not _sync_thread_stop.is_set():
|
||
try:
|
||
store = get_platform_store()
|
||
running = store.running_compute_tasks()
|
||
if running:
|
||
asyncio.run(_poll_once())
|
||
except Exception: # noqa: BLE001 - keep polling loop alive
|
||
pass
|
||
time.sleep(interval)
|
||
|
||
|
||
async def _poll_once() -> None:
|
||
from app.modules.compute_gateway.sync import poll_compute_jobs_once
|
||
await poll_compute_jobs_once()
|
||
|
||
|
||
def start_compute_sync_worker() -> None:
|
||
"""启动后台轮询线程（幂等，多次调用安全）。"""
|
||
global _sync_thread
|
||
if _sync_thread is not None and _sync_thread.is_alive():
|
||
return
|
||
_sync_thread_stop.clear()
|
||
_sync_thread = threading.Thread(target=_compute_sync_loop, daemon=True)
|
||
_sync_thread.start()
|
||
|
||
|
||
def stop_compute_sync_worker() -> None:
|
||
"""停止后台轮询线程。"""
|
||
_sync_thread_stop.set()
|
||
|
||
|
||
# ── preset / config ────────────────────────────────────────────────────
|
||
def apply_presets(payload: dict[str, Any]) -> dict[str, Any]:
|
||
"""根据 preset 字段补全缺失的超参；preset=custom 时不覆盖。"""
|
||
payload = dict(payload)
|
||
preset = payload.get("preset", "standard")
|
||
if preset in PRESETS and payload.get("preset") != "custom":
|
||
for key, value in PRESETS[preset].items():
|
||
payload.setdefault(key, value)
|
||
return payload
|
||
|
||
|
||
def build_training_config(payload: dict[str, Any]) -> dict[str, Any]:
|
||
"""把前端创建/启动载荷标准化为执行器可消费的 config。"""
|
||
payload = apply_presets(dict(payload))
|
||
gpus = payload.get("gpus") or [0]
|
||
return {
|
||
"name": payload.get("name", ""),
|
||
"description": payload.get("description", ""),
|
||
"train_type": payload.get("train_type", "SFT"),
|
||
"train_method": payload.get("train_method", "lora"),
|
||
"template": payload.get("template", "qwen"),
|
||
"base_model": payload.get("base_model", "") or payload.get("base_model_id", ""),
|
||
"train_dataset_id": payload.get("train_dataset_id", ""),
|
||
"eval_dataset_id": payload.get("eval_dataset_id", ""),
|
||
"auto_merge": bool(payload.get("auto_merge", False)),
|
||
"output_model_name": payload.get("output_model_name", ""),
|
||
"gpus": gpus,
|
||
"num_gpus": payload.get("num_gpus", len(gpus)),
|
||
"batch_size": payload.get("batch_size", 2),
|
||
"learning_rate": payload.get("learning_rate", 0.0002),
|
||
"n_epochs": payload.get("n_epochs", 3),
|
||
"save_steps": payload.get("save_steps", 50),
|
||
"lr_scheduler_type": payload.get("lr_scheduler_type", "cosine"),
|
||
"max_length": payload.get("max_length", 2048),
|
||
"warmup_ratio": payload.get("warmup_ratio", 0.03),
|
||
"weight_decay": payload.get("weight_decay", 0.01),
|
||
"lora_rank": payload.get("lora_rank", 8),
|
||
"lora_alpha": payload.get("lora_alpha", 16),
|
||
"lora_dropout": payload.get("lora_dropout", 0.05),
|
||
"resume_from": payload.get("resume_from"),
|
||
}
|
||
|
||
|
||
def launch_training(task_id: str) -> None:
|
||
"""在后台线程启动真实训练（本机 subprocess fallback，当算力节点不可用时使用）。
|
||
|
||
架构原则：GPU 计算应派发到算力服务进程执行。
|
||
当 platform_store.start_task 检测到在线算力节点时，会走 _dispatch_to_compute 派发路径；
|
||
仅当无可用算力节点且非 simulator 模式时，降级到本机 runner（违反 §1.1，待移除）。
|
||
"""
|
||
threading.Thread(target=runner.run_training, args=(task_id,), daemon=True).start()
|
||
|
||
|
||
def pause(task_id: str) -> bool:
|
||
return runner.pause(task_id)
|
||
|
||
|
||
def resume(task_id: str) -> bool:
|
||
return runner.resume(task_id)
|
||
|
||
|
||
def cancel(task_id: str) -> bool:
|
||
return runner.cancel(task_id)
|