bot_runner.py

"""Bot subprocess runner.

In multi-tenant mode each bot runs as a subprocess inside the hosting
HF Space.  We spawn it on a free port, watch via ``psutil``, and offer
start/stop/restart + status queries.  Cross-space bots (legacy ones that
live on dedicated HF Spaces) are not handled here -- see ``bot_files.py``.

Resource enforcement
--------------------
Docker SDK on HF Spaces does not allow nested cgroups, so true
per-process CPU/RAM limits from inside the container aren't possible.
We enforce two soft policies:

- **RAM**: ``psutil.Process(pid).memory_info().rss`` polled every second;
  if RSS exceeds ``bot.ram_mb`` by 20 % for 30 seconds, the bot gets a
  ``SIGTERM``.
- **CPU**: best-effort ``os.sched_setaffinity`` to restrict the bot's
  process to a subset of cores; ``psutil.cpu_percent`` sampled and
  exposed as ``status['cpu_percent']`` for the UI. We also feed the
  watchdog a *hard* limit (``user_hard_cpu_limit``) so a single bot
  that spikes above its slice for too long gets SIGTERMed.

Both are advisory -- a runaway bot will impact neighbours -- but in
practice bots respect memory limits because Python raises ``MemoryError``
naturally and the soft-kill watchdog is a backstop.
"""
from __future__ import annotations

import asyncio
import logging
import os
import shutil
import signal
import subprocess
import sys
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional, Sequence

import psutil

from allocator import user_hard_cpu_limit
from config import PANEL_CONFIG
from models import BotInstance, BotStatus

log = logging.getLogger("hosting-panel.runner")


# ---------------------------------------------------------------------------
# Process record (in-memory; mirrors the BotInstance row)
# ---------------------------------------------------------------------------

@dataclass
class BotProcessRecord:
    pid: int
    started_at: float
    port: int
    proc: subprocess.Popen
    last_rss_mb: float = 0.0
    last_cpu_pct: float = 0.0
    last_kill_at: Optional[float] = None   # set when watchdog trips
    bytes_log_tail: int = 0               # offset in log file for streaming


# Module-level process registry, keyed by bot.id.
_PROCESSES: dict[int, BotProcessRecord] = {}
_REGISTRY_LOCK = asyncio.Lock()


def get_record(bot_id: int) -> Optional[BotProcessRecord]:
    return _PROCESSES.get(bot_id)


# ---------------------------------------------------------------------------
# Filesystem layout
# ---------------------------------------------------------------------------

def bot_workdir(bot_id: int) -> Path:
    """Where the bot's source + runtime artefacts live on disk.

    We keep them LOCAL to the Space; the shared HF Dataset is a backup
    mirror so other Spaces can re-materialise the same files.
    """
    d = PANEL_CONFIG.data_dir / "bots" / str(bot_id)
    d.mkdir(parents=True, exist_ok=True)
    return d


def bot_log_path(bot_id: int) -> Path:
    return bot_workdir(bot_id) / "run.log"


# ---------------------------------------------------------------------------
# Build entrypoint command
# ---------------------------------------------------------------------------

def _bot_command(bot: BotInstance, port: int) -> list[str]:
    """Build the argv we will exec for this bot.

    Two options:

    1. The bot ships a ``run.sh`` (or ``run.py``) at repo root -- run it.
    2. Otherwise fall back to ``python -m panel_bot_runner <port>`` which
       we ship as a generic uvicorn host for ``app:app`` modules found
       inside the source tree.
    """
    wd = bot_workdir(bot.id)
    if (wd / "run.sh").exists():
        return ["/bin/bash", str(wd / "run.sh")]
    if (wd / "run.py").exists():
        return [sys.executable, str(wd / "run.py")]
    if (wd / "app.py").exists():
        return [
            sys.executable, "-m", "uvicorn",
            "app:app", "--host", "0.0.0.0",
            "--port", str(port),
        ]
    # Last resort: keep process alive on the allocated port so health checks pass.
    return [
        sys.executable, "-c",
        "import http.server, socketserver, sys;"
        "class H(http.server.BaseHTTPRequestHandler):"
        " def do_GET(s): s.send_response(200); s.end_headers(); s.wfile.write(b'ok');"
        " def log_message(*a): pass;"
        f"HTTP_PORT={port};"
        "with socketserver.TCPServer(('0.0.0.0', HTTP_PORT), H) as s: s.serve_forever()",
    ]


# ---------------------------------------------------------------------------
# Start / stop / restart
# ---------------------------------------------------------------------------

