Codex
269 lines
10 KiB
Python
269 lines
10 KiB
Python
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 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["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)
|
|
return {
|
|
"available": True,
|
|
"agent_online": bool(last_seen and datetime.now(UTC) - last_seen <= ONLINE_WINDOW),
|
|
"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 _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()
|