tools/space.py

#!/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())