926 lines
33 KiB
Python
Executable file
926 lines
33 KiB
Python
Executable file
from __future__ import annotations
|
||
|
||
import argparse
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
import os
|
||
import re
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
import threading
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
|
||
from fastapi.responses import HTMLResponse
|
||
from fastapi.staticfiles import StaticFiles
|
||
import uvicorn
|
||
|
||
from yolo_webui.config import TrainingConfig
|
||
from yolo_webui.export_runner import ExportConfig
|
||
from yolo_webui.trainer import TrainingRunner
|
||
|
||
# Set up logging
|
||
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
|
||
logger = logging.getLogger("yolo_webui")
|
||
SESSION_NAME_PATTERN = re.compile(r"^[A-Za-z0-9_-]{1,64}$")
|
||
|
||
try:
|
||
from ultralytics import settings as _ultralytics_settings
|
||
_workspace_root = Path.cwd().resolve()
|
||
_ultralytics_settings.update({
|
||
"runs_dir": str((_workspace_root / "runs").resolve()),
|
||
"datasets_dir": str((_workspace_root / "datasets").resolve()),
|
||
"weights_dir": str((_workspace_root / "models").resolve()),
|
||
})
|
||
except Exception as _exc:
|
||
logger.debug(f"Ultralytics settings init skipped: {_exc}")
|
||
|
||
|
||
@dataclass
|
||
class LiveState:
|
||
status: str = "idle" # idle, preparing, training, stopping, succeeded, cancelled, failed
|
||
epoch: int = 0
|
||
total_epochs: int = 0
|
||
logs: list[str] = field(default_factory=list)
|
||
metrics: list[dict[str, Any]] = field(default_factory=list)
|
||
output_dir: str | None = None
|
||
stop_requested: bool = False
|
||
last_event_kind: str | None = None
|
||
|
||
def reset(self) -> None:
|
||
self.status = "idle"
|
||
self.epoch = 0
|
||
self.total_epochs = 0
|
||
self.logs = []
|
||
self.metrics = []
|
||
self.output_dir = None
|
||
self.stop_requested = False
|
||
self.last_event_kind = None
|
||
|
||
def snapshot(self) -> dict[str, Any]:
|
||
"""Return a detached snapshot safe to serialize after releasing a lock."""
|
||
return {
|
||
"status": self.status,
|
||
"epoch": self.epoch,
|
||
"total_epochs": self.total_epochs,
|
||
"logs": list(self.logs),
|
||
"metrics": [dict(metric) for metric in self.metrics],
|
||
"output_dir": self.output_dir,
|
||
}
|
||
|
||
|
||
class TrainingManager:
|
||
"""Manages the background training subprocess and WebSocket clients."""
|
||
|
||
def __init__(self) -> None:
|
||
self.state = LiveState()
|
||
self.runner = TrainingRunner()
|
||
self.active_websockets: set[WebSocket] = set()
|
||
self._lock = threading.Lock()
|
||
self._thread: threading.Thread | None = None
|
||
self._event_loop: asyncio.AbstractEventLoop | None = None
|
||
self._broadcast_lock: asyncio.Lock | None = None
|
||
|
||
def add_websocket(self, websocket: WebSocket) -> None:
|
||
loop = asyncio.get_running_loop()
|
||
with self._lock:
|
||
self.active_websockets.add(websocket)
|
||
if self._event_loop is not loop:
|
||
self._event_loop = loop
|
||
self._broadcast_lock = asyncio.Lock()
|
||
|
||
def remove_websocket(self, websocket: WebSocket) -> None:
|
||
with self._lock:
|
||
self.active_websockets.discard(websocket)
|
||
|
||
def broadcast(self, data: dict[str, Any]) -> None:
|
||
payload = json.dumps(data)
|
||
with self._lock:
|
||
loop = self._event_loop
|
||
has_sockets = bool(self.active_websockets)
|
||
|
||
if not has_sockets:
|
||
return
|
||
if loop is None or loop.is_closed():
|
||
logger.warning("WebSocket event loop is unavailable; broadcast was skipped")
|
||
return
|
||
|
||
coroutine = self._send_payload(payload)
|
||
try:
|
||
try:
|
||
running_loop = asyncio.get_running_loop()
|
||
except RuntimeError:
|
||
running_loop = None
|
||
|
||
if running_loop is loop:
|
||
future = loop.create_task(coroutine)
|
||
else:
|
||
future = asyncio.run_coroutine_threadsafe(coroutine, loop)
|
||
future.add_done_callback(self._log_broadcast_failure)
|
||
except Exception:
|
||
coroutine.close()
|
||
logger.exception("Failed to schedule WebSocket broadcast")
|
||
|
||
async def _send_payload(self, payload: str) -> None:
|
||
broadcast_lock = self._broadcast_lock
|
||
if broadcast_lock is None:
|
||
return
|
||
|
||
async with broadcast_lock:
|
||
with self._lock:
|
||
sockets = list(self.active_websockets)
|
||
|
||
failed: list[WebSocket] = []
|
||
for websocket in sockets:
|
||
try:
|
||
await websocket.send_text(payload)
|
||
except Exception:
|
||
failed.append(websocket)
|
||
logger.warning("Dropping failed WebSocket client", exc_info=True)
|
||
|
||
if failed:
|
||
with self._lock:
|
||
for websocket in failed:
|
||
self.active_websockets.discard(websocket)
|
||
|
||
@staticmethod
|
||
def _log_broadcast_failure(future: Any) -> None:
|
||
if future.cancelled():
|
||
return
|
||
error = future.exception()
|
||
if error is not None:
|
||
logger.error(
|
||
"WebSocket broadcast failed",
|
||
exc_info=(type(error), error, error.__traceback__),
|
||
)
|
||
|
||
def add_log(self, text: str, level: str = "info") -> None:
|
||
log_entry = f"__LOG_LEVEL_{level.upper()}__:{text}"
|
||
with self._lock:
|
||
if level == "progress" and self.state.logs and self.state.logs[-1].startswith("__LOG_LEVEL_PROGRESS__"):
|
||
self.state.logs[-1] = log_entry
|
||
else:
|
||
self.state.logs.append(log_entry)
|
||
self.broadcast({"type": "log", "message": text, "level": level})
|
||
|
||
def start_training(self, config: TrainingConfig) -> None:
|
||
with self._lock:
|
||
if self.state.status in ("preparing", "training", "stopping") or (
|
||
self._thread is not None and self._thread.is_alive()
|
||
):
|
||
raise ValueError("Обучение уже выполняется.")
|
||
|
||
self.state.reset()
|
||
self.state.status = "preparing"
|
||
self.runner.prepare_run()
|
||
self._thread = threading.Thread(target=self._run_subprocess, args=(config,), daemon=True)
|
||
self._thread.start()
|
||
|
||
self.broadcast({"type": "status", "status": self.state.status})
|
||
self.add_log(f"Запуск: задача={config.task}, модель={config.model}, датасет={config.dataset}", "started")
|
||
if config.mlflow.enabled:
|
||
self.add_log(f"MLflow: {config.mlflow.tracking_uri} · эксперимент {config.mlflow.experiment_name}", "info")
|
||
if config.augmentation.enabled:
|
||
self.add_log(f"Аугментация: enabled=True, mosaic={config.augmentation.mosaic}, mixup={config.augmentation.mixup}", "info")
|
||
|
||
def stop_training(self) -> None:
|
||
with self._lock:
|
||
if self.state.status not in ("preparing", "training"):
|
||
return
|
||
self.state.status = "stopping"
|
||
self.state.stop_requested = True
|
||
self.runner.request_stop()
|
||
|
||
self.broadcast({"type": "status", "status": self.state.status})
|
||
self.add_log("Запрошена остановка обучения...", "warning")
|
||
|
||
def _handle_subprocess_line(self, line_str: str, is_progress: bool = False) -> None:
|
||
line_str = line_str.strip()
|
||
if not line_str:
|
||
return
|
||
|
||
if line_str == "__YOLO_WEBUI_READY__":
|
||
self.runner.mark_subprocess_ready()
|
||
elif line_str.startswith("__YOLO_WEBUI_EVENT__:"):
|
||
try:
|
||
event_data = json.loads(line_str[len("__YOLO_WEBUI_EVENT__:") :])
|
||
kind = event_data["kind"]
|
||
message = event_data["message"]
|
||
epoch = event_data["epoch"]
|
||
total = event_data["total_epochs"]
|
||
|
||
metrics_dict = {}
|
||
if kind == "epoch":
|
||
with self._lock:
|
||
self.state.epoch = epoch
|
||
self.state.total_epochs = total
|
||
if " · " in message:
|
||
parts = message.split(" · ")[1:]
|
||
for p in parts:
|
||
if "=" in p:
|
||
k, v = p.split("=", 1)
|
||
try:
|
||
metrics_dict[k.strip()] = float(v.strip())
|
||
except ValueError:
|
||
pass
|
||
if metrics_dict:
|
||
metrics_dict["epoch"] = epoch
|
||
with self._lock:
|
||
self.state.metrics.append(metrics_dict)
|
||
|
||
status_update = None
|
||
with self._lock:
|
||
self.state.last_event_kind = kind
|
||
if kind == "started" and self.state.status == "preparing":
|
||
self.state.status = "training"
|
||
status_update = self.state.status
|
||
|
||
if status_update is not None:
|
||
self.broadcast({"type": "status", "status": status_update})
|
||
|
||
self.add_log(message, "progress" if is_progress else kind)
|
||
self.broadcast({
|
||
"type": "progress",
|
||
"epoch": epoch,
|
||
"total_epochs": total,
|
||
"metrics": metrics_dict,
|
||
"message": message
|
||
})
|
||
except Exception as e:
|
||
logger.error(f"Error parsing event: {e}")
|
||
elif line_str.startswith("__YOLO_WEBUI_RESULT__:"):
|
||
with self._lock:
|
||
self.state.output_dir = line_str[len("__YOLO_WEBUI_RESULT__:") :]
|
||
else:
|
||
self.add_log(line_str, "info")
|
||
|
||
def _run_subprocess(self, config: TrainingConfig) -> None:
|
||
temp_config_path = None
|
||
process = None
|
||
try:
|
||
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
|
||
json.dump(config.to_dict(), f)
|
||
temp_config_path = f.name
|
||
|
||
# Run python with -u to disable block buffering for real-time progress output
|
||
cmd = [sys.executable, "-u", "-m", "yolo_webui.subprocess_runner", temp_config_path]
|
||
process = subprocess.Popen(
|
||
cmd,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.STDOUT,
|
||
text=True,
|
||
bufsize=1,
|
||
)
|
||
self.runner.set_subprocess(process, ready=False)
|
||
|
||
buffer = ""
|
||
while True:
|
||
char = process.stdout.read(1)
|
||
if not char:
|
||
if buffer:
|
||
self._handle_subprocess_line(buffer, is_progress=False)
|
||
break
|
||
|
||
if char in ("\r", "\n"):
|
||
if buffer:
|
||
self._handle_subprocess_line(buffer, is_progress=(char == "\r"))
|
||
buffer = ""
|
||
else:
|
||
buffer += char
|
||
|
||
process.wait()
|
||
rc = process.returncode
|
||
self.runner.clear_subprocess()
|
||
self._finalize_process_result(rc)
|
||
|
||
except Exception as exc:
|
||
logger.exception("Error in training process thread:")
|
||
if process is not None:
|
||
try:
|
||
if process.poll() is None:
|
||
process.kill()
|
||
process.wait(timeout=5)
|
||
except Exception:
|
||
pass
|
||
self.runner.clear_subprocess()
|
||
|
||
with self._lock:
|
||
self.state.status = "failed"
|
||
self.add_log(f"Внутренняя ошибка менеджера: {exc}", "error")
|
||
finally:
|
||
if temp_config_path and os.path.exists(temp_config_path):
|
||
try:
|
||
os.unlink(temp_config_path)
|
||
except Exception:
|
||
pass
|
||
self.runner.clear_subprocess()
|
||
with self._lock:
|
||
final_status = self.state.status
|
||
output_dir = self.state.output_dir
|
||
self.broadcast({"type": "status", "status": final_status, "output_dir": output_dir})
|
||
|
||
def _finalize_process_result(self, return_code: int) -> None:
|
||
with self._lock:
|
||
stopped = self.state.stop_requested
|
||
last_event_kind = self.state.last_event_kind
|
||
|
||
was_cancelled = stopped and (
|
||
return_code == 0
|
||
or last_event_kind == "cancelled"
|
||
or self.runner.force_stop_triggered
|
||
)
|
||
|
||
if was_cancelled:
|
||
status = "cancelled"
|
||
message = "Обучение остановлено пользователем."
|
||
level = "warning"
|
||
elif return_code == 0:
|
||
status = "succeeded"
|
||
message = "Обучение успешно завершено."
|
||
level = "success"
|
||
else:
|
||
status = "failed"
|
||
message = "Процесс обучения завершился с ошибкой. Проверьте логи выше."
|
||
level = "error"
|
||
|
||
with self._lock:
|
||
self.state.status = status
|
||
self.add_log(message, level)
|
||
|
||
|
||
manager = TrainingManager()
|
||
app = FastAPI(title="YOLO Train Studio Web")
|
||
|
||
# Serve UI static folder
|
||
static_dir = Path(__file__).parent / "static"
|
||
if static_dir.exists():
|
||
app.mount("/static", StaticFiles(directory=static_dir), name="static")
|
||
|
||
|
||
@app.get("/", response_class=HTMLResponse)
|
||
async def get_index():
|
||
index_file = static_dir / "index.html"
|
||
if not index_file.exists():
|
||
return HTMLResponse(
|
||
content="<h1>YOLO Train Studio Web</h1><p>Static assets are missing. Place index.html under static/.</p>",
|
||
status_code=404,
|
||
)
|
||
return HTMLResponse(content=index_file.read_text(encoding="utf-8"))
|
||
|
||
|
||
def get_sessions_dir() -> Path:
|
||
path = Path("runs") / "sessions"
|
||
path.mkdir(parents=True, exist_ok=True)
|
||
return path
|
||
|
||
|
||
def write_json_atomic(path: Path, data: dict[str, Any]) -> None:
|
||
"""Replace a JSON file atomically so concurrent readers never see partial data."""
|
||
temp_path: Path | None = None
|
||
try:
|
||
with tempfile.NamedTemporaryFile(
|
||
"w",
|
||
prefix=f".{path.name}.",
|
||
suffix=".tmp",
|
||
dir=path.parent,
|
||
delete=False,
|
||
encoding="utf-8",
|
||
) as temp_file:
|
||
temp_path = Path(temp_file.name)
|
||
json.dump(data, temp_file, ensure_ascii=False, indent=2)
|
||
os.replace(temp_path, path)
|
||
temp_path = None
|
||
finally:
|
||
if temp_path is not None:
|
||
try:
|
||
temp_path.unlink()
|
||
except FileNotFoundError:
|
||
pass
|
||
|
||
|
||
def parse_training_config(config_data: dict[str, Any]) -> TrainingConfig:
|
||
config = TrainingConfig.from_dict(config_data)
|
||
config.validate()
|
||
return config
|
||
|
||
|
||
def parse_export_config(config_data: dict[str, Any]) -> dict[str, Any]:
|
||
config = ExportConfig.from_mapping(config_data)
|
||
return {
|
||
"model": str(config.model),
|
||
"format": config.export_format,
|
||
"imgsz": config.imgsz,
|
||
"half": config.half,
|
||
"int8": config.int8,
|
||
"dynamic": config.dynamic,
|
||
"simplify": config.simplify,
|
||
"batch": config.batch,
|
||
"workspace": config.workspace,
|
||
}
|
||
|
||
|
||
def get_session_path(name: str, *, allow_last_run: bool = True) -> Path:
|
||
if SESSION_NAME_PATTERN.fullmatch(name) is None:
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail="Имя сессии может содержать только латинские буквы, цифры, '_' и '-'.",
|
||
)
|
||
if not allow_last_run and name == "last_run":
|
||
raise HTTPException(status_code=400, detail="Имя 'last_run' зарезервировано.")
|
||
return get_sessions_dir() / f"{name}.json"
|
||
|
||
|
||
@app.get("/api/config/defaults")
|
||
async def get_defaults():
|
||
# Return defaults by instantiating with dummy paths and serializing
|
||
defaults = TrainingConfig(dataset="coco8.yaml", model="yolo11n.pt")
|
||
return defaults.to_dict()
|
||
|
||
|
||
@app.get("/api/sessions")
|
||
async def list_sessions():
|
||
sessions_dir = get_sessions_dir()
|
||
files = sessions_dir.glob("*.json")
|
||
names = [f.stem for f in files if f.name != "last_run.json"]
|
||
return sorted(names)
|
||
|
||
|
||
@app.get("/api/datasets")
|
||
async def list_datasets():
|
||
datasets_dir = Path("datasets")
|
||
if not datasets_dir.exists():
|
||
return []
|
||
|
||
items = []
|
||
try:
|
||
for path in datasets_dir.iterdir():
|
||
if path.is_dir() and not path.name.startswith("."):
|
||
items.append({
|
||
"name": path.name,
|
||
"path": str(path.absolute()),
|
||
"type": "directory"
|
||
})
|
||
elif path.is_file() and path.suffix.lower() in (".yaml", ".yml"):
|
||
items.append({
|
||
"name": path.name,
|
||
"path": str(path.absolute()),
|
||
"type": "yaml"
|
||
})
|
||
except Exception as e:
|
||
logger.error(f"Failed to list datasets: {e}")
|
||
|
||
return sorted(items, key=lambda x: x["name"])
|
||
|
||
|
||
@app.get("/api/models")
|
||
async def list_models():
|
||
items = []
|
||
|
||
# Check models directory
|
||
models_dir = Path("models")
|
||
if models_dir.exists():
|
||
try:
|
||
for path in models_dir.iterdir():
|
||
if path.is_file() and path.suffix.lower() in (".pt", ".pth", ".yaml", ".yml"):
|
||
items.append({
|
||
"name": path.name,
|
||
"path": str(path.absolute()),
|
||
"source": "models"
|
||
})
|
||
except Exception as e:
|
||
logger.error(f"Failed to list models in models/: {e}")
|
||
|
||
# Check runs directory for .pt and .pth files
|
||
runs_dir = Path("runs")
|
||
if runs_dir.exists():
|
||
try:
|
||
for path in runs_dir.rglob("*.pt"):
|
||
items.append({
|
||
"name": path.relative_to(runs_dir).as_posix(),
|
||
"path": str(path.absolute()),
|
||
"source": "runs"
|
||
})
|
||
for path in runs_dir.rglob("*.pth"):
|
||
items.append({
|
||
"name": path.relative_to(runs_dir).as_posix(),
|
||
"path": str(path.absolute()),
|
||
"source": "runs"
|
||
})
|
||
except Exception as e:
|
||
logger.error(f"Failed to list models in runs/: {e}")
|
||
|
||
return sorted(items, key=lambda x: x["name"])
|
||
|
||
|
||
@app.get("/api/sessions/{name}")
|
||
async def load_session(name: str):
|
||
file_path = get_session_path(name)
|
||
if not file_path.exists():
|
||
raise HTTPException(status_code=404, detail="Сессия не найдена.")
|
||
try:
|
||
with file_path.open("r", encoding="utf-8") as f:
|
||
return json.load(f)
|
||
except Exception as exc:
|
||
raise HTTPException(status_code=500, detail=f"Не удалось загрузить сессию: {exc}")
|
||
|
||
|
||
@app.post("/api/sessions/{name}")
|
||
async def save_session(name: str, config_data: dict[str, Any]):
|
||
file_path = get_session_path(name, allow_last_run=False)
|
||
try:
|
||
config = parse_training_config(config_data)
|
||
except ValueError as exc:
|
||
raise HTTPException(status_code=400, detail=str(exc))
|
||
except Exception as exc:
|
||
raise HTTPException(status_code=400, detail=f"Некорректная конфигурация: {exc}")
|
||
try:
|
||
write_json_atomic(file_path, config.to_dict())
|
||
return {"message": "Сессия успешно сохранена."}
|
||
except Exception as exc:
|
||
raise HTTPException(status_code=500, detail=f"Не удалось сохранить сессию: {exc}")
|
||
|
||
|
||
@app.delete("/api/sessions/{name}")
|
||
async def delete_session(name: str):
|
||
file_path = get_session_path(name, allow_last_run=False)
|
||
if not file_path.exists():
|
||
raise HTTPException(status_code=404, detail="Сессия не найдена.")
|
||
try:
|
||
file_path.unlink()
|
||
return {"message": "Сессия удалена."}
|
||
except Exception as exc:
|
||
raise HTTPException(status_code=500, detail=f"Не удалось удалить сессию: {exc}")
|
||
|
||
|
||
@app.get("/api/train/status")
|
||
async def get_status():
|
||
with manager._lock:
|
||
return manager.state.snapshot()
|
||
|
||
|
||
@app.post("/api/train/start")
|
||
async def start_training(config_data: dict[str, Any]):
|
||
try:
|
||
config = parse_training_config(config_data)
|
||
except ValueError as exc:
|
||
raise HTTPException(status_code=400, detail=str(exc))
|
||
except Exception as exc:
|
||
raise HTTPException(status_code=400, detail=f"Некорректная конфигурация: {exc}")
|
||
|
||
try:
|
||
manager.start_training(config)
|
||
except ValueError as exc:
|
||
raise HTTPException(status_code=400, detail=str(exc))
|
||
|
||
# Only a configuration accepted for execution becomes the last run.
|
||
try:
|
||
sessions_dir = get_sessions_dir()
|
||
last_run_path = sessions_dir / "last_run.json"
|
||
write_json_atomic(last_run_path, config.to_dict())
|
||
except Exception as exc:
|
||
logger.error(f"Failed to auto-save last run: {exc}")
|
||
|
||
return {"message": "Обучение запущено."}
|
||
|
||
|
||
@app.post("/api/train/stop")
|
||
async def stop_training():
|
||
manager.stop_training()
|
||
return {"message": "Запрос на остановку отправлен."}
|
||
|
||
|
||
class ExportManager:
|
||
def __init__(self) -> None:
|
||
self.state = LiveState()
|
||
self.active_websockets: set[WebSocket] = set()
|
||
self._lock = threading.Lock()
|
||
self._thread: threading.Thread | None = None
|
||
self._event_loop: asyncio.AbstractEventLoop | None = None
|
||
self._broadcast_lock: asyncio.Lock | None = None
|
||
self._process: subprocess.Popen | None = None
|
||
|
||
def add_websocket(self, websocket: WebSocket) -> None:
|
||
loop = asyncio.get_running_loop()
|
||
with self._lock:
|
||
self.active_websockets.add(websocket)
|
||
if self._event_loop is not loop:
|
||
self._event_loop = loop
|
||
self._broadcast_lock = asyncio.Lock()
|
||
|
||
def remove_websocket(self, websocket: WebSocket) -> None:
|
||
with self._lock:
|
||
self.active_websockets.discard(websocket)
|
||
|
||
def broadcast(self, data: dict[str, Any]) -> None:
|
||
payload = json.dumps(data)
|
||
with self._lock:
|
||
loop = self._event_loop
|
||
has_sockets = bool(self.active_websockets)
|
||
|
||
if not has_sockets:
|
||
return
|
||
if loop is None or loop.is_closed():
|
||
return
|
||
|
||
coroutine = self._send_payload(payload)
|
||
try:
|
||
try:
|
||
running_loop = asyncio.get_running_loop()
|
||
except RuntimeError:
|
||
running_loop = None
|
||
|
||
if running_loop is loop:
|
||
future = loop.create_task(coroutine)
|
||
else:
|
||
future = asyncio.run_coroutine_threadsafe(coroutine, loop)
|
||
future.add_done_callback(TrainingManager._log_broadcast_failure)
|
||
except Exception:
|
||
coroutine.close()
|
||
logger.exception("Failed to schedule export WebSocket broadcast")
|
||
|
||
async def _send_payload(self, payload: str) -> None:
|
||
broadcast_lock = self._broadcast_lock
|
||
if broadcast_lock is None:
|
||
return
|
||
|
||
async with broadcast_lock:
|
||
with self._lock:
|
||
sockets = list(self.active_websockets)
|
||
|
||
failed: list[WebSocket] = []
|
||
for websocket in sockets:
|
||
try:
|
||
await websocket.send_text(payload)
|
||
except Exception:
|
||
failed.append(websocket)
|
||
|
||
if failed:
|
||
with self._lock:
|
||
for websocket in failed:
|
||
self.active_websockets.discard(websocket)
|
||
|
||
def add_log(self, text: str, level: str = "info") -> None:
|
||
log_entry = f"__LOG_LEVEL_{level.upper()}__:{text}"
|
||
with self._lock:
|
||
self.state.logs.append(log_entry)
|
||
self.broadcast({"type": "log", "message": text, "level": level})
|
||
|
||
def start_export(self, config_data: dict[str, Any]) -> None:
|
||
with self._lock:
|
||
if self.state.status in ("preparing", "exporting", "stopping") or (
|
||
self._thread is not None and self._thread.is_alive()
|
||
):
|
||
raise ValueError("Экспорт уже выполняется.")
|
||
|
||
self.state.reset()
|
||
self.state.status = "preparing"
|
||
self._thread = threading.Thread(target=self._run_subprocess, args=(config_data,), daemon=True)
|
||
self._thread.start()
|
||
|
||
self.broadcast({"type": "status", "status": self.state.status})
|
||
self.add_log(f"Запуск экспорта: {config_data.get('model')} -> {config_data.get('format')}", "started")
|
||
|
||
def stop_export(self) -> None:
|
||
with self._lock:
|
||
if self.state.status not in ("preparing", "exporting"):
|
||
return
|
||
|
||
process = self._process
|
||
if process is not None:
|
||
if process.poll() is not None:
|
||
# The subprocess has already completed; let its real return code
|
||
# determine the final status instead of relabeling it cancelled.
|
||
return
|
||
try:
|
||
process.terminate()
|
||
except Exception:
|
||
logger.warning("Failed to terminate export subprocess", exc_info=True)
|
||
return
|
||
|
||
self.state.status = "stopping"
|
||
self.state.stop_requested = True
|
||
|
||
self.broadcast({"type": "status", "status": self.state.status})
|
||
self.add_log("Запрошена остановка экспорта...", "warning")
|
||
|
||
def _run_subprocess(self, config_data: dict[str, Any]) -> None:
|
||
temp_config_path = None
|
||
process: subprocess.Popen[str] | None = None
|
||
try:
|
||
with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8") as f:
|
||
json.dump(config_data, f)
|
||
temp_config_path = f.name
|
||
|
||
cmd = [sys.executable, "-u", "-m", "yolo_webui.export_runner", temp_config_path]
|
||
with self._lock:
|
||
if self.state.stop_requested:
|
||
cancelled_before_start = True
|
||
else:
|
||
cancelled_before_start = False
|
||
process = subprocess.Popen(
|
||
cmd,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.STDOUT,
|
||
text=True,
|
||
bufsize=1,
|
||
)
|
||
self._process = process
|
||
|
||
if cancelled_before_start:
|
||
self._finalize_process_result(0)
|
||
return
|
||
|
||
if process is None or process.stdout is None:
|
||
raise RuntimeError("Не удалось открыть stdout процесса экспорта.")
|
||
|
||
buffer = ""
|
||
while True:
|
||
char = process.stdout.read(1)
|
||
if not char:
|
||
if buffer:
|
||
self._handle_subprocess_line(buffer)
|
||
break
|
||
|
||
if char in ("\r", "\n"):
|
||
if buffer:
|
||
self._handle_subprocess_line(buffer)
|
||
buffer = ""
|
||
else:
|
||
buffer += char
|
||
|
||
process.wait()
|
||
rc = process.returncode
|
||
|
||
with self._lock:
|
||
if self._process is process:
|
||
self._process = None
|
||
|
||
self._finalize_process_result(rc)
|
||
|
||
except Exception as exc:
|
||
logger.exception("Error in export process thread:")
|
||
if process is not None:
|
||
try:
|
||
if process.poll() is None:
|
||
process.kill()
|
||
process.wait(timeout=5)
|
||
except Exception:
|
||
pass
|
||
|
||
with self._lock:
|
||
if self._process is process:
|
||
self._process = None
|
||
self.state.status = "failed"
|
||
self.add_log(f"Внутренняя ошибка менеджера: {exc}", "error")
|
||
finally:
|
||
if temp_config_path and os.path.exists(temp_config_path):
|
||
try:
|
||
os.unlink(temp_config_path)
|
||
except Exception:
|
||
pass
|
||
with self._lock:
|
||
final_status = self.state.status
|
||
output_dir = self.state.output_dir
|
||
self.broadcast({"type": "status", "status": final_status, "output_dir": output_dir})
|
||
|
||
def _handle_subprocess_line(self, line_str: str) -> None:
|
||
line_str = line_str.strip()
|
||
if not line_str:
|
||
return
|
||
|
||
if line_str == "__YOLO_WEBUI_READY__":
|
||
status_changed = False
|
||
with self._lock:
|
||
if self.state.status == "preparing":
|
||
self.state.status = "exporting"
|
||
status_changed = True
|
||
if status_changed:
|
||
self.broadcast({"type": "status", "status": "exporting"})
|
||
elif line_str.startswith("__YOLO_WEBUI_RESULT__:"):
|
||
with self._lock:
|
||
self.state.output_dir = line_str[len("__YOLO_WEBUI_RESULT__:") :]
|
||
else:
|
||
self.add_log(line_str, "info")
|
||
|
||
def _finalize_process_result(self, return_code: int) -> None:
|
||
with self._lock:
|
||
stopped = self.state.stop_requested
|
||
|
||
if stopped:
|
||
status = "cancelled"
|
||
message = "Экспорт остановлен пользователем."
|
||
level = "warning"
|
||
elif return_code == 0:
|
||
status = "succeeded"
|
||
message = "Экспорт успешно завершен."
|
||
level = "success"
|
||
else:
|
||
status = "failed"
|
||
message = "Процесс экспорта завершился с ошибкой."
|
||
level = "error"
|
||
|
||
with self._lock:
|
||
self.state.status = status
|
||
self.add_log(message, level)
|
||
|
||
export_manager = ExportManager()
|
||
|
||
@app.get("/api/export/status")
|
||
async def get_export_status():
|
||
with export_manager._lock:
|
||
snapshot = export_manager.state.snapshot()
|
||
return {
|
||
"status": snapshot["status"],
|
||
"output_dir": snapshot["output_dir"],
|
||
"logs": snapshot["logs"],
|
||
}
|
||
|
||
@app.post("/api/export/start")
|
||
async def start_export(config_data: dict[str, Any]):
|
||
try:
|
||
validated_data = parse_export_config(config_data)
|
||
export_manager.start_export(validated_data)
|
||
return {"message": "Экспорт запущен."}
|
||
except ValueError as exc:
|
||
raise HTTPException(status_code=400, detail=str(exc))
|
||
|
||
@app.post("/api/export/stop")
|
||
async def stop_export():
|
||
export_manager.stop_export()
|
||
return {"message": "Запрос на остановку экспорта отправлен."}
|
||
|
||
@app.websocket("/api/export/ws")
|
||
async def export_websocket_endpoint(websocket: WebSocket):
|
||
await websocket.accept()
|
||
export_manager.add_websocket(websocket)
|
||
|
||
with export_manager._lock:
|
||
snapshot = export_manager.state.snapshot()
|
||
state_dict = {
|
||
"type": "init",
|
||
"status": snapshot["status"],
|
||
"output_dir": snapshot["output_dir"],
|
||
"logs": [log.split(":", 1) for log in snapshot["logs"] if ":" in log],
|
||
}
|
||
try:
|
||
await websocket.send_text(json.dumps(state_dict))
|
||
while True:
|
||
await websocket.receive_text()
|
||
except WebSocketDisconnect:
|
||
pass
|
||
except Exception:
|
||
pass
|
||
finally:
|
||
export_manager.remove_websocket(websocket)
|
||
|
||
|
||
@app.websocket("/api/ws")
|
||
async def websocket_endpoint(websocket: WebSocket):
|
||
await websocket.accept()
|
||
manager.add_websocket(websocket)
|
||
|
||
# Send current state upon connection
|
||
with manager._lock:
|
||
snapshot = manager.state.snapshot()
|
||
state_dict = {
|
||
"type": "init",
|
||
"status": snapshot["status"],
|
||
"epoch": snapshot["epoch"],
|
||
"total_epochs": snapshot["total_epochs"],
|
||
"output_dir": snapshot["output_dir"],
|
||
"metrics": snapshot["metrics"],
|
||
# We format log items for the UI
|
||
"logs": [log.split(":", 1) for log in snapshot["logs"] if ":" in log],
|
||
}
|
||
try:
|
||
await websocket.send_text(json.dumps(state_dict))
|
||
while True:
|
||
# Keep connection alive; discard incoming messages
|
||
await websocket.receive_text()
|
||
except WebSocketDisconnect:
|
||
pass
|
||
except Exception:
|
||
logger.warning("WebSocket connection failed", exc_info=True)
|
||
finally:
|
||
manager.remove_websocket(websocket)
|
||
|
||
|
||
def main() -> None:
|
||
parser = argparse.ArgumentParser(description="YOLO Train Studio Web UI")
|
||
parser.add_argument("--host", default="127.0.0.1", help="Host address to bind to")
|
||
parser.add_argument("--port", type=int, default=8000, help="Port to bind to")
|
||
args = parser.parse_args()
|
||
|
||
# Headless matplotlib
|
||
os.environ["MPLBACKEND"] = "Agg"
|
||
|
||
logger.info(f"Starting server on http://{args.host}:{args.port}")
|
||
uvicorn.run(app, host=args.host, port=args.port)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|