app.py

"""
NII Relay — a tiny HF Space that acts as a message relay between you (in the
browser) and shell daemons running inside no-SSH containers on the cluster.

Flow
----
1. A daemon inside a container registers with a CLIENT_ID and polls this Space
   every few seconds (via the `gradio_client` API endpoints below).
2. You pick a container + bash session in the web UI and submit a shell command.
3. The next poll hands the command to the daemon, which marks it ACKNOWLEDGED,
   runs it in a persistent bash session, and posts stdout + exit code back.
4. The UI shows the output.

State is in-memory only. If the Space restarts, daemons simply re-register on
their next poll — but any not-yet-delivered commands are lost. That is fine for
interactive debugging.

Auth: the Space is PRIVATE, so reaching it at all requires an HF token with
access — that token is the auth boundary. There is no separate app-level secret.
"""

import os
import re
import time
import uuid
import shutil
import tempfile
import pathlib
import threading
from collections import defaultdict

import gradio as gr

CLIENT_STALE_AFTER = 30                                # seconds w/o a poll => "offline"
MAX_OUTPUT_CHARS = 200_000                             # cap stored output per command
MAX_HISTORY_PER_CLIENT = 200                           # ring-buffer of commands
MAX_FILE_BYTES = 100 * 1024 * 1024                     # 100 MB cap on transfers
FILE_TTL = 3600                                        # delete stored files older than 1h

FILES_DIR = pathlib.Path(tempfile.gettempdir()) / "relay_files"
FILES_DIR.mkdir(exist_ok=True)

_HEX32 = re.compile(r"\A[0-9a-f]{32}\Z")


def _valid_id(s) -> bool:
    return isinstance(s, str) and bool(_HEX32.match(s))

LOCK = threading.Lock()
CLIENTS: dict[str, dict] = {}                          # cid -> {last_seen, meta}
COMMANDS: dict[str, dict] = {}                         # command_id -> record
BY_CLIENT: dict[str, list[str]] = defaultdict(list)    # cid -> [command_id, ...]
PUT_FILES: dict[str, dict] = {}                        # file_id -> {path, filename} (server->client payloads)


def _now() -> float:
    return time.time()


def _human(n: float) -> str:
    for unit in ("B", "KB", "MB", "GB"):
        if n < 1024 or unit == "GB":
            return f"{n:.0f}{unit}" if unit == "B" else f"{n:.1f}{unit}"
        n /= 1024


def _touch_client(cid: str, meta: str = ""):
    c = CLIENTS.setdefault(cid, {"first_seen": _now(), "meta": meta})
    c["last_seen"] = _now()
    if meta:
        c["meta"] = meta


def _unlink(path: str):
    try:
        os.unlink(path)
    except OSError:
        pass


def _drop_command(cmd_id: str):
    """Remove a command record and any file it owns on disk. Call under LOCK."""
    rec = COMMANDS.pop(cmd_id, None)
    if not rec:
        return
    if rec.get("stored_path"):
        _unlink(rec["stored_path"])
    fid = rec.get("file_id")
    if fid:
        info = PUT_FILES.pop(fid, None)
        if info:
            _unlink(info["path"])


def _append_and_trim(cid: str, cmd_id: str):
    """Append to a client's history and ring-buffer it, deleting evicted files.
    Call under LOCK. Used by every append site so growth is bounded uniformly."""
    hist = BY_CLIENT[cid]
    hist.append(cmd_id)
    while len(hist) > MAX_HISTORY_PER_CLIENT:
        _drop_command(hist.pop(0))


def _file_reaper():
    """Backstop against disk growth on the ephemeral Space: delete stored files
    older than FILE_TTL and prune PUT_FILES whose file is gone."""
    while True:
        time.sleep(600)
        cutoff = _now() - FILE_TTL
        try:
            for p in FILES_DIR.iterdir():
                try:
                    if p.is_file() and p.stat().st_mtime < cutoff:
                        p.unlink()
                except OSError:
                    pass
        except OSError:
            pass
        with LOCK:
            for fid in list(PUT_FILES):
                if not os.path.isfile(PUT_FILES[fid]["path"]):
                    PUT_FILES.pop(fid, None)


