| from __future__ import annotations |
|
|
| import math |
| import re |
| import time |
| from dataclasses import dataclass |
| from datetime import datetime |
| from typing import Any |
|
|
| from adam.models import Job, SystemSnapshot |
|
|
|
|
| @dataclass(slots=True) |
| class AtlasDecision: |
| severity: str |
| message: str |
| action: str = "none" |
|
|
|
|
| class AtlasSupervisor: |
| """Stateful, conservative runtime guard for active training jobs.""" |
|
|
| def __init__(self, config: Any | None = None) -> None: |
| getter = config.get if config is not None else lambda _key, default: default |
| self.warning_temp = float(getter("atlas_warning_temperature", 82)) |
| self.critical_temp = float(getter("atlas_critical_temperature", 90)) |
| self.warning_disk_gb = float(getter("atlas_warning_disk_gb", 10)) |
| self.critical_disk_gb = float(getter("atlas_critical_disk_gb", 1)) |
| self.stall_minutes = float(getter("atlas_stall_minutes", 30)) |
| self._state: dict[str, dict[str, Any]] = {} |
|
|
| @staticmethod |
| def _is_training(job: Job) -> bool: |
| return 0 <= job.current_step < len(job.plan.steps) and job.plan.steps[job.current_step].tool_id.endswith("_trainer") |
|
|
| def observe(self, job: Job, snapshot: SystemSnapshot, *, now: float | None = None) -> AtlasDecision: |
| if not self._is_training(job): |
| return AtlasDecision("info", "ATLAS is standing by; the active step is not training.") |
| moment = time.monotonic() if now is None else now |
| state = self._state.setdefault(job.id, { |
| "progress": job.progress, "changed_at": moment, "hot_samples": 0, |
| "disk_samples": 0, "log_index": 0, |
| }) |
| if job.progress != state["progress"]: |
| state["progress"] = job.progress |
| state["changed_at"] = moment |
|
|
| new_logs = job.logs[int(state["log_index"]):] |
| state["log_index"] = len(job.logs) |
| recent = "\n".join(new_logs) |
| if re.search(r"\bloss\s*[:=]?\s*(?:nan|[+-]?inf)\b", recent, re.I): |
| return AtlasDecision("critical", "Non-finite loss detected. ATLAS paused the job for review.", "pause") |
|
|
| temperature = snapshot.gpu_temperature |
| state["hot_samples"] = state["hot_samples"] + 1 if temperature is not None and temperature >= self.critical_temp else 0 |
| if state["hot_samples"] >= 3: |
| return AtlasDecision("critical", f"GPU temperature remained at {temperature:.0f}°C. ATLAS paused the job.", "pause") |
|
|
| free_disk = max(0.0, snapshot.disk_total_gb - snapshot.disk_used_gb) |
| state["disk_samples"] = state["disk_samples"] + 1 if snapshot.disk_total_gb and free_disk <= self.critical_disk_gb else 0 |
| if state["disk_samples"] >= 2: |
| return AtlasDecision("critical", f"Only {free_disk:.1f} GB remains on the output drive. ATLAS paused the job.", "pause") |
|
|
| stalled_minutes = (moment - state["changed_at"]) / 60 |
| if stalled_minutes >= self.stall_minutes and snapshot.gpu_percent < 5: |
| return AtlasDecision("warning", f"No recorded progress and little GPU activity for {stalled_minutes:.0f} minutes. Check the trainer.") |
| if temperature is not None and temperature >= self.warning_temp: |
| return AtlasDecision("warning", f"GPU temperature is elevated at {temperature:.0f}°C; ATLAS is watching it closely.") |
| if snapshot.disk_total_gb and free_disk <= self.warning_disk_gb: |
| return AtlasDecision("warning", f"Output drive space is getting low ({free_disk:.1f} GB free).") |
| if snapshot.memory_percent >= 95: |
| return AtlasDecision("warning", f"System memory usage is very high at {snapshot.memory_percent:.0f}%.") |
|
|
| try: |
| started = datetime.fromisoformat(job.started_at) if job.started_at else None |
| elapsed_minutes = (datetime.now(started.tzinfo) - started).total_seconds() / 60 if started else 0 |
| except ValueError: |
| elapsed_minutes = 0 |
| expected = float(job.plan.orion_review.get("estimated_high_minutes", 0) or 0) |
| if expected and elapsed_minutes > expected * 2 and job.progress < 90: |
| return AtlasDecision("warning", "Runtime is now more than twice ORION's broad estimate. The job is still running.") |
| return AtlasDecision("healthy", "Training behavior is within the current safety limits.") |
|
|
| def forget(self, job_id: str) -> None: |
| self._state.pop(job_id, None) |
|
|
