""" 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", )