refactor: move server/ to tools/server/

tools/server/routes_train.py

@@ -0,0 +1,131 @@
+# ai_mouse/server/routes_train.py
+"""Training and status routes."""
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from pathlib import Path
+from typing import AsyncGenerator
+
+from fastapi import APIRouter
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+
+from .deps import get_data_dir
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter()
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _paths() -> tuple[Path, Path, Path]:
+    data_dir = get_data_dir()
+    return (
+        data_dir / "traces.jsonl",
+        data_dir / "models_v2",
+        data_dir / "models_v2_pretrained",
+    )
+
+
+def _trace_count() -> int:
+    traces_path, _, _ = _paths()
+    if not traces_path.exists():
+        return 0
+    return sum(
+        1
+        for line in traces_path.read_text(encoding="utf-8").splitlines()
+        if line.strip()
+    )
+
+
+def _model_trained() -> bool:
+    _, models_dir, _ = _paths()
+    return (models_dir / "flow_model.pt").exists()
+
+
+# ---------------------------------------------------------------------------
+# Request models
+# ---------------------------------------------------------------------------
+
+
+class TrainRequest(BaseModel):
+    epochs: int = 200
+    data_path: str | None = None
+    output_dir: str | None = None
+
+
+# ---------------------------------------------------------------------------
+# Routes
+# ---------------------------------------------------------------------------
+
+
+@router.get("/status")
+def get_status() -> dict:
+    return {"trace_count": _trace_count(), "model_trained": _model_trained()}
+
+
+async def _train_sse_generator(req: TrainRequest) -> AsyncGenerator[str, None]:
+    """Run training in a thread via asyncio.to_thread, yield SSE events via asyncio.Queue."""
+    queue: asyncio.Queue[dict] = asyncio.Queue()
+
+    def callback(msg: dict) -> None:
+        queue.put_nowait(msg)
+
+    async def run_training_async() -> None:
+        from tools.trainer import train
+
+        traces_path, models_dir, pretrained_dir = _paths()
+        data_path = Path(req.data_path) if req.data_path else traces_path
+        output_dir = Path(req.output_dir) if req.output_dir else models_dir
+
+        # Auto-detect pretrained checkpoint and switch to fine-tune mode
+        resume_from: Path | None = None
+        effective_lr = 3e-4
+        if (pretrained_dir / "flow_model.pt").exists():
+            resume_from = pretrained_dir
+            effective_lr = 1e-5  # fine-tune lr
+            logger.info("Detected pretrained checkpoint, fine-tuning at lr=%g", effective_lr)
+            queue.put_nowait({
+                "info": f"Detected pretrained checkpoint at {pretrained_dir.name}, "
+                        f"running fine-tune at lr={effective_lr}",
+            })
+
+        try:
+            await asyncio.to_thread(
+                train,
+                data_path=data_path,
+                output_dir=output_dir,
+                epochs=req.epochs,
+                lr=effective_lr,
+                progress_callback=callback,
+                resume_from=resume_from,
+            )
+        except Exception as exc:  # noqa: BLE001
+            queue.put_nowait({"error": str(exc)})
+
+    task = asyncio.create_task(run_training_async())
+
+    while True:
+        msg = await queue.get()
+        yield f"data: {json.dumps(msg)}\n\n"
+        if msg.get("done") or msg.get("error"):
+            break
+
+    await task
+
+
+@router.post("/train")
+async def train_model(req: TrainRequest) -> StreamingResponse:
+    return StreamingResponse(
+        _train_sse_generator(req),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "X-Accel-Buffering": "no",
+        },
+    )