#!/usr/bin/env python3 """Talk to the deployed Space from a terminal: status, logs, and test generations. Removes the copy-paste loop of pulling container logs out of the web UI. python tools/space.py status # runtime stage + hardware python tools/space.py logs # our [nv-generate] lines only python tools/space.py logs --all # everything, unfiltered python tools/space.py logs --build # build log instead of run log python tools/space.py wait # block until RUNNING (or fail) python tools/space.py timings # just the phase/banner lines python tools/space.py generate --model mr # fire one generation, report timing Auth comes from $HF_TOKEN (read access is enough for status and logs). `generate` consumes real ZeroGPU quota, so it asks before running unless -y. """ from __future__ import annotations import argparse import json import os import subprocess import sys import time import urllib.error import urllib.request SPACE = os.environ.get("NV_SPACE_ID", "nvidia/nv-generate") API = f"https://huggingface.co/api/spaces/{SPACE}" APP_URL = os.environ.get("NV_SPACE_URL", f"https://{SPACE.replace('/', '-')}.hf.space") # Startup can take minutes: ~20 GB of weights plus the mask archive extraction. READY_MARKER = "DONE - READY FOR USE" DEGRADED_MARKER = "DEGRADED - READY WITH FAILURES" def _token() -> str: tok = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") if not tok: sys.exit("set HF_TOKEN (read scope is enough for status/logs)") return tok def _get(url: str, timeout: int = 30) -> bytes: req = urllib.request.Request(url, headers={"Authorization": f"Bearer {_token()}"}) try: with urllib.request.urlopen(req, timeout=timeout) as r: return r.read() except urllib.error.HTTPError as e: sys.exit(f"HTTP {e.code} from {url}\n{e.read().decode('utf-8', 'replace')[:400]}") def runtime() -> dict: return json.loads(_get(API)).get("runtime", {}) def log_lines(build: bool = False, timeout: int = 25) -> list[tuple[str, str]]: """Return [(timestamp, text)] from the log endpoint. The endpoint is an SSE stream that replays the backlog and then stays open forever, so the read has to be bounded by wall time rather than EOF. curl's --max-time does exactly that and behaves consistently here; urllib's socket timeouts did not reliably surface partial stream data. """ url = f"{API}/logs/{'build' if build else 'run'}" proc = subprocess.run( ["curl", "-s", "--max-time", str(timeout), "-H", f"Authorization: Bearer {_token()}", url], capture_output=True, ) out = [] for line in proc.stdout.decode("utf-8", "replace").splitlines(): if not line.startswith("data: "): continue try: rec = json.loads(line[6:]) except json.JSONDecodeError: continue out.append((rec.get("timestamp", ""), rec.get("data", "").rstrip("\n"))) return out def cmd_status(args) -> int: r = runtime() print(f"space : {SPACE}") print(f"stage : {r.get('stage')}") print(f"hardware : {r.get('hardware', {}).get('current')}") print(f"url : {APP_URL}") return 0 def cmd_logs(args) -> int: for ts, text in log_lines(build=args.build, timeout=args.timeout): if args.all or "[nv-generate]" in text or "Error" in text or "Traceback" in text: print(f"{ts[11:19] if len(ts) > 19 else ts:9} {text}") return 0 def cmd_timings(args) -> int: keep = ("[nv-generate]",) for ts, text in log_lines(timeout=args.timeout): if any(k in text for k in keep): print(f"{ts[11:19] if len(ts) > 19 else ts:9} {text}") return 0 def cmd_wait(args) -> int: """Poll until the app reports ready, or the stage goes terminal.""" deadline = time.time() + args.timeout last = None while time.time() < deadline: stage = runtime().get("stage") if stage != last: print(f"[{time.strftime('%H:%M:%S')}] stage={stage}", flush=True) last = stage if stage in ("RUNTIME_ERROR", "BUILD_ERROR", "CONFIG_ERROR"): print("space failed to start; last log lines:") for ts, text in log_lines()[-25:]: print(" ", text) return 1 if stage == "RUNNING": for _ts, text in log_lines(): if READY_MARKER in text or DEGRADED_MARKER in text: print(f"ready: {text.split('] ', 1)[-1]}") return 0 time.sleep(args.interval) print("timed out waiting for the space to become ready") return 1 def cmd_generate(args) -> int: """Trigger one generation through the Gradio API and report its timing.""" try: from gradio_client import Client except ImportError: sys.exit("pip install gradio_client") if not args.yes: reply = input( f"This runs a real generation on {SPACE} and consumes ZeroGPU quota " f"(CT is billed 2x on xlarge). Continue? [y/N] " ) if reply.strip().lower() not in ("y", "yes"): return 1 before = {t for t, _ in log_lines()} client = Client(APP_URL, hf_token=_token(), verbose=False) print(f"calling {args.api_name} ...", flush=True) t0 = time.perf_counter() result = client.predict(api_name=args.api_name) wall = time.perf_counter() - t0 print(f"returned in {wall:.1f}s; result type {type(result).__name__}") for ts, text in log_lines(): if ts not in before and "[nv-generate]" in text: print(" ", text) return 0 def main() -> int: ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) sub = ap.add_subparsers(dest="cmd", required=True) sub.add_parser("status").set_defaults(fn=cmd_status) p = sub.add_parser("logs") p.add_argument("--all", action="store_true", help="do not filter to our lines") p.add_argument("--build", action="store_true", help="build log instead of run log") p.add_argument("--timeout", type=int, default=30) p.set_defaults(fn=cmd_logs) p = sub.add_parser("timings") p.add_argument("--timeout", type=int, default=30) p.set_defaults(fn=cmd_timings) p = sub.add_parser("wait") p.add_argument("--timeout", type=int, default=900) p.add_argument("--interval", type=int, default=15) p.set_defaults(fn=cmd_wait) p = sub.add_parser("generate") p.add_argument("--api-name", default="/_generate") p.add_argument("-y", "--yes", action="store_true", help="skip the quota prompt") p.set_defaults(fn=cmd_generate) args = ap.parse_args() return args.fn(args) if __name__ == "__main__": raise SystemExit(main())