from __future__ import annotations import base64 import hashlib import json import os import uuid from datetime import UTC from datetime import datetime from datetime import timedelta from pathlib import Path from threading import Lock from typing import Any ALLOWED_TRAINING_ARTIFACTS = { "lstm_forecaster.json", "torch_retrain_guard.json", "torch_threshold_calibration.json", } RUNNING_TIMEOUT = timedelta(hours=12) ONLINE_WINDOW = timedelta(minutes=3) class TrainingCoordinator: def __init__(self, runtime_dir: Path) -> None: self.runtime_dir = runtime_dir self.state_path = runtime_dir / "training_coordination.json" self.upload_root = runtime_dir / ".training_uploads" self._lock = Lock() def status(self) -> dict[str, Any]: with self._lock: state = self._load_state() self._expire_stale_jobs(state) self._save_state(state) return self._public_status(state) def request_retrain(self, payload: dict[str, Any] | None = None) -> dict[str, Any]: payload = payload or {} with self._lock: state = self._load_state() self._expire_stale_jobs(state) existing = self._active_job(state) if existing is not None: self._save_state(state) return {"queued": False, "reason": "active_job_exists", "job": existing, "status": self._public_status(state)} now = _now() job = { "id": str(uuid.uuid4()), "status": "pending", "requested_at": now, "requested_by": str(payload.get("source") or "api"), "parameters": _safe_parameters(payload.get("parameters")), "message": "", "artifacts": [], } state.setdefault("jobs", []).append(job) self._trim_jobs(state) self._save_state(state) return {"queued": True, "job": job, "status": self._public_status(state)} def heartbeat(self, payload: dict[str, Any] | None = None) -> dict[str, Any]: payload = payload or {} with self._lock: state = self._load_state() worker = self._worker_from_payload(payload) state["worker"] = worker self._save_state(state) return {"ok": True, "worker": worker, "status": self._public_status(state)} def claim(self, payload: dict[str, Any] | None = None) -> dict[str, Any]: payload = payload or {} with self._lock: state = self._load_state() self._expire_stale_jobs(state) worker = self._worker_from_payload(payload) state["worker"] = worker job = self._oldest_pending_job(state) if job is None: self._save_state(state) return {"claimed": False, "job": None, "status": self._public_status(state)} now = _now() job["status"] = "running" job["claimed_at"] = now job["claimed_by"] = worker["id"] job["worker"] = worker self._save_state(state) return {"claimed": True, "job": job, "status": self._public_status(state)} def save_artifact_chunk(self, job_id: str, payload: dict[str, Any]) -> dict[str, Any]: name = Path(str(payload.get("name") or "")).name if name not in ALLOWED_TRAINING_ARTIFACTS: raise ValueError(f"artifact is not allowed: {name}") index = int(payload.get("index", -1)) total = int(payload.get("total", 0)) sha256 = str(payload.get("sha256") or "").strip().lower() if index < 0 or total <= 0 or index >= total: raise ValueError("invalid artifact chunk index") if not sha256: raise ValueError("artifact sha256 is required") try: chunk = base64.b64decode(str(payload.get("data_base64") or ""), validate=True) except (ValueError, TypeError) as exc: raise ValueError("invalid artifact chunk payload") from exc chunk_dir = self.upload_root / job_id / name chunk_dir.mkdir(parents=True, exist_ok=True) (chunk_dir / f"{index:06d}.part").write_bytes(chunk) if not all((chunk_dir / f"{part:06d}.part").is_file() for part in range(total)): return {"complete": False, "received": index + 1, "total": total} target_tmp = self.runtime_dir / f".{name}.{job_id}.tmp" digest = hashlib.sha256() with target_tmp.open("wb") as output: for part in range(total): data = (chunk_dir / f"{part:06d}.part").read_bytes() digest.update(data) output.write(data) if digest.hexdigest().lower() != sha256: target_tmp.unlink(missing_ok=True) raise ValueError("artifact sha256 mismatch") self.runtime_dir.mkdir(parents=True, exist_ok=True) os.replace(target_tmp, self.runtime_dir / name) _remove_tree(chunk_dir) with self._lock: state = self._load_state() job = self._job_by_id(state, job_id) if job is not None: artifacts = job.setdefault("artifacts", []) artifacts = [item for item in artifacts if item.get("name") != name] artifacts.append({"name": name, "sha256": sha256, "uploaded_at": _now()}) job["artifacts"] = artifacts self._save_state(state) return {"complete": True, "name": name, "sha256": sha256} def progress(self, job_id: str, payload: dict[str, Any] | None = None) -> dict[str, Any]: payload = payload or {} with self._lock: state = self._load_state() job = self._job_by_id(state, job_id) if job is None: raise ValueError(f"training job not found: {job_id}") if isinstance(payload.get("worker"), dict): state["worker"] = self._worker_from_payload(payload["worker"]) job["status"] = str(payload.get("status") or job.get("status") or "running") job["phase"] = str(payload.get("phase") or job.get("phase") or "running") job["message"] = str(payload.get("message") or job.get("message") or "") job["progress_percent"] = _coerce_percent(payload.get("progress_percent"), job.get("progress_percent", 0)) job["updated_at"] = _now() if isinstance(payload.get("details"), dict): job["details"] = payload["details"] self._save_state(state) return {"ok": True, "job": job, "status": self._public_status(state)} def complete(self, job_id: str, payload: dict[str, Any] | None = None) -> dict[str, Any]: payload = payload or {} with self._lock: state = self._load_state() job = self._job_by_id(state, job_id) if job is None: raise ValueError(f"training job not found: {job_id}") success = bool(payload.get("success", payload.get("status") == "completed")) job["status"] = "completed" if success else "failed" job["phase"] = "completed" if success else "failed" job["progress_percent"] = 100 if success else _coerce_percent(payload.get("progress_percent"), job.get("progress_percent", 0)) job["completed_at"] = _now() job["message"] = str(payload.get("message") or "") if isinstance(payload.get("summary"), dict): job["summary"] = payload["summary"] self._save_state(state) return {"ok": True, "job": job, "status": self._public_status(state)} def _load_state(self) -> dict[str, Any]: try: data = json.loads(self.state_path.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError): data = {} if not isinstance(data, dict): data = {} data.setdefault("jobs", []) return data def _save_state(self, state: dict[str, Any]) -> None: self.runtime_dir.mkdir(parents=True, exist_ok=True) tmp = self.state_path.with_suffix(".tmp") tmp.write_text(json.dumps(state, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") os.replace(tmp, self.state_path) def _worker_from_payload(self, payload: dict[str, Any]) -> dict[str, Any]: return { "id": str(payload.get("worker_id") or payload.get("id") or "windows-training-host"), "name": str(payload.get("name") or "DESKTOP-TMFDL0H"), "path": str(payload.get("path") or "C:\\Repos\\TradeBot"), "version": str(payload.get("version") or "1"), "last_seen_at": _now(), } def _public_status(self, state: dict[str, Any]) -> dict[str, Any]: worker = state.get("worker") if isinstance(state.get("worker"), dict) else {} last_seen = _parse_time(str(worker.get("last_seen_at") or "")) active = self._active_job(state) latest = _latest_job(state) recently_seen = bool(last_seen and datetime.now(UTC) - last_seen <= ONLINE_WINDOW) agent_busy = bool( active and active.get("status") == "running" and worker and active.get("claimed_by") == worker.get("id") ) return { "available": True, "agent_online": recently_seen or agent_busy, "agent_recently_seen": recently_seen, "agent_busy": agent_busy, "worker": worker, "active_job": active, "latest_job": latest, "pending_jobs": sum(1 for job in state.get("jobs", []) if job.get("status") == "pending"), } def _active_job(self, state: dict[str, Any]) -> dict[str, Any] | None: for job in reversed(state.get("jobs", [])): if job.get("status") in {"pending", "running"}: return job return None def _oldest_pending_job(self, state: dict[str, Any]) -> dict[str, Any] | None: for job in state.get("jobs", []): if job.get("status") == "pending": return job return None def _job_by_id(self, state: dict[str, Any], job_id: str) -> dict[str, Any] | None: for job in state.get("jobs", []): if job.get("id") == job_id: return job return None def _expire_stale_jobs(self, state: dict[str, Any]) -> None: now = datetime.now(UTC) for job in state.get("jobs", []): if job.get("status") != "running": continue claimed_at = _parse_time(str(job.get("claimed_at") or "")) if claimed_at and now - claimed_at > RUNNING_TIMEOUT: job["status"] = "failed" job["completed_at"] = _now() job["message"] = "training worker timeout" def _trim_jobs(self, state: dict[str, Any]) -> None: jobs = state.get("jobs", []) if isinstance(jobs, list) and len(jobs) > 30: state["jobs"] = jobs[-30:] def _safe_parameters(value: Any) -> dict[str, Any]: if not isinstance(value, dict): return {} allowed = {"symbols", "limit", "lookbacks", "architectures", "hidden_sizes", "layers", "dropouts", "epochs"} return {key: value[key] for key in allowed if key in value} def _latest_job(state: dict[str, Any]) -> dict[str, Any] | None: jobs = state.get("jobs", []) if not jobs: return None latest = jobs[-1] return latest if isinstance(latest, dict) else None def _now() -> str: return datetime.now(UTC).isoformat(timespec="seconds") def _parse_time(value: str) -> datetime | None: if not value: return None try: return datetime.fromisoformat(value.replace("Z", "+00:00")) except ValueError: return None def _coerce_percent(value: Any, default: Any = 0) -> int: try: number = int(float(value)) except (TypeError, ValueError): try: number = int(float(default)) except (TypeError, ValueError): number = 0 return max(0, min(number, 100)) def _remove_tree(path: Path) -> None: if not path.exists(): return for child in path.iterdir(): if child.is_dir(): _remove_tree(child) else: child.unlink(missing_ok=True) path.rmdir()