Spaces:
Sleeping
Sleeping
| # SPDX-License-Identifier: Apache-2.0 | |
| # Copyright (c) Axolotl AI | |
| # Licensed under the Apache License, Version 2.0 | |
| """Minimal Hugging Face Spaces ZeroGPU probe surface. | |
| This module is intentionally small and dependency-light: importing it does not | |
| import Gradio, spaces, or torch. The prototype Space imports those packages | |
| lazily only when the Space app is built or the GPU probe is invoked. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import time | |
| from collections.abc import Callable | |
| from typing import Any | |
| DEFAULT_PROBE_DURATION_S = 10 | |
| MAX_ZEROGPU_DURATION_S = 120 | |
| def coerce_duration(value: Any = None) -> int: | |
| """Clamp a requested ZeroGPU duration to the supported prototype range.""" | |
| if value is None or value == "": | |
| value = os.environ.get("HATCHERY_ZEROGPU_PROBE_DURATION", DEFAULT_PROBE_DURATION_S) | |
| try: | |
| duration = int(float(value)) | |
| except (TypeError, ValueError): | |
| duration = DEFAULT_PROBE_DURATION_S | |
| return max(1, min(MAX_ZEROGPU_DURATION_S, duration)) | |
| def probe_duration(duration_s: Any = None) -> int: | |
| """Callable duration hook compatible with ``@spaces.GPU(duration=...)``.""" | |
| return coerce_duration(duration_s) | |
| def _import_spaces() -> Any | None: | |
| try: | |
| import spaces | |
| except ImportError: | |
| return None | |
| return spaces | |
| def _decorate_gpu( | |
| fn: Callable[[Any], dict[str, Any]], | |
| *, | |
| spaces_mod: Any | None = None, | |
| ) -> Callable[[Any], dict[str, Any]]: | |
| spaces_mod = _import_spaces() if spaces_mod is None else spaces_mod | |
| if spaces_mod is None: | |
| return fn | |
| return spaces_mod.GPU(duration=probe_duration)(fn) | |
| def make_gpu_probe(*, spaces_mod: Any | None = None) -> Callable[[Any], dict[str, Any]]: | |
| """Build the ZeroGPU-decorated probe function. | |
| The returned callable imports torch inside the decorated function body. That | |
| preserves the ZeroGPU rule that GPU work happens inside ``@spaces.GPU``. | |
| """ | |
| def _probe(duration_s: Any = None) -> dict[str, Any]: | |
| requested_duration_s = coerce_duration(duration_s) | |
| started = time.perf_counter() | |
| import torch | |
| cuda_available = bool(torch.cuda.is_available()) | |
| device_count = int(torch.cuda.device_count()) if cuda_available else 0 | |
| device_name = torch.cuda.get_device_name(0) if cuda_available else None | |
| allocation_ok = False | |
| if cuda_available: | |
| torch.zeros(1, device="cuda") | |
| torch.cuda.synchronize() | |
| allocation_ok = True | |
| return { | |
| "ok": True, | |
| "requested_duration_s": requested_duration_s, | |
| "elapsed_s": round(time.perf_counter() - started, 4), | |
| "torch_version": getattr(torch, "__version__", "unknown"), | |
| "cuda_available": cuda_available, | |
| "cuda_device_count": device_count, | |
| "cuda_device_name": device_name, | |
| "cuda_allocation_ok": allocation_ok, | |
| } | |
| return _decorate_gpu(_probe, spaces_mod=spaces_mod) | |
| def build_fastapi_app( | |
| gpu_probe: Callable[[Any], dict[str, Any]] | None = None, | |
| *, | |
| title: str = "Hatchery ZeroGPU Probe API", | |
| ) -> Any: | |
| """Build the FastAPI probe app mounted by the prototype Gradio Space.""" | |
| from fastapi import FastAPI, Query | |
| probe = gpu_probe or make_gpu_probe() | |
| app = FastAPI(title=title) | |
| def health() -> dict[str, Any]: | |
| return { | |
| "ok": True, | |
| "service": "hatchery-zerogpu-probe", | |
| "max_duration_s": MAX_ZEROGPU_DURATION_S, | |
| } | |
| def zerogpu_probe( | |
| duration_s: float | None = Query(default=None, ge=1, le=MAX_ZEROGPU_DURATION_S), | |
| ) -> dict[str, Any]: | |
| return probe(duration_s) | |
| return app | |
| def build_demo() -> Any: | |
| """Build the Gradio UI and mount the FastAPI probe app below ``/v1``.""" | |
| import gradio as gr | |
| probe = make_gpu_probe() | |
| with gr.Blocks(title="Hatchery ZeroGPU Probe") as demo: | |
| gr.Markdown("# Hatchery ZeroGPU Probe") | |
| gr.Markdown( | |
| "Run a tiny GPU allocation probe from Gradio, then test the mounted " | |
| "FastAPI route at `/v1/zerogpu/probe`." | |
| ) | |
| duration = gr.Slider( | |
| minimum=1, | |
| maximum=MAX_ZEROGPU_DURATION_S, | |
| value=DEFAULT_PROBE_DURATION_S, | |
| step=1, | |
| label="Requested duration seconds", | |
| ) | |
| output = gr.JSON(label="Probe result") | |
| run = gr.Button("Run GPU probe", variant="primary") | |
| run.click(fn=probe, inputs=[duration], outputs=[output], api_name="zerogpu_probe") | |
| api = build_fastapi_app(probe) | |
| demo.app.mount("/v1", api) | |
| demo.app.mount("/api/v1", api) | |
| return demo | |