173 lines
4.8 KiB
Python
173 lines
4.8 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from pathlib import Path
|
|
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from fastapi.responses import FileResponse, StreamingResponse
|
|
|
|
from . import db
|
|
from .catalog import get_catalog
|
|
from .config import settings
|
|
from .jobs import cancel_job, create_job
|
|
from .modules.system.service import disk_usage, get_conda_envs, get_gpus, scan_results
|
|
from .modules.weights.service import load_manifest, sync_weights, verify_weights
|
|
from .paths import ensure_inside
|
|
from .schemas import JobCreate, ProfileCreate, WeightSyncRequest
|
|
|
|
app = FastAPI(title="Seg Data Server Net", version="0.1.0")
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
@app.on_event("startup")
|
|
def startup() -> None:
|
|
db.init_db()
|
|
settings.log_dir.mkdir(parents=True, exist_ok=True)
|
|
settings.weights_root.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
@app.get("/api/health")
|
|
def health() -> dict:
|
|
return {
|
|
"ok": True,
|
|
"source_root": str(settings.source_root),
|
|
"project_root": str(settings.project_root),
|
|
"disk": disk_usage(),
|
|
}
|
|
|
|
|
|
@app.get("/api/system/gpus")
|
|
def api_gpus() -> dict:
|
|
return get_gpus()
|
|
|
|
|
|
@app.get("/api/system/envs")
|
|
def api_envs() -> dict:
|
|
return get_conda_envs()
|
|
|
|
|
|
@app.get("/api/catalog")
|
|
def api_catalog() -> dict:
|
|
return get_catalog()
|
|
|
|
|
|
@app.get("/api/profiles")
|
|
def api_profiles(kind: str | None = None) -> list[dict]:
|
|
return db.list_profiles(kind)
|
|
|
|
|
|
@app.post("/api/profiles")
|
|
def api_save_profile(profile: ProfileCreate) -> dict:
|
|
return db.upsert_profile(profile.name, profile.kind, profile.data)
|
|
|
|
|
|
@app.post("/api/jobs")
|
|
def api_create_job(request: JobCreate) -> dict:
|
|
try:
|
|
return create_job(request)
|
|
except Exception as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
|
|
|
|
@app.get("/api/jobs")
|
|
def api_jobs(limit: int = 100) -> list[dict]:
|
|
return db.list_jobs(limit)
|
|
|
|
|
|
@app.get("/api/jobs/{job_id}")
|
|
def api_job(job_id: str) -> dict:
|
|
job = db.get_job(job_id)
|
|
if not job:
|
|
raise HTTPException(status_code=404, detail="job not found")
|
|
job["log_tail"] = db.log_tail(job["log_path"])
|
|
return job
|
|
|
|
|
|
@app.post("/api/jobs/{job_id}/cancel")
|
|
def api_cancel_job(job_id: str) -> dict:
|
|
job = cancel_job(job_id)
|
|
if not job:
|
|
raise HTTPException(status_code=404, detail="job not found")
|
|
return job
|
|
|
|
|
|
@app.get("/api/jobs/{job_id}/events")
|
|
async def api_job_events(job_id: str):
|
|
async def stream():
|
|
last_size = 0
|
|
while True:
|
|
job = db.get_job(job_id)
|
|
if not job:
|
|
yield "event: error\ndata: job not found\n\n"
|
|
return
|
|
path = Path(job["log_path"])
|
|
chunk = ""
|
|
if path.exists():
|
|
size = path.stat().st_size
|
|
if size > last_size:
|
|
with path.open("rb") as handle:
|
|
handle.seek(last_size)
|
|
chunk = handle.read(size - last_size).decode("utf-8", errors="replace")
|
|
last_size = size
|
|
payload = {"job": job, "chunk": chunk}
|
|
yield f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"
|
|
if job["status"] in {"success", "failed", "cancelled"}:
|
|
return
|
|
await asyncio.sleep(1)
|
|
|
|
return StreamingResponse(stream(), media_type="text/event-stream")
|
|
|
|
|
|
@app.get("/api/results")
|
|
def api_results() -> list[dict]:
|
|
return scan_results()
|
|
|
|
|
|
@app.get("/api/artifacts/{artifact_path:path}")
|
|
def api_artifact(artifact_path: str):
|
|
candidate = Path(artifact_path)
|
|
if not candidate.is_absolute():
|
|
candidate = settings.source_root / candidate
|
|
try:
|
|
resolved = candidate.resolve()
|
|
allowed = False
|
|
for root in (settings.source_root, settings.project_root):
|
|
try:
|
|
ensure_inside(resolved, root)
|
|
allowed = True
|
|
break
|
|
except Exception:
|
|
continue
|
|
if not allowed:
|
|
raise ValueError("artifact path is outside allowed roots")
|
|
except Exception as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
if not resolved.exists() or not resolved.is_file():
|
|
raise HTTPException(status_code=404, detail="artifact not found")
|
|
return FileResponse(resolved)
|
|
|
|
|
|
@app.get("/api/weights")
|
|
def api_weights() -> dict:
|
|
return load_manifest()
|
|
|
|
|
|
@app.post("/api/weights/sync")
|
|
def api_weight_sync(request: WeightSyncRequest) -> dict:
|
|
return sync_weights(request.mode, request.hash_files, request.skip_existing)
|
|
|
|
|
|
@app.post("/api/weights/verify")
|
|
def api_weight_verify() -> dict:
|
|
return verify_weights()
|
|
|