async def start_bot(bot: BotInstance) -> BotProcessRecord:
    """Spawn the bot's subprocess.  Raises if already running."""
    if bot.id in _PROCESSES:
        return _PROCESSES[bot.id]
    if not bot.port:
        raise RuntimeError("Bot has no allocated port; call allocator first")

    wd = bot_workdir(bot.id)
    env = os.environ.copy()
    env["PANEL_BOT_PORT"] = str(bot.port)
    env["PANEL_BOT_ID"] = str(bot.id)
    log_path = bot_log_path(bot.id)
    log_fh = open(log_path, "ab", buffering=0)

    cmd = _bot_command(bot, bot.port)
    log.info(
        "Starting bot %s on port %d: %s",
        bot.slug, bot.port, " ".join(cmd[:3]) + "...",
    )
    proc = subprocess.Popen(
        cmd,
        cwd=str(wd),
        env=env,
        stdout=log_fh,
        stderr=subprocess.STDOUT,
        # New process group so we can kill children too.
        start_new_session=True,
    )
    # Best-effort CPU pinning so the kernel scheduler doesn't let one
    # bot borrow CPU slices meant for its neighbour. ``sched_setaffinity``
    # silently no-ops in unprivileged HF Spaces (which is fine: the
    # watchdog's per-core CPU sampling still catches runaway loops).
    apply_cpu_affinity(proc.pid, cpu_cores_for_bot(bot.id))
    rec = BotProcessRecord(
        pid=proc.pid,
        started_at=time.time(),
        port=bot.port,
        proc=proc,
        bytes_log_tail=log_path.stat().st_size,
    )
    async with _REGISTRY_LOCK:
        _PROCESSES[bot.id] = rec
    bot.pid = proc.pid
    return rec


async def stop_bot(bot: BotInstance, timeout: float = 10.0) -> bool:
    """SIGTERM then SIGKILL.  Returns True if it exited gracefully."""
    rec = _PROCESSES.pop(bot.id, None)
    if not rec:
        bot.pid = None
        return True

    try:
        os.killpg(os.getpgid(rec.proc.pid), signal.SIGTERM)
        try:
            rec.proc.wait(timeout=timeout)
            log.info("Bot %s stopped gracefully (pid=%d)", bot.slug, rec.pid)
            return True
        except subprocess.TimeoutExpired:
            log.warning("Bot %s did not exit; sending SIGKILL (pid=%d)",
                        bot.slug, rec.pid)
            os.killpg(os.getpgid(rec.proc.pid), signal.SIGKILL)
            rec.proc.wait(timeout=2.0)
            return False
    except ProcessLookupError:
        return True
    finally:
        bot.pid = None


async def restart_bot(bot: BotInstance) -> BotProcessRecord:
    await stop_bot(bot)
    return await start_bot(bot)


# ---------------------------------------------------------------------------
# Status (used by /api/bots/{id}/live and the UI dashboard)
# ---------------------------------------------------------------------------

def bot_status(bot: BotInstance) -> dict:
    rec = _PROCESSES.get(bot.id)
    if not rec:
        return {
            "running": False,
            "pid": None,
            "uptime_sec": 0,
            "cpu_percent": 0.0,
            "rss_mb": 0.0,
            "port": bot.port,
        }
    uptime = time.time() - rec.started_at
    cpu_pct = 0.0
    rss_mb = rec.last_rss_mb
    try:
        p = psutil.Process(rec.pid)
        with p.oneshot():
            cpu_pct = p.cpu_percent(interval=None)
            rss_mb = p.memory_info().rss / (1024 * 1024)
        rec.last_cpu_pct = cpu_pct
        rec.last_rss_mb = rss_mb
    except (psutil.NoSuchProcess, psutil.AccessDenied):
        pass
    alive = rec.proc.poll() is None
    return {
        "running": alive,
        "pid": rec.pid,
        "uptime_sec": round(uptime, 1),
        "cpu_percent": round(cpu_pct, 2),
        "rss_mb": round(rss_mb, 2),
        "port": rec.port,
    }


# ---------------------------------------------------------------------------
# Watchdog: kill bots that exceed RAM quota
# ---------------------------------------------------------------------------

