from __future__ import annotations import json import logging import math import re import threading import time from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Any from PySide6.QtCore import QCoreApplication, QObject, QThread, QTimer, Signal from adam.executor import ToolAdjustmentRequested, ToolCancelled, ToolExecutionError, ToolExecutor from adam.assets import AssetRegistry from adam.atlas import AtlasSupervisor from adam.experiment_tracker import ExperimentStore from adam.models import ExecutionPlan, Job, JobStatus, StepStatus, utc_now from adam.nova import evaluate_job_output from adam.training_assistant import append_preflight_summary def _safe_int(value: Any) -> int: try: return max(0, int(value or 0)) except (TypeError, ValueError): return 0 class JobWorker(QThread): event = Signal(object) def __init__(self, job: Job, executor: ToolExecutor) -> None: super().__init__() self.job = job self.executor = executor self.cancel_event = threading.Event() self.run_event = threading.Event() self.run_event.set() self.adjustment_event = threading.Event() self.adjustment_request: dict[str, Any] = {} def pause(self) -> None: self.run_event.clear() def resume(self) -> None: self.run_event.set() def cancel(self) -> None: self.cancel_event.set() self.run_event.set() def request_adjustment(self, updates: dict[str, Any]) -> None: self.adjustment_request.clear() self.adjustment_request.update(updates) self.adjustment_event.set() self.run_event.set() def run(self) -> None: total_steps = len(self.job.plan.steps) try: for index, step in enumerate(self.job.plan.steps): if self.cancel_event.is_set(): raise ToolCancelled("Job cancelled by user.") self.event.emit( { "type": "step_started", "index": index, "message": step.title, } ) preview_state = {"epoch": 0, "path": ""} last_progress_emit = {"time": 0.0, "overall": -1, "message": ""} progress_samples: list[dict[str, Any]] = [] def on_progress(percent: int, message: str, step_index: int = index, **details: Any) -> None: overall = int(((step_index + percent / 100) / total_steps) * 100) now = time.monotonic() progress_eta = self._estimate_step_eta(details, progress_samples, now) changed = overall != last_progress_emit["overall"] or message != last_progress_emit["message"] terminal = percent >= 100 or overall >= 100 if not terminal and (not changed or now - last_progress_emit["time"] < 0.25): return last_progress_emit.update({"time": now, "overall": overall, "message": message}) self.event.emit( { "type": "progress", "step_percent": percent, "overall": overall, "message": message, **progress_eta, } ) self._discover_external_preview(step, message, preview_state) def on_log(message: str) -> None: self.event.emit({"type": "log", "message": message}) def on_preview(preview: dict[str, Any]) -> None: preview_state["epoch"] = int(preview.get("epoch", 0) or 0) preview_state["path"] = str(preview.get("path", "")) self.event.emit({"type": "preview", **preview}) result = self.executor.execute( step.tool_id, step.arguments, job_id=self.job.id, cancel_event=self.cancel_event, run_event=self.run_event, progress_callback=on_progress, log_callback=on_log, preview_callback=on_preview, adjustment_event=self.adjustment_event, adjustment_request=self.adjustment_request, ) self.event.emit( { "type": "step_finished", "index": index, "result": result, } ) self.event.emit({"type": "completed"}) except ToolCancelled as exc: self.event.emit({"type": "cancelled", "message": str(exc)}) except ToolAdjustmentRequested as exc: self.event.emit({"type": "adjustment_ready", "message": str(exc), **exc.details}) except Exception as exc: self.event.emit( { "type": "failed", "message": str(exc), "exception": type(exc).__name__, } ) def _discover_external_preview( self, step: Any, message: str, state: dict[str, Any] ) -> None: """Discover conventional preview files from any registered trainer.""" if not step.tool_id.endswith("_trainer") or not step.arguments.get("preview_enabled", False): return match = re.search(r"\bepoch\s+(\d+)(?:\s+(?:of|/)|/)?", message, re.I) if not match: return epoch = int(match.group(1)) interval = max(1, int(step.arguments.get("preview_every", 5) or 5)) if epoch % interval or epoch == int(state.get("epoch", 0)): return output = Path(str(step.arguments.get("output_dir", ""))).expanduser() if not output.is_dir(): return try: candidates = [ path for path in output.rglob("*") if path.is_file() and path.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp", ".bmp"} and any(token in path.name.lower() for token in ("preview", "sample", "epoch")) ] latest = max(candidates, key=lambda path: path.stat().st_mtime) if candidates else None except OSError: latest = None if latest and str(latest) != state.get("path"): state["epoch"] = epoch state["path"] = str(latest) self.event.emit({ "type": "preview", "path": str(latest), "epoch": epoch, "next_epoch": min(int(step.arguments.get("epochs", epoch + interval)), epoch + interval), "prompt": str(step.arguments.get("preview_prompt", "")), "seed": step.arguments.get("preview_seed"), "steps": int(step.arguments.get("preview_steps", 0) or 0), }) @staticmethod def _estimate_step_eta( details: dict[str, Any], samples: list[dict[str, Any]], now: float, ) -> dict[str, Any]: """Estimate remaining runtime from real step cadence instead of percent alone.""" current = _safe_int(details.get("current_step", details.get("step", details.get("current")))) total = _safe_int(details.get("total_steps", details.get("total"))) unit = str(details.get("unit", "step") or "step") epoch = _safe_int(details.get("epoch")) total_epochs = _safe_int(details.get("total_epochs")) if details.get("reset_eta"): samples.clear() if (not current or not total) and epoch and total_epochs: current, total, unit = epoch, total_epochs, "epoch" payload: dict[str, Any] = { "progress_current": current, "progress_total": total, "progress_unit": unit, } if not current or not total or current >= total: return payload last = samples[-1] if samples else None if last and ( int(last.get("total", total)) != total or str(last.get("unit", unit)) != unit or current < last["current"] ): # A different tqdm operation or a restarted counter needs a fresh # cadence; carrying the old rate creates wildly incorrect ETAs. samples.clear() last = None if last and current <= last["current"]: return payload samples.append({"time": now, "current": float(current), "total": float(total), "unit": unit}) del samples[:-25] if len(samples) < 2: return payload first = samples[0] elapsed = now - first["time"] completed = current - int(first["current"]) if completed <= 0 or elapsed <= 0: return payload lifetime_seconds_per_unit = elapsed / completed recent = samples[-8:] recent_first = recent[0] recent_completed = current - int(recent_first["current"]) recent_elapsed = now - recent_first["time"] recent_seconds_per_unit = ( recent_elapsed / recent_completed if recent_completed > 0 and recent_elapsed > 0 else lifetime_seconds_per_unit ) seconds_per_unit = (recent_seconds_per_unit * 0.65) + (lifetime_seconds_per_unit * 0.35) remaining = max(0, int(round((total - current) * seconds_per_unit))) if remaining: payload["eta_seconds"] = remaining payload["progress_rate"] = 1 / seconds_per_unit if seconds_per_unit > 0 else 0.0 payload["estimated_completion_at"] = ( datetime.now(timezone.utc) + timedelta(seconds=remaining) ).isoformat() return payload class JobManager(QObject): job_created = Signal(object) job_updated = Signal(object) log_added = Signal(str, str) active_changed = Signal(object) notification = Signal(str, str) def __init__( self, root: Path, executor: ToolExecutor, logger: logging.Logger, config: Any | None = None, ) -> None: super().__init__() self.root = root self.executor = executor self.logger = logger self.config = config if config is not None else {} self.jobs_path = root / "data" / "jobs.json" self.assets = AssetRegistry(root) self.atlas = AtlasSupervisor(config) self.experiments = ExperimentStore(root) self.jobs: list[Job] = [] self._queue: list[str] = [] self._worker: JobWorker | None = None self._active_job: Job | None = None self._last_snapshot = None self._pending_update_job_ids: set[str] = set() self._pending_update_timer_active = False self._load() self._schedule_timer = QTimer(self) self._schedule_timer.timeout.connect(self._release_due_scheduled) if QCoreApplication.instance() is not None: self._schedule_timer.start(15_000) QTimer.singleShot(0, self._release_due_scheduled) if self._queue: QTimer.singleShot(0, self._start_next) @property def active_job(self) -> Job | None: return self._active_job @property def pending_jobs(self) -> list[Job]: return [ job for job in self.jobs if job.status in {JobStatus.AWAITING_CONFIRMATION, JobStatus.SCHEDULED, JobStatus.QUEUED} ] def submit(self, plan: ExecutionPlan, scheduled_for: str | None = None) -> Job: # Every entry point must review a plan before its queue state is chosen. # UI and Remote may prepare it earlier to keep filesystem work off Qt. append_preflight_summary(plan, self.config) is_future = self._is_future(scheduled_for) status = ( JobStatus.AWAITING_CONFIRMATION if plan.requires_confirmation else JobStatus.SCHEDULED if is_future else JobStatus.QUEUED ) job = Job(plan=plan, status=status, scheduled_for=scheduled_for if is_future else None) self.jobs.insert(0, job) self._append_log(job, f"Plan created: {plan.summary}") if plan.requires_confirmation: self._append_log(job, "Waiting for user confirmation.") if is_future: self._append_log(job, f"Requested start time: {self._display_time(scheduled_for)}.") elif is_future: self._append_log(job, f"Scheduled for {self._display_time(scheduled_for)}.") else: self._queue.append(job.id) self._save() self.job_created.emit(job) self.job_updated.emit(job) if not plan.requires_confirmation: self._start_next() return job def confirm(self, job_id: str) -> None: job = self.get(job_id) if job.status != JobStatus.AWAITING_CONFIRMATION: return job.status = JobStatus.SCHEDULED if self._is_future(job.scheduled_for) else JobStatus.QUEUED self._append_log(job, "Plan approved by user.") if job.status == JobStatus.SCHEDULED: self._append_log(job, f"Training will become eligible at {self._display_time(job.scheduled_for)}.") else: self._queue.append(job.id) self._save() self.job_updated.emit(job) self._start_next() def reject(self, job_id: str) -> None: job = self.get(job_id) if job.status != JobStatus.AWAITING_CONFIRMATION: return job.status = JobStatus.CANCELLED job.ended_at = utc_now() self._append_log(job, "Plan cancelled before execution.") self._save() self.job_updated.emit(job) def pause(self, job_id: str) -> None: job = self.get(job_id) if job is self._active_job and job.status == JobStatus.RUNNING and self._worker: self._worker.pause() job.status = JobStatus.PAUSED self._append_log(job, "Job paused.") self._save() self.job_updated.emit(job) def resume(self, job_id: str) -> None: job = self.get(job_id) if job is self._active_job and job.status == JobStatus.PAUSED and self._worker: self._worker.resume() job.status = JobStatus.RUNNING self._append_log(job, "Job resumed.") self._save() self.job_updated.emit(job) def cancel(self, job_id: str) -> None: job = self.get(job_id) if job is self._active_job and self._worker: self._append_log(job, "Cancellation requested.") self._worker.cancel() self._save() self.job_updated.emit(job) return if job.id in self._queue: self._queue.remove(job.id) if job.status in { JobStatus.SCHEDULED, JobStatus.QUEUED, JobStatus.AWAITING_CONFIRMATION, JobStatus.DRAFT, }: job.status = JobStatus.CANCELLED job.ended_at = utc_now() self._append_log(job, "Job cancelled.") self._save() self.job_updated.emit(job) def request_training_adjustment(self, job_id: str, updates: dict[str, Any]) -> None: """Apply safe DDPM settings after the current epoch and resume automatically.""" job = self.get(job_id) if job is not self._active_job or job.status not in {JobStatus.RUNNING, JobStatus.PAUSED} or not self._worker: raise ValueError("Only the active training job can be adjusted.") if not (0 <= job.current_step < len(job.plan.steps)): raise ValueError("The active training step is unavailable.") step = job.plan.steps[job.current_step] if step.tool_id != "ddpm_trainer": raise ValueError("Safe epoch-boundary adjustment currently supports DDPM training.") if step.arguments.get("progressive_stages"): raise ValueError( "Change batch settings before starting a progressive run; each stage manages its own saved handoff." ) allowed = {"batch_size", "training_intensity", "gradient_accumulation_steps"} cleaned = {key: int(value) for key, value in updates.items() if key in allowed} if not cleaned or not 1 <= cleaned.get("batch_size", 1) <= 64 \ or not 10 <= cleaned.get("training_intensity", 100) <= 100 \ or not 1 <= cleaned.get("gradient_accumulation_steps", 1) <= 64: raise ValueError("The requested training settings are outside ADAM's safe range.") previous = {key: step.arguments.get(key) for key in cleaned} if all(previous[key] == value for key, value in cleaned.items()): raise ValueError("Those settings are already active.") job.status = JobStatus.RUNNING self._append_log(job, f"Adjustment queued for the end of this epoch: {cleaned}.") self._worker.request_adjustment(cleaned) self._save() self.job_updated.emit(job) def safer_vram_retry(self, job_id: str) -> Job: """Create a checkpoint-aware DDPM retry with a smaller physical batch.""" original = self.get(job_id) if original.status != JobStatus.FAILED or not self._looks_like_vram_failure(original): raise ValueError("This job did not fail with a recognizable VRAM error.") plan = ExecutionPlan.from_dict(original.to_dict()["plan"]) start_index = max(0, min(original.current_step, len(plan.steps) - 1)) plan.steps = plan.steps[start_index:] step = plan.steps[0] old_batch = max(1, int(step.arguments.get("batch_size", 1))) if old_batch <= 1: raise ValueError("Batch size is already 1; lower resolution or enable other memory-saving options.") original_epochs = max(1, int(step.arguments.get("epochs", 1))) resume_note = self._prepare_ddpm_resume(step.arguments, step.tool_id) remaining_epochs = max(1, int(step.arguments.get("epochs", original_epochs))) completed_epochs = max(0, original_epochs - remaining_epochs) if resume_note else 0 new_batch = max(1, old_batch // 2) old_accumulation = max(1, int(step.arguments.get("gradient_accumulation_steps", 1))) step.arguments["batch_size"] = new_batch step.arguments["gradient_accumulation_steps"] = min(64, old_accumulation * max(1, math.ceil(old_batch / new_batch))) if completed_epochs: step.arguments["completed_epochs"] = completed_epochs for item in plan.steps: item.status = StepStatus.PENDING plan.id = original.plan.id + "-vram-retry" plan.created_at = utc_now() plan.requires_confirmation = True plan.confirmation_reason = "VRAM recovery reduced the physical batch and preserved the effective batch with gradient accumulation." retry = self.submit(plan) self._append_log(retry, f"VRAM recovery changed batch {old_batch} → {new_batch} and gradient accumulation {old_accumulation} → {step.arguments['gradient_accumulation_steps']}.") if resume_note: self._append_log(retry, resume_note) return retry @staticmethod def _looks_like_vram_failure(job: Job) -> bool: text = "\n".join([job.error or "", *job.logs[-100:]]).lower() return any(token in text for token in ("out of memory", "cuda oom", "cuda error: out of memory")) @staticmethod def _is_future(value: str | None) -> bool: if not value: return False try: scheduled = datetime.fromisoformat(value) if scheduled.tzinfo is None: scheduled = scheduled.astimezone() return scheduled.astimezone(timezone.utc) > datetime.now(timezone.utc) except (TypeError, ValueError): return False @staticmethod def _display_time(value: str | None) -> str: try: return datetime.fromisoformat(str(value)).astimezone().strftime("%b %d at %I:%M %p") except ValueError: return str(value or "the requested time") def _release_due_scheduled(self) -> None: released: list[Job] = [] for job in reversed(self.jobs): if job.status == JobStatus.SCHEDULED and not self._is_future(job.scheduled_for): job.status = JobStatus.QUEUED self._queue.append(job.id) self._append_log(job, "Scheduled start time reached; waiting for the training slot.") released.append(job) if not released: return self._save() for job in released: self.job_updated.emit(job) self._start_next() def get(self, job_id: str) -> Job: for job in self.jobs: if job.id == job_id: return job raise KeyError(f"Unknown job: {job_id}") def retry(self, job_id: str) -> Job: """Create an approval-gated retry, resuming interrupted DDPM work when possible.""" original = self.get(job_id) plan = ExecutionPlan.from_dict(original.to_dict()["plan"]) resume_note = "" if original.status == JobStatus.INTERRUPTED: start_index = max(0, min(original.current_step, len(plan.steps) - 1)) plan.steps = plan.steps[start_index:] if plan.steps: step = plan.steps[0] resume_note = self._prepare_ddpm_resume(step.arguments, step.tool_id) for step in plan.steps: step.status = StepStatus.PENDING plan.id = original.plan.id + "-retry" plan.created_at = utc_now() plan.requires_confirmation = True plan.confirmation_reason = ( "This retries a previous job. Review paths, checkpoints, and settings " "because files or available resources may have changed." ) retry = self.submit(plan) if resume_note: self._append_log(retry, resume_note) return retry @staticmethod def _prepare_ddpm_resume(arguments: dict[str, Any], tool_id: str) -> str: """Attach the newest complete Accelerate checkpoint to a DDPM retry.""" if tool_id != "ddpm_trainer": return "" output = Path(str(arguments.get("output_dir", ""))).expanduser() dataset = Path(str(arguments.get("dataset_dir", ""))).expanduser() try: checkpoints = sorted( (path for path in output.glob("checkpoint-*") if path.is_dir() and (path / "unet" / "diffusion_pytorch_model.safetensors").is_file() and (path / "optimizer.bin").is_file() and (path / "scheduler.bin").is_file()), key=lambda path: int(path.name.rsplit("-", 1)[-1]), ) image_count = sum(1 for item in dataset.iterdir() if item.is_file() and item.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp", ".bmp"}) batch_size = max(1, int(arguments.get("batch_size", 1))) completed_epochs = int(checkpoints[-1].name.rsplit("-", 1)[-1]) // max(1, math.ceil(image_count / batch_size)) remaining_epochs = int(arguments.get("epochs", 0)) - completed_epochs except (IndexError, OSError, ValueError, TypeError): return "" if remaining_epochs <= 0: return "" arguments["resume_from"] = str(checkpoints[-1]) arguments["epochs"] = remaining_epochs return f"Resuming interrupted DDPM training from {checkpoints[-1].name} (about epoch {completed_epochs}; {remaining_epochs} epochs remaining)." def end_task(self, job_id: str) -> bool: """Acknowledge an interrupted job and leave it inactive in history. This is intentionally limited to interrupted jobs: active work must still use ``cancel`` so its worker receives the cancellation signal. """ job = self.get(job_id) if job.status != JobStatus.INTERRUPTED: return False job.status = JobStatus.CANCELLED job.ended_at = utc_now() job.logs = [ line for line in job.logs if not line.startswith("[startup] Previous session ended before this job.") ] self._append_log(job, "Interrupted job ended by user; no retry is pending.") self._save() self.job_updated.emit(job) return True def remove_completed_or_failed(self) -> int: """Remove completed and failed history records without touching output files.""" removable = {JobStatus.FINISHED, JobStatus.FAILED} before = len(self.jobs) self.jobs = [job for job in self.jobs if job.status not in removable] removed = before - len(self.jobs) if removed: self._save() return removed def _start_next(self) -> None: if self._worker is not None and self._worker.isRunning(): return while self._queue: job_id = self._queue.pop(0) job = self.get(job_id) if job.status != JobStatus.QUEUED: continue self._active_job = job job.status = JobStatus.RUNNING job.started_at = utc_now() self._append_log(job, f"Job {job.id} started.") self._worker = JobWorker(job, self.executor) self._worker.event.connect(self._handle_event) self._worker.finished.connect(self._worker_finished) self._save() self.active_changed.emit(job) self.job_updated.emit(job) self.notification.emit("Training started" if self._has_training(job) else "Job started", job.plan.project_name) self._worker.start() return self._active_job = None self.active_changed.emit(None) def _handle_event(self, event: dict[str, Any]) -> None: job = self._active_job if job is None: return event_type = event.get("type") if event_type == "step_started": index = int(event["index"]) job.current_step = index job.plan.steps[index].status = StepStatus.RUNNING if job.plan.steps[index].tool_id.endswith("_trainer"): job.preview_path = None job.preview_epoch = 0 job.preview_next_epoch = 0 job.preview_prompt = "" job.preview_seed = None job.preview_steps = 0 job.preview_kind = "training" job.preview_current = 0 job.preview_total = 0 job.preview_image_index = 0 job.preview_image_count = 0 job.eta_seconds = None job.estimated_completion_at = None job.progress_current = 0 job.progress_total = 0 job.progress_rate = 0.0 job.progress_unit = "step" self._append_log(job, f"Starting: {job.plan.steps[index].title}") elif event_type == "progress": job.progress = int(event["overall"]) job.eta_seconds = _safe_int(event.get("eta_seconds")) or None job.estimated_completion_at = str(event.get("estimated_completion_at") or "") or None job.progress_current = _safe_int(event.get("progress_current")) job.progress_total = _safe_int(event.get("progress_total")) job.progress_rate = float(event.get("progress_rate", 0.0) or 0.0) job.progress_unit = str(event.get("progress_unit", "step") or "step") message = str(event["message"]) if message and (not job.logs or message not in job.logs[-1]): self._append_log(job, message) self._schedule_job_update(job) return elif event_type == "log": self._append_log(job, str(event["message"])) self._schedule_job_update(job) return elif event_type == "preview": job.preview_path = str(event.get("path", "")) or None job.preview_epoch = int(event.get("epoch", 0) or 0) job.preview_next_epoch = int(event.get("next_epoch", 0) or 0) job.preview_prompt = str(event.get("prompt", "")) seed = event.get("seed") job.preview_seed = int(seed) if seed is not None else None job.preview_steps = int(event.get("steps", 0) or 0) job.preview_kind = str(event.get("kind", "training")) job.preview_current = int(event.get("current", 0) or 0) job.preview_total = int(event.get("total", 0) or 0) job.preview_image_index = int(event.get("image_index", 0) or 0) job.preview_image_count = int(event.get("image_count", 0) or 0) label = "Denoising" if job.preview_kind == "generation" else "Training" position = f" step {job.preview_current}" if job.preview_current else f" epoch {job.preview_epoch}" self._append_log(job, f"{label} preview updated at{position}.") self._schedule_job_update(job) return elif event_type == "step_finished": index = int(event["index"]) job.plan.steps[index].status = StepStatus.FINISHED result = event.get("result") or {} self.assets.ingest_result(result) if result.get("output_folder"): job.output_folder = str(result["output_folder"]) if job.plan.steps[index].tool_id.endswith("_trainer"): evaluation = evaluate_job_output(job) if evaluation: evaluation["step"] = job.plan.steps[index].title evaluation["model_name"] = str(result.get("model_name", "")) reports = list(job.nova_report.get("evaluations", [])) reports.append(evaluation) job.nova_report = { "agent": "NOVA", "evaluations": reports, "latest": evaluation, } self._append_log( job, f"NOVA — {evaluation['status']}: {evaluation['summary']}", ) self._append_log(job, f"Finished: {job.plan.steps[index].title}") elif event_type == "completed": job.status = JobStatus.FINISHED job.progress = 100 job.eta_seconds = 0 job.estimated_completion_at = utc_now() job.ended_at = utc_now() self._append_log(job, "Job finished successfully.") self._record_experiment(job) self.notification.emit("Job complete", job.plan.project_name) elif event_type == "cancelled": job.status = JobStatus.CANCELLED job.ended_at = utc_now() self._append_log(job, str(event.get("message", "Job cancelled."))) for step in job.plan.steps: if step.status == StepStatus.RUNNING: step.status = StepStatus.SKIPPED self._record_experiment(job) self.notification.emit("Job cancelled", job.plan.project_name) elif event_type == "adjustment_ready": index = max(0, min(job.current_step, len(job.plan.steps) - 1)) remaining_steps = job.plan.steps[index:] step = remaining_steps[0] updates = dict(event.get("updates") or {}) step.arguments.update(updates) checkpoint = str(event.get("checkpoint", "")) completed_epochs = max(0, int(event.get("completed_epochs", 0) or 0)) total_epochs = max(1, int(step.arguments.get("epochs", 1))) step.arguments["epochs"] = max(1, total_epochs - completed_epochs) step.arguments["completed_epochs"] = completed_epochs if checkpoint: step.arguments["resume_from"] = checkpoint for pending in remaining_steps: pending.status = StepStatus.PENDING job.plan.steps = remaining_steps job.current_step = -1 job.status = JobStatus.QUEUED job.eta_seconds = None job.estimated_completion_at = None self._queue.insert(0, job.id) self._append_log(job, f"Epoch {completed_epochs} checkpoint is complete. Restarting with {updates}.") self.notification.emit("Training settings ready", "Restarting from the completed epoch checkpoint.") elif event_type == "failed": job.status = JobStatus.FAILED job.ended_at = utc_now() job.error = str(event.get("message", "Unknown error")) if 0 <= job.current_step < len(job.plan.steps): job.plan.steps[job.current_step].status = StepStatus.FAILED self._append_log(job, f"Failed: {job.error}") self.logger.error( "Job %s failed (%s): %s", job.id, event.get("exception"), job.error, ) self._record_experiment(job) self.notification.emit("Job failed", job.error) self._save() self.job_updated.emit(job) def _schedule_job_update(self, job: Job) -> None: self._pending_update_job_ids.add(job.id) if QCoreApplication.instance() is None: self._flush_pending_job_updates() return if self._pending_update_timer_active: return self._pending_update_timer_active = True QTimer.singleShot(300, self._flush_pending_job_updates) def _flush_pending_job_updates(self) -> None: if not self._pending_update_job_ids: self._pending_update_timer_active = False return pending_ids = list(self._pending_update_job_ids) self._pending_update_job_ids.clear() self._pending_update_timer_active = False self._save() for job_id in pending_ids: try: self.job_updated.emit(self.get(job_id)) except KeyError: continue def _worker_finished(self) -> None: self._worker = None self._active_job = None self.active_changed.emit(None) self._start_next() def supervise(self, snapshot: Any) -> None: """Let ATLAS inspect the active training run and apply critical pauses.""" self._last_snapshot = snapshot job = self._active_job if job is None or job.status != JobStatus.RUNNING: return decision = self.atlas.observe(job, snapshot) previous = str(job.atlas_report.get("message", "")) unchanged = ( previous == decision.message and job.atlas_report.get("severity") == decision.severity and job.atlas_report.get("action") == decision.action ) if unchanged: return job.atlas_report = { "agent": "ATLAS", "severity": decision.severity, "message": decision.message, "action": decision.action, "updated_at": utc_now(), } if decision.message != previous and decision.severity in {"warning", "critical"}: self._append_log(job, f"ATLAS {decision.severity.upper()} — {decision.message}") self.notification.emit(f"ATLAS {decision.severity}", decision.message) self._save() self.job_updated.emit(job) if decision.action == "pause" and job.status == JobStatus.RUNNING: self.pause(job.id) def _record_experiment(self, job: Job) -> None: if not self._has_training(job): return try: self.experiments.record_job(job, self._last_snapshot) except Exception as exc: self.logger.warning("Experiment tracking failed for %s: %s", job.id, exc) def _append_log(self, job: Job, message: str) -> None: timestamp = datetime.now().strftime("%H:%M:%S") line = f"[{timestamp}] {message}" job.logs.append(line) job.logs = job.logs[-1000:] self.log_added.emit(job.id, line) self.logger.info("Job %s | %s", job.id, message) @staticmethod def _has_training(job: Job) -> bool: return any(step.tool_id.endswith("trainer") for step in job.plan.steps) def _load(self) -> None: if not self.jobs_path.exists(): return try: payload = json.loads(self.jobs_path.read_text(encoding="utf-8")) self.jobs = [Job.from_dict(item) for item in payload.get("jobs", [])] except (OSError, ValueError, TypeError, KeyError, json.JSONDecodeError): self.jobs = [] return for job in self.jobs: if job.status in {JobStatus.RUNNING, JobStatus.PAUSED}: job.status = JobStatus.INTERRUPTED job.ended_at = utc_now() job.logs.append( "[startup] Previous session ended before this job. " "Review it before retrying." ) elif job.status == JobStatus.QUEUED: self._queue.append(job.id) if not any("Queued job restored" in line for line in job.logs[-5:]): job.logs.append("[startup] Queued job restored and will run when ADAM is ready.") self._save() def _save(self) -> None: self.jobs_path.parent.mkdir(parents=True, exist_ok=True) temporary = self.jobs_path.with_suffix(".tmp") temporary.write_text( json.dumps({"jobs": [job.to_dict() for job in self.jobs]}, indent=2), encoding="utf-8", ) temporary.replace(self.jobs_path) def shutdown(self) -> None: if self._worker is not None and self._worker.isRunning(): self._worker.cancel() self._worker.wait(2500) if self._active_job and self._active_job.status in { JobStatus.RUNNING, JobStatus.PAUSED, }: self._active_job.status = JobStatus.CANCELLED self._active_job.ended_at = utc_now() self._append_log(self._active_job, "ADAM closed; the active job was stopped.") self._save()