Spaces:
Paused
Paused
Leave new observed-boundary and KW jobs to compatible Jetstream2 workers
Browse files
app.py
CHANGED
|
@@ -1,841 +1,855 @@
|
|
| 1 |
-
"""CREST_gpu_worker β ZeroGPU event worker (V27 P1.5 + P1.6 sessions). Since
|
| 2 |
-
the HPC worker was released (2026-10-16) it is the only 2-D engine: it claims
|
| 3 |
-
events/queue/ jobs, solves them in <=120 s GPU slices within a daily GPU
|
| 4 |
-
budget (authenticated calls draw on the account's ZeroGPU quota), and β when
|
| 5 |
-
PUBLISH=1 β publishes exactly like the HPC worker did.
|
| 6 |
-
|
| 7 |
-
P1.6: solver state is a resident crestimap.session.EventSession per episode
|
| 8 |
-
(depth AND momentum held in RAM between bundles), so returning to an episode
|
| 9 |
-
this Space already solved is a minutes-long catch-up visit from the last
|
| 10 |
-
obs-supported junction, not an hours-long anchored re-solve. Sessions for
|
| 11 |
-
episodes this worker has visited get scheduling priority over new events.
|
| 12 |
-
|
| 13 |
-
Queue protocol: identical to crestimap/worker.py (claim/heartbeat/stale
|
| 14 |
-
takeover/failed.json). 30 m (dem_res "1") jobs only; never coarsens β an
|
| 15 |
-
oversized or 10 m job is simply not claimed.
|
| 16 |
-
|
| 17 |
-
ZeroGPU note: @spaces.GPU functions run in a forked child with the GPU;
|
| 18 |
-
arguments/returns are pickled, and parent-side mutations inside the child are
|
| 19 |
-
lost. gpu_chunk() therefore RETURNS session.state_tuple() and the parent
|
| 20 |
-
reapplies it. Files the child writes (depth frames) land on the shared fs.
|
| 21 |
-
|
| 22 |
-
Never prints or logs the HF token.
|
| 23 |
-
"""
|
| 24 |
-
from __future__ import annotations
|
| 25 |
-
|
| 26 |
-
import datetime
|
| 27 |
-
import io
|
| 28 |
-
import json
|
| 29 |
-
import os
|
| 30 |
-
import shutil
|
| 31 |
-
import subprocess
|
| 32 |
-
import sys
|
| 33 |
-
import tarfile
|
| 34 |
-
import tempfile
|
| 35 |
-
import threading
|
| 36 |
-
import time
|
| 37 |
-
import traceback
|
| 38 |
-
|
| 39 |
-
# ---------------------------------------------------------------- bootstrap
|
| 40 |
-
APP = os.path.dirname(os.path.abspath(__file__))
|
| 41 |
-
FORK = "https://github.com/mchen15ouedu/CREST-iMAP.git"
|
| 42 |
-
CREST_AI = "https://huggingface.co/spaces/vincewin/CREST_demo"
|
| 43 |
-
# Solver commit, pinned like the CPU Spaces' Dockerfile CRESTIMAP_REF and the
|
| 44 |
-
# Jetstream2 kit's FORK_REF: a push to v2 must never reach this worker on a
|
| 45 |
-
# restart without review. c5184759 = draining-time outflow limiter +
|
| 46 |
-
# interflow store, approved 2026-10-20. Full SHA only (a branch would float).
|
| 47 |
-
FORK_REF = os.environ.get("CRESTIMAP_REF",
|
| 48 |
-
"c518475956c933a9c0fc46962b3825685a54b181")
|
| 49 |
-
|
| 50 |
-
|
| 51 |
-
def _git(*args, cwd=None):
|
| 52 |
-
return subprocess.run(["git", *args], cwd=cwd, check=True,
|
| 53 |
-
capture_output=True, text=True).stdout.strip()
|
| 54 |
-
|
| 55 |
-
|
| 56 |
-
def _bootstrap():
|
| 57 |
-
"""Blobless sparse clones: crestimap/ from the fork at FORK_REF, hf_data/
|
| 58 |
-
from the CREST_demo Space (eventsim/eventstore reuse β nothing
|
| 59 |
-
reimplemented).
|
| 60 |
-
|
| 61 |
-
crestimap is imported straight from the clone (a detached sparse checkout
|
| 62 |
-
of FORK_REF), so provenance.stamp() records the pinned SHA on every
|
| 63 |
-
manifest; a copied-out package has no .git and stamped only
|
| 64 |
-
"2.0.0.dev0"."""
|
| 65 |
-
fork = os.path.join(APP, "_fork")
|
| 66 |
-
if not os.path.isdir(os.path.join(fork, "crestimap")):
|
| 67 |
-
shutil.rmtree(fork, ignore_errors=True)
|
| 68 |
-
# all branches are fetched (blobless), so a pinned commit stays
|
| 69 |
-
# reachable whichever branch carries it
|
| 70 |
-
_git("clone", "--filter=blob:none", "--no-checkout", FORK, fork)
|
| 71 |
-
_git("sparse-checkout", "init", "--cone", cwd=fork)
|
| 72 |
-
_git("sparse-checkout", "set", "crestimap", cwd=fork)
|
| 73 |
-
_git("checkout", "--detach", FORK_REF, cwd=fork)
|
| 74 |
-
want = _git("rev-parse", "--verify", f"{FORK_REF}^{{commit}}", cwd=fork)
|
| 75 |
-
if _git("rev-parse", "HEAD", cwd=fork) != want:
|
| 76 |
-
raise RuntimeError(f"crestimap clone is not at the pinned {FORK_REF}")
|
| 77 |
-
if fork not in sys.path:
|
| 78 |
-
sys.path.insert(0, fork)
|
| 79 |
-
ai = os.path.join(APP, "_crest_ai")
|
| 80 |
-
if not os.path.isdir(os.path.join(ai, "hf_data")):
|
| 81 |
-
subprocess.run(["git", "clone", "--filter=blob:none", "--no-checkout",
|
| 82 |
-
CREST_AI, ai], check=True)
|
| 83 |
-
subprocess.run(["git", "checkout", "HEAD", "--", "hf_data"],
|
| 84 |
-
cwd=ai, check=True)
|
| 85 |
-
if ai not in sys.path:
|
| 86 |
-
sys.path.insert(0, ai)
|
| 87 |
-
|
| 88 |
-
|
| 89 |
-
_bootstrap()
|
| 90 |
-
try:
|
| 91 |
-
_FORK_SHA = subprocess.run(
|
| 92 |
-
["git", "-C", os.path.join(APP, "_fork"), "rev-parse", "HEAD"],
|
| 93 |
-
capture_output=True, text=True, timeout=15).stdout.strip()[:12] or None
|
| 94 |
-
except Exception:
|
| 95 |
-
_FORK_SHA = None
|
| 96 |
-
|
| 97 |
-
import gradio as gr # noqa: E402
|
| 98 |
-
|
| 99 |
-
try:
|
| 100 |
-
import spaces # noqa: E402
|
| 101 |
-
HAS_ZERO = True
|
| 102 |
-
except Exception: # local dev
|
| 103 |
-
HAS_ZERO = False
|
| 104 |
-
|
| 105 |
-
TS = "%Y-%m-%dT%H:%M:%SZ"
|
| 106 |
-
TMIN = "%Y-%m-%dT%H:%MZ"
|
| 107 |
-
QPREFIX = "events/queue"
|
| 108 |
-
REPO = os.environ.get("CREST_FEEDBACK_REPO", "vincewin/CREST_data")
|
| 109 |
-
# The HPC worker was released 2026-10-16; this Space is now the only 2-D
|
| 110 |
-
# engine, so there is no reason to leave a job unclaimed for anyone else.
|
| 111 |
-
ZERO_DELAY_S = float(os.environ.get("ZERO_DELAY_S", "0"))
|
| 112 |
-
# GPU calls must be AUTHENTICATED to get the account's ZeroGPU quota (PRO:
|
| 113 |
-
# 40 min/day, fixed 24 h window from first use). Self-calls over 127.0.0.1
|
| 114 |
-
# carry no X-IP-Token and land in a shared anonymous "Space app" pool that
|
| 115 |
-
# gave ~0 GPU-s/day from 2026-09-23 to 10-16 ("Space app has reached its GPU
|
| 116 |
-
# limit" on every claim). ZERO_AUTH=1 self-drives through the public Space
|
| 117 |
-
# URL with HF_TOKEN instead.
|
| 118 |
-
ZERO_AUTH = os.environ.get("ZERO_AUTH", "1") == "1"
|
| 119 |
-
# Our own daily cap, below the 40-min PRO quota: the quota is shared with
|
| 120 |
-
# every other ZeroGPU use under the same account, and past it ZeroGPU bills
|
| 121 |
-
# prepaid credits automatically ($1 per 10 min). 0 disables the cap.
|
| 122 |
-
DAILY_BUDGET_S = float(os.environ.get("ZERO_DAILY_BUDGET_S", "1800"))
|
| 123 |
-
GPU_CALL_S = 120.0 # declared @spaces.GPU duration (admission needs it)
|
| 124 |
-
if 0 < DAILY_BUDGET_S < GPU_CALL_S:
|
| 125 |
-
DAILY_BUDGET_S = GPU_CALL_S # a smaller cap could never run
|
| 126 |
-
# a NEW event (full anchored solve) is only claimed with this much budget
|
| 127 |
-
# left; otherwise it would be claimed, paused ~a day, and solved stale
|
| 128 |
-
MIN_NEW_S = float(os.environ.get("ZERO_MIN_NEW_S", "600"))
|
| 129 |
-
# "No GPU was available after 60s" is platform capacity, not this job's
|
| 130 |
-
# fault and not quota: wait and retry the same chunk.
|
| 131 |
-
CAPACITY_BACKOFF_S = float(os.environ.get("ZERO_CAPACITY_BACKOFF_S", "120"))
|
| 132 |
-
# Backup -> active drain under pressure (user directive 2026-08-14): with
|
| 133 |
-
# more than this many jobs pending, skip the polite deferral and claim the
|
| 134 |
-
# next unclaimed event immediately β one solver per event, so this Space
|
| 135 |
-
# always works a DIFFERENT event than the HPC.
|
| 136 |
-
PRESSURE_N = int(os.environ.get("ZERO_PRESSURE_N", "4"))
|
| 137 |
-
STALE_S = float(os.environ.get("ZERO_STALE_S", "600"))
|
| 138 |
-
HB_S = float(os.environ.get("ZERO_HB_S", "240"))
|
| 139 |
-
POLL_S = float(os.environ.get("ZERO_POLL_S", "60"))
|
| 140 |
-
# ZeroGPU quota exhaustion is NOT a poison job: release the claim (the HPC
|
| 141 |
-
# or a later quota window takes the bundle) and stop claiming for a while.
|
| 142 |
-
# Without this (seen 2026-09-11) the backup rung claimed a bundle, hit
|
| 143 |
-
# "Space app has reached its GPU limit", and CONSUMED it as failed.
|
| 144 |
-
QUOTA_BACKOFF_S = float(os.environ.get("ZERO_QUOTA_BACKOFF_S", "3600"))
|
| 145 |
-
MAX_CELLS = int(os.environ.get("ZERO_MAX_CELLS", "2000000"))
|
| 146 |
-
PUBLISH = os.environ.get("PUBLISH", "0") == "1"
|
| 147 |
-
CHUNK0_SIM_S = float(os.environ.get("ZERO_CHUNK0_SIM_S", "300"))
|
| 148 |
-
GPU_TARGET_S = float(os.environ.get("ZERO_GPU_TARGET_S", "80"))
|
| 149 |
-
DEM_RES_DEG = {"1": 1.0 / 3600.0, "13": 1.0 / 10800.0}
|
| 150 |
-
IDENT = f"zerogpu:{os.environ.get('SPACE_ID', 'local')}:{os.getpid()}"
|
| 151 |
-
os.environ.setdefault("EVENT_ENGINE_IDENT", IDENT) # provenance stamp
|
| 152 |
-
|
| 153 |
-
_tok = os.environ.get("HF_TOKEN", "").strip()
|
| 154 |
-
if _tok:
|
| 155 |
-
os.environ["HF_TOKEN"] = _tok # eventstore reads env
|
| 156 |
-
|
| 157 |
-
MAX_SESSIONS = int(os.environ.get("ZERO_MAX_SESSIONS", "6"))
|
| 158 |
-
# > one daily budget pause (~23.5 h), so parked sessions survive the pause and
|
| 159 |
-
# the next bundle is a cheap catch-up rather than a full re-solve
|
| 160 |
-
EVICT_IDLE_H = float(os.environ.get("ZERO_EVICT_IDLE_H", "36"))
|
| 161 |
-
# cadence of the server-side tick loop (see _self_drive); matches the old
|
| 162 |
-
# gr.Timer(5) so idle polling behaviour is unchanged
|
| 163 |
-
SELF_TICK_S = float(os.environ.get("ZERO_SELF_TICK_S", "5"))
|
| 164 |
-
|
| 165 |
-
LOG: list[str] = []
|
| 166 |
-
W = {"phase": "idle", "job": None, "runner": None, "chunk_sim_s": CHUNK0_SIM_S,
|
| 167 |
-
"last_poll": 0.0, "last_hb": 0.0, "queued": "", "workdir": None,
|
| 168 |
-
"done": set(), "solved": [], "t_wall0": 0.0, "last_tick": 0.0,
|
| 169 |
-
# daily GPU budget: [(start, seconds), ...] of GPU calls in a SLIDING
|
| 170 |
-
# 24 h window, so use stays under the cap in any 24 h (hence in any of
|
| 171 |
-
# ZeroGPU's own windows); RAM-only β keep the account's prepaid credit
|
| 172 |
-
# balance at $0 so an over-run after a restart becomes a quota error
|
| 173 |
-
"gpu_calls": [], "capacity_until": 0.0, "quota_until": 0.0,
|
| 174 |
-
"paused_note": "", "auth_suspect": False,
|
| 175 |
-
# /tick calls by outcome (see tick_api); "accepted" must keep rising
|
| 176 |
-
"tick_api": {"accepted": 0, "rejected": 0, "empty": 0,
|
| 177 |
-
"last_reject_log": 0.0},
|
| 178 |
-
# what the last _scan saw (status "scan"): time, queue ids, backlog,
|
| 179 |
-
# candidates and per-reason skip counts
|
| 180 |
-
"scan": {},
|
| 181 |
-
# P1.6 resident sessions: ev -> {"session", "bbox", "visits", "last",
|
| 182 |
-
# "root"}. RAM-only; a Space rebuild recovers via the bundle's
|
| 183 |
-
# init_depth (state-hierarchy level 2).
|
| 184 |
-
"sessions": {}}
|
| 185 |
-
_tick_lock = threading.Lock()
|
| 186 |
-
# secret the self-drive loop passes to /tick (in-process only; never logged)
|
| 187 |
-
_SELF_KEY = __import__("secrets").token_hex(16)
|
| 188 |
-
|
| 189 |
-
|
| 190 |
-
def log(msg: str):
|
| 191 |
-
LOG.append(f"[{datetime.datetime.utcnow():%H:%M:%S}Z] {msg}")
|
| 192 |
-
del LOG[:-300]
|
| 193 |
-
print(LOG[-1], flush=True)
|
| 194 |
-
|
| 195 |
-
|
| 196 |
-
def _api():
|
| 197 |
-
from huggingface_hub import HfApi
|
| 198 |
-
return HfApi(token=_tok or None)
|
| 199 |
-
|
| 200 |
-
|
| 201 |
-
def _read_json(path: str):
|
| 202 |
-
from huggingface_hub import hf_hub_download
|
| 203 |
-
try:
|
| 204 |
-
p = hf_hub_download(REPO, path, repo_type="dataset", token=_tok or None,
|
| 205 |
-
force_download=True)
|
| 206 |
-
with open(p, encoding="utf-8") as fp:
|
| 207 |
-
return json.load(fp)
|
| 208 |
-
except Exception:
|
| 209 |
-
return None
|
| 210 |
-
|
| 211 |
-
|
| 212 |
-
def _commit(ops, message: str) -> bool:
|
| 213 |
-
try:
|
| 214 |
-
_api().create_commit(repo_id=REPO, repo_type="dataset", operations=ops,
|
| 215 |
-
commit_message=message)
|
| 216 |
-
return True
|
| 217 |
-
except Exception as e:
|
| 218 |
-
log(f"commit failed ({message}): {type(e).__name__}")
|
| 219 |
-
return False
|
| 220 |
-
|
| 221 |
-
|
| 222 |
-
def _commit_claim(ev: str, queued: str) -> bool:
|
| 223 |
-
from huggingface_hub import CommitOperationAdd
|
| 224 |
-
body = json.dumps({"worker": IDENT,
|
| 225 |
-
"hb": datetime.datetime.utcnow().strftime(TS),
|
| 226 |
-
"queued": queued})
|
| 227 |
-
return _commit([CommitOperationAdd(f"{QPREFIX}/{ev}.claim",
|
| 228 |
-
io.BytesIO(body.encode()))],
|
| 229 |
-
f"claim {ev} [{IDENT}]")
|
| 230 |
-
|
| 231 |
-
|
| 232 |
-
def _used_s(now: float) -> float:
|
| 233 |
-
"""GPU seconds booked in the last 24 h (old calls are forgotten)."""
|
| 234 |
-
W["gpu_calls"] = [(t, s) for t, s in W["gpu_calls"] if now - t < 86400.0]
|
| 235 |
-
return sum(s for _, s in W["gpu_calls"])
|
| 236 |
-
|
| 237 |
-
|
| 238 |
-
def _budget_left(now: float) -> float:
|
| 239 |
-
"""Seconds of daily GPU budget left (inf when uncapped)."""
|
| 240 |
-
if DAILY_BUDGET_S <= 0:
|
| 241 |
-
return float("inf")
|
| 242 |
-
return DAILY_BUDGET_S - _used_s(now)
|
| 243 |
-
|
| 244 |
-
|
| 245 |
-
def _budget_free_at(now: float, need: float) -> float:
|
| 246 |
-
"""When at least `need` seconds of budget will be free again."""
|
| 247 |
-
left = _budget_left(now)
|
| 248 |
-
for t, s in sorted(W["gpu_calls"]):
|
| 249 |
-
if left >= need:
|
| 250 |
-
break
|
| 251 |
-
left += s
|
| 252 |
-
now = t + 86400.0
|
| 253 |
-
return now
|
| 254 |
-
|
| 255 |
-
|
| 256 |
-
def _charge(seconds: float, started: float):
|
| 257 |
-
"""Book a GPU call that began at `started` and lasted `seconds` (wall,
|
| 258 |
-
incl. fork/pickle overhead β conservative)."""
|
| 259 |
-
if DAILY_BUDGET_S > 0:
|
| 260 |
-
W["gpu_calls"].append((started, max(0.0, seconds)))
|
| 261 |
-
|
| 262 |
-
|
| 263 |
-
def _gpu_wait_reason(now: float) -> str | None:
|
| 264 |
-
"""Why no GPU call may start now (None = go)."""
|
| 265 |
-
if now < W["quota_until"]:
|
| 266 |
-
return f"ZeroGPU quota exhausted until {_hhmm(W['quota_until'])}"
|
| 267 |
-
if now < W["capacity_until"]:
|
| 268 |
-
return "no ZeroGPU capacity β retrying shortly"
|
| 269 |
-
if _budget_left(now) < GPU_CALL_S:
|
| 270 |
-
return (f"daily GPU budget used ({_used_s(now) / 60:.0f} of "
|
| 271 |
-
f"{DAILY_BUDGET_S / 60:.0f} min in 24 h) β resumes "
|
| 272 |
-
f"{_hhmm(_budget_free_at(now, GPU_CALL_S))}")
|
| 273 |
-
return None
|
| 274 |
-
|
| 275 |
-
|
| 276 |
-
def _hhmm(t: float) -> str:
|
| 277 |
-
return datetime.datetime.utcfromtimestamp(t).strftime("%m-%d %H:%MZ")
|
| 278 |
-
|
| 279 |
-
|
| 280 |
-
def _quota_retry_s(exc) -> float | None:
|
| 281 |
-
"""Seconds until ZeroGPU says to retry ('... try again in H:MM:SS'), or
|
| 282 |
-
None when the error names no retry time (e.g. the anonymous pool's
|
| 283 |
-
'Space app has reached its GPU limit')."""
|
| 284 |
-
import re
|
| 285 |
-
m = re.search(r"try again in (\d+):(\d{2}):(\d{2})", str(exc), re.I)
|
| 286 |
-
if m:
|
| 287 |
-
h, mi, s = (int(x) for x in m.groups())
|
| 288 |
-
return float(h * 3600 + mi * 60 + s) + 60.0
|
| 289 |
-
return None
|
| 290 |
-
|
| 291 |
-
|
| 292 |
-
def _auth_suspect(exc) -> bool:
|
| 293 |
-
"""The error text of the ANONYMOUS pool: our calls were not attributed
|
| 294 |
-
to the account (ZERO_AUTH did not take effect)."""
|
| 295 |
-
m = f"{exc}".lower()
|
| 296 |
-
return "space app has reached its gpu limit" in m or "unlogged" in m
|
| 297 |
-
|
| 298 |
-
|
| 299 |
-
def _is_capacity_error(exc) -> bool:
|
| 300 |
-
return "no gpu was available" in f"{exc}".lower()
|
| 301 |
-
|
| 302 |
-
|
| 303 |
-
def _claim_fresh(claim) -> bool:
|
| 304 |
-
if not claim:
|
| 305 |
-
return False
|
| 306 |
-
try:
|
| 307 |
-
hb = datetime.datetime.strptime(claim.get("hb", ""), TS)
|
| 308 |
-
except ValueError:
|
| 309 |
-
return False
|
| 310 |
-
return (datetime.datetime.utcnow() - hb).total_seconds() < STALE_S
|
| 311 |
-
|
| 312 |
-
|
| 313 |
-
def _scan():
|
| 314 |
-
"""Oldest BACKUP-eligible job: bundle present, 30 m, inside cell budget,
|
| 315 |
-
unclaimed (or stale), and queued more than ZERO_DELAY_S ago (the HPC gets
|
| 316 |
-
first shot β this Space is ladder rung 2, not a race participant)."""
|
| 317 |
-
try:
|
| 318 |
-
files = _api().list_repo_files(REPO, repo_type="dataset")
|
| 319 |
-
except Exception as e:
|
| 320 |
-
log(f"list_repo_files failed: {type(e).__name__}")
|
| 321 |
-
return None
|
| 322 |
-
qf = [f for f in files if f.startswith(QPREFIX + "/")]
|
| 323 |
-
ids = sorted({os.path.basename(f)[:-5] for f in qf
|
| 324 |
-
if f.endswith(".json") and not f.endswith(".failed.json")})
|
| 325 |
-
now = datetime.datetime.utcnow()
|
| 326 |
-
# queue pressure: backlog > PRESSURE_N promotes this Space from deferred
|
| 327 |
-
# backup to active drain (claim immediately; claims keep events disjoint)
|
| 328 |
-
backlog = sum(1 for ev in ids if f"{QPREFIX}/{ev}.forcing.tar.gz" in qf)
|
| 329 |
-
delay = 0.0 if backlog > PRESSURE_N else ZERO_DELAY_S
|
| 330 |
-
if backlog > PRESSURE_N and W.get("pressure") != backlog:
|
| 331 |
-
W["pressure"] = backlog
|
| 332 |
-
log(f"queue pressure: {backlog} pending jobs > {PRESSURE_N} β "
|
| 333 |
-
f"claiming without deferral")
|
| 334 |
-
cands = []
|
| 335 |
-
skip = W["scan"]["skipped"] = {}
|
| 336 |
-
W["scan"].update(at=now.strftime(TS), ids=len(ids), backlog=backlog)
|
| 337 |
-
|
| 338 |
-
def _skip(why):
|
| 339 |
-
skip[why] = skip.get(why, 0) + 1
|
| 340 |
-
|
| 341 |
-
for ev in ids:
|
| 342 |
-
if f"{QPREFIX}/{ev}.forcing.tar.gz" not in qf:
|
| 343 |
-
_skip("no_forcing")
|
| 344 |
-
continue
|
| 345 |
-
spec = _read_json(f"{QPREFIX}/{ev}.json")
|
| 346 |
-
if not spec:
|
| 347 |
-
_skip("spec_unreadable")
|
| 348 |
-
continue
|
| 349 |
-
|
| 350 |
-
|
| 351 |
-
|
| 352 |
-
|
| 353 |
-
|
| 354 |
-
|
| 355 |
-
|
| 356 |
-
|
| 357 |
-
|
| 358 |
-
|
| 359 |
-
|
| 360 |
-
if
|
| 361 |
-
_skip("
|
| 362 |
-
continue
|
| 363 |
-
if
|
| 364 |
-
|
| 365 |
-
|
| 366 |
-
|
| 367 |
-
_skip("
|
| 368 |
-
continue
|
| 369 |
-
|
| 370 |
-
|
| 371 |
-
|
| 372 |
-
|
| 373 |
-
|
| 374 |
-
|
| 375 |
-
|
| 376 |
-
|
| 377 |
-
|
| 378 |
-
|
| 379 |
-
|
| 380 |
-
|
| 381 |
-
|
| 382 |
-
|
| 383 |
-
|
| 384 |
-
|
| 385 |
-
|
| 386 |
-
|
| 387 |
-
|
| 388 |
-
|
| 389 |
-
|
| 390 |
-
|
| 391 |
-
|
| 392 |
-
|
| 393 |
-
|
| 394 |
-
|
| 395 |
-
|
| 396 |
-
|
| 397 |
-
|
| 398 |
-
|
| 399 |
-
|
| 400 |
-
|
| 401 |
-
|
| 402 |
-
|
| 403 |
-
|
| 404 |
-
|
| 405 |
-
|
| 406 |
-
|
| 407 |
-
|
| 408 |
-
|
| 409 |
-
|
| 410 |
-
|
| 411 |
-
|
| 412 |
-
|
| 413 |
-
|
| 414 |
-
|
| 415 |
-
|
| 416 |
-
|
| 417 |
-
|
| 418 |
-
|
| 419 |
-
|
| 420 |
-
|
| 421 |
-
|
| 422 |
-
|
| 423 |
-
|
| 424 |
-
|
| 425 |
-
|
| 426 |
-
|
| 427 |
-
|
| 428 |
-
|
| 429 |
-
|
| 430 |
-
|
| 431 |
-
|
| 432 |
-
|
| 433 |
-
|
| 434 |
-
|
| 435 |
-
|
| 436 |
-
|
| 437 |
-
|
| 438 |
-
|
| 439 |
-
|
| 440 |
-
|
| 441 |
-
|
| 442 |
-
|
| 443 |
-
|
| 444 |
-
|
| 445 |
-
|
| 446 |
-
|
| 447 |
-
|
| 448 |
-
|
| 449 |
-
|
| 450 |
-
|
| 451 |
-
|
| 452 |
-
|
| 453 |
-
|
| 454 |
-
|
| 455 |
-
|
| 456 |
-
|
| 457 |
-
|
| 458 |
-
|
| 459 |
-
|
| 460 |
-
|
| 461 |
-
|
| 462 |
-
|
| 463 |
-
|
| 464 |
-
|
| 465 |
-
|
| 466 |
-
|
| 467 |
-
|
| 468 |
-
|
| 469 |
-
|
| 470 |
-
|
| 471 |
-
|
| 472 |
-
|
| 473 |
-
|
| 474 |
-
|
| 475 |
-
|
| 476 |
-
|
| 477 |
-
|
| 478 |
-
|
| 479 |
-
|
| 480 |
-
|
| 481 |
-
|
| 482 |
-
|
| 483 |
-
|
| 484 |
-
|
| 485 |
-
|
| 486 |
-
|
| 487 |
-
|
| 488 |
-
|
| 489 |
-
|
| 490 |
-
|
| 491 |
-
|
| 492 |
-
|
| 493 |
-
|
| 494 |
-
|
| 495 |
-
|
| 496 |
-
|
| 497 |
-
|
| 498 |
-
|
| 499 |
-
|
| 500 |
-
|
| 501 |
-
|
| 502 |
-
|
| 503 |
-
|
| 504 |
-
|
| 505 |
-
|
| 506 |
-
|
| 507 |
-
|
| 508 |
-
|
| 509 |
-
|
| 510 |
-
|
| 511 |
-
|
| 512 |
-
|
| 513 |
-
|
| 514 |
-
|
| 515 |
-
|
| 516 |
-
|
| 517 |
-
|
| 518 |
-
|
| 519 |
-
|
| 520 |
-
|
| 521 |
-
|
| 522 |
-
|
| 523 |
-
|
| 524 |
-
|
| 525 |
-
|
| 526 |
-
|
| 527 |
-
if
|
| 528 |
-
|
| 529 |
-
|
| 530 |
-
|
| 531 |
-
|
| 532 |
-
|
| 533 |
-
|
| 534 |
-
|
| 535 |
-
|
| 536 |
-
|
| 537 |
-
|
| 538 |
-
|
| 539 |
-
|
| 540 |
-
|
| 541 |
-
|
| 542 |
-
|
| 543 |
-
|
| 544 |
-
|
| 545 |
-
|
| 546 |
-
W
|
| 547 |
-
|
| 548 |
-
|
| 549 |
-
|
| 550 |
-
|
| 551 |
-
|
| 552 |
-
|
| 553 |
-
|
| 554 |
-
|
| 555 |
-
|
| 556 |
-
|
| 557 |
-
|
| 558 |
-
|
| 559 |
-
|
| 560 |
-
|
| 561 |
-
|
| 562 |
-
|
| 563 |
-
if
|
| 564 |
-
|
| 565 |
-
|
| 566 |
-
|
| 567 |
-
|
| 568 |
-
|
| 569 |
-
|
| 570 |
-
|
| 571 |
-
|
| 572 |
-
|
| 573 |
-
|
| 574 |
-
|
| 575 |
-
|
| 576 |
-
|
| 577 |
-
|
| 578 |
-
|
| 579 |
-
|
| 580 |
-
|
| 581 |
-
|
| 582 |
-
|
| 583 |
-
|
| 584 |
-
|
| 585 |
-
|
| 586 |
-
|
| 587 |
-
|
| 588 |
-
|
| 589 |
-
|
| 590 |
-
|
| 591 |
-
|
| 592 |
-
|
| 593 |
-
|
| 594 |
-
|
| 595 |
-
|
| 596 |
-
|
| 597 |
-
|
| 598 |
-
|
| 599 |
-
|
| 600 |
-
|
| 601 |
-
|
| 602 |
-
|
| 603 |
-
|
| 604 |
-
|
| 605 |
-
|
| 606 |
-
|
| 607 |
-
|
| 608 |
-
|
| 609 |
-
|
| 610 |
-
|
| 611 |
-
|
| 612 |
-
|
| 613 |
-
|
| 614 |
-
|
| 615 |
-
|
| 616 |
-
|
| 617 |
-
|
| 618 |
-
|
| 619 |
-
|
| 620 |
-
|
| 621 |
-
|
| 622 |
-
|
| 623 |
-
|
| 624 |
-
|
| 625 |
-
|
| 626 |
-
|
| 627 |
-
|
| 628 |
-
|
| 629 |
-
|
| 630 |
-
|
| 631 |
-
|
| 632 |
-
|
| 633 |
-
|
| 634 |
-
|
| 635 |
-
|
| 636 |
-
|
| 637 |
-
|
| 638 |
-
|
| 639 |
-
|
| 640 |
-
|
| 641 |
-
|
| 642 |
-
|
| 643 |
-
|
| 644 |
-
|
| 645 |
-
|
| 646 |
-
|
| 647 |
-
|
| 648 |
-
|
| 649 |
-
|
| 650 |
-
|
| 651 |
-
|
| 652 |
-
|
| 653 |
-
|
| 654 |
-
|
| 655 |
-
|
| 656 |
-
|
| 657 |
-
|
| 658 |
-
|
| 659 |
-
|
| 660 |
-
|
| 661 |
-
|
| 662 |
-
|
| 663 |
-
|
| 664 |
-
|
| 665 |
-
|
| 666 |
-
|
| 667 |
-
|
| 668 |
-
|
| 669 |
-
|
| 670 |
-
|
| 671 |
-
|
| 672 |
-
|
| 673 |
-
|
| 674 |
-
|
| 675 |
-
|
| 676 |
-
|
| 677 |
-
|
| 678 |
-
|
| 679 |
-
|
| 680 |
-
|
| 681 |
-
|
| 682 |
-
|
| 683 |
-
|
| 684 |
-
|
| 685 |
-
|
| 686 |
-
|
| 687 |
-
|
| 688 |
-
|
| 689 |
-
|
| 690 |
-
|
| 691 |
-
|
| 692 |
-
|
| 693 |
-
|
| 694 |
-
|
| 695 |
-
|
| 696 |
-
|
| 697 |
-
|
| 698 |
-
|
| 699 |
-
|
| 700 |
-
|
| 701 |
-
|
| 702 |
-
|
| 703 |
-
|
| 704 |
-
|
| 705 |
-
|
| 706 |
-
|
| 707 |
-
|
| 708 |
-
|
| 709 |
-
|
| 710 |
-
|
| 711 |
-
|
| 712 |
-
|
| 713 |
-
|
| 714 |
-
|
| 715 |
-
|
| 716 |
-
|
| 717 |
-
|
| 718 |
-
|
| 719 |
-
|
| 720 |
-
|
| 721 |
-
|
| 722 |
-
|
| 723 |
-
|
| 724 |
-
|
| 725 |
-
|
| 726 |
-
|
| 727 |
-
|
| 728 |
-
|
| 729 |
-
|
| 730 |
-
|
| 731 |
-
|
| 732 |
-
|
| 733 |
-
|
| 734 |
-
|
| 735 |
-
|
| 736 |
-
|
| 737 |
-
|
| 738 |
-
|
| 739 |
-
|
| 740 |
-
|
| 741 |
-
"
|
| 742 |
-
"
|
| 743 |
-
|
| 744 |
-
"
|
| 745 |
-
|
| 746 |
-
|
| 747 |
-
|
| 748 |
-
|
| 749 |
-
|
| 750 |
-
|
| 751 |
-
|
| 752 |
-
|
| 753 |
-
|
| 754 |
-
|
| 755 |
-
|
| 756 |
-
|
| 757 |
-
|
| 758 |
-
|
| 759 |
-
|
| 760 |
-
|
| 761 |
-
|
| 762 |
-
|
| 763 |
-
|
| 764 |
-
|
| 765 |
-
|
| 766 |
-
|
| 767 |
-
|
| 768 |
-
|
| 769 |
-
|
| 770 |
-
|
| 771 |
-
|
| 772 |
-
|
| 773 |
-
|
| 774 |
-
|
| 775 |
-
|
| 776 |
-
|
| 777 |
-
|
| 778 |
-
|
| 779 |
-
|
| 780 |
-
|
| 781 |
-
|
| 782 |
-
|
| 783 |
-
|
| 784 |
-
|
| 785 |
-
|
| 786 |
-
|
| 787 |
-
|
| 788 |
-
|
| 789 |
-
|
| 790 |
-
|
| 791 |
-
|
| 792 |
-
|
| 793 |
-
|
| 794 |
-
|
| 795 |
-
|
| 796 |
-
|
| 797 |
-
|
| 798 |
-
|
| 799 |
-
|
| 800 |
-
|
| 801 |
-
|
| 802 |
-
|
| 803 |
-
|
| 804 |
-
|
| 805 |
-
|
| 806 |
-
|
| 807 |
-
|
| 808 |
-
|
| 809 |
-
|
| 810 |
-
|
| 811 |
-
|
| 812 |
-
|
| 813 |
-
|
| 814 |
-
|
| 815 |
-
|
| 816 |
-
|
| 817 |
-
|
| 818 |
-
|
| 819 |
-
|
| 820 |
-
|
| 821 |
-
|
| 822 |
-
|
| 823 |
-
|
| 824 |
-
|
| 825 |
-
|
| 826 |
-
|
| 827 |
-
|
| 828 |
-
|
| 829 |
-
|
| 830 |
-
|
| 831 |
-
|
| 832 |
-
|
| 833 |
-
|
| 834 |
-
|
| 835 |
-
|
| 836 |
-
|
| 837 |
-
|
| 838 |
-
|
| 839 |
-
|
| 840 |
-
|
| 841 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""CREST_gpu_worker β ZeroGPU event worker (V27 P1.5 + P1.6 sessions). Since
|
| 2 |
+
the HPC worker was released (2026-10-16) it is the only 2-D engine: it claims
|
| 3 |
+
events/queue/ jobs, solves them in <=120 s GPU slices within a daily GPU
|
| 4 |
+
budget (authenticated calls draw on the account's ZeroGPU quota), and β when
|
| 5 |
+
PUBLISH=1 β publishes exactly like the HPC worker did.
|
| 6 |
+
|
| 7 |
+
P1.6: solver state is a resident crestimap.session.EventSession per episode
|
| 8 |
+
(depth AND momentum held in RAM between bundles), so returning to an episode
|
| 9 |
+
this Space already solved is a minutes-long catch-up visit from the last
|
| 10 |
+
obs-supported junction, not an hours-long anchored re-solve. Sessions for
|
| 11 |
+
episodes this worker has visited get scheduling priority over new events.
|
| 12 |
+
|
| 13 |
+
Queue protocol: identical to crestimap/worker.py (claim/heartbeat/stale
|
| 14 |
+
takeover/failed.json). 30 m (dem_res "1") jobs only; never coarsens β an
|
| 15 |
+
oversized or 10 m job is simply not claimed.
|
| 16 |
+
|
| 17 |
+
ZeroGPU note: @spaces.GPU functions run in a forked child with the GPU;
|
| 18 |
+
arguments/returns are pickled, and parent-side mutations inside the child are
|
| 19 |
+
lost. gpu_chunk() therefore RETURNS session.state_tuple() and the parent
|
| 20 |
+
reapplies it. Files the child writes (depth frames) land on the shared fs.
|
| 21 |
+
|
| 22 |
+
Never prints or logs the HF token.
|
| 23 |
+
"""
|
| 24 |
+
from __future__ import annotations
|
| 25 |
+
|
| 26 |
+
import datetime
|
| 27 |
+
import io
|
| 28 |
+
import json
|
| 29 |
+
import os
|
| 30 |
+
import shutil
|
| 31 |
+
import subprocess
|
| 32 |
+
import sys
|
| 33 |
+
import tarfile
|
| 34 |
+
import tempfile
|
| 35 |
+
import threading
|
| 36 |
+
import time
|
| 37 |
+
import traceback
|
| 38 |
+
|
| 39 |
+
# ---------------------------------------------------------------- bootstrap
|
| 40 |
+
APP = os.path.dirname(os.path.abspath(__file__))
|
| 41 |
+
FORK = "https://github.com/mchen15ouedu/CREST-iMAP.git"
|
| 42 |
+
CREST_AI = "https://huggingface.co/spaces/vincewin/CREST_demo"
|
| 43 |
+
# Solver commit, pinned like the CPU Spaces' Dockerfile CRESTIMAP_REF and the
|
| 44 |
+
# Jetstream2 kit's FORK_REF: a push to v2 must never reach this worker on a
|
| 45 |
+
# restart without review. c5184759 = draining-time outflow limiter +
|
| 46 |
+
# interflow store, approved 2026-10-20. Full SHA only (a branch would float).
|
| 47 |
+
FORK_REF = os.environ.get("CRESTIMAP_REF",
|
| 48 |
+
"c518475956c933a9c0fc46962b3825685a54b181")
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
def _git(*args, cwd=None):
|
| 52 |
+
return subprocess.run(["git", *args], cwd=cwd, check=True,
|
| 53 |
+
capture_output=True, text=True).stdout.strip()
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def _bootstrap():
|
| 57 |
+
"""Blobless sparse clones: crestimap/ from the fork at FORK_REF, hf_data/
|
| 58 |
+
from the CREST_demo Space (eventsim/eventstore reuse β nothing
|
| 59 |
+
reimplemented).
|
| 60 |
+
|
| 61 |
+
crestimap is imported straight from the clone (a detached sparse checkout
|
| 62 |
+
of FORK_REF), so provenance.stamp() records the pinned SHA on every
|
| 63 |
+
manifest; a copied-out package has no .git and stamped only
|
| 64 |
+
"2.0.0.dev0"."""
|
| 65 |
+
fork = os.path.join(APP, "_fork")
|
| 66 |
+
if not os.path.isdir(os.path.join(fork, "crestimap")):
|
| 67 |
+
shutil.rmtree(fork, ignore_errors=True)
|
| 68 |
+
# all branches are fetched (blobless), so a pinned commit stays
|
| 69 |
+
# reachable whichever branch carries it
|
| 70 |
+
_git("clone", "--filter=blob:none", "--no-checkout", FORK, fork)
|
| 71 |
+
_git("sparse-checkout", "init", "--cone", cwd=fork)
|
| 72 |
+
_git("sparse-checkout", "set", "crestimap", cwd=fork)
|
| 73 |
+
_git("checkout", "--detach", FORK_REF, cwd=fork)
|
| 74 |
+
want = _git("rev-parse", "--verify", f"{FORK_REF}^{{commit}}", cwd=fork)
|
| 75 |
+
if _git("rev-parse", "HEAD", cwd=fork) != want:
|
| 76 |
+
raise RuntimeError(f"crestimap clone is not at the pinned {FORK_REF}")
|
| 77 |
+
if fork not in sys.path:
|
| 78 |
+
sys.path.insert(0, fork)
|
| 79 |
+
ai = os.path.join(APP, "_crest_ai")
|
| 80 |
+
if not os.path.isdir(os.path.join(ai, "hf_data")):
|
| 81 |
+
subprocess.run(["git", "clone", "--filter=blob:none", "--no-checkout",
|
| 82 |
+
CREST_AI, ai], check=True)
|
| 83 |
+
subprocess.run(["git", "checkout", "HEAD", "--", "hf_data"],
|
| 84 |
+
cwd=ai, check=True)
|
| 85 |
+
if ai not in sys.path:
|
| 86 |
+
sys.path.insert(0, ai)
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
_bootstrap()
|
| 90 |
+
try:
|
| 91 |
+
_FORK_SHA = subprocess.run(
|
| 92 |
+
["git", "-C", os.path.join(APP, "_fork"), "rev-parse", "HEAD"],
|
| 93 |
+
capture_output=True, text=True, timeout=15).stdout.strip()[:12] or None
|
| 94 |
+
except Exception:
|
| 95 |
+
_FORK_SHA = None
|
| 96 |
+
|
| 97 |
+
import gradio as gr # noqa: E402
|
| 98 |
+
|
| 99 |
+
try:
|
| 100 |
+
import spaces # noqa: E402
|
| 101 |
+
HAS_ZERO = True
|
| 102 |
+
except Exception: # local dev
|
| 103 |
+
HAS_ZERO = False
|
| 104 |
+
|
| 105 |
+
TS = "%Y-%m-%dT%H:%M:%SZ"
|
| 106 |
+
TMIN = "%Y-%m-%dT%H:%MZ"
|
| 107 |
+
QPREFIX = "events/queue"
|
| 108 |
+
REPO = os.environ.get("CREST_FEEDBACK_REPO", "vincewin/CREST_data")
|
| 109 |
+
# The HPC worker was released 2026-10-16; this Space is now the only 2-D
|
| 110 |
+
# engine, so there is no reason to leave a job unclaimed for anyone else.
|
| 111 |
+
ZERO_DELAY_S = float(os.environ.get("ZERO_DELAY_S", "0"))
|
| 112 |
+
# GPU calls must be AUTHENTICATED to get the account's ZeroGPU quota (PRO:
|
| 113 |
+
# 40 min/day, fixed 24 h window from first use). Self-calls over 127.0.0.1
|
| 114 |
+
# carry no X-IP-Token and land in a shared anonymous "Space app" pool that
|
| 115 |
+
# gave ~0 GPU-s/day from 2026-09-23 to 10-16 ("Space app has reached its GPU
|
| 116 |
+
# limit" on every claim). ZERO_AUTH=1 self-drives through the public Space
|
| 117 |
+
# URL with HF_TOKEN instead.
|
| 118 |
+
ZERO_AUTH = os.environ.get("ZERO_AUTH", "1") == "1"
|
| 119 |
+
# Our own daily cap, below the 40-min PRO quota: the quota is shared with
|
| 120 |
+
# every other ZeroGPU use under the same account, and past it ZeroGPU bills
|
| 121 |
+
# prepaid credits automatically ($1 per 10 min). 0 disables the cap.
|
| 122 |
+
DAILY_BUDGET_S = float(os.environ.get("ZERO_DAILY_BUDGET_S", "1800"))
|
| 123 |
+
GPU_CALL_S = 120.0 # declared @spaces.GPU duration (admission needs it)
|
| 124 |
+
if 0 < DAILY_BUDGET_S < GPU_CALL_S:
|
| 125 |
+
DAILY_BUDGET_S = GPU_CALL_S # a smaller cap could never run
|
| 126 |
+
# a NEW event (full anchored solve) is only claimed with this much budget
|
| 127 |
+
# left; otherwise it would be claimed, paused ~a day, and solved stale
|
| 128 |
+
MIN_NEW_S = float(os.environ.get("ZERO_MIN_NEW_S", "600"))
|
| 129 |
+
# "No GPU was available after 60s" is platform capacity, not this job's
|
| 130 |
+
# fault and not quota: wait and retry the same chunk.
|
| 131 |
+
CAPACITY_BACKOFF_S = float(os.environ.get("ZERO_CAPACITY_BACKOFF_S", "120"))
|
| 132 |
+
# Backup -> active drain under pressure (user directive 2026-08-14): with
|
| 133 |
+
# more than this many jobs pending, skip the polite deferral and claim the
|
| 134 |
+
# next unclaimed event immediately β one solver per event, so this Space
|
| 135 |
+
# always works a DIFFERENT event than the HPC.
|
| 136 |
+
PRESSURE_N = int(os.environ.get("ZERO_PRESSURE_N", "4"))
|
| 137 |
+
STALE_S = float(os.environ.get("ZERO_STALE_S", "600"))
|
| 138 |
+
HB_S = float(os.environ.get("ZERO_HB_S", "240"))
|
| 139 |
+
POLL_S = float(os.environ.get("ZERO_POLL_S", "60"))
|
| 140 |
+
# ZeroGPU quota exhaustion is NOT a poison job: release the claim (the HPC
|
| 141 |
+
# or a later quota window takes the bundle) and stop claiming for a while.
|
| 142 |
+
# Without this (seen 2026-09-11) the backup rung claimed a bundle, hit
|
| 143 |
+
# "Space app has reached its GPU limit", and CONSUMED it as failed.
|
| 144 |
+
QUOTA_BACKOFF_S = float(os.environ.get("ZERO_QUOTA_BACKOFF_S", "3600"))
|
| 145 |
+
MAX_CELLS = int(os.environ.get("ZERO_MAX_CELLS", "2000000"))
|
| 146 |
+
PUBLISH = os.environ.get("PUBLISH", "0") == "1"
|
| 147 |
+
CHUNK0_SIM_S = float(os.environ.get("ZERO_CHUNK0_SIM_S", "300"))
|
| 148 |
+
GPU_TARGET_S = float(os.environ.get("ZERO_GPU_TARGET_S", "80"))
|
| 149 |
+
DEM_RES_DEG = {"1": 1.0 / 3600.0, "13": 1.0 / 10800.0}
|
| 150 |
+
IDENT = f"zerogpu:{os.environ.get('SPACE_ID', 'local')}:{os.getpid()}"
|
| 151 |
+
os.environ.setdefault("EVENT_ENGINE_IDENT", IDENT) # provenance stamp
|
| 152 |
+
|
| 153 |
+
_tok = os.environ.get("HF_TOKEN", "").strip()
|
| 154 |
+
if _tok:
|
| 155 |
+
os.environ["HF_TOKEN"] = _tok # eventstore reads env
|
| 156 |
+
|
| 157 |
+
MAX_SESSIONS = int(os.environ.get("ZERO_MAX_SESSIONS", "6"))
|
| 158 |
+
# > one daily budget pause (~23.5 h), so parked sessions survive the pause and
|
| 159 |
+
# the next bundle is a cheap catch-up rather than a full re-solve
|
| 160 |
+
EVICT_IDLE_H = float(os.environ.get("ZERO_EVICT_IDLE_H", "36"))
|
| 161 |
+
# cadence of the server-side tick loop (see _self_drive); matches the old
|
| 162 |
+
# gr.Timer(5) so idle polling behaviour is unchanged
|
| 163 |
+
SELF_TICK_S = float(os.environ.get("ZERO_SELF_TICK_S", "5"))
|
| 164 |
+
|
| 165 |
+
LOG: list[str] = []
|
| 166 |
+
W = {"phase": "idle", "job": None, "runner": None, "chunk_sim_s": CHUNK0_SIM_S,
|
| 167 |
+
"last_poll": 0.0, "last_hb": 0.0, "queued": "", "workdir": None,
|
| 168 |
+
"done": set(), "solved": [], "t_wall0": 0.0, "last_tick": 0.0,
|
| 169 |
+
# daily GPU budget: [(start, seconds), ...] of GPU calls in a SLIDING
|
| 170 |
+
# 24 h window, so use stays under the cap in any 24 h (hence in any of
|
| 171 |
+
# ZeroGPU's own windows); RAM-only β keep the account's prepaid credit
|
| 172 |
+
# balance at $0 so an over-run after a restart becomes a quota error
|
| 173 |
+
"gpu_calls": [], "capacity_until": 0.0, "quota_until": 0.0,
|
| 174 |
+
"paused_note": "", "auth_suspect": False,
|
| 175 |
+
# /tick calls by outcome (see tick_api); "accepted" must keep rising
|
| 176 |
+
"tick_api": {"accepted": 0, "rejected": 0, "empty": 0,
|
| 177 |
+
"last_reject_log": 0.0},
|
| 178 |
+
# what the last _scan saw (status "scan"): time, queue ids, backlog,
|
| 179 |
+
# candidates and per-reason skip counts
|
| 180 |
+
"scan": {},
|
| 181 |
+
# P1.6 resident sessions: ev -> {"session", "bbox", "visits", "last",
|
| 182 |
+
# "root"}. RAM-only; a Space rebuild recovers via the bundle's
|
| 183 |
+
# init_depth (state-hierarchy level 2).
|
| 184 |
+
"sessions": {}}
|
| 185 |
+
_tick_lock = threading.Lock()
|
| 186 |
+
# secret the self-drive loop passes to /tick (in-process only; never logged)
|
| 187 |
+
_SELF_KEY = __import__("secrets").token_hex(16)
|
| 188 |
+
|
| 189 |
+
|
| 190 |
+
def log(msg: str):
|
| 191 |
+
LOG.append(f"[{datetime.datetime.utcnow():%H:%M:%S}Z] {msg}")
|
| 192 |
+
del LOG[:-300]
|
| 193 |
+
print(LOG[-1], flush=True)
|
| 194 |
+
|
| 195 |
+
|
| 196 |
+
def _api():
|
| 197 |
+
from huggingface_hub import HfApi
|
| 198 |
+
return HfApi(token=_tok or None)
|
| 199 |
+
|
| 200 |
+
|
| 201 |
+
def _read_json(path: str):
|
| 202 |
+
from huggingface_hub import hf_hub_download
|
| 203 |
+
try:
|
| 204 |
+
p = hf_hub_download(REPO, path, repo_type="dataset", token=_tok or None,
|
| 205 |
+
force_download=True)
|
| 206 |
+
with open(p, encoding="utf-8") as fp:
|
| 207 |
+
return json.load(fp)
|
| 208 |
+
except Exception:
|
| 209 |
+
return None
|
| 210 |
+
|
| 211 |
+
|
| 212 |
+
def _commit(ops, message: str) -> bool:
|
| 213 |
+
try:
|
| 214 |
+
_api().create_commit(repo_id=REPO, repo_type="dataset", operations=ops,
|
| 215 |
+
commit_message=message)
|
| 216 |
+
return True
|
| 217 |
+
except Exception as e:
|
| 218 |
+
log(f"commit failed ({message}): {type(e).__name__}")
|
| 219 |
+
return False
|
| 220 |
+
|
| 221 |
+
|
| 222 |
+
def _commit_claim(ev: str, queued: str) -> bool:
|
| 223 |
+
from huggingface_hub import CommitOperationAdd
|
| 224 |
+
body = json.dumps({"worker": IDENT,
|
| 225 |
+
"hb": datetime.datetime.utcnow().strftime(TS),
|
| 226 |
+
"queued": queued})
|
| 227 |
+
return _commit([CommitOperationAdd(f"{QPREFIX}/{ev}.claim",
|
| 228 |
+
io.BytesIO(body.encode()))],
|
| 229 |
+
f"claim {ev} [{IDENT}]")
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
def _used_s(now: float) -> float:
|
| 233 |
+
"""GPU seconds booked in the last 24 h (old calls are forgotten)."""
|
| 234 |
+
W["gpu_calls"] = [(t, s) for t, s in W["gpu_calls"] if now - t < 86400.0]
|
| 235 |
+
return sum(s for _, s in W["gpu_calls"])
|
| 236 |
+
|
| 237 |
+
|
| 238 |
+
def _budget_left(now: float) -> float:
|
| 239 |
+
"""Seconds of daily GPU budget left (inf when uncapped)."""
|
| 240 |
+
if DAILY_BUDGET_S <= 0:
|
| 241 |
+
return float("inf")
|
| 242 |
+
return DAILY_BUDGET_S - _used_s(now)
|
| 243 |
+
|
| 244 |
+
|
| 245 |
+
def _budget_free_at(now: float, need: float) -> float:
|
| 246 |
+
"""When at least `need` seconds of budget will be free again."""
|
| 247 |
+
left = _budget_left(now)
|
| 248 |
+
for t, s in sorted(W["gpu_calls"]):
|
| 249 |
+
if left >= need:
|
| 250 |
+
break
|
| 251 |
+
left += s
|
| 252 |
+
now = t + 86400.0
|
| 253 |
+
return now
|
| 254 |
+
|
| 255 |
+
|
| 256 |
+
def _charge(seconds: float, started: float):
|
| 257 |
+
"""Book a GPU call that began at `started` and lasted `seconds` (wall,
|
| 258 |
+
incl. fork/pickle overhead β conservative)."""
|
| 259 |
+
if DAILY_BUDGET_S > 0:
|
| 260 |
+
W["gpu_calls"].append((started, max(0.0, seconds)))
|
| 261 |
+
|
| 262 |
+
|
| 263 |
+
def _gpu_wait_reason(now: float) -> str | None:
|
| 264 |
+
"""Why no GPU call may start now (None = go)."""
|
| 265 |
+
if now < W["quota_until"]:
|
| 266 |
+
return f"ZeroGPU quota exhausted until {_hhmm(W['quota_until'])}"
|
| 267 |
+
if now < W["capacity_until"]:
|
| 268 |
+
return "no ZeroGPU capacity β retrying shortly"
|
| 269 |
+
if _budget_left(now) < GPU_CALL_S:
|
| 270 |
+
return (f"daily GPU budget used ({_used_s(now) / 60:.0f} of "
|
| 271 |
+
f"{DAILY_BUDGET_S / 60:.0f} min in 24 h) β resumes "
|
| 272 |
+
f"{_hhmm(_budget_free_at(now, GPU_CALL_S))}")
|
| 273 |
+
return None
|
| 274 |
+
|
| 275 |
+
|
| 276 |
+
def _hhmm(t: float) -> str:
|
| 277 |
+
return datetime.datetime.utcfromtimestamp(t).strftime("%m-%d %H:%MZ")
|
| 278 |
+
|
| 279 |
+
|
| 280 |
+
def _quota_retry_s(exc) -> float | None:
|
| 281 |
+
"""Seconds until ZeroGPU says to retry ('... try again in H:MM:SS'), or
|
| 282 |
+
None when the error names no retry time (e.g. the anonymous pool's
|
| 283 |
+
'Space app has reached its GPU limit')."""
|
| 284 |
+
import re
|
| 285 |
+
m = re.search(r"try again in (\d+):(\d{2}):(\d{2})", str(exc), re.I)
|
| 286 |
+
if m:
|
| 287 |
+
h, mi, s = (int(x) for x in m.groups())
|
| 288 |
+
return float(h * 3600 + mi * 60 + s) + 60.0
|
| 289 |
+
return None
|
| 290 |
+
|
| 291 |
+
|
| 292 |
+
def _auth_suspect(exc) -> bool:
|
| 293 |
+
"""The error text of the ANONYMOUS pool: our calls were not attributed
|
| 294 |
+
to the account (ZERO_AUTH did not take effect)."""
|
| 295 |
+
m = f"{exc}".lower()
|
| 296 |
+
return "space app has reached its gpu limit" in m or "unlogged" in m
|
| 297 |
+
|
| 298 |
+
|
| 299 |
+
def _is_capacity_error(exc) -> bool:
|
| 300 |
+
return "no gpu was available" in f"{exc}".lower()
|
| 301 |
+
|
| 302 |
+
|
| 303 |
+
def _claim_fresh(claim) -> bool:
|
| 304 |
+
if not claim:
|
| 305 |
+
return False
|
| 306 |
+
try:
|
| 307 |
+
hb = datetime.datetime.strptime(claim.get("hb", ""), TS)
|
| 308 |
+
except ValueError:
|
| 309 |
+
return False
|
| 310 |
+
return (datetime.datetime.utcnow() - hb).total_seconds() < STALE_S
|
| 311 |
+
|
| 312 |
+
|
| 313 |
+
def _scan():
|
| 314 |
+
"""Oldest BACKUP-eligible job: bundle present, 30 m, inside cell budget,
|
| 315 |
+
unclaimed (or stale), and queued more than ZERO_DELAY_S ago (the HPC gets
|
| 316 |
+
first shot β this Space is ladder rung 2, not a race participant)."""
|
| 317 |
+
try:
|
| 318 |
+
files = _api().list_repo_files(REPO, repo_type="dataset")
|
| 319 |
+
except Exception as e:
|
| 320 |
+
log(f"list_repo_files failed: {type(e).__name__}")
|
| 321 |
+
return None
|
| 322 |
+
qf = [f for f in files if f.startswith(QPREFIX + "/")]
|
| 323 |
+
ids = sorted({os.path.basename(f)[:-5] for f in qf
|
| 324 |
+
if f.endswith(".json") and not f.endswith(".failed.json")})
|
| 325 |
+
now = datetime.datetime.utcnow()
|
| 326 |
+
# queue pressure: backlog > PRESSURE_N promotes this Space from deferred
|
| 327 |
+
# backup to active drain (claim immediately; claims keep events disjoint)
|
| 328 |
+
backlog = sum(1 for ev in ids if f"{QPREFIX}/{ev}.forcing.tar.gz" in qf)
|
| 329 |
+
delay = 0.0 if backlog > PRESSURE_N else ZERO_DELAY_S
|
| 330 |
+
if backlog > PRESSURE_N and W.get("pressure") != backlog:
|
| 331 |
+
W["pressure"] = backlog
|
| 332 |
+
log(f"queue pressure: {backlog} pending jobs > {PRESSURE_N} β "
|
| 333 |
+
f"claiming without deferral")
|
| 334 |
+
cands = []
|
| 335 |
+
skip = W["scan"]["skipped"] = {}
|
| 336 |
+
W["scan"].update(at=now.strftime(TS), ids=len(ids), backlog=backlog)
|
| 337 |
+
|
| 338 |
+
def _skip(why):
|
| 339 |
+
skip[why] = skip.get(why, 0) + 1
|
| 340 |
+
|
| 341 |
+
for ev in ids:
|
| 342 |
+
if f"{QPREFIX}/{ev}.forcing.tar.gz" not in qf:
|
| 343 |
+
_skip("no_forcing")
|
| 344 |
+
continue
|
| 345 |
+
spec = _read_json(f"{QPREFIX}/{ev}.json")
|
| 346 |
+
if not spec:
|
| 347 |
+
_skip("spec_unreadable")
|
| 348 |
+
continue
|
| 349 |
+
# This backup image implements only the original forcing contract.
|
| 350 |
+
# New coupling jobs must remain available to a compatible HPC worker.
|
| 351 |
+
try:
|
| 352 |
+
protocol = int(spec.get("worker_protocol", spec.get("coupling_protocol", 1)))
|
| 353 |
+
except (TypeError, ValueError):
|
| 354 |
+
_skip("unsupported_protocol")
|
| 355 |
+
continue
|
| 356 |
+
if protocol > 1:
|
| 357 |
+
_skip("unsupported_protocol")
|
| 358 |
+
continue
|
| 359 |
+
required_engine = spec.get("required_engine_prefix")
|
| 360 |
+
if required_engine and not IDENT.startswith(required_engine):
|
| 361 |
+
_skip("engine_mismatch")
|
| 362 |
+
continue
|
| 363 |
+
if (ev, spec.get("queued", "")) in W["done"]:
|
| 364 |
+
_skip("done")
|
| 365 |
+
continue
|
| 366 |
+
if str(spec.get("dem_res", "1")) != "1":
|
| 367 |
+
_skip("not_30m")
|
| 368 |
+
continue # 30 m only on this rung
|
| 369 |
+
try:
|
| 370 |
+
age = (now - datetime.datetime.strptime(spec["queued"], TS)
|
| 371 |
+
).total_seconds()
|
| 372 |
+
except Exception:
|
| 373 |
+
age = 0.0
|
| 374 |
+
if age < delay:
|
| 375 |
+
_skip("deferred")
|
| 376 |
+
continue
|
| 377 |
+
if not spec.get("domain"):
|
| 378 |
+
# pre-basin bundle (rectangular-window era): never solved as a
|
| 379 |
+
# rectangle β the runner re-enqueues it with a HUC12 basin domain
|
| 380 |
+
log(f"skip {ev}: no basin domain in spec β awaiting re-enqueue")
|
| 381 |
+
_skip("no_domain")
|
| 382 |
+
continue
|
| 383 |
+
w, s, e, n = spec["bbox_basin"]
|
| 384 |
+
res = DEM_RES_DEG["1"]
|
| 385 |
+
cells = int((spec.get("domain") or {}).get("n_bbox") or
|
| 386 |
+
int(round((n - s) / res)) * int(round((e - w) / res)))
|
| 387 |
+
if cells > MAX_CELLS:
|
| 388 |
+
log(f"skip {ev}: {cells / 1e6:.1f} M cells > "
|
| 389 |
+
f"{MAX_CELLS / 1e6:.1f} M budget (never coarsen)")
|
| 390 |
+
_skip("over_budget")
|
| 391 |
+
continue
|
| 392 |
+
if f"{QPREFIX}/{ev}.claim" in qf:
|
| 393 |
+
claim = _read_json(f"{QPREFIX}/{ev}.claim")
|
| 394 |
+
if _claim_fresh(claim) and claim.get("worker") != IDENT:
|
| 395 |
+
_skip("claimed")
|
| 396 |
+
continue
|
| 397 |
+
cands.append({"id": ev, "spec": spec})
|
| 398 |
+
W["scan"]["candidates"] = len(cands)
|
| 399 |
+
# P1.6: episodes with a resident session first β a catch-up visit is
|
| 400 |
+
# minutes; new events (full anchored solve) fill the remaining time
|
| 401 |
+
cands.sort(key=lambda j: (j["id"] not in W["sessions"],
|
| 402 |
+
j["spec"].get("queued", "")))
|
| 403 |
+
return cands[0] if cands else None
|
| 404 |
+
|
| 405 |
+
|
| 406 |
+
# ------------------------------------------------------------------ GPU work
|
| 407 |
+
def _gpu_chunk_impl(session, sim_s: float, device: str):
|
| 408 |
+
done = session.chunk(sim_s, device=device)
|
| 409 |
+
return done, session.state_tuple()
|
| 410 |
+
|
| 411 |
+
|
| 412 |
+
if HAS_ZERO:
|
| 413 |
+
@spaces.GPU(duration=120)
|
| 414 |
+
def gpu_chunk(session, sim_s: float):
|
| 415 |
+
return _gpu_chunk_impl(session, sim_s, "cuda")
|
| 416 |
+
else:
|
| 417 |
+
def gpu_chunk(session, sim_s: float):
|
| 418 |
+
return _gpu_chunk_impl(session, sim_s, "cpu")
|
| 419 |
+
|
| 420 |
+
|
| 421 |
+
# ------------------------------------------------------------------ lifecycle
|
| 422 |
+
def _start_job(job):
|
| 423 |
+
from crestimap import EventConfig
|
| 424 |
+
from crestimap.session import EventSession
|
| 425 |
+
|
| 426 |
+
ev, spec = job["id"], job["spec"]
|
| 427 |
+
# CLAIM-LESS SHADOW: an unvalidated backup must never block the ladder.
|
| 428 |
+
# In "on" mode the Space (and the HPC's scan) treat any fresh claim as
|
| 429 |
+
# "a worker is on it" β a shadow claim would make the Space wait 45 min
|
| 430 |
+
# for a result we are forbidden to publish, and stop the HPC claiming.
|
| 431 |
+
# So we only write claims when we're allowed to deliver (PUBLISH=1).
|
| 432 |
+
if PUBLISH:
|
| 433 |
+
if not _commit_claim(ev, spec.get("queued", "")):
|
| 434 |
+
return False
|
| 435 |
+
time.sleep(3)
|
| 436 |
+
claim = _read_json(f"{QPREFIX}/{ev}.claim")
|
| 437 |
+
if not claim or claim.get("worker") != IDENT:
|
| 438 |
+
log(f"lost claim race on {ev}")
|
| 439 |
+
return False
|
| 440 |
+
|
| 441 |
+
rec = W["sessions"].get(ev)
|
| 442 |
+
if rec and (rec["bbox"] != spec["bbox_basin"]
|
| 443 |
+
or rec.get("cold") != spec.get("cold")):
|
| 444 |
+
log(f"{ev}: {'cold re-run requested' if rec.get('cold') != spec.get('cold') else 'bbox changed'}"
|
| 445 |
+
f" β dropping resident session")
|
| 446 |
+
shutil.rmtree(rec["root"], ignore_errors=True)
|
| 447 |
+
W["sessions"].pop(ev, None)
|
| 448 |
+
rec = None
|
| 449 |
+
if rec is None:
|
| 450 |
+
rec = {"session": None, "root": tempfile.mkdtemp(prefix=f"sess_{ev}_"),
|
| 451 |
+
"bbox": spec["bbox_basin"], "cold": spec.get("cold"),
|
| 452 |
+
"visits": 0, "last": time.time()}
|
| 453 |
+
W["sessions"][ev] = rec # registered now so _fail_job cleans it up
|
| 454 |
+
n = rec["visits"] + 1
|
| 455 |
+
forcing = os.path.join(rec["root"], f"forcing_{n}")
|
| 456 |
+
out_dir = os.path.join(rec["root"], f"out_{n}")
|
| 457 |
+
os.makedirs(forcing, exist_ok=True)
|
| 458 |
+
from huggingface_hub import hf_hub_download
|
| 459 |
+
tar = hf_hub_download(REPO, f"{QPREFIX}/{ev}.forcing.tar.gz",
|
| 460 |
+
repo_type="dataset", token=_tok or None,
|
| 461 |
+
force_download=True)
|
| 462 |
+
with tarfile.open(tar) as tf:
|
| 463 |
+
tf.extractall(forcing)
|
| 464 |
+
log(f"{ev}: forcing extracted ({len(os.listdir(forcing))} grids)")
|
| 465 |
+
|
| 466 |
+
t0 = datetime.datetime.strptime(spec["t0"], TMIN)
|
| 467 |
+
t_end = datetime.datetime.strptime(spec["t_end"], TMIN)
|
| 468 |
+
if rec["session"] is None:
|
| 469 |
+
cfg = EventConfig(
|
| 470 |
+
event_id=ev, bbox=tuple(spec["bbox_basin"]),
|
| 471 |
+
t0=t0, t_end=t_end,
|
| 472 |
+
sim_start=datetime.datetime.strptime(spec["sim_start"], TMIN),
|
| 473 |
+
ef5_output_dir=forcing, out_dir=out_dir,
|
| 474 |
+
model=spec.get("model", "crest"), dem_res="1",
|
| 475 |
+
dem_cache=os.path.join(tempfile.gettempdir(), "dem_cache"),
|
| 476 |
+
max_cells=MAX_CELLS, trigger=spec.get("trigger"),
|
| 477 |
+
device="cpu", progress=log) # module-level fn: fork-picklable
|
| 478 |
+
session = EventSession(cfg, resume=True) # CPU: DEM + pre-wet/resume
|
| 479 |
+
if max(session.grid.dx, session.grid.dy) > 45.0: # ~30 m native
|
| 480 |
+
raise RuntimeError(f"grid dx {session.grid.dx:.0f} m β load_dem "
|
| 481 |
+
f"coarsened; refusing (resolution-first)")
|
| 482 |
+
rec["session"] = session
|
| 483 |
+
session = rec["session"]
|
| 484 |
+
session.begin_visit(forcing, out_dir, t0, t_end)
|
| 485 |
+
ny, nx = session.grid.z.shape
|
| 486 |
+
W.update(job=job, runner=session, workdir=None, phase="solving",
|
| 487 |
+
queued=spec.get("queued", ""), chunk_sim_s=CHUNK0_SIM_S,
|
| 488 |
+
t_wall0=time.time(), last_hb=time.time())
|
| 489 |
+
log(f"{ev}: visit {n} begins ({ny}x{nx}, "
|
| 490 |
+
f"{(session.visit['t_end_rel'] - session.t) / 3600:.1f} sim-h from "
|
| 491 |
+
f"{session.state_time():%m-%d %H:%M})")
|
| 492 |
+
return True
|
| 493 |
+
|
| 494 |
+
|
| 495 |
+
def _finish_job():
|
| 496 |
+
job, session = W["job"], W["runner"]
|
| 497 |
+
ev, spec = job["id"], job["spec"]
|
| 498 |
+
queued = W["queued"]
|
| 499 |
+
rec = W["sessions"][ev]
|
| 500 |
+
out_dir = session.visit["out_dir"]
|
| 501 |
+
n = rec["visits"] + 1
|
| 502 |
+
from huggingface_hub import CommitOperationDelete
|
| 503 |
+
try:
|
| 504 |
+
manifest = session.end_visit() # rolls back to the junction
|
| 505 |
+
manifest["gauge"] = spec["gauge"]["id"]
|
| 506 |
+
manifest["hydro"] = spec.get("hydro") or []
|
| 507 |
+
manifest["status"] = "active"
|
| 508 |
+
manifest["basin"] = spec.get("basin")
|
| 509 |
+
manifest["huc12s"] = spec.get("huc12s")
|
| 510 |
+
manifest["domain"] = {**{k: v for k, v in (spec.get("domain") or {}).items()
|
| 511 |
+
if k != "hucs"},
|
| 512 |
+
**(manifest.get("domain") or {})}
|
| 513 |
+
from hf_data import eventsim, eventstore
|
| 514 |
+
eventsim._update_archive(out_dir, manifest, log)
|
| 515 |
+
if not manifest.get("archive_frames"):
|
| 516 |
+
manifest["status"] = "ended"
|
| 517 |
+
manifest["dry"] = True
|
| 518 |
+
eventsim._make_pngs(out_dir, manifest)
|
| 519 |
+
wall = time.time() - W["t_wall0"]
|
| 520 |
+
log(f"{ev}: visit {n} done β {len(manifest['frames'])} frames, "
|
| 521 |
+
f"wall {wall / 60:.1f} min")
|
| 522 |
+
if PUBLISH:
|
| 523 |
+
ok = eventstore.publish_event(out_dir, manifest)
|
| 524 |
+
log(f"{ev}: publish {'OK' if ok else 'FAILED'}")
|
| 525 |
+
ops = [CommitOperationDelete(f"{QPREFIX}/{ev}.claim")]
|
| 526 |
+
cur = _read_json(f"{QPREFIX}/{ev}.json")
|
| 527 |
+
if ok and cur and cur.get("queued", "") == queued:
|
| 528 |
+
ops += [CommitOperationDelete(f"{QPREFIX}/{ev}.json"),
|
| 529 |
+
CommitOperationDelete(f"{QPREFIX}/{ev}.forcing.tar.gz")]
|
| 530 |
+
_commit(ops, f"zero-worker done {ev}")
|
| 531 |
+
else:
|
| 532 |
+
log(f"{ev}: shadow done β no claim was held, nothing to release")
|
| 533 |
+
W["solved"].append({"id": ev, "visit": n,
|
| 534 |
+
"frames": len(manifest["frames"]),
|
| 535 |
+
"wall_min": round(wall / 60, 1),
|
| 536 |
+
"published": PUBLISH})
|
| 537 |
+
rec["visits"] = n
|
| 538 |
+
rec["last"] = time.time()
|
| 539 |
+
# bound disk: only the just-published visit's dirs remain
|
| 540 |
+
for stale in os.listdir(rec["root"]):
|
| 541 |
+
if stale not in (f"forcing_{n}", f"out_{n}"):
|
| 542 |
+
shutil.rmtree(os.path.join(rec["root"], stale),
|
| 543 |
+
ignore_errors=True)
|
| 544 |
+
finally:
|
| 545 |
+
W["done"].add((ev, queued))
|
| 546 |
+
W.update(job=None, runner=None, workdir=None, phase="idle")
|
| 547 |
+
|
| 548 |
+
|
| 549 |
+
def _is_quota_error(exc) -> bool:
|
| 550 |
+
m = f"{type(exc).__name__}: {exc}".lower()
|
| 551 |
+
return "gpu limit" in m or "quota" in m
|
| 552 |
+
|
| 553 |
+
|
| 554 |
+
def _fail_job(exc):
|
| 555 |
+
job = W["job"]
|
| 556 |
+
ev, queued = job["id"], W["queued"]
|
| 557 |
+
tb = "\n".join(traceback.format_exc().strip().splitlines()[-8:])
|
| 558 |
+
if _is_quota_error(exc):
|
| 559 |
+
# not this job's fault: hand it back untouched and back off
|
| 560 |
+
W["quota_until"] = time.time() + QUOTA_BACKOFF_S
|
| 561 |
+
log(f"{ev}: ZeroGPU quota exhausted β releasing the claim, no "
|
| 562 |
+
f"claims for {QUOTA_BACKOFF_S / 60:.0f} min ({exc})")
|
| 563 |
+
if PUBLISH:
|
| 564 |
+
from huggingface_hub import CommitOperationDelete
|
| 565 |
+
_commit([CommitOperationDelete(f"{QPREFIX}/{ev}.claim")],
|
| 566 |
+
f"zero-worker quota β released {ev}")
|
| 567 |
+
rec = W["sessions"].pop(ev, None) # mid-visit state is untrustworthy
|
| 568 |
+
if rec:
|
| 569 |
+
shutil.rmtree(rec["root"], ignore_errors=True)
|
| 570 |
+
W.update(job=None, runner=None, workdir=None, phase="idle")
|
| 571 |
+
return # NOT added to done: retryable
|
| 572 |
+
log(f"{ev}: FAILED β {type(exc).__name__}: {exc}")
|
| 573 |
+
if PUBLISH: # shadow held no claim
|
| 574 |
+
from huggingface_hub import CommitOperationAdd, CommitOperationDelete
|
| 575 |
+
ops = [CommitOperationDelete(f"{QPREFIX}/{ev}.claim")]
|
| 576 |
+
cur = _read_json(f"{QPREFIX}/{ev}.json")
|
| 577 |
+
if cur and cur.get("queued", "") == queued:
|
| 578 |
+
body = json.dumps({"worker": IDENT,
|
| 579 |
+
"failed": datetime.datetime.utcnow().strftime(TS),
|
| 580 |
+
"error": f"{type(exc).__name__}: {exc}"[:2000],
|
| 581 |
+
"traceback_tail": tb[:4000]})
|
| 582 |
+
ops += [CommitOperationAdd(f"{QPREFIX}/{ev}.failed.json",
|
| 583 |
+
io.BytesIO(body.encode())),
|
| 584 |
+
CommitOperationDelete(f"{QPREFIX}/{ev}.json"),
|
| 585 |
+
CommitOperationDelete(f"{QPREFIX}/{ev}.forcing.tar.gz")]
|
| 586 |
+
_commit(ops, f"zero-worker failed {ev}")
|
| 587 |
+
W["done"].add((ev, queued))
|
| 588 |
+
# mid-visit state is not trustworthy β drop the session; the next
|
| 589 |
+
# bundle rebuilds via init_depth (state-hierarchy level 2)
|
| 590 |
+
rec = W["sessions"].pop(ev, None)
|
| 591 |
+
if rec:
|
| 592 |
+
shutil.rmtree(rec["root"], ignore_errors=True)
|
| 593 |
+
W.update(job=None, runner=None, workdir=None, phase="idle")
|
| 594 |
+
|
| 595 |
+
|
| 596 |
+
def _evict_sessions():
|
| 597 |
+
"""Idle and LRU eviction of parked sessions (never the active one)."""
|
| 598 |
+
now = time.time()
|
| 599 |
+
active = (W["job"] or {}).get("id")
|
| 600 |
+
for ev in list(W["sessions"]):
|
| 601 |
+
if ev == active:
|
| 602 |
+
continue
|
| 603 |
+
idle_h = (now - W["sessions"][ev]["last"]) / 3600.0
|
| 604 |
+
if idle_h > EVICT_IDLE_H:
|
| 605 |
+
log(f"evicting session {ev} (idle {idle_h:.0f} h)")
|
| 606 |
+
shutil.rmtree(W["sessions"][ev]["root"], ignore_errors=True)
|
| 607 |
+
del W["sessions"][ev]
|
| 608 |
+
while len(W["sessions"]) > MAX_SESSIONS:
|
| 609 |
+
parked = [e for e in W["sessions"] if e != active]
|
| 610 |
+
if not parked:
|
| 611 |
+
break
|
| 612 |
+
ev = min(parked, key=lambda e: W["sessions"][e]["last"])
|
| 613 |
+
log(f"evicting session {ev} (LRU, cap {MAX_SESSIONS})")
|
| 614 |
+
shutil.rmtree(W["sessions"][ev]["root"], ignore_errors=True)
|
| 615 |
+
del W["sessions"][ev]
|
| 616 |
+
|
| 617 |
+
|
| 618 |
+
def view():
|
| 619 |
+
"""What a browser sees: status + log only. It never advances the state
|
| 620 |
+
machine β a visitor's request would run @spaces.GPU under the VISITOR's
|
| 621 |
+
ZeroGPU identity, and an anonymous visitor's quota error could stall the
|
| 622 |
+
only 2-D engine for a day."""
|
| 623 |
+
return _status(), "\n".join(LOG[-40:])
|
| 624 |
+
|
| 625 |
+
|
| 626 |
+
def tick_api(key: str = ""):
|
| 627 |
+
"""The /tick endpoint: advances the worker only for the self-drive loop,
|
| 628 |
+
which passes this process's secret; anyone else just gets the status.
|
| 629 |
+
|
| 630 |
+
Every call is counted (status "tick_api"): from 10-19 to 10-20 the worker
|
| 631 |
+
sat idle next to an eligible job with no log line at all, which means
|
| 632 |
+
either the key never arrived or tick() found nothing. These counters tell
|
| 633 |
+
the two apart from outside."""
|
| 634 |
+
c = W["tick_api"]
|
| 635 |
+
if not key:
|
| 636 |
+
c["empty"] += 1
|
| 637 |
+
return view()
|
| 638 |
+
if key != _SELF_KEY:
|
| 639 |
+
c["rejected"] += 1
|
| 640 |
+
if time.time() - c["last_reject_log"] > 3600:
|
| 641 |
+
c["last_reject_log"] = time.time()
|
| 642 |
+
log(f"tick_api: rejected a key (len {len(key)}; "
|
| 643 |
+
f"{c['rejected']} so far)")
|
| 644 |
+
return view()
|
| 645 |
+
c["accepted"] += 1
|
| 646 |
+
if c["accepted"] == 1:
|
| 647 |
+
log("tick_api: self-drive key accepted β worker advancing")
|
| 648 |
+
return tick()
|
| 649 |
+
|
| 650 |
+
|
| 651 |
+
def tick():
|
| 652 |
+
"""One state-machine step (single-flight); called via tick_api by the
|
| 653 |
+
self-drive loop."""
|
| 654 |
+
if not _tick_lock.acquire(blocking=False):
|
| 655 |
+
return _status(), "\n".join(LOG[-40:])
|
| 656 |
+
try:
|
| 657 |
+
now = time.time()
|
| 658 |
+
wait = _gpu_wait_reason(now)
|
| 659 |
+
if (wait or "") != W["paused_note"]:
|
| 660 |
+
W["paused_note"] = wait or ""
|
| 661 |
+
log(f"paused: {wait}" if wait else "GPU available β resuming")
|
| 662 |
+
if W["phase"] == "idle":
|
| 663 |
+
# never claim a job we couldn't start computing on right away
|
| 664 |
+
if now - W["last_poll"] >= POLL_S and wait is None:
|
| 665 |
+
W["last_poll"] = now
|
| 666 |
+
_evict_sessions()
|
| 667 |
+
job = _scan()
|
| 668 |
+
if (job and job["id"] not in W["sessions"]
|
| 669 |
+
and _budget_left(now) < MIN_NEW_S):
|
| 670 |
+
job = None # a new event waits for a fuller budget
|
| 671 |
+
if job:
|
| 672 |
+
kind = ("catch-up" if job["id"] in W["sessions"]
|
| 673 |
+
else "new")
|
| 674 |
+
log(f"claiming {job['id']} (backup rung, {kind}, "
|
| 675 |
+
f"queued {job['spec'].get('queued')})")
|
| 676 |
+
try:
|
| 677 |
+
_start_job(job)
|
| 678 |
+
except Exception as e:
|
| 679 |
+
W["job"] = job
|
| 680 |
+
W["queued"] = job["spec"].get("queued", "")
|
| 681 |
+
_fail_job(e)
|
| 682 |
+
elif W["phase"] == "solving":
|
| 683 |
+
ev = W["job"]["id"]
|
| 684 |
+
if PUBLISH and now - W["last_hb"] >= HB_S:
|
| 685 |
+
_commit_claim(ev, W["queued"])
|
| 686 |
+
W["last_hb"] = now
|
| 687 |
+
session = W["runner"]
|
| 688 |
+
if wait is not None:
|
| 689 |
+
# paused mid-visit: the session keeps its state in RAM and
|
| 690 |
+
# the claim stays fresh (heartbeat above); resume later
|
| 691 |
+
return _status(), "\n".join(LOG[-40:])
|
| 692 |
+
t0 = time.time()
|
| 693 |
+
try:
|
| 694 |
+
done, state = gpu_chunk(session, W["chunk_sim_s"])
|
| 695 |
+
session.set_state(state)
|
| 696 |
+
except Exception as e:
|
| 697 |
+
# admission failures leave the session untouched (the chunk
|
| 698 |
+
# never ran): hold the job and retry the same chunk later β
|
| 699 |
+
# but only when ZeroGPU names a retry time; a quota error
|
| 700 |
+
# without one (the anonymous pool: our identity was not
|
| 701 |
+
# recognised) takes the old release-and-back-off path
|
| 702 |
+
retry = _quota_retry_s(e) if _is_quota_error(e) else None
|
| 703 |
+
if _is_capacity_error(e):
|
| 704 |
+
W["capacity_until"] = time.time() + CAPACITY_BACKOFF_S
|
| 705 |
+
log(f"{ev}: no ZeroGPU capacity β retry in "
|
| 706 |
+
f"{CAPACITY_BACKOFF_S:.0f} s ({e})")
|
| 707 |
+
elif retry is not None:
|
| 708 |
+
W["quota_until"] = time.time() + retry
|
| 709 |
+
log(f"{ev}: ZeroGPU quota exhausted β holding the job "
|
| 710 |
+
f"until {_hhmm(W['quota_until'])} ({e})")
|
| 711 |
+
else:
|
| 712 |
+
if _auth_suspect(e):
|
| 713 |
+
W["auth_suspect"] = True
|
| 714 |
+
log(f"{ev}: anonymous-pool quota error β GPU calls are "
|
| 715 |
+
f"NOT attributed to the account (ZERO_AUTH failed?)")
|
| 716 |
+
elif not _is_quota_error(e):
|
| 717 |
+
_charge(time.time() - t0, t0) # it ran, then failed
|
| 718 |
+
_fail_job(e)
|
| 719 |
+
return _status(), "\n".join(LOG[-40:])
|
| 720 |
+
_charge(time.time() - t0, t0)
|
| 721 |
+
wall = max(time.time() - t0, 1e-3)
|
| 722 |
+
rate = W["chunk_sim_s"] / wall # sim-s per wall-s
|
| 723 |
+
W["chunk_sim_s"] = float(min(max(GPU_TARGET_S * rate, 60.0), 14400.0))
|
| 724 |
+
total_h = session.visit["t_end_rel"] / 3600.0
|
| 725 |
+
log(f"{ev}: sim t={session.t / 3600:.2f}/{total_h:.0f} h "
|
| 726 |
+
f"(chunk {wall:.0f} s wall, next {W['chunk_sim_s'] / 60:.0f} sim-min)")
|
| 727 |
+
if done:
|
| 728 |
+
try:
|
| 729 |
+
_finish_job()
|
| 730 |
+
except Exception as e:
|
| 731 |
+
_fail_job(e)
|
| 732 |
+
return _status(), "\n".join(LOG[-40:])
|
| 733 |
+
finally:
|
| 734 |
+
_tick_lock.release()
|
| 735 |
+
|
| 736 |
+
|
| 737 |
+
def _status():
|
| 738 |
+
s = W["runner"]
|
| 739 |
+
return {"ident": IDENT, "phase": W["phase"], "publish": PUBLISH,
|
| 740 |
+
"job": (W["job"] or {}).get("id"),
|
| 741 |
+
"sim_h": round(s.t / 3600, 2) if s else None,
|
| 742 |
+
"total_h": (round(s.visit["t_end_rel"] / 3600, 1)
|
| 743 |
+
if s and s.visit else None),
|
| 744 |
+
"chunk_sim_min": round(W["chunk_sim_s"] / 60, 1),
|
| 745 |
+
# seconds since the server-side loop last completed a tick β the
|
| 746 |
+
# liveness signal: if this grows without bound the worker is idle
|
| 747 |
+
# even though the container reads RUNNING (the 2026-08-24 failure)
|
| 748 |
+
"tick_age_s": (round(time.time() - W["last_tick"], 1)
|
| 749 |
+
if W["last_tick"] else None),
|
| 750 |
+
"sessions": {ev: {"visits": r["visits"],
|
| 751 |
+
"held_at": (r["session"].state_time()
|
| 752 |
+
.strftime("%m-%d %H:%MZ")
|
| 753 |
+
if r["session"] else None)}
|
| 754 |
+
for ev, r in W["sessions"].items()},
|
| 755 |
+
"solved": W["solved"][-5:], "zero_gpu": HAS_ZERO,
|
| 756 |
+
"tick_api": {k: W["tick_api"][k]
|
| 757 |
+
for k in ("accepted", "rejected", "empty")},
|
| 758 |
+
"scan": W["scan"],
|
| 759 |
+
"fork_sha": _FORK_SHA,
|
| 760 |
+
"gpu_budget": {"used_min_24h": round(_used_s(time.time()) / 60, 1),
|
| 761 |
+
"budget_min": round(DAILY_BUDGET_S / 60, 1),
|
| 762 |
+
"authenticated": bool(ZERO_AUTH and _tok),
|
| 763 |
+
"auth_suspect": W["auth_suspect"],
|
| 764 |
+
"paused": W["paused_note"] or None}}
|
| 765 |
+
|
| 766 |
+
|
| 767 |
+
def _keep_awake():
|
| 768 |
+
import requests
|
| 769 |
+
host = os.environ.get("SPACE_HOST", "")
|
| 770 |
+
while host:
|
| 771 |
+
time.sleep(900)
|
| 772 |
+
try:
|
| 773 |
+
requests.get(f"https://{host}", timeout=30)
|
| 774 |
+
except Exception:
|
| 775 |
+
pass
|
| 776 |
+
|
| 777 |
+
|
| 778 |
+
def _self_drive():
|
| 779 |
+
"""Server-side tick loop β the worker's actual engine.
|
| 780 |
+
|
| 781 |
+
gr.Timer is a FRONTEND timer: it only fires while a browser session has
|
| 782 |
+
the Space open, and demo.load only fires on a page load. _keep_awake's
|
| 783 |
+
plain GET keeps the container off sleep but opens no Gradio session, so
|
| 784 |
+
with nobody watching this worker sat RUNNING and completely idle β found
|
| 785 |
+
2026-08-24: 59 min of silence between boot and the first externally
|
| 786 |
+
triggered /tick, and every past ZeroGPU solve coincided with someone
|
| 787 |
+
polling it. The fleet/ungauged workers self-drive with a plain thread;
|
| 788 |
+
this brings the GPU worker in line.
|
| 789 |
+
|
| 790 |
+
The tick is driven through the Space's own /tick API over localhost
|
| 791 |
+
rather than by calling tick() inline, so the state machine advances on
|
| 792 |
+
exactly the request path it was built for (queue single-flight, and the
|
| 793 |
+
@spaces.GPU allocation happens inside a normal Gradio request)."""
|
| 794 |
+
from gradio_client import Client
|
| 795 |
+
base = f"http://127.0.0.1:{os.environ.get('GRADIO_SERVER_PORT', '7860')}"
|
| 796 |
+
space = os.environ.get("SPACE_ID", "")
|
| 797 |
+
if ZERO_AUTH and _tok and space:
|
| 798 |
+
# through HF's proxy with the account token: the proxy attaches the
|
| 799 |
+
# caller identity (X-IP-Token), so @spaces.GPU draws on the account's
|
| 800 |
+
# ZeroGPU quota instead of the anonymous per-Space pool
|
| 801 |
+
src, kw, how = space, {"hf_token": _tok}, "authenticated (account quota)"
|
| 802 |
+
else:
|
| 803 |
+
src, kw, how = base, {}, "localhost, unauthenticated (anonymous pool)"
|
| 804 |
+
client = None
|
| 805 |
+
while True:
|
| 806 |
+
try:
|
| 807 |
+
if client is None:
|
| 808 |
+
# token positionally: gradio_client 1.x names it hf_token,
|
| 809 |
+
# 2.x names it token
|
| 810 |
+
client = (Client(src, kw["hf_token"], verbose=False) if kw
|
| 811 |
+
else Client(src, verbose=False))
|
| 812 |
+
log(f"self-drive: tick loop connected β {how}, "
|
| 813 |
+
f"{SELF_TICK_S:.0f} s cadence")
|
| 814 |
+
# bounded wait: a half-open proxy connection must not hang the
|
| 815 |
+
# only engine forever (a healthy tick is <= ~3 min)
|
| 816 |
+
client.submit(_SELF_KEY, api_name="/tick").result(timeout=900)
|
| 817 |
+
W["last_tick"] = time.time()
|
| 818 |
+
except Exception as e: # server still booting, etc.
|
| 819 |
+
try:
|
| 820 |
+
if client is not None:
|
| 821 |
+
client.close()
|
| 822 |
+
except Exception:
|
| 823 |
+
pass
|
| 824 |
+
client = None
|
| 825 |
+
log(f"self-drive: {type(e).__name__}: {str(e)[:200]} β retrying")
|
| 826 |
+
time.sleep(15)
|
| 827 |
+
continue
|
| 828 |
+
time.sleep(SELF_TICK_S)
|
| 829 |
+
|
| 830 |
+
|
| 831 |
+
threading.Thread(target=_keep_awake, daemon=True).start()
|
| 832 |
+
threading.Thread(target=_self_drive, daemon=True).start()
|
| 833 |
+
log(f"boot: {IDENT} publish={PUBLISH} zero={HAS_ZERO} "
|
| 834 |
+
f"delay={ZERO_DELAY_S:.0f}s budget={MAX_CELLS / 1e6:.1f}M cells "
|
| 835 |
+
f"gpu/day={DAILY_BUDGET_S / 60:.0f} min auth={bool(ZERO_AUTH and _tok)} "
|
| 836 |
+
f"crestimap={_FORK_SHA}")
|
| 837 |
+
|
| 838 |
+
with gr.Blocks(title="CREST GPU worker (ZeroGPU)") as demo:
|
| 839 |
+
gr.Markdown("## CREST event worker β ZeroGPU\n"
|
| 840 |
+
"The 2-D inundation engine since the HPC worker was released "
|
| 841 |
+
"(2026-10-16). Works within a daily GPU budget; shadow mode "
|
| 842 |
+
"unless PUBLISH=1.")
|
| 843 |
+
st = gr.JSON(label="worker state")
|
| 844 |
+
lg = gr.Textbox(label="log", lines=18, max_lines=18)
|
| 845 |
+
# the page only DISPLAYS (view); the worker advances through /tick,
|
| 846 |
+
# which needs the in-process key the self-drive loop sends
|
| 847 |
+
t = gr.Timer(5)
|
| 848 |
+
t.tick(view, outputs=[st, lg], api_name=False)
|
| 849 |
+
demo.load(view, outputs=[st, lg], api_name=False)
|
| 850 |
+
key_in = gr.Textbox(visible=False)
|
| 851 |
+
tick_btn = gr.Button(visible=False)
|
| 852 |
+
tick_btn.click(tick_api, inputs=[key_in], outputs=[st, lg],
|
| 853 |
+
api_name="tick", concurrency_limit=1)
|
| 854 |
+
|
| 855 |
+
demo.queue().launch()
|