adam/atlas.py

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)