Add poller daemon (client.py)

daemon/client.py

@@ -0,0 +1,210 @@
+#!/usr/bin/env python3
+"""
+NII Relay daemon — runs inside a no-SSH container, polls the HF Space relay, and
+executes shell commands in persistent named bash sessions.
+
+Run it (foreground, or under nohup / tmux / systemd):
+
+    export RELAY_SPACE="your-username/nii-relay"   # the Space repo id
+    export RELAY_SECRET="<same secret as the Space>"
+    export HF_TOKEN="hf_..."                       # read token; needed if Space is private
+    export CLIENT_ID="$(hostname)"                 # how it shows up in the UI (optional)
+    python client.py
+
+Each command carries a `session` name. Commands with the same session name share
+one long-lived bash process, so `cd`, exports, and activated venvs persist.
+A new session name spins up a fresh shell.
+"""
+
+import os
+import sys
+import time
+import signal
+import threading
+import subprocess
+import uuid
+import queue as queuelib
+
+from gradio_client import Client
+
+SPACE = os.environ.get("RELAY_SPACE", "")
+SECRET = os.environ.get("RELAY_SECRET", "")
+HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")
+CLIENT_ID = os.environ.get("CLIENT_ID") or os.uname().nodename
+POLL_INTERVAL = float(os.environ.get("POLL_INTERVAL", "5"))
+DEFAULT_TIMEOUT = int(os.environ.get("CMD_TIMEOUT", "120"))
+MAX_OUTPUT_CHARS = 200_000
+
+_log_lock = threading.Lock()
+
+
+def log(*a):
+    with _log_lock:
+        print(f"[{time.strftime('%H:%M:%S')}]", *a, flush=True)
+
+
+# --------------------------------------------------------------------------- #
+# Persistent bash session                                                      #
+# --------------------------------------------------------------------------- #
+class Shell:
+    """A long-lived bash process. Commands are run serially; output is captured
+    up to a sentinel line that also carries the exit code."""
+
+    def __init__(self, name: str):
+        self.name = name
+        self._q: "queuelib.Queue[str]" = queuelib.Queue()
+        self._start()
+
+    def _start(self):
+        self.proc = subprocess.Popen(
+            ["bash", "--norc", "--noprofile"],
+            stdin=subprocess.PIPE, stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT, text=True, bufsize=1,
+            start_new_session=True,  # own process group, so we can SIGINT it
+            env={**os.environ, "PS1": "", "TERM": "dumb"},
+        )
+        self._reader = threading.Thread(target=self._pump, daemon=True)
+        self._reader.start()
+
+    def _pump(self):
+        for line in self.proc.stdout:
+            self._q.put(line)
+        self._q.put(None)  # EOF sentinel
+
+    def _drain_pending(self):
+        while True:
+            try:
+                self._q.get_nowait()
+            except queuelib.Empty:
+                return
+
+    def run(self, command: str, timeout: int) -> tuple[str, int]:
+        if self.proc.poll() is not None:  # shell died (e.g. `exit`) — respawn
+            log(f"session {self.name}: shell exited, restarting")
+            self._start()
+
+        self._drain_pending()
+        marker = "__NII_END_" + uuid.uuid4().hex + "__"
+        # Run the command, then print a sentinel with the exit status on its own line.
+        script = f"{command}\nprintf '\\n%s %s\\n' '{marker}' \"$?\"\n"
+        try:
+            self.proc.stdin.write(script)
+            self.proc.stdin.flush()
+        except BrokenPipeError:
+            return ("(shell pipe broken)", 127)
+
+        out_lines: list[str] = []
+        exit_code = 0
+        deadline = time.monotonic() + max(1, timeout)
+        interrupted = False
+        while True:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0 and not interrupted:
+                # Interrupt the running command but keep the shell alive.
+                try:
+                    os.killpg(self.proc.pid, signal.SIGINT)
+                except ProcessLookupError:
+                    pass
+                out_lines.append(f"\n[relay: timed out after {timeout}s, sent SIGINT]\n")
+                interrupted = True
+                deadline = time.monotonic() + 5  # grace to collect the sentinel
+                continue
+            try:
+                line = self._q.get(timeout=max(0.1, remaining if not interrupted else 0.5))
+            except queuelib.Empty:
+                if interrupted:
+                    break
+                continue
+            if line is None:  # shell EOF
+                if interrupted:
+                    exit_code = 124  # conventional timeout exit code
+                else:
+                    out_lines.append("\n[relay: shell closed]\n")
+                    exit_code = 137
+                break
+            stripped = line.rstrip("\n")
+            if stripped.startswith(marker):
+                parts = stripped.split()
+                try:
+                    exit_code = int(parts[-1])
+                except (ValueError, IndexError):
+                    exit_code = 0
+                break
+            out_lines.append(line)
+
+        output = "".join(out_lines)
+        # Trim the single leading newline our printf added before the sentinel.
+        if output.endswith("\n"):
+            output = output[:-1]
+        if len(output) > MAX_OUTPUT_CHARS:
+            output = output[:MAX_OUTPUT_CHARS] + "\n[relay: output truncated]"
+        return (output, exit_code)
+
+
+SHELLS: dict[str, Shell] = {}
+
+
+def get_shell(name: str) -> Shell:
+    sh = SHELLS.get(name)
+    if sh is None:
+        sh = SHELLS[name] = Shell(name)
+        log(f"opened bash session: {name}")
+    return sh
+
+
+# --------------------------------------------------------------------------- #
+# Relay client                                                                 #
+# --------------------------------------------------------------------------- #
+def connect() -> Client:
+    while True:
+        try:
+            log(f"connecting to Space {SPACE} …")
+            client = Client(SPACE, hf_token=HF_TOKEN, verbose=False)
+            client.predict(CLIENT_ID, SECRET, _meta(), api_name="/register")
+            log(f"registered as '{CLIENT_ID}'")
+            return client
+        except Exception as e:
+            log(f"connect failed: {e!r} — retrying in 10s")
+            time.sleep(10)
+
+
+def _meta() -> str:
+    return f"{os.uname().sysname} {os.uname().release} pid={os.getpid()}"
+
+
+def handle(client: Client, cmd: dict):
+    cid = cmd["id"]
+    session = cmd.get("session") or "main"
+    command = cmd["command"]
+    timeout = int(cmd.get("timeout") or DEFAULT_TIMEOUT)
+    log(f"[{session}] $ {command}")
+    try:
+        output, code = get_shell(session).run(command, timeout)
+    except Exception as e:
+        output, code = (f"[relay: daemon error: {e!r}]", 1)
+    client.predict(cid, CLIENT_ID, SECRET, code, output, session, api_name="/result")
+    log(f"[{session}] → exit {code} ({len(output)} bytes)")
+
+
+def main():
+    if not SPACE:
+        sys.exit("set RELAY_SPACE to the Space repo id, e.g. your-username/nii-relay")
+    log(f"daemon starting — client_id='{CLIENT_ID}', poll every {POLL_INTERVAL}s")
+    client = connect()
+    while True:
+        try:
+            res = client.predict(CLIENT_ID, SECRET, api_name="/poll")
+            commands = (res or {}).get("commands", [])
+            for cmd in commands:
+                handle(client, cmd)
+        except KeyboardInterrupt:
+            log("bye")
+            return
+        except Exception as e:
+            log(f"poll error: {e!r} — reconnecting")
+            client = connect()
+        time.sleep(POLL_INTERVAL)
+
+
+if __name__ == "__main__":
+    main()

daemon/requirements.txt

@@ -0,0 +1 @@
+gradio_client>=1.3