monitor_gpu_logger.py

"""
Remote GPU Logger Service
Concurrently monitors remote GPU servers via SSH and logs metrics (VRAM, utilization,
power, and active process user) into a local SQLite database.
"""

import os
import sys
import time
import sqlite3
import argparse
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

# Add repo root to import config
sys.path.insert(0, str(Path(__file__).resolve().parent))
from config import get_config

def init_db(db_file: str):
    """Initialize database and tables if they do not exist."""
    db_dir = os.path.dirname(os.path.abspath(db_file))
    if db_dir and not os.path.exists(db_dir):
        os.makedirs(db_dir, exist_ok=True)
        
    conn = sqlite3.connect(db_file)
    c = conn.cursor()
    
    # 1. Real-time status table (only stores the latest snapshot per GPU)
    c.execute('''CREATE TABLE IF NOT EXISTS gpu_status_latest
                 (ip TEXT, gpu_index TEXT, gpu_name TEXT,
                  memory_used INTEGER, memory_total INTEGER, 
                  utilization INTEGER, power_draw TEXT, power_limit TEXT, top_user TEXT,
                  update_time TEXT,
                  PRIMARY KEY (ip, gpu_index))''')
    
    # 2. Historical records table (stores time-series entries)
    c.execute('''CREATE TABLE IF NOT EXISTS gpu_history
                 (timestamp TEXT, ip TEXT, gpu_index TEXT, gpu_name TEXT,
                  memory_used INTEGER, memory_total INTEGER, 
                  utilization INTEGER, power_draw TEXT, power_limit TEXT, top_user TEXT)''')
    
    conn.commit()
    conn.close()

def connect_ssh(server: dict, timeout: int = 4):
    """Establish SSH connection using paramiko with key or password auth."""
    import paramiko
    
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    
    host = server.get("host", server.get("ip"))
    port = int(server.get("port", 22))
    user = server.get("username", server.get("user"))
    auth_type = server.get("auth_type", "key").lower()
    key_path = server.get("key_path", "")
    password = server.get("password", None)
    
    if auth_type == "key" and key_path:
        key_file = os.path.expanduser(key_path)
        if not os.path.exists(key_file):
            raise FileNotFoundError(f"SSH private key not found at: {key_file}")
            
        # Try RSAKey, then fallback to general key loader
        try:
            pkey = paramiko.RSAKey.from_private_key_file(key_file)
        except Exception:
            try:
                pkey = paramiko.Ed25519Key.from_private_key_file(key_file)
            except Exception:
                pkey = paramiko.AutoAddPolicy()
                # Let paramiko handle key types
                client.connect(host, port=port, username=user, key_filename=key_file, timeout=timeout)
                return client
                
        client.connect(host, port=port, username=user, pkey=pkey, timeout=timeout)
    else:
        # Password authentication
        client.connect(host, port=port, username=user, password=password, timeout=timeout)
        
    return client

def fetch_server_data(server: dict, busy_threshold_mb: int = 1024, timeout: int = 4):
    """Connect to a single remote server via SSH and collect GPU & process metrics."""
    data_points = []
    current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    host = server.get("host", server.get("ip"))
    server_id = server.get("id", host)
    
    try:
        ssh = connect_ssh(server, timeout=timeout)
        
        # 1. Fetch GPU base metrics (including UUID for process correlation)
        smi_cmd = "nvidia-smi --query-gpu=index,uuid,name,memory.used,memory.total,utilization.gpu,power.draw,power.limit --format=csv,noheader,nounits"
        stdin, stdout, stderr = ssh.exec_command(smi_cmd, timeout=timeout)
        gpu_raw = stdout.read().decode('utf-8', errors='ignore').strip()
        
        # 2. Fetch compute app processes by gpu_uuid
        proc_cmd = "nvidia-smi --query-compute-apps=pid,used_memory,gpu_uuid --format=csv,noheader,nounits"
        stdin, stdout, stderr = ssh.exec_command(proc_cmd, timeout=timeout)
        proc_raw = stdout.read().decode('utf-8', errors='ignore').strip()
        
        # Build GPU UUID -> Username mapping
        uuid_user_map = {}
        if proc_raw:
            for line in proc_raw.split('\n'):
                line = line.strip()
                if not line:
                    continue
                parts = [x.strip() for x in line.split(',')]
                if len(parts) >= 3:
                    pid = parts[0]
                    try:
                        mem_used_proc = float(parts[1])
                    except ValueError:
                        mem_used_proc = 0.0
                    g_uuid = parts[2]
                    
                    if mem_used_proc > busy_threshold_mb:
                        found_user = f"PID:{pid}"
                        try:
                            # Query process owner
                            _, out_u, _ = ssh.exec_command(f"ps -p {pid} -o user=", timeout=timeout)
                            user_out = out_u.read().decode('utf-8', errors='ignore').strip()
                            if user_out:
                                found_user = user_out
                        except Exception:
                            pass
                        uuid_user_map[g_uuid] = found_user
        
        ssh.close()

        # Parse GPU data lines
        if gpu_raw:
            for line in gpu_raw.split('\n'):
                parts = [x.strip() for x in line.split(',')]
                if len(parts) >= 8:
                    try:
                        g_idx = parts[0]
                        g_uuid = parts[1]
                        # Clean common vendor prefixes for cleaner display
                        raw_name = parts[2]
                        g_name = raw_name.replace("NVIDIA ", "").replace("GeForce ", "").replace("RTX ", "")
                        mem_used = int(float(parts[3]))
                        mem_total = int(float(parts[4]))
                        util = int(float(parts[5]))
                        power_draw = str(int(float(parts[6]))) 
                        power_limit = str(int(float(parts[7])))
                        
                        # Match user from UUID map
                        user_from_map = uuid_user_map.get(g_uuid)
                        if user_from_map:
                            top_user = user_from_map
                        elif mem_used > busy_threshold_mb:
                            top_user = "Busy"
                        else:
                            top_user = "Idle"
                        
                        data_points.append({
                            "ip": host,
                            "server_id": server_id,
                            "gpu_index": g_idx,
                            "gpu_name": g_name,
                            "memory_used": mem_used,
                            "memory_total": mem_total,
                            "utilization": util,
                            "power_draw": power_draw,
                            "power_limit": power_limit,
                            "top_user": top_user,
                            "update_time": current_time
                        })
                    except Exception as e:
                        print(f"[{datetime.now().strftime('%H:%M:%S')}] [WARN] Parse Error on {host}: {e}")

    except Exception as e:
        print(f"[{datetime.now().strftime('%H:%M:%S')}] [WARN] Connection / Fetch failed on {host} ({server_id}): {e}")
        
    return data_points