# --------------------------------------------------------------------------- #
# Machine API (called by the daemon via gradio_client)                         #
# --------------------------------------------------------------------------- #
def api_register(client_id: str, meta: str):
    client_id = (client_id or "").strip()
    if not client_id:
        raise gr.Error("empty client_id")
    with LOCK:
        _touch_client(client_id, meta or "")
    return {"ok": True, "ts": _now()}


def api_poll(client_id: str):
    """Return all pending commands for this client and flip them to acknowledged."""
    client_id = (client_id or "").strip()
    out = []
    with LOCK:
        _touch_client(client_id)
        for cmd_id in list(BY_CLIENT.get(client_id, [])):
            rec = COMMANDS.get(cmd_id)
            if rec is None:
                continue
            if rec["status"] == "pending":
                rec["status"] = "acknowledged"
                rec["acknowledged_at"] = _now()
                out.append({
                    "id": rec["id"],
                    "type": rec.get("type", "exec"),
                    "session": rec.get("session", "main"),
                    "command": rec.get("command", ""),
                    "timeout": rec.get("timeout", 120),
                    # file-transfer fields (None for plain exec commands)
                    "file_id": rec.get("file_id"),
                    "dest": rec.get("dest"),
                    "remote_path": rec.get("remote_path"),
                    "filename": rec.get("filename"),
                })
    return {"commands": out, "ts": _now()}


def api_result(command_id: str, client_id: str,
               exit_code: float, stdout: str, session: str):
    if not _valid_id(command_id):
        raise gr.Error("bad command id")
    with LOCK:
        rec = COMMANDS.get(command_id)
        if rec is None:
            # Space probably restarted; record it anyway so output isn't lost.
            rec = COMMANDS[command_id] = {
                "id": command_id, "client_id": client_id, "session": session,
                "command": "(unknown — relay restarted)", "created": _now(),
                "status": "pending",
            }
            _append_and_trim(client_id, command_id)
        rec["status"] = "done"
        rec["finished_at"] = _now()
        rec["exit_code"] = int(exit_code)
        rec["stdout"] = (stdout or "")[:MAX_OUTPUT_CHARS]
        _touch_client(client_id)
    return {"ok": True}


def api_fetch_file(file_id: str):
    """Daemon calls this to download a server->client (push) payload."""
    if not _valid_id(file_id):
        raise gr.Error("bad file id")
    info = PUT_FILES.get(file_id)
    if not info or not os.path.isfile(info["path"]):
        raise gr.Error("unknown or expired file_id")
    return info["path"]


def api_upload_file(command_id: str, client_id: str, fileobj, error: str):
    """Daemon calls this to post the result of a client->server (pull) request.
    `fileobj` is the uploaded file path on success, or None when `error` is set."""
    if not _valid_id(command_id):
        raise gr.Error("bad command id")

    # Do filesystem work BEFORE taking the lock (a 100 MB copy must not stall
    # every concurrent poll/UI handler). Enforce the size cap server-side too —
    # don't trust the client.
    saved = None
    fname = sz = None
    too_big = False
    if fileobj:
        sz = os.path.getsize(fileobj)
        if sz > MAX_FILE_BYTES:
            too_big = True
        else:
            fname = os.path.basename(fileobj)
            saved = FILES_DIR / f"{command_id}__{fname}"
            shutil.copy(fileobj, saved)

    with LOCK:
        rec = COMMANDS.get(command_id)
        if rec is None:  # relay restarted — keep the file anyway
            rec = COMMANDS[command_id] = {
                "id": command_id, "client_id": client_id, "session": "file",
                "type": "get", "command": "(file fetch — relay restarted)",
                "created": _now(), "status": "pending",
            }
            _append_and_trim(client_id, command_id)
        rec["status"] = "done"
        rec["finished_at"] = _now()
        _touch_client(client_id)
        if too_big:
            rec.update(exit_code=1,
                       stdout=f"file exceeds the {_human(MAX_FILE_BYTES)} limit")
        elif saved is not None:
            rec.update(stored_path=str(saved), filename=fname, size=sz,
                       exit_code=0, stdout=f"fetched {fname} ({_human(sz)})")
        else:
            rec.update(exit_code=1, stdout=(error or "fetch failed"))
    return {"ok": True}


