Spaces:
Running on Zero
Running on Zero
| """One-shot, caller-authenticated single-Space CLI. No inference on import. | |
| Examples: | |
| python client.py presets | |
| python client.py health | |
| python client.py plan --input request.json | |
| python client.py run --input request.json --run-dir runs/unique-name \ | |
| --expected-revision <approved-generator-commit-sha> | |
| Run requires HF_INFERENCE_TOKEN (or --token-env NAME) and the operator's | |
| HF_POOL_ZEROGPU_OVERAGE_DISABLED=1 attestation. This client does not change or | |
| verify the account's billing settings. No SDK credential lookup is allowed. | |
| A claimed directory is never reused, including after errors. An unknown hop | |
| may still be running remotely. Reconcile its event_id/session_hash manually; | |
| there is deliberately no resume, retry, cancellation, or token rotation. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| from concurrent.futures import Future | |
| from contextlib import contextmanager, redirect_stderr, redirect_stdout, suppress | |
| from dataclasses import dataclass | |
| import hashlib | |
| import json | |
| import math | |
| import os | |
| from pathlib import Path | |
| import re | |
| import signal | |
| import stat | |
| import sys | |
| import tempfile | |
| import threading | |
| import time | |
| from urllib.request import HTTPRedirectHandler, Request, build_opener | |
| GENERATOR = "https://owenisas-workflow-studio.hf.space" | |
| GENERATOR_REPO = "owenisas/workflow-studio" | |
| class AuthenticationRequired(ValueError): | |
| pass | |
| class OverageAttestationRequired(ValueError): | |
| pass | |
| class OwnerTokenForbidden(ValueError): | |
| pass | |
| class RevisionMismatch(ValueError): | |
| pass | |
| class RuntimeNotReady(ValueError): | |
| pass | |
| class IntegrityError(ValueError): | |
| pass | |
| class InputError(ValueError): | |
| pass | |
| class CLIArgumentError(ValueError): | |
| pass | |
| class TerminationRequested(KeyboardInterrupt): | |
| pass | |
| def _termination_guard(): | |
| if threading.current_thread() is not threading.main_thread(): | |
| yield | |
| return | |
| previous = signal.getsignal(signal.SIGTERM) | |
| def terminate(signum, frame): | |
| raise TerminationRequested() | |
| signal.signal(signal.SIGTERM, terminate) | |
| try: | |
| yield | |
| finally: | |
| signal.signal(signal.SIGTERM, previous) | |
| def _unique_object(pairs): | |
| result = {} | |
| for key, value in pairs: | |
| if key in result: | |
| raise InputError() | |
| result[key] = value | |
| return result | |
| def _bad_constant(value): | |
| raise InputError() | |
| def _local_path(value, base): | |
| if not isinstance(value, str) or not value or "://" in value or "\x00" in value: | |
| raise InputError() | |
| path = Path(value) | |
| if not path.is_absolute(): | |
| path = base / path | |
| if path.is_symlink() or not path.is_file(): | |
| raise InputError() | |
| return path.resolve(strict=True) | |
| class RequestInput: | |
| preset_id: str | |
| prompt: str | |
| canvas: str = "" | |
| duration: float = 5 | |
| seed: int = 42 | |
| image: Path | None = None | |
| last_image: Path | None = None | |
| def plan_args(self): | |
| return (self.preset_id, self.prompt, self.canvas, self.duration, self.seed, | |
| self.image is not None, self.last_image is not None) | |
| def load_input(filename): | |
| path = _local_path(str(filename), Path.cwd()) | |
| with _open_regular(path) as handle: | |
| raw = handle.read(1024 * 1024 + 1) | |
| if len(raw) > 1024 * 1024: | |
| raise InputError() | |
| data = json.loads(raw, object_pairs_hook=_unique_object, parse_constant=_bad_constant) | |
| allowed = {"preset_id", "prompt", "canvas", "duration", "seed", "image", "last_image"} | |
| if not isinstance(data, dict) or set(data) - allowed: | |
| raise InputError() | |
| for field in ("preset_id", "prompt"): | |
| if not isinstance(data.get(field), str) or not data[field].strip(): | |
| raise InputError() | |
| if (len(data["prompt"]) > 12000 or data["prompt"] != data["prompt"].strip() | |
| or not isinstance(data.get("canvas", ""), str)): | |
| raise InputError() | |
| duration, seed = data.get("duration", 5), data.get("seed", 42) | |
| if type(duration) not in (int, float) or not 2 <= duration <= 8 or not math.isfinite(duration): | |
| raise InputError() | |
| if type(seed) is not int or not 0 <= seed <= 2147483647: | |
| raise InputError() | |
| for field in ("image", "last_image"): | |
| if data.get(field) is not None: | |
| data[field] = _local_path(data[field], path.parent) | |
| return RequestInput(**data) | |
| def _open_regular(path): | |
| # Nonblocking prevents a swapped FIFO from hanging before fstat(). | |
| fd = os.open(path, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK) | |
| try: | |
| if not stat.S_ISREG(os.fstat(fd).st_mode): | |
| raise InputError() | |
| return os.fdopen(fd, "rb") | |
| except BaseException: | |
| os.close(fd) | |
| raise | |
| def _fsync_dir(path): | |
| fd = os.open(path, os.O_RDONLY) | |
| try: | |
| os.fsync(fd) | |
| finally: | |
| os.close(fd) | |
| def _sanitize(value, token=""): | |
| if isinstance(value, str): | |
| if token: | |
| value = value.replace(token, "[REDACTED]") | |
| return re.sub(r"hf_[A-Za-z0-9_]+", "[REDACTED]", value) | |
| if isinstance(value, dict): | |
| return {_sanitize(str(k), token): _sanitize(v, token) for k, v in value.items() | |
| if not re.search(r"token|authorization|password|secret|signature", str(k), re.I)} | |
| if isinstance(value, (list, tuple)): | |
| return [_sanitize(v, token) for v in value] | |
| if value is None or type(value) in (bool, int): | |
| return value | |
| if type(value) is float and math.isfinite(value): | |
| return value | |
| raise ValueError() | |
| def _atomic_json(path, data): | |
| # Both initial publication and replacements have private modes and fsync. | |
| fd, temporary = tempfile.mkstemp(prefix=".journal-", dir=path.parent) | |
| try: | |
| with os.fdopen(fd, "w", encoding="utf-8") as handle: | |
| os.fchmod(handle.fileno(), 0o600) | |
| json.dump(data, handle, ensure_ascii=True, allow_nan=False, sort_keys=True) | |
| handle.write("\n") | |
| handle.flush() | |
| os.fsync(handle.fileno()) | |
| os.replace(temporary, path) | |
| _fsync_dir(path.parent) | |
| finally: | |
| with suppress(FileNotFoundError): | |
| os.unlink(temporary) | |
| class Journal: | |
| def __init__(self, directory, token, revision): | |
| # Resolve the parent only: a final symlink must fail the exclusive mkdir. | |
| directory = Path(directory) | |
| self.directory = directory.parent.resolve(strict=True) / directory.name | |
| self.directory.mkdir(mode=0o700, exist_ok=False) | |
| self.directory.chmod(0o700) | |
| _fsync_dir(self.directory.parent) | |
| self.token = token | |
| self.data = { | |
| "schema_version": 2, "state": "preflight", "checkpoint": 0, | |
| "generator": GENERATOR, "generator_revision": revision, | |
| "overage_disabled_operator_attestation": True, | |
| "hops": {"workflow": {"state": "not_submitted", "event_id": None, "session_hash": None}}, | |
| } | |
| self.save() | |
| def save(self): | |
| self.data["checkpoint"] += 1 | |
| self.data["updated_at_unix"] = time.time() | |
| _atomic_json(self.directory / "journal.json", _sanitize(self.data, self.token)) | |
| def _copy_private(source, destination, expected=None): | |
| """Snapshot regular bytes, detect concurrent source writes, then publish.""" | |
| fd, temporary = tempfile.mkstemp(prefix=".copy-", dir=destination.parent) | |
| hasher = hashlib.sha256() | |
| size = 0 | |
| try: | |
| with os.fdopen(fd, "wb") as target, _open_regular(source) as original: | |
| before = os.fstat(original.fileno()) | |
| while chunk := original.read(1024 * 1024): | |
| hasher.update(chunk) | |
| size += len(chunk) | |
| target.write(chunk) | |
| after = os.fstat(original.fileno()) | |
| stamp = lambda s: (s.st_dev, s.st_ino, s.st_size, s.st_mtime_ns, s.st_ctime_ns) | |
| if stamp(before) != stamp(after) or size != before.st_size: | |
| raise IntegrityError() | |
| digest = hasher.hexdigest() | |
| if expected is not None and digest != expected: | |
| raise IntegrityError() | |
| os.fchmod(target.fileno(), 0o600) | |
| target.flush() | |
| os.fsync(target.fileno()) | |
| os.replace(temporary, destination) | |
| _fsync_dir(destination.parent) | |
| return {"path": destination.name, "sha256": digest, "bytes": size} | |
| finally: | |
| with suppress(FileNotFoundError): | |
| os.unlink(temporary) | |
| def _verify_file(path, expected): | |
| with _open_regular(path) as handle: | |
| actual = hashlib.file_digest(handle, "sha256").hexdigest() | |
| if actual != expected: | |
| raise IntegrityError() | |
| class _NoRedirect(HTTPRedirectHandler): | |
| def redirect_request(self, req, fp, code, msg, headers, newurl): | |
| # Never forward caller authorization to an unexpected destination. | |
| return None | |
| def _hub_json(route, token): | |
| request = Request("https://huggingface.co/api/" + route, | |
| headers={"Authorization": "Bearer " + token, "Accept": "application/json"}) | |
| with build_opener(_NoRedirect).open(request, timeout=20) as response: | |
| return json.load(response) | |
| def fetch_identity(token): | |
| return _hub_json("whoami-v2", token) | |
| def fetch_metadata(repo, token): | |
| if repo != GENERATOR_REPO: | |
| raise ValueError() | |
| return _hub_json("spaces/" + repo, token) | |
| def _client_factory(src, **options): | |
| from gradio_client import Client | |
| # 2.5.0 uses an OR with this env even when analytics_enabled=False. | |
| os.environ["GRADIO_ANALYTICS_ENABLED"] = "False" | |
| return Client(src, **options) | |
| def _file_wrapper(path): | |
| from gradio_client import handle_file | |
| return handle_file(str(path)) | |
| def _new_client(factory, src, token, downloads: bool | Path = False): | |
| # token=False is intentional: None permits saved SDK owner credentials. | |
| return factory(src, token=token, max_workers=8, verbose=False, | |
| download_files=downloads, analytics_enabled=False, | |
| httpx_kwargs={"timeout": 20}) | |
| def _within_deadline(call, deadline, clock, cleanup=None): | |
| """Bound a read-only preflight operation, never an inference submission. | |
| SDK predict/initialization may wait beyond an HTTP request timeout. The | |
| daemon performs only this one read; it cannot advance to a GPU submission | |
| or update the journal after the caller times out. Late-created clients close. | |
| """ | |
| if clock() >= deadline: | |
| raise TimeoutError() | |
| future = Future() | |
| def invoke(): | |
| try: | |
| future.set_result(call()) | |
| except BaseException as error: | |
| future.set_exception(error) | |
| threading.Thread(target=invoke, name="workflow-preflight", daemon=True).start() | |
| try: | |
| remaining = deadline - clock() | |
| if remaining <= 0: | |
| raise TimeoutError() | |
| return future.result(timeout=remaining) | |
| except BaseException: | |
| if cleanup is not None: | |
| def close_late(done): | |
| with suppress(Exception): | |
| if not done.cancelled() and done.exception() is None: | |
| cleanup(done.result()) | |
| future.add_done_callback(close_late) | |
| raise | |
| def _capture(hop, client, job=None): | |
| session = getattr(client, "session_hash", None) | |
| if isinstance(session, str): | |
| hop["session_hash"] = session | |
| event = getattr(getattr(job, "communicator", None), "event_id", None) | |
| if isinstance(event, str): | |
| hop["event_id"] = event | |
| def _submit_once(client, name, args, api, journal, deadline, clock): | |
| if clock() >= deadline: | |
| raise TimeoutError() | |
| hop = journal.data["hops"][name] | |
| # An abrupt process death after this fsync is ambiguous, never resubmittable. | |
| hop["state"] = "unknown" | |
| journal.data["state"] = "unknown" | |
| journal.data["active_hop"] = name | |
| _capture(hop, client) | |
| journal.save() | |
| job = None | |
| try: | |
| job = client.submit(*args, api_name=api) | |
| while True: | |
| _capture(hop, client, job) | |
| journal.save() | |
| remaining = deadline - clock() | |
| if remaining <= 0: | |
| raise TimeoutError() | |
| try: | |
| result = job.future.result(timeout=min(2.0, remaining)) | |
| except TimeoutError: | |
| if job.future.done(): | |
| raise | |
| continue | |
| _capture(hop, client, job) | |
| hop["state"] = "complete" | |
| journal.data["state"] = name + "_complete" | |
| journal.save() | |
| return result | |
| finally: | |
| # Event IDs often arrive while waiting, including just before errors. | |
| _capture(hop, client, job) | |
| journal.save() | |
| def _artifact(result, downloads): | |
| if not isinstance(result, (tuple, list)) or len(result) != 2 or not isinstance(result[1], dict): | |
| raise IntegrityError() | |
| value = result[0] | |
| # gr.Video may deserialize into (video_path, subtitle_path). | |
| if isinstance(value, (list, tuple)) and len(value) == 2: | |
| value = value[0] | |
| if isinstance(value, dict): | |
| value = value.get("path") or value.get("video") | |
| if isinstance(value, dict): | |
| value = value.get("path") | |
| if isinstance(value, Path): | |
| value = str(value) | |
| path = _local_path(value, downloads) | |
| if not path.is_relative_to(downloads.resolve()): | |
| raise IntegrityError() | |
| return path, result[1] | |
| def _sha(value): | |
| if not isinstance(value, str) or re.fullmatch(r"[0-9a-f]{64}", value) is None: | |
| raise IntegrityError() | |
| return value | |
| def _plan(client, request): | |
| result = client.predict(*request.plan_args(), api_name="/plan") | |
| if (not isinstance(result, dict) or not isinstance(result.get("canvas"), str) | |
| or not result["canvas"] or not isinstance(result.get("adapter"), str)): | |
| raise InputError() | |
| prompt = result.get("prompt") | |
| if not isinstance(prompt, str) or not prompt or prompt != prompt.strip() or len(prompt) > 12000: | |
| raise InputError() | |
| for field in ("height", "width", "num_frames", "steps"): | |
| if type(result.get(field)) is not int or result[field] <= 0: | |
| raise InputError() | |
| return result | |
| def _run(args, request, token, factory, metadata, identity, wrap, clock): | |
| journal = Journal(args.run_dir, token, args.expected_revision) | |
| clients = [] | |
| deadline = clock() + args.max_wall_seconds | |
| try: | |
| staged = {} | |
| for name in ("image", "last_image"): | |
| source = getattr(request, name) | |
| if source is not None: | |
| suffix = source.suffix if re.fullmatch(r"\.[A-Za-z0-9]{1,10}", source.suffix) else ".bin" | |
| target = journal.directory / (name + suffix) | |
| staged[name] = _copy_private(source, target) | |
| journal.data["inputs"] = staged | |
| journal.data["request"] = {"preset_id": request.preset_id, "prompt": request.prompt, | |
| "canvas": request.canvas, "duration": request.duration, "seed": request.seed} | |
| journal.save() | |
| caller = _within_deadline(lambda: identity(token), deadline, clock).get("name") | |
| if not isinstance(caller, str) or not caller: | |
| raise AuthenticationRequired() | |
| if caller.casefold() == "owenisas": | |
| raise OwnerTokenForbidden() | |
| journal.data["caller"] = caller | |
| if _within_deadline(lambda: metadata(GENERATOR_REPO, token), deadline, clock).get("sha") != args.expected_revision: | |
| raise RevisionMismatch() | |
| downloads = journal.directory / "workflow-downloads" | |
| downloads.mkdir(mode=0o700) | |
| generator = _within_deadline( | |
| lambda: _new_client(factory, GENERATOR, token, downloads), | |
| deadline, clock, cleanup=lambda client: client.close()) | |
| clients.append(generator) | |
| health = _within_deadline(lambda: generator.predict(api_name="/health"), deadline, clock) | |
| # Models are staged and loaded inside GPU calls, not held resident here. | |
| if (not isinstance(health, dict) or type(health.get("schema_version")) is not int | |
| or health["schema_version"] != 2 or health.get("status") != "ready" | |
| or health.get("blockers") != [] or health.get("runtime_initialized") is not True | |
| or health.get("conditioning_transport") != "in_process"): | |
| raise RuntimeNotReady() | |
| journal.data["health"] = health | |
| plan = _within_deadline(lambda: _plan(generator, request), deadline, clock) | |
| journal.data["plan"] = plan | |
| journal.save() | |
| def image(name): | |
| if name not in staged: | |
| return None | |
| path = journal.directory / staged[name]["path"] | |
| _verify_file(path, staged[name]["sha256"]) | |
| return wrap(path) | |
| first, last = image("image"), image("last_image") | |
| result = _submit_once(generator, "workflow", | |
| (request.preset_id, plan["prompt"], first, last, plan["canvas"], | |
| request.duration, request.seed), | |
| "/generate", journal, deadline, clock) | |
| path, report = _artifact(result, downloads) | |
| expected = _sha(report.get("output_sha256")) | |
| output = _copy_private(path, journal.directory / "final.mp4", expected) | |
| journal.data["result"] = _sanitize(report, token) | |
| journal.data["output"] = output | |
| journal.data["state"] = "succeeded" | |
| journal.save() | |
| return {"state": "succeeded", "output": str(journal.directory / "final.mp4"), | |
| "result": journal.data["result"]} | |
| except BaseException as error: | |
| journal.data["state"] = ("unknown" if any(h["state"] == "unknown" | |
| for h in journal.data["hops"].values()) else "failed") | |
| journal.data["error"] = type(error).__name__ | |
| journal.save() | |
| raise | |
| finally: | |
| for client in clients: | |
| with suppress(Exception): | |
| client.close() | |
| class _Parser(argparse.ArgumentParser): | |
| def error(self, message): | |
| raise CLIArgumentError() | |
| def _parser(): | |
| parser = _Parser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) | |
| commands = parser.add_subparsers(dest="command", required=True) | |
| commands.add_parser("presets", help="List public presets without credentials") | |
| commands.add_parser("health", help="Read public runtime readiness without credentials") | |
| plan = commands.add_parser("plan", help="Validate and plan; no files uploaded") | |
| plan.add_argument("--input", required=True, type=Path) | |
| run = commands.add_parser("run", help="One-shot authenticated single-Space inference") | |
| run.add_argument("--input", required=True, type=Path) | |
| run.add_argument("--run-dir", required=True, type=Path) | |
| run.add_argument("--expected-revision", required=True) | |
| run.add_argument("--token-env", default="HF_INFERENCE_TOKEN") | |
| run.add_argument("--max-wall-seconds", default=900.0, type=float) | |
| return parser | |
| def main(argv=None, *, env=None, client_factory=None, metadata_fetcher=None, | |
| identity_fetcher=None, file_wrapper=None, clock=time.monotonic): | |
| """Dependencies are injectable for offline tests; return an exit status.""" | |
| env = os.environ if env is None else env | |
| token = "" | |
| try: | |
| args = _parser().parse_args(argv) | |
| request = load_input(args.input) if args.command in ("plan", "run") else None | |
| factory = client_factory or _client_factory | |
| if args.command == "run": | |
| if re.fullmatch(r"[0-9a-f]{40}", args.expected_revision) is None: | |
| raise InputError() | |
| if not math.isfinite(args.max_wall_seconds) or not 0 < args.max_wall_seconds <= 86400: | |
| raise InputError() | |
| if re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", args.token_env) is None: | |
| raise InputError() | |
| token = env.get(args.token_env, "") | |
| if not isinstance(token, str) or not token or any(c.isspace() for c in token): | |
| raise AuthenticationRequired() | |
| if env.get("HF_POOL_ZEROGPU_OVERAGE_DISABLED") != "1": | |
| raise OverageAttestationRequired() | |
| # Do not relay SDK prints, provider errors, or warnings to CLI output. | |
| with _termination_guard(), open(os.devnull, "w") as sink, redirect_stdout(sink), redirect_stderr(sink): | |
| if args.command == "run": | |
| result = _run(args, request, token, factory, metadata_fetcher or fetch_metadata, | |
| identity_fetcher or fetch_identity, file_wrapper or _file_wrapper, clock) | |
| else: | |
| client = _new_client(factory, GENERATOR, False) | |
| try: | |
| result = (_plan(client, request) if args.command == "plan" | |
| else client.predict(api_name="/" + args.command)) | |
| finally: | |
| with suppress(Exception): | |
| client.close() | |
| print(json.dumps(_sanitize(result, token), sort_keys=True, allow_nan=False)) | |
| return 0 | |
| except SystemExit as error: | |
| # argparse help is the only non-JSON successful output. | |
| return int(error.code or 0) | |
| except (Exception, KeyboardInterrupt) as error: | |
| print(json.dumps({"state": "error", "error": _sanitize(type(error).__name__, token)}), file=sys.stderr) | |
| if isinstance(error, TerminationRequested): | |
| return 143 | |
| return 130 if isinstance(error, KeyboardInterrupt) else 1 | |
| if __name__ == "__main__": | |
| code = main() | |
| sys.stdout.flush() | |
| sys.stderr.flush() | |
| # ThreadPoolExecutor's atexit join can outlive an unknown remote job. The | |
| # CLI must exit within its polling bound without cancelling that job. | |
| # All durable writes above are fsynced before returning from main(). | |
| os._exit(code) | |