def main_loop(config_path: str = None):
    cfg = get_config()
    db_file = cfg["database"]["resolved_path"]
    monitor_cfg = cfg.get("monitor", {})
    servers = cfg.get("servers", [])
    
    refresh_sec = monitor_cfg.get("refresh_interval_sec", 5)
    history_interval_sec = monitor_cfg.get("history_log_interval_sec", 60)
    busy_threshold = monitor_cfg.get("busy_memory_threshold_mb", 1024)
    ssh_timeout = monitor_cfg.get("ssh_timeout_sec", 4)
    max_workers = monitor_cfg.get("max_workers", max(1, len(servers)))
    
    print("=" * 60)
    print("🚀 Remote GPU Logger Service Started")
    print(f"📁 Target Database: {db_file}")
    print(f"🖥️ Monitored Servers: {len(servers)}")
    print(f"⏱️ Refresh interval: {refresh_sec}s | History archive interval: {history_interval_sec}s")
    print("=" * 60)
    
    init_db(db_file)
    last_history_time = 0
    
    while True:
        start_ts = time.time()
        all_data = []
        
        # Concurrent polling across all nodes
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            tasks = [executor.submit(fetch_server_data, s, busy_threshold, ssh_timeout) for s in servers]
            for t in tasks:
                res = t.result()
                if res:
                    all_data.extend(res)
        
        if all_data:
            try:
                conn = sqlite3.connect(db_file)
                c = conn.cursor()
                
                # 1. Update latest snapshot table
                latest_tuples = [
                    (d['ip'], d['gpu_index'], d['gpu_name'], d['memory_used'], d['memory_total'], 
                     d['utilization'], d['power_draw'], d['power_limit'], d['top_user'], d['update_time']) 
                    for d in all_data
                ]
                c.executemany("REPLACE INTO gpu_status_latest VALUES (?,?,?,?,?,?,?,?,?,?)", latest_tuples)
                
                # 2. Archive to history table periodically
                now_ts = time.time()
                if now_ts - last_history_time >= history_interval_sec:
                    history_tuples = [
                        (d['update_time'], d['ip'], d['gpu_index'], d['gpu_name'], 
                         d['memory_used'], d['memory_total'], d['utilization'], 
                         d['power_draw'], d['power_limit'], d['top_user']) 
                        for d in all_data
                    ]
                    c.executemany("INSERT INTO gpu_history VALUES (?,?,?,?,?,?,?,?,?,?)", history_tuples)
                    conn.commit()
                    last_history_time = now_ts
                    print(f"[{datetime.now().strftime('%H:%M:%S')}] [INFO] Snapshot updated + History archived ({len(history_tuples)} GPUs)")
                else:
                    conn.commit()
                
                conn.close()
            except Exception as e:
                print(f"[{datetime.now().strftime('%H:%M:%S')}] [ERROR] Database write error: {e}")
        else:
            # When no servers could be reached
            pass
        
        elapsed = time.time() - start_ts
        sleep_time = max(0, refresh_sec - elapsed)
        time.sleep(sleep_time)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Remote GPU Cluster Logger Service")
    parser.add_argument("--config", default=None, help="Path to custom config.json / config.yaml")
    args = parser.parse_args()
    
    if args.config:
        from config import load_config
        load_config(args.config)
        
    main_loop(args.config)