# --------------------------------------------------------------------------- #
# Human UI helpers                                                             #
# --------------------------------------------------------------------------- #
def _client_choices() -> list[str]:
    with LOCK:
        items = sorted(CLIENTS.items(), key=lambda kv: kv[1].get("last_seen", 0), reverse=True)
        labels = []
        for cid, c in items:
            age = _now() - c.get("last_seen", 0)
            dot = "🟢" if age < CLIENT_STALE_AFTER else "🔴"
            labels.append(f"{dot} {cid}")
        return labels


def _strip_dot(label: str) -> str:
    return label.split(" ", 1)[1] if label and label[:1] in "🟢🔴" else (label or "")


def ui_send(client_label: str, session: str, command: str, timeout: int):
    cid = _strip_dot(client_label)
    if not cid:
        return "⚠️ pick a container first", command
    if not command.strip():
        return "⚠️ empty command", command
    session = (session or "main").strip()
    cmd_id = uuid.uuid4().hex
    with LOCK:
        COMMANDS[cmd_id] = {
            "id": cmd_id, "client_id": cid, "session": session,
            "command": command, "timeout": int(timeout),
            "status": "pending", "created": _now(),
        }
        _append_and_trim(cid, cmd_id)
    return f"➡️ sent to `{cid}` [{session}]", ""


def ui_broadcast(session: str, command: str, timeout: int):
    """Queue the same command on the given session for every online container."""
    if not command.strip():
        return "⚠️ empty command", command
    session = (session or "main").strip()
    now = _now()
    sent = []
    with LOCK:
        online = [cid for cid, c in CLIENTS.items()
                  if now - c.get("last_seen", 0) < CLIENT_STALE_AFTER]
        for cid in online:
            cmd_id = uuid.uuid4().hex
            COMMANDS[cmd_id] = {
                "id": cmd_id, "client_id": cid, "session": session,
                "command": command, "timeout": int(timeout),
                "status": "pending", "created": now,
            }
            _append_and_trim(cid, cmd_id)
            sent.append(cid)
    if not sent:
        return "⚠️ no online containers to broadcast to", command
    return (f"📡 broadcast to {len(sent)} container(s) [{session}]: "
            + ", ".join(f"`{c}`" for c in sent), "")


def ui_put_file(client_label: str, fileobj, dest: str):
    """Queue a push: send an uploaded file to the selected container."""
    cid = _strip_dot(client_label)
    if not cid:
        return "⚠️ pick a container first", None
    if not fileobj:
        return "⚠️ choose a file to send", None
    sz = os.path.getsize(fileobj)
    if sz > MAX_FILE_BYTES:
        return f"⚠️ file is {_human(sz)} — over the {_human(MAX_FILE_BYTES)} limit", None
    fname = os.path.basename(fileobj)
    file_id = uuid.uuid4().hex
    saved = FILES_DIR / f"{file_id}__{fname}"
    shutil.copy(fileobj, saved)
    dest = (dest or "").strip() or fname
    cmd_id = uuid.uuid4().hex
    with LOCK:
        PUT_FILES[file_id] = {"path": str(saved), "filename": fname}
        COMMANDS[cmd_id] = {
            "id": cmd_id, "client_id": cid, "session": "file", "type": "put",
            "file_id": file_id, "dest": dest, "filename": fname, "size": sz,
            "command": f"📤 send {fname} ({_human(sz)}) → {dest}",
            "timeout": 300, "status": "pending", "created": _now(),
        }
        _append_and_trim(cid, cmd_id)
    return f"➡️ sending `{fname}` to `{cid}`:`{dest}`", None


