from __future__ import annotations
import base64
import importlib.util
import importlib.metadata
import html
import json
import os
import platform
import subprocess
import sys
import time
from functools import lru_cache
from pathlib import Path
import gradio as gr
try:
import spaces as hf_spaces
except Exception: # Local dev environments usually do not have HF ZeroGPU.
hf_spaces = None
HERE = Path(__file__).resolve().parent
DEMO_WEBMAP = HERE / "demo_webmap"
MANIFEST_PATH = DEMO_WEBMAP / "manifest.json"
LIVE_ROOT = Path(os.getenv("STORMSCOPE_LIVE_ROOT", "/tmp/stormscope-live"))
LIVE_TIMEOUT_SEC = int(os.getenv("STORMSCOPE_LIVE_TIMEOUT_SEC", "900"))
PYDEPS_ROOT = Path(os.getenv("STORMSCOPE_PYDEPS_ROOT", str(LIVE_ROOT / "pydeps")))
EARTH2GRID_REF = os.getenv("STORMSCOPE_EARTH2GRID_REF", "8fdff5a78d324f8d25afe224915301b3169bffe2")
def gpu_decorator(*args, **kwargs):
if hf_spaces is not None and hasattr(hf_spaces, "GPU"):
return hf_spaces.GPU(*args, **kwargs)
def passthrough(func):
return func
return passthrough
def data_uri(path: Path) -> str:
encoded = base64.b64encode(path.read_bytes()).decode("ascii")
return f"data:image/png;base64,{encoded}"
def embed_manifest(webmap_dir: Path, *, cached: bool = False) -> dict:
manifest_path = webmap_dir / "manifest.json"
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
for product in manifest["products"].values():
for frame in product["frames"]:
frame_path = webmap_dir / frame["url"]
frame["url"] = data_uri(frame_path)
frame.pop("source_npz", None)
manifest["source"] = "bundled-demo" if cached else "live"
return manifest
@lru_cache(maxsize=1)
def embedded_demo_manifest() -> dict:
return embed_manifest(DEMO_WEBMAP, cached=True)
def map_html(manifest: dict | None = None) -> str:
manifest = manifest or embedded_demo_manifest()
manifest_json = json.dumps(manifest)
document = f"""
"""
return (
''
)
def live_enabled() -> bool:
return os.getenv("STORMSCOPE_ENABLE_LIVE", "0").lower() in {"1", "true", "yes", "on"}
def package_available() -> bool:
return importlib.util.find_spec("rustwx_stormscope") is not None
def env_flag(name: str, default: str = "0") -> bool:
return os.getenv(name, default).lower() in {"1", "true", "yes", "on"}
def ensure_pydeps_path() -> None:
if PYDEPS_ROOT.exists():
path = str(PYDEPS_ROOT)
if path not in sys.path:
sys.path.insert(0, path)
def command_probe(cmd: list[str]) -> str:
try:
proc = subprocess.run(
cmd,
cwd=HERE,
text=True,
capture_output=True,
timeout=45,
)
except Exception as exc:
return "+ " + " ".join(cmd) + f"\nfailed: {type(exc).__name__}: {exc}"
stdout = proc.stdout.strip()
stderr = proc.stderr.strip()
lines = [
"+ " + " ".join(cmd),
f"exit code: {proc.returncode}",
]
if stdout:
lines.append(stdout[:1200])
if stderr:
lines.append(stderr[:1200])
return "\n".join(lines)
def bootstrap_earth2grid() -> str:
ensure_pydeps_path()
if importlib.util.find_spec("earth2grid") is not None:
return "earth2grid bootstrap: already importable"
if not env_flag("STORMSCOPE_BOOTSTRAP_EARTH2GRID"):
return "earth2grid bootstrap: disabled (set STORMSCOPE_BOOTSTRAP_EARTH2GRID=1)"
PYDEPS_ROOT.mkdir(parents=True, exist_ok=True)
requirement = f"earth2grid @ git+https://github.com/NVlabs/earth2grid.git@{EARTH2GRID_REF}"
env = os.environ.copy()
env["PYTHONPATH"] = str(PYDEPS_ROOT) + os.pathsep + env.get("PYTHONPATH", "")
cmd = [
sys.executable,
"-m",
"pip",
"install",
"--no-cache-dir",
"--no-build-isolation",
"--target",
str(PYDEPS_ROOT),
"--upgrade",
requirement,
]
try:
log = run_command(cmd, cwd=HERE, env=env)
except Exception as exc:
return f"earth2grid bootstrap failed: {type(exc).__name__}: {exc}"
importlib.invalidate_caches()
ensure_pydeps_path()
if importlib.util.find_spec("earth2grid") is None:
return log + "\n\nearth2grid bootstrap finished but module is still not importable"
return log + "\n\nearth2grid bootstrap: ok"
def module_probe(module_name: str, package_name: str | None = None) -> list[str]:
package_name = package_name or module_name
lines = [f"[{module_name}]"]
spec = importlib.util.find_spec(module_name)
lines.append(f"spec: {spec is not None}")
try:
version = importlib.metadata.version(package_name)
lines.append(f"version: {version}")
except Exception as exc:
lines.append(f"version: unavailable ({type(exc).__name__})")
if spec is None:
return lines
started = time.monotonic()
try:
__import__(module_name)
lines.append(f"import: ok ({time.monotonic() - started:.2f}s)")
except Exception as exc:
lines.append(f"import: failed ({type(exc).__name__}: {exc})")
return lines
def package_probe() -> str:
lines = [
"StormScope package probe",
f"python: {platform.python_version()}",
f"rustwx_stormscope importable: {package_available()}",
]
if not package_available():
lines.append("")
lines.append("The Space payload does not include the release package source.")
lines.append("Push with `scripts/push-hf-space.py --profile package-probe` to test package import and CLI help.")
return "\n".join(lines)
lines.append("")
lines.append(command_probe([sys.executable, "-m", "rustwx_stormscope.run", "--help"]))
lines.append("")
lines.append(command_probe([sys.executable, "-m", "rustwx_stormscope.webmap_preview", "--help"]))
return "\n".join(lines)
@gpu_decorator(duration=120)
def weather_stack_probe() -> str:
lines = [
"Weather model live-inference stack probe",
f"python: {platform.python_version()}",
f"platform: {platform.platform()}",
f"STORMSCOPE_ENABLE_LIVE: {os.getenv('STORMSCOPE_ENABLE_LIVE', '0')}",
f"cwd: {HERE}",
"",
]
try:
import torch
lines.extend(
[
"[torch/cuda]",
f"torch: {torch.__version__}",
f"cuda available: {torch.cuda.is_available()}",
]
)
if torch.cuda.is_available():
lines.extend(
[
f"gpu: {torch.cuda.get_device_name(0)}",
f"capability: {torch.cuda.get_device_capability(0)}",
f"arch list: {torch.cuda.get_arch_list()}",
f"bf16 supported: {torch.cuda.is_bf16_supported()}",
f"device memory GB: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f}",
]
)
except Exception as exc:
lines.append(f"[torch/cuda]\nfailed: {type(exc).__name__}: {exc}")
lines.append("")
lines.append("[runtime bootstrap]")
lines.append(bootstrap_earth2grid())
lines.append("")
for module_name, package_name in [
("rustwx_stormscope", "rustwx-stormscope"),
("earth2studio", "earth2studio"),
("earth2grid", "earth2grid"),
("physicsnemo", "nvidia-physicsnemo"),
("natten", "natten"),
("torchao", "torchao"),
("h5netcdf", "h5netcdf"),
("xarray", "xarray"),
("scipy", "scipy"),
]:
lines.extend(module_probe(module_name, package_name))
lines.append("")
lines.append("[earth2studio StormScope import]")
started = time.monotonic()
try:
from earth2studio.models.px.stormscope import StormScopeBase, StormScopeGOES, StormScopeMRMS
lines.append(f"import: ok ({time.monotonic() - started:.2f}s)")
lines.append(f"classes: {StormScopeBase.__name__}, {StormScopeGOES.__name__}, {StormScopeMRMS.__name__}")
except Exception as exc:
lines.append(f"import: failed ({type(exc).__name__}: {exc})")
lines.append("")
lines.append("[nvidia-smi]")
lines.append(command_probe(["nvidia-smi"]))
return "\n".join(lines)
def live_status() -> str:
lines = [
"Live inference status",
f"STORMSCOPE_ENABLE_LIVE: {os.getenv('STORMSCOPE_ENABLE_LIVE', '0')}",
f"rustwx_stormscope importable: {package_available()}",
f"live root: {LIVE_ROOT}",
f"timeout seconds: {LIVE_TIMEOUT_SEC}",
]
if not live_enabled():
lines.append("")
lines.append("Live inference is disabled for this public demo deployment.")
lines.append("Set STORMSCOPE_ENABLE_LIVE=1 only after the Space image has the full StormScope dependency stack.")
elif not package_available():
lines.append("")
lines.append("Live inference is enabled, but rustwx_stormscope is not installed in this Space runtime.")
else:
lines.append("")
lines.append("Live inference is enabled and the package is importable.")
return "\n".join(lines)
def subprocess_env() -> dict:
env = os.environ.copy()
cache_root = Path(env.get("STORMSCOPE_CACHE_DIR", str(LIVE_ROOT / "cache")))
env.setdefault("RUSTWX_STORMSCOPE_CACHE_DIR", str(cache_root))
env.setdefault("EARTH2STUDIO_CACHE", str(cache_root / "earth2studio"))
env.setdefault("RUSTWX_STORMSCOPE_GOES_CACHE", str(cache_root / "goes_nc"))
env.setdefault("RUSTWX_STORMSCOPE_MRMS_CACHE", str(cache_root / "mrms_grib2"))
env.setdefault("TORCHINDUCTOR_CACHE_DIR", str(cache_root / "torchinductor"))
env.setdefault("TRITON_CACHE_DIR", str(cache_root / "triton"))
env["PYTHONPATH"] = str(PYDEPS_ROOT) + os.pathsep + env.get("PYTHONPATH", "")
return env
def run_command(cmd: list[str], *, cwd: Path, env: dict) -> str:
started = time.monotonic()
proc = subprocess.run(
cmd,
cwd=cwd,
env=env,
text=True,
capture_output=True,
timeout=LIVE_TIMEOUT_SEC,
)
elapsed = time.monotonic() - started
lines = [
"+ " + " ".join(cmd),
f"exit code: {proc.returncode}",
f"elapsed: {elapsed:.1f}s",
]
if proc.stdout.strip():
lines.append("")
lines.append(proc.stdout.strip()[-4000:])
if proc.stderr.strip():
lines.append("")
lines.append(proc.stderr.strip()[-4000:])
if proc.returncode != 0:
raise RuntimeError("\n".join(lines))
return "\n".join(lines)
def latest_meta(outdir: Path) -> Path:
metas = sorted(outdir.glob("stormscope_nowcast_6km_*_meta.json"), key=lambda path: path.stat().st_mtime)
if not metas:
raise FileNotFoundError(f"no StormScope metadata written in {outdir}")
return metas[-1]
def live_duration(init_time: str, frames: int, num_steps: int, backend: str, mxfp8: bool) -> int:
del init_time, backend, mxfp8
return min(240, max(90, int(frames) * max(30, int(num_steps) * 4)))
@gpu_decorator(duration=live_duration)
def run_live_nowcast(init_time: str, frames: int, num_steps: int, backend: str, mxfp8: bool) -> tuple[str, str]:
if not live_enabled():
return map_html(), live_status()
if not package_available():
return map_html(), live_status()
ensure_pydeps_path()
frames = max(1, min(int(frames), 3))
num_steps = max(4, min(int(num_steps), 24))
backend = backend if backend in {"earth2studio", "rust"} else "earth2studio"
init_time = (init_time or "").strip()
run_id = time.strftime("%Y%m%d_%H%M%S")
run_root = LIVE_ROOT / "runs" / run_id
outdir = run_root / "out"
webmap_dir = run_root / "webmap"
outdir.mkdir(parents=True, exist_ok=True)
webmap_dir.mkdir(parents=True, exist_ok=True)
env = subprocess_env()
for path in [
env["RUSTWX_STORMSCOPE_CACHE_DIR"],
env["EARTH2STUDIO_CACHE"],
env["RUSTWX_STORMSCOPE_GOES_CACHE"],
env["RUSTWX_STORMSCOPE_MRMS_CACHE"],
]:
Path(path).mkdir(parents=True, exist_ok=True)
run_cmd = [
sys.executable,
"-m",
"rustwx_stormscope.run",
"--mode",
"nowcast",
"--res",
"6km",
"--steps",
str(frames),
"--num-steps",
str(num_steps),
"--no-compile",
"--backend",
backend,
"--outdir",
str(outdir),
]
if init_time and init_time.lower() not in {"latest", "auto"}:
run_cmd.extend(["--date", init_time])
if mxfp8:
run_cmd.append("--mxfp8")
preview_cmd = [
sys.executable,
"-m",
"rustwx_stormscope.webmap_preview",
"--outdir",
str(outdir),
"--output-dir",
str(webmap_dir),
"--products",
"refc,ir",
"--view",
"central_plains",
"--max-frames",
str(frames),
"--force",
]
logs = [live_status(), ""]
try:
logs.append(bootstrap_earth2grid())
logs.append("")
logs.append(run_command(run_cmd, cwd=HERE, env=env))
meta = latest_meta(outdir)
preview_cmd.extend(["--meta", str(meta)])
logs.append("")
logs.append(run_command(preview_cmd, cwd=HERE, env=env))
manifest = embed_manifest(webmap_dir)
return map_html(manifest), "\n".join(logs)
except subprocess.TimeoutExpired as exc:
return map_html(), "\n".join(logs + [f"live inference timed out after {exc.timeout}s"])
except Exception as exc:
return map_html(), "\n".join(logs + [f"live inference failed: {type(exc).__name__}: {exc}"])
@gpu_decorator(duration=45)
def zero_gpu_probe() -> str:
lines = [
"ZeroGPU runtime probe",
f"python: {platform.python_version()}",
f"platform: {platform.platform()}",
f"STORMSCOPE_ENABLE_LIVE: {os.getenv('STORMSCOPE_ENABLE_LIVE', '0')}",
]
try:
import torch
lines.append(f"torch: {torch.__version__}")
lines.append(f"cuda available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
lines.append(f"gpu: {torch.cuda.get_device_name(0)}")
lines.append(f"capability: {torch.cuda.get_device_capability(0)}")
lines.append(f"arch list: {torch.cuda.get_arch_list()}")
except Exception as exc:
lines.append(f"torch probe failed: {type(exc).__name__}: {exc}")
lines.append("")
lines.append(f"Live StormScope inference enabled: {live_enabled()}")
lines.append("The production fast path targets a validated CUDA/PyTorch/NATTEN/torchao stack with MXFP8.")
lines.append("This probe reports the GPU Hugging Face assigns at launch time; hardware can vary.")
lines.append("This Space proves the public web-map product and hosted UX; live inference needs a separately validated HF dependency path.")
lines.append("")
lines.append(package_probe())
lines.append("")
lines.append(live_status())
return "\n".join(lines)
with gr.Blocks(title="StormScope Fast Demo") as demo:
gr.Markdown(
"""
# StormScope Fast Demo
StormScope 10-minute nowcast frames rendered as transparent overlays on an
OpenStreetMap basemap. This is the hosted visualization shape for the
`rustwx-stormscope` release package.
"""
)
map_output = gr.HTML(map_html())
gr.Markdown("## Live Server Probe")
with gr.Row():
probe = gr.Button("Check ZeroGPU Runtime", variant="secondary")
package_button = gr.Button("Check Package CLI", variant="secondary")
weather_button = gr.Button("Check Weather Stack", variant="secondary")
probe_output = gr.Textbox(label="Runtime probe", lines=12)
probe.click(zero_gpu_probe, outputs=probe_output)
package_button.click(package_probe, outputs=probe_output)
weather_button.click(weather_stack_probe, outputs=probe_output)
gr.Markdown("## Live Nowcast")
with gr.Row():
init_time = gr.Textbox(label="Init time", value="2026-05-30T22:40:00", scale=2)
frames = gr.Slider(label="Frames", minimum=1, maximum=3, step=1, value=1, scale=1)
sampler_steps = gr.Slider(label="Sampler steps", minimum=4, maximum=24, step=1, value=8, scale=1)
with gr.Row():
backend = gr.Dropdown(label="Data backend", choices=["earth2studio", "rust"], value="earth2studio")
mxfp8 = gr.Checkbox(label="MXFP8", value=False)
live_button = gr.Button("Run Live Nowcast", variant="primary")
live_output = gr.Textbox(label="Live inference log", value=live_status(), lines=16)
live_button.click(
run_live_nowcast,
inputs=[init_time, frames, sampler_steps, backend, mxfp8],
outputs=[map_output, live_output],
)
gr.Markdown(
"""
The bundled cycle is pre-rendered from StormScope output. Full live inference is
guarded by `STORMSCOPE_ENABLE_LIVE=1` and uses the packaged CLI when the full
CUDA/PyTorch/NATTEN/torchao stack, model weights, and live satellite fetch path
are available on ZeroGPU.
"""
)
if __name__ == "__main__":
demo.launch()