| import os |
| import json |
| import time |
| import threading |
| import subprocess |
| from datetime import datetime |
| from abc import ABC, abstractmethod |
|
|
|
|
| class BaseAgent(ABC): |
| def __init__(self, name, status_file="agents/agent_status.json"): |
| self.name = name |
| self.status_file = status_file |
| self.status = "idle" |
| self.progress = 0 |
| self.message = "" |
| self.start_time = None |
| self.gpu_available = self._check_gpu() |
| self._lock = threading.Lock() |
| self._running = False |
| os.makedirs(os.path.dirname(status_file), exist_ok=True) |
| self._write_status() |
|
|
| def _check_gpu(self): |
| try: |
| result = subprocess.run( |
| ["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader"], |
| capture_output=True, text=True, timeout=5 |
| ) |
| if result.returncode == 0 and result.stdout.strip(): |
| return result.stdout.strip() |
| except Exception: |
| pass |
| return None |
|
|
| def get_gpu_stats(self): |
| try: |
| result = subprocess.run( |
| ["nvidia-smi", "--query-gpu=temperature.gpu,memory.used,memory.total,utilization.gpu", |
| "--format=csv,noheader,nounits"], |
| capture_output=True, text=True, timeout=5 |
| ) |
| if result.returncode == 0: |
| parts = result.stdout.strip().split(",") |
| return { |
| "temp": int(parts[0].strip()), |
| "mem_used": int(parts[1].strip()), |
| "mem_total": int(parts[2].strip()), |
| "util": int(parts[3].strip()) |
| } |
| except Exception: |
| pass |
| return None |
|
|
| def set_status(self, status, progress=None, message=""): |
| with self._lock: |
| self.status = status |
| if progress is not None: |
| self.progress = progress |
| if message: |
| self.message = message |
| if status == "running" and self.start_time is None: |
| self.start_time = time.time() |
| elif status in ("completed", "failed", "idle"): |
| self.start_time = None |
| self._write_status() |
|
|
| def _write_status(self): |
| try: |
| all_status = {} |
| if os.path.exists(self.status_file): |
| with open(self.status_file, "r") as f: |
| all_status = json.load(f) |
|
|
| gpu = self.get_gpu_stats() |
| all_status[self.name] = { |
| "status": self.status, |
| "progress": self.progress, |
| "message": self.message, |
| "gpu_available": self.gpu_available is not None, |
| "gpu_stats": gpu, |
| "start_time": self.start_time, |
| "updated_at": datetime.now().isoformat() |
| } |
|
|
| with open(self.status_file, "w") as f: |
| json.dump(all_status, f, indent=2) |
| except Exception: |
| pass |
|
|
| def log(self, msg): |
| ts = datetime.now().strftime("%H:%M:%S") |
| print(f"[{ts}] [{self.name}] {msg}") |
| self.message = msg |
| self._write_status() |
|
|
| def wait_for_gpu(self, max_temp=80, max_mem_pct=90): |
| self.log(f"Waiting for GPU (max temp: {max_temp}C, max mem: {max_mem_pct}%)...") |
| while True: |
| stats = self.get_gpu_stats() |
| if stats is None: |
| self.log("No GPU stats available, proceeding anyway") |
| return True |
| temp_ok = stats["temp"] < max_temp |
| mem_pct = (stats["mem_used"] / max(1, stats["mem_total"])) * 100 |
| mem_ok = mem_pct < max_mem_pct |
| if temp_ok and mem_ok: |
| self.log(f"GPU ready: {stats['temp']}C, {mem_pct:.0f}% mem used") |
| return True |
| self.log(f"GPU busy: {stats['temp']}C, {mem_pct:.0f}% mem, {stats['util']}% util - waiting 10s") |
| time.sleep(10) |
|
|
| @abstractmethod |
| def run(self, **kwargs): |
| pass |
|
|
| def start(self, **kwargs): |
| self._running = True |
| self.set_status("running", 0, "Starting...") |
| try: |
| self.run(**kwargs) |
| self.set_status("completed", 100, "Done") |
| except Exception as e: |
| self.set_status("failed", self.progress, f"Error: {str(e)}") |
| self.log(f"FAILED: {e}") |
| raise |
| finally: |
| self._running = False |
|
|
| def stop(self): |
| self._running = False |
| self.set_status("stopped", self.progress, "Stopped by user") |
|
|