def ui_get_file(client_label: str, remote_path: str):
    """Queue a pull: ask the container for a file at `remote_path`."""
    cid = _strip_dot(client_label)
    if not cid:
        return "⚠️ pick a container first"
    remote_path = (remote_path or "").strip()
    if not remote_path:
        return "⚠️ enter a path on the container"
    cmd_id = uuid.uuid4().hex
    with LOCK:
        COMMANDS[cmd_id] = {
            "id": cmd_id, "client_id": cid, "session": "file", "type": "get",
            "remote_path": remote_path, "command": f"📥 fetch {remote_path}",
            "timeout": 300, "status": "pending", "created": _now(),
        }
        _append_and_trim(cid, cmd_id)
    return f"➡️ requested `{remote_path}` from `{cid}`"


def _fetched_files(client_label: str) -> list[str]:
    cid = _strip_dot(client_label)
    paths = []
    with LOCK:
        for cmd_id in BY_CLIENT.get(cid, []):
            r = COMMANDS.get(cmd_id)
            if r and r.get("type") == "get" and r.get("stored_path") \
                    and os.path.isfile(r["stored_path"]):
                paths.append(r["stored_path"])
    return paths


_STATUS_ICON = {"pending": "⏳", "acknowledged": "📨", "done": "✅"}


def ui_history(client_label: str, n: int = 12) -> str:
    cid = _strip_dot(client_label)
    if not cid:
        return "_Select a container to see its command history._"
    with LOCK:
        ids = list(BY_CLIENT.get(cid, []))[-int(n):][::-1]
        recs = [dict(COMMANDS[i]) for i in ids if i in COMMANDS]
    if not recs:
        return f"_No commands for `{cid}` yet._"
    blocks = []
    for r in recs:
        icon = _STATUS_ICON.get(r["status"], "•")
        head = f"{icon} **[{r.get('session','main')}]** `{r['command']}`"
        if r["status"] == "done":
            ec = r.get("exit_code", "?")
            out = r.get("stdout", "") or "(no output)"
            head += f"  → exit {ec}"
            blocks.append(head + f"\n```text\n{out}\n```")
        else:
            blocks.append(head + f"\n_{r['status']}…_")
    return "\n\n".join(blocks)


def _keep_selection(choices: list[str], client_label: str) -> str | None:
    cur = _strip_dot(client_label)
    for ch in choices:
        if _strip_dot(ch) == cur:
            return ch
    return choices[0] if choices else None


def ui_refresh(client_label: str):
    choices = _client_choices()
    keep = _keep_selection(choices, client_label)
    return gr.update(choices=choices, value=keep), ui_history(keep or "")


def ui_tick(client_label: str, last_choices: list[str], last_hist: str,
            last_files: list[str]):
    """Timer handler that only emits updates when something actually changed,
    so the dropdown, output, and downloads don't repaint (flicker) every tick."""
    choices = _client_choices()
    if choices == last_choices:
        dd_out = gr.skip()                       # dropdown unchanged -> leave it alone
        label_for_hist = client_label
    else:
        keep = _keep_selection(choices, client_label)
        dd_out = gr.update(choices=choices, value=keep)
        label_for_hist = keep or ""
    hist = ui_history(label_for_hist)
    hist_out = gr.skip() if hist == last_hist else hist
    files = _fetched_files(label_for_hist)
    files_out = gr.skip() if files == last_files else gr.update(value=files)
    return dd_out, hist_out, files_out, choices, hist, files