async def resource_watchdog_loop(check_interval: float = 5.0) -> None:
    """Background loop: every 5s sample RSS (and CPU) for every running bot.

    RAM: if RSS exceeds ``ram_mb`` by 20 % for 30 seconds, SIGTERM.
    CPU: if the bot's CPU% sustained for 30 s above its *hard* limit
    (the slightly-underbooked value the allocator publishes to
    ``bot_runner.user_hard_cpu_limit``), SIGTERM. The number is
    internal-only; the user still sees their full quota in the UI.
    """
    log.info("resource watchdog running every %.0fs", check_interval)
    while True:
        try:
            await asyncio.sleep(check_interval)
            for bot_id, rec in list(_PROCESSES.items()):
                if rec.proc.poll() is not None:
                    # Process exited on its own; flip the DB status to
                    # ERROR so the UI doesn't keep showing a stale
                    # "running" badge for a dead subprocess.  We read
                    # the tail of run.log to put a useful message in
                    # the deployment timeline.
                    exit_code = rec.proc.returncode
                    _PROCESSES.pop(bot_id, None)
                    try:
                        from database import async_session
                        from models import BotInstance, BotStatus, DeploymentLog
                        async with async_session() as db:
                            bot = await db.get(BotInstance, bot_id)
                            if bot is not None and bot.status == BotStatus.RUNNING:
                                bot.status = BotStatus.ERROR
                                # Best-effort: pull the last few lines of
                                # the subprocess log to surface in the
                                # deployment timeline.
                                tail = ""
                                try:
                                    lp = bot_log_path(bot_id)
                                    if lp.exists():
                                        text = lp.read_text(encoding="utf-8", errors="replace")
                                        lines = text.splitlines()[-10:]
                                        tail = "\n".join(lines).strip()
                                except Exception:
                                    pass
                                msg = f"Subprocess exited unexpectedly (code={exit_code})."
                                if tail:
                                    msg += f"\n--- last log lines ---\n{tail}"
                                db.add(DeploymentLog(
                                    bot_id=bot_id, action="watchdog", status="failed",
                                    message=msg[:4000],
                                ))
                                await db.commit()
                                log.warning("bot %d marked ERROR: %s", bot_id, msg[:200])
                    except Exception as exc:  # noqa: BLE001
                        log.warning("watchdog: failed to mark bot %d ERROR: %r", bot_id, exc)
                    continue
                try:
                    p = psutil.Process(rec.pid)
                    rss = p.memory_info().rss / (1024 * 1024)
                    cpu_pct = p.cpu_percent(interval=None)
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    continue

                limit_mb = _RECORD_QUOTA.get(bot_id, 0)
                hard_cpu = _RECORD_HARD_CPU.get(bot_id, 0.0)

                # RAM enforcement
                ram_breach = limit_mb > 0 and rss > limit_mb * 1.20
                # CPU enforcement (Linux reports process CPU% relative to
                # one core, so 100% == 1.0 cores. We allow 30 % margin to
                # avoid spurious kills on brief spikes.)
                cpu_breach = hard_cpu > 0 and cpu_pct > (hard_cpu * 100.0) * 1.30

                if ram_breach or cpu_breach:
                    if rec.last_kill_at is None:
                        rec.last_kill_at = time.time()
                    elif time.time() - rec.last_kill_at > 30:
                        log.warning(
                            "Bot %d resource breach (rss=%.0f/%.0f MB, cpu=%.1f/%.2f cores) "
                            "for >30s; SIGTERM",
                            bot_id, rss, limit_mb, cpu_pct / 100.0, hard_cpu,
                        )
                        try:
                            os.killpg(os.getpgid(rec.pid), signal.SIGTERM)
                        except ProcessLookupError:
                            pass
                else:
                    rec.last_kill_at = None
        except asyncio.CancelledError:
            return
        except Exception as exc:
            log.warning("resource watchdog error: %s", exc)


# ---------------------------------------------------------------------------
# CPU-affinity helpers
# ---------------------------------------------------------------------------

def _host_cpu_count() -> int:
    """Number of CPUs the host (HF Space container) actually exposes."""
    try:
        aff = psutil.Process().cpu_affinity()
        if aff:
            return max(1, len(aff))
        n = psutil.cpu_count(logical=True)
        return max(1, n or 2)
    except Exception:
        return 2


def cpu_cores_for_bot(bot_id: int, host_cpus: int | None = None) -> list[int]:
    """Pick the subset of host CPUs that ``bot_id`` should be pinned to.

    With 5 users per cpu-basic Space (2 cores) we round-robin bots across
    cores so any two on the same one share via Linux CFS, and 3+1 leaves
    headroom for orchestrator traffic. ``host_cpus`` is the caller's cached
    count to avoid re-querying the kernel per bot.
    """
    n = max(1, host_cpus or _host_cpu_count())
    return [bot_id % n]


def apply_cpu_affinity(pid: int, cores) -> None:
    """Best-effort ``sched_setaffinity`` for the spawned subprocess.

    HF Spaces run on a cgroupv2 slice; the panel process itself is locked
    to the container's CPU set. Pinning each bot to a small subset keeps
    the kernel scheduler's per-core accounting honest when we add the
    watchdog's per-core CPU sampling on top of it.
    """
    if not cores:
        return
    try:
        os.sched_setaffinity(pid, set(int(c) for c in cores))
    except (PermissionError, OSError, NotImplementedError):
        # Unprivileged HF Space or non-POSIX; not fatal.
        pass


# Per-bot quota (set by the route right after assign_bot).
# We track *both* RAM and hard-CPU here so the watchdog can enforce
# against the user's actual hard limit (internal-only), not the soft
# quota that's shown in the UI.
_RECORD_QUOTA: dict[int, int] = {}
_RECORD_HARD_CPU: dict[int, float] = {}


def set_record_quota(bot_id: int, ram_mb: int, hard_cpu_cores: float | None = None) -> None:
    _RECORD_QUOTA[bot_id] = ram_mb
    if hard_cpu_cores is not None:
        _RECORD_HARD_CPU[bot_id] = float(hard_cpu_cores)


def clear_record_quota(bot_id: int) -> None:
    _RECORD_QUOTA.pop(bot_id, None)
    _RECORD_HARD_CPU.pop(bot_id, None)