a13labs.infra/ansible/roles/auto_suspend/files/auto_suspend.py

#!/usr/bin/env python3
"""
Auto-suspend daemon — checks system idleness and suspends when idle.

Uses an exponential moving average (EMA) of CPU idle % to smooth out
transient spikes. GPU utilization is tracked as a binary busy/idle value
per cycle: if any GPU exceeds the threshold, the idle counter resets.
The system suspends only after idle_timeout_cycles consecutive idle checks.

State is persisted to /var/lib/auto_suspend/state.json and survives
across timer invocations. State is fully reset on suspend resume.
"""

import argparse
import json
import logging
import logging.handlers
import os
import signal
import subprocess
import sys
import tempfile
from pathlib import Path

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

DEFAULT_STATE_DIR = "/var/lib/auto_suspend"
DEFAULT_LOG_FILE = "/var/log/auto_suspend.log"
DEFAULT_ALPHA = 0.3
DEFAULT_MIN_SAMPLES = 3
DEFAULT_IDLE_TIMEOUT_CYCLES = 4
DEFAULT_CPU_IDLE_THRESHOLD = 90.0
DEFAULT_GPU_UTIL_THRESHOLD = 5.0

STATE_FILE_NAME = "state.json"

# ---------------------------------------------------------------------------
# Logging setup
# ---------------------------------------------------------------------------

def setup_logging(log_file, debug=False):
    """Configure logging to both file and stderr."""
    logger = logging.getLogger("auto_suspend")
    logger.setLevel(logging.DEBUG if debug else logging.INFO)

    formatter = logging.Formatter(
        "%(asctime)s [%(levelname)s] %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )

    # File handler with rotation (1 MB, 3 backups)
    log_path = Path(log_file)
    log_path.parent.mkdir(parents=True, exist_ok=True)
    file_handler = logging.handlers.RotatingFileHandler(
        str(log_path),
        maxBytes=1 * 1024 * 1024,
        backupCount=3,
    )
    file_handler.setLevel(logging.DEBUG)
    file_handler.setFormatter(formatter)
    logger.addHandler(file_handler)

    # Stderr handler (visible in journalctl -u auto-suspend.service)
    stderr_handler = logging.StreamHandler(sys.stderr)
    stderr_handler.setLevel(logging.DEBUG)
    stderr_handler.setFormatter(formatter)
    logger.addHandler(stderr_handler)

    return logger

# ---------------------------------------------------------------------------
# State persistence
# ---------------------------------------------------------------------------

def fresh_state():
    """Return a fresh state dict with all values at defaults."""
    return {
        "cpu_idle_n": None,
        "cpu_idle_n1": None,
        "cpu_avg": None,
        "gpu_idle_counter": 0,
        "idle_counter": 0,
        "sample_count": 0,
    }


def load_state(state_dir, logger):
    """Load state from disk, or return fresh state on any error."""
    state_path = Path(state_dir) / STATE_FILE_NAME
    if not state_path.exists():
        logger.debug("No state file found at %s — starting fresh", state_path)
        return fresh_state()

    try:
        with open(state_path) as f:
            state = json.load(f)
        # Validate required keys
        for key in fresh_state():
            if key not in state:
                state[key] = fresh_state()[key]
        logger.debug("Loaded state: sample_count=%d, idle_counter=%d, gpu_idle=%d",
                      state["sample_count"], state["idle_counter"],
                      state.get("gpu_idle_counter", 0))
        return state
    except (json.JSONDecodeError, IOError, OSError) as exc:
        logger.warning("Failed to load state (%s) — starting fresh", exc)
        return fresh_state()


def save_state(state_dir, state, logger):
    """Atomically write state to disk (tempfile + os.replace)."""
    state_path = Path(state_dir) / STATE_FILE_NAME
    state_path.parent.mkdir(parents=True, exist_ok=True)

    fd, tmp = tempfile.mkstemp(
        dir=str(state_path.parent), suffix=".tmp", prefix=".state_"
    )
    try:
        with os.fdopen(fd, "w") as f:
            json.dump(state, f, indent=2)
            f.write("\n")
        os.replace(tmp, str(state_path))
        logger.debug("State saved: sample_count=%d, idle_counter=%d, gpu_idle=%d",
                      state["sample_count"], state["idle_counter"],
                      state.get("gpu_idle_counter", 0))
    except BaseException:
        try:
            os.unlink(tmp)
        except OSError:
            pass
        raise

# ---------------------------------------------------------------------------
# Measurement functions
# ---------------------------------------------------------------------------

def measure_cpu(logger):
    """
    Measure CPU idle percentage using mpstat.

    Runs ``mpstat 1 1`` and extracts the last field of the "Average:" line,
    which is ``%idle``.  Returns a float, or None on failure.

    Uses ``LC_ALL=C`` to force English locale output regardless of system
    locale (avoids issues like Portuguese ``Média:`` with comma decimals).
    """
    env = dict(os.environ)
    env["LC_ALL"] = "C"
    try:
        result = subprocess.run(
            ["mpstat", "1", "1"],
            capture_output=True, text=True, timeout=15,
            env=env,
        )
        if result.returncode != 0:
            logger.warning("mpstat returned non-zero: %s", result.stderr.strip())
            return None

        for line in result.stdout.strip().splitlines():
            line = line.strip()
            if line.startswith("Average:") or line.startswith("all"):
                fields = line.split()
                # Find %idle column index from header
                idle_pct = float(fields[-1])
                logger.debug("CPU idle: %.1f%%", idle_pct)
                return idle_pct

        logger.warning("mpstat output did not contain expected data line")
        return None
    except subprocess.TimeoutExpired:
        logger.warning("mpstat timed out")
        return None
    except (ValueError, IndexError) as exc:
        logger.warning("Failed to parse mpstat output: %s", exc)
        return None


def measure_gpu_nvidia(logger):
    """
    Measure GPU utilization via nvidia-smi.

    Returns the **maximum** utilization across all GPUs (any busy GPU means
    the system is not idle). Returns None if nvidia-smi is unavailable.
    """
    try:
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=utilization.gpu",
             "--format=csv,noheader,nounits"],
            capture_output=True, text=True, timeout=15,
        )
        if result.returncode != 0:
            # nvidia-smi may fail if no NVIDIA GPU is present
            logger.debug("nvidia-smi returned %d — no NVIDIA GPU", result.returncode)
            return None

        utils = []
        for line in result.stdout.strip().splitlines():
            line = line.strip()
            if line:
                utils.append(float(line))

        if not utils:
            return None

        max_util = max(utils)
        logger.info("NVIDIA GPU utilization: max=%.1f%%", max_util)
        return max_util
    except FileNotFoundError:
        logger.debug("nvidia-smi not found")
        return None
    except subprocess.TimeoutExpired:
        logger.warning("nvidia-smi timed out")
        return None
    except (ValueError, IndexError) as exc:
        logger.warning("Failed to parse nvidia-smi output: %s", exc)
        return None


def measure_gpu_amd(logger):
    """
    Measure GPU utilization via amd-smi.

    Parses the ``gfx.value`` field from ``amd-smi monitor --json``.
    Returns the utilization of the first GPU, or None if amd-smi is unavailable.
    """
    try:
        result = subprocess.run(
            ["amd-smi", "monitor", "--json"],
            capture_output=True, text=True, timeout=15,
        )
        if result.returncode != 0:
            logger.debug("amd-smi returned %d — no AMD GPU", result.returncode)
            return None

        raw = result.stdout.strip()

        data = json.loads(result.stdout)
        if not data or not isinstance(data, list):
            logger.debug("amd-smi JSON output empty or invalid")
            return None

        gpu = data[0]
        gfx = gpu.get("gfx")
        if isinstance(gfx, dict):
            gfx_val = gfx.get("value")
        else:
            gfx_val = gfx
        if gfx_val is None:
            logger.debug("amd-smi JSON has no 'gfx.value' field; fields: %s",
                        list(gpu.keys()))
            return None
        util = float(gfx_val)
        logger.info("AMD GPU utilization: %.1f%% (gfx=%s)", util, gfx_val)
        return util
    except FileNotFoundError:
        logger.debug("amd-smi not found")
        return None
    except subprocess.TimeoutExpired:
        logger.warning("amd-smi timed out")
        return None
    except (json.JSONDecodeError, ValueError, TypeError) as exc:
        logger.warning("Failed to parse amd-smi output: %s", exc)
        return None


def get_active_users(logger):
    """
    Check for logged-in users.

    Returns True if ``who`` reports any sessions (including SSH/pts).
    """
    try:
        result = subprocess.run(
            ["who"],
            capture_output=True, text=True, timeout=5,
        )
        users = [line for line in result.stdout.strip().splitlines() if line.strip()]
        if users:
            logger.debug("Active users: %s", len(users))
            return True
        return False
    except subprocess.TimeoutExpired:
        logger.warning("who timed out")
        return True  # conservative: assume user present

# ---------------------------------------------------------------------------
# EMA update
# ---------------------------------------------------------------------------

def update_averages(state, cpu_n, gpu_n, gpu_threshold, alpha, logger):
    """
    Update CPU EMA averages and GPU idle counter.

    GPU is binary per-cycle: if utilization > threshold, the GPU idle
    counter resets. If idle (below threshold), the counter increments.

    Args:
        state: current state dict (modified in place for averages)
        cpu_n: current CPU idle %
        gpu_n: current GPU utilization % (None = no GPU, treated as idle)
        gpu_threshold: utilization threshold above which GPU is "busy"
        alpha: EMA weight for current sample
        logger: logger instance

    Returns:
        dict with updated state (new keys: cpu_avg, cpu_idle_n, cpu_idle_n1,
        gpu_idle_counter, sample_count)
    """
    new_state = dict(state)
    new_state["sample_count"] = state["sample_count"] + 1

    # Shift current → previous
    new_state["cpu_idle_n1"] = state["cpu_idle_n"]

    # Update current measurements
    new_state["cpu_idle_n"] = cpu_n

    # EMA for CPU idle average
    if state["cpu_avg"] is not None and cpu_n is not None:
        new_state["cpu_avg"] = round(alpha * cpu_n + (1 - alpha) * state["cpu_avg"], 2)
    elif cpu_n is not None:
        # First valid sample seeds the average
        new_state["cpu_avg"] = cpu_n

    # GPU idle counter — binary per-cycle (no EMA)
    gpu_busy = gpu_n is not None and gpu_n > gpu_threshold
    if gpu_busy:
        new_state["gpu_idle_counter"] = 0
        logger.debug("GPU busy (%.1f%% >= %.0f%%) — reset gpu_idle_counter",
                     gpu_n, gpu_threshold)
    else:
        # gpu_n is None (no GPU) or below threshold — counts as idle
        new_state["gpu_idle_counter"] = state.get("gpu_idle_counter", 0) + 1

    logger.debug(
        "Updated averages: cpu_avg=%.1f%%, gpu_idle_counter=%d, sample_count=%d",
        new_state["cpu_avg"] if new_state["cpu_avg"] is not None else 0,
        new_state["gpu_idle_counter"],
        new_state["sample_count"],
    )
    return new_state

# ---------------------------------------------------------------------------
# Evaluation
# ---------------------------------------------------------------------------

def get_why_not_idle(state, has_users, cpu_idle_threshold=DEFAULT_CPU_IDLE_THRESHOLD,
                     gpu_util_threshold=DEFAULT_GPU_UTIL_THRESHOLD, min_samples=DEFAULT_MIN_SAMPLES,
                     gpu_idle_counter=None, logger=None):
    """Return a human-readable reason why the system is not idle."""
    reasons = []
    if has_users:
        reasons.append("users logged in")
    if state["cpu_avg"] is not None:
        if state["cpu_avg"] <= cpu_idle_threshold:
            reasons.append(
                f"CPU idle avg {state['cpu_avg']:.1f}% <= threshold {cpu_idle_threshold:.0f}%"
            )
    if gpu_idle_counter is not None and gpu_idle_counter == 0 and state["sample_count"] >= min_samples:
        reasons.append("GPU busy (idle_counter=0)")
    if not reasons:
        if state["sample_count"] < min_samples:
            reasons.append(f"warming up ({state['sample_count']}/{min_samples} samples)")
        else:
            reasons.append("unknown")
    return "; ".join(reasons)


def evaluate(state, has_users, gpu_idle_counter, idle_threshold, gpu_threshold,
             min_samples, logger):
    """
    Evaluate whether the system should be considered idle.

    GPU is binary per-cycle: busy if any GPU > threshold, otherwise idle.
    Returns True if all conditions are met for one idle cycle.
    """
    # Need enough samples before trusting the average
    if state["sample_count"] < min_samples:
        logger.debug(
            "Skipping evaluation: warmup (%d/%d samples)",
            state["sample_count"], min_samples,
        )
        return False

    # Check active users
    if has_users:
        logger.debug("Not idle: users logged in")
        return False

    # Check CPU idle average
    if state["cpu_avg"] is None or state["cpu_avg"] <= idle_threshold:
        logger.debug(
            "Not idle: CPU avg %.1f%% <= threshold %.0f%%",
            state["cpu_avg"] or 0, idle_threshold,
        )
        return False

    # GPU is binary per-cycle: busy if counter was reset this cycle
    if gpu_idle_counter == 0 and state["sample_count"] >= min_samples:
        logger.debug(
            "Not idle: GPU busy this cycle (idle_counter=0)"
        )
        return False

    return True

# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(description="Auto-suspend based on system idle state")
    parser.add_argument("--alpha", type=float, default=DEFAULT_ALPHA,
                        help="EMA weight for current sample (default: %(default)s)")
    parser.add_argument("--min-samples", type=int, default=DEFAULT_MIN_SAMPLES,
                        help="Warmup samples before trusting averages (default: %(default)s)")
    parser.add_argument("--idle-timeout-cycles", type=int, default=DEFAULT_IDLE_TIMEOUT_CYCLES,
                        help="Consecutive idle cycles needed (default: %(default)s)")
    parser.add_argument("--cpu-idle-threshold", type=float, default=DEFAULT_CPU_IDLE_THRESHOLD,
                        help="CPU idle %% threshold (default: %(default)s)")
    parser.add_argument("--gpu-util-threshold", type=float, default=DEFAULT_GPU_UTIL_THRESHOLD,
                        help="GPU utilization %% threshold (default: %(default)s)")
    parser.add_argument("--state-dir", default=DEFAULT_STATE_DIR,
                        help="Directory for state file (default: %(default)s)")
    parser.add_argument("--log-file", default=DEFAULT_LOG_FILE,
                        help="Log file path (default: %(default)s)")
    parser.add_argument("--debug", action="store_true",
                        help="Enable debug logging")
    parser.add_argument("--dry-run", action="store_true",
                        help="Run check without persisting state or suspending")
    args = parser.parse_args()

    logger = setup_logging(args.log_file, args.debug)

    logger.info("Auto-suspend check started")
    logger.debug(
        "Config: alpha=%.2f, min_samples=%d, idle_cycles=%d, "
        "cpu_threshold=%.0f%%, gpu_threshold=%.0f%%",
        args.alpha, args.min_samples, args.idle_timeout_cycles,
        args.cpu_idle_threshold, args.gpu_util_threshold,
    )

    # --- Load state ---
    state = load_state(args.state_dir, logger)

    # --- Measure ---
    cpu_idle = measure_cpu(logger)
    gpu_nvidia = measure_gpu_nvidia(logger)
    gpu_amd = measure_gpu_amd(logger)

    # Combine GPU metrics: max across all available GPUs; None if none found
    gpu_values = [v for v in [gpu_nvidia, gpu_amd] if v is not None]
    gpu_util = max(gpu_values) if gpu_values else None

    logger.debug(
        "Measurements: cpu_idle=%.1f%%, gpu_util=%.1f%%",
        cpu_idle if cpu_idle is not None else 0,
        gpu_util if gpu_util is not None else 0,
    )

    # --- Update averages ---
    state = update_averages(
        state, cpu_idle, gpu_util, args.gpu_util_threshold, args.alpha, logger,
    )

    # --- Evaluate ---
    has_users = get_active_users(logger)

    is_idle = evaluate(
        state, has_users, state["gpu_idle_counter"],
        idle_threshold=args.cpu_idle_threshold,
        gpu_threshold=args.gpu_util_threshold,
        min_samples=args.min_samples,
        logger=logger,
    )

    if is_idle:
        state["idle_counter"] = state.get("idle_counter", 0) + 1
        logger.info(
            "System idle (cpu_avg=%.1f%%, gpu_idle=%d, "
            "idle_counter=%d/%d)",
            state["cpu_avg"] or 0,
            state["gpu_idle_counter"],
            state["idle_counter"],
            args.idle_timeout_cycles,
        )

        if state["idle_counter"] >= args.idle_timeout_cycles:
            if args.dry_run:
                print("Would suspend: idle timeout reached (idle_counter=%d/%d)" % (
                    state["idle_counter"], args.idle_timeout_cycles))
            else:
                logger.info("Idle timeout reached — suspending...")
                save_state(args.state_dir, state, logger)

                try:
                    result = subprocess.run(
                        ["shutdown", "now"],
                        capture_output=True, text=True, timeout=10,
                    )
                    if result.returncode != 0:
                        logger.error("shutdown failed: %s", result.stderr.strip())
                    else:
                        logger.info("shutdown succeeded")
                except subprocess.TimeoutExpired:
                    logger.error("shutdown timed out")
                except FileNotFoundError:
                    logger.error("shutdown not found — cannot suspend")
                except BaseException as exc:
                    logger.error("Failed to suspend: %s", exc)

                # Reset state on resume
                save_state(args.state_dir, fresh_state(), logger)
            return
    else:
        if not args.dry_run:
            state["idle_counter"] = 0
        reason = get_why_not_idle(
            state, has_users, args.cpu_idle_threshold, args.gpu_util_threshold,
            args.min_samples, state.get("gpu_idle_counter"), logger)
        logger.info("Not idle (%s) — reset counter", reason)

    # --- Save and exit ---
    if not args.dry_run:
        save_state(args.state_dir, state, logger)
    logger.info("Auto-suspend check completed")


def signal_handler(signum, frame):
    """Handle SIGTERM/SIGINT gracefully."""
    sys.exit(0)


if __name__ == "__main__":
    signal.signal(signal.SIGTERM, signal_handler)
    signal.signal(signal.SIGINT, signal_handler)
    main()