# --------------------------------------------------------------------------- #
# Layout                                                                        #
# --------------------------------------------------------------------------- #
with gr.Blocks(title="NII Relay", analytics_enabled=False) as demo:
    gr.Markdown(
        "# 🛰️ NII Relay\n"
        "Send shell commands to no-SSH containers via a polling daemon. "
        "Pick a container, name a bash session (reuse a name to keep its `cd`/env, "
        "or type a new name to open a fresh session), and run."
    )

    with gr.Row():
        client_dd = gr.Dropdown(label="Container", choices=_client_choices(),
                                interactive=True, allow_custom_value=True, scale=3)
        refresh_btn = gr.Button("↻ Refresh", scale=1)

    with gr.Row():
        session_tb = gr.Textbox(label="bash session", value="main", scale=2)
        timeout_nb = gr.Number(label="timeout (s)", value=120, precision=0, scale=1)

    command_tb = gr.Textbox(label="command", lines=2,
                            placeholder="nvidia-smi    |    cd /mnt/data && ls -la    |    tail -n 50 train.log")
    with gr.Row():
        send_btn = gr.Button("Run ▶", variant="primary")
        broadcast_btn = gr.Button("📡 Broadcast to all online", variant="secondary")
        send_status = gr.Markdown("")

    history_md = gr.Markdown("_Select a container to see its command history._")

    with gr.Accordion("📁 File transfer (≤ 100 MB)", open=False):
        with gr.Row():
            with gr.Column():
                gr.Markdown("**Send to container** (push)")
                put_file = gr.File(label="file", type="filepath")
                put_dest = gr.Textbox(label="destination path on container",
                                      placeholder="/mnt/data/  or  /mnt/data/foo.bin")
                put_btn = gr.Button("Send ⬆", variant="primary")
                put_status = gr.Markdown("")
            with gr.Column():
                gr.Markdown("**Fetch from container** (pull)")
                get_path = gr.Textbox(label="path on container",
                                      placeholder="/mnt/data/train.log")
                get_btn = gr.Button("Fetch ⬇", variant="primary")
                get_status = gr.Markdown("")
                downloads = gr.Files(label="fetched files (click to download)",
                                     interactive=False)

    # remember what we last pushed, so the timer can skip no-op repaints
    last_choices = gr.State([])
    last_hist = gr.State("")
    last_files = gr.State([])

    timer = gr.Timer(3.0)

    # wiring
    send_btn.click(ui_send, [client_dd, session_tb, command_tb, timeout_nb],
                   [send_status, command_tb])
    command_tb.submit(ui_send, [client_dd, session_tb, command_tb, timeout_nb],
                      [send_status, command_tb])
    broadcast_btn.click(ui_broadcast, [session_tb, command_tb, timeout_nb],
                        [send_status, command_tb])
    refresh_btn.click(ui_refresh, [client_dd], [client_dd, history_md])
    client_dd.change(ui_history, [client_dd], history_md)
    client_dd.change(_fetched_files, [client_dd], downloads)
    put_btn.click(ui_put_file, [client_dd, put_file, put_dest], [put_status, put_file])
    get_btn.click(ui_get_file, [client_dd, get_path], get_status)
    timer.tick(ui_tick, [client_dd, last_choices, last_hist, last_files],
               [client_dd, history_md, downloads, last_choices, last_hist, last_files])

    # --- hidden machine API endpoints (called by the daemon) ---
    # Auth is the private Space + HF token (gating who can reach these at all);
    # there is no separate app-level secret.
    with gr.Row(visible=False):
        a_cid = gr.Textbox()
        a_meta = gr.Textbox()
        a_cmdid = gr.Textbox()
        a_session = gr.Textbox()
        a_exit = gr.Number()
        a_stdout = gr.Textbox()
        a_fileid = gr.Textbox()
        a_err = gr.Textbox()
        a_file_in = gr.File()
        a_file_out = gr.File()
        a_out = gr.JSON()
        b_reg = gr.Button()
        b_poll = gr.Button()
        b_res = gr.Button()
        b_fetch = gr.Button()
        b_upload = gr.Button()
        b_reg.click(api_register, [a_cid, a_meta], a_out, api_name="register")
        b_poll.click(api_poll, [a_cid], a_out, api_name="poll")
        b_res.click(api_result, [a_cmdid, a_cid, a_exit, a_stdout, a_session],
                    a_out, api_name="result")
        b_fetch.click(api_fetch_file, [a_fileid], a_file_out, api_name="fetch_file")
        b_upload.click(api_upload_file, [a_cmdid, a_cid, a_file_in, a_err],
                       a_out, api_name="upload_result")


threading.Thread(target=_file_reaper, daemon=True).start()

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=16).launch(
        server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)),
        theme=gr.themes.Soft(), max_file_size="150mb",
    )