File size: 4,405 Bytes
e0265b9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | 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)
|