File size: 5,314 Bytes
e0265b9 c61c435 e0265b9 c61c435 e0265b9 c61c435 e0265b9 c61c435 e0265b9 c61c435 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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | from __future__ import annotations
import math
import re
import shutil
import time
from pathlib import Path
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)
disk_known = bool(snapshot.disk_total_gb)
disk_label = "monitored drive"
output = job.plan.steps[job.current_step].arguments.get("output_dir") or job.output_folder
disk_error = False
if output:
disk_label = "output drive"
try:
target = Path(str(output)).expanduser().resolve()
while not target.exists() and target != target.parent:
target = target.parent
usage = shutil.disk_usage(target)
free_disk = usage.free / (1024 ** 3)
disk_known = True
except (OSError, ValueError):
disk_known = False
disk_error = True
state["disk_samples"] = state["disk_samples"] + 1 if disk_known 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 {disk_label}. 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 disk_error:
return AtlasDecision("warning", "ATLAS could not check the output drive's free space. Check that the output location is available.")
if disk_known and free_disk <= self.warning_disk_gb:
return AtlasDecision("warning", f"Space on the {disk_label} 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)
|