vincewin commited on
Commit
3271f17
Β·
verified Β·
1 Parent(s): 71616ca

Leave new observed-boundary and KW jobs to compatible Jetstream2 workers

Browse files
Files changed (1) hide show
  1. app.py +855 -841
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
- if (ev, spec.get("queued", "")) in W["done"]:
350
- _skip("done")
351
- continue
352
- if str(spec.get("dem_res", "1")) != "1":
353
- _skip("not_30m")
354
- continue # 30 m only on this rung
355
- try:
356
- age = (now - datetime.datetime.strptime(spec["queued"], TS)
357
- ).total_seconds()
358
- except Exception:
359
- age = 0.0
360
- if age < delay:
361
- _skip("deferred")
362
- continue
363
- if not spec.get("domain"):
364
- # pre-basin bundle (rectangular-window era): never solved as a
365
- # rectangle β€” the runner re-enqueues it with a HUC12 basin domain
366
- log(f"skip {ev}: no basin domain in spec β€” awaiting re-enqueue")
367
- _skip("no_domain")
368
- continue
369
- w, s, e, n = spec["bbox_basin"]
370
- res = DEM_RES_DEG["1"]
371
- cells = int((spec.get("domain") or {}).get("n_bbox") or
372
- int(round((n - s) / res)) * int(round((e - w) / res)))
373
- if cells > MAX_CELLS:
374
- log(f"skip {ev}: {cells / 1e6:.1f} M cells > "
375
- f"{MAX_CELLS / 1e6:.1f} M budget (never coarsen)")
376
- _skip("over_budget")
377
- continue
378
- if f"{QPREFIX}/{ev}.claim" in qf:
379
- claim = _read_json(f"{QPREFIX}/{ev}.claim")
380
- if _claim_fresh(claim) and claim.get("worker") != IDENT:
381
- _skip("claimed")
382
- continue
383
- cands.append({"id": ev, "spec": spec})
384
- W["scan"]["candidates"] = len(cands)
385
- # P1.6: episodes with a resident session first β€” a catch-up visit is
386
- # minutes; new events (full anchored solve) fill the remaining time
387
- cands.sort(key=lambda j: (j["id"] not in W["sessions"],
388
- j["spec"].get("queued", "")))
389
- return cands[0] if cands else None
390
-
391
-
392
- # ------------------------------------------------------------------ GPU work
393
- def _gpu_chunk_impl(session, sim_s: float, device: str):
394
- done = session.chunk(sim_s, device=device)
395
- return done, session.state_tuple()
396
-
397
-
398
- if HAS_ZERO:
399
- @spaces.GPU(duration=120)
400
- def gpu_chunk(session, sim_s: float):
401
- return _gpu_chunk_impl(session, sim_s, "cuda")
402
- else:
403
- def gpu_chunk(session, sim_s: float):
404
- return _gpu_chunk_impl(session, sim_s, "cpu")
405
-
406
-
407
- # ------------------------------------------------------------------ lifecycle
408
- def _start_job(job):
409
- from crestimap import EventConfig
410
- from crestimap.session import EventSession
411
-
412
- ev, spec = job["id"], job["spec"]
413
- # CLAIM-LESS SHADOW: an unvalidated backup must never block the ladder.
414
- # In "on" mode the Space (and the HPC's scan) treat any fresh claim as
415
- # "a worker is on it" β€” a shadow claim would make the Space wait 45 min
416
- # for a result we are forbidden to publish, and stop the HPC claiming.
417
- # So we only write claims when we're allowed to deliver (PUBLISH=1).
418
- if PUBLISH:
419
- if not _commit_claim(ev, spec.get("queued", "")):
420
- return False
421
- time.sleep(3)
422
- claim = _read_json(f"{QPREFIX}/{ev}.claim")
423
- if not claim or claim.get("worker") != IDENT:
424
- log(f"lost claim race on {ev}")
425
- return False
426
-
427
- rec = W["sessions"].get(ev)
428
- if rec and (rec["bbox"] != spec["bbox_basin"]
429
- or rec.get("cold") != spec.get("cold")):
430
- log(f"{ev}: {'cold re-run requested' if rec.get('cold') != spec.get('cold') else 'bbox changed'}"
431
- f" β€” dropping resident session")
432
- shutil.rmtree(rec["root"], ignore_errors=True)
433
- W["sessions"].pop(ev, None)
434
- rec = None
435
- if rec is None:
436
- rec = {"session": None, "root": tempfile.mkdtemp(prefix=f"sess_{ev}_"),
437
- "bbox": spec["bbox_basin"], "cold": spec.get("cold"),
438
- "visits": 0, "last": time.time()}
439
- W["sessions"][ev] = rec # registered now so _fail_job cleans it up
440
- n = rec["visits"] + 1
441
- forcing = os.path.join(rec["root"], f"forcing_{n}")
442
- out_dir = os.path.join(rec["root"], f"out_{n}")
443
- os.makedirs(forcing, exist_ok=True)
444
- from huggingface_hub import hf_hub_download
445
- tar = hf_hub_download(REPO, f"{QPREFIX}/{ev}.forcing.tar.gz",
446
- repo_type="dataset", token=_tok or None,
447
- force_download=True)
448
- with tarfile.open(tar) as tf:
449
- tf.extractall(forcing)
450
- log(f"{ev}: forcing extracted ({len(os.listdir(forcing))} grids)")
451
-
452
- t0 = datetime.datetime.strptime(spec["t0"], TMIN)
453
- t_end = datetime.datetime.strptime(spec["t_end"], TMIN)
454
- if rec["session"] is None:
455
- cfg = EventConfig(
456
- event_id=ev, bbox=tuple(spec["bbox_basin"]),
457
- t0=t0, t_end=t_end,
458
- sim_start=datetime.datetime.strptime(spec["sim_start"], TMIN),
459
- ef5_output_dir=forcing, out_dir=out_dir,
460
- model=spec.get("model", "crest"), dem_res="1",
461
- dem_cache=os.path.join(tempfile.gettempdir(), "dem_cache"),
462
- max_cells=MAX_CELLS, trigger=spec.get("trigger"),
463
- device="cpu", progress=log) # module-level fn: fork-picklable
464
- session = EventSession(cfg, resume=True) # CPU: DEM + pre-wet/resume
465
- if max(session.grid.dx, session.grid.dy) > 45.0: # ~30 m native
466
- raise RuntimeError(f"grid dx {session.grid.dx:.0f} m β€” load_dem "
467
- f"coarsened; refusing (resolution-first)")
468
- rec["session"] = session
469
- session = rec["session"]
470
- session.begin_visit(forcing, out_dir, t0, t_end)
471
- ny, nx = session.grid.z.shape
472
- W.update(job=job, runner=session, workdir=None, phase="solving",
473
- queued=spec.get("queued", ""), chunk_sim_s=CHUNK0_SIM_S,
474
- t_wall0=time.time(), last_hb=time.time())
475
- log(f"{ev}: visit {n} begins ({ny}x{nx}, "
476
- f"{(session.visit['t_end_rel'] - session.t) / 3600:.1f} sim-h from "
477
- f"{session.state_time():%m-%d %H:%M})")
478
- return True
479
-
480
-
481
- def _finish_job():
482
- job, session = W["job"], W["runner"]
483
- ev, spec = job["id"], job["spec"]
484
- queued = W["queued"]
485
- rec = W["sessions"][ev]
486
- out_dir = session.visit["out_dir"]
487
- n = rec["visits"] + 1
488
- from huggingface_hub import CommitOperationDelete
489
- try:
490
- manifest = session.end_visit() # rolls back to the junction
491
- manifest["gauge"] = spec["gauge"]["id"]
492
- manifest["hydro"] = spec.get("hydro") or []
493
- manifest["status"] = "active"
494
- manifest["basin"] = spec.get("basin")
495
- manifest["huc12s"] = spec.get("huc12s")
496
- manifest["domain"] = {**{k: v for k, v in (spec.get("domain") or {}).items()
497
- if k != "hucs"},
498
- **(manifest.get("domain") or {})}
499
- from hf_data import eventsim, eventstore
500
- eventsim._update_archive(out_dir, manifest, log)
501
- if not manifest.get("archive_frames"):
502
- manifest["status"] = "ended"
503
- manifest["dry"] = True
504
- eventsim._make_pngs(out_dir, manifest)
505
- wall = time.time() - W["t_wall0"]
506
- log(f"{ev}: visit {n} done β€” {len(manifest['frames'])} frames, "
507
- f"wall {wall / 60:.1f} min")
508
- if PUBLISH:
509
- ok = eventstore.publish_event(out_dir, manifest)
510
- log(f"{ev}: publish {'OK' if ok else 'FAILED'}")
511
- ops = [CommitOperationDelete(f"{QPREFIX}/{ev}.claim")]
512
- cur = _read_json(f"{QPREFIX}/{ev}.json")
513
- if ok and cur and cur.get("queued", "") == queued:
514
- ops += [CommitOperationDelete(f"{QPREFIX}/{ev}.json"),
515
- CommitOperationDelete(f"{QPREFIX}/{ev}.forcing.tar.gz")]
516
- _commit(ops, f"zero-worker done {ev}")
517
- else:
518
- log(f"{ev}: shadow done β€” no claim was held, nothing to release")
519
- W["solved"].append({"id": ev, "visit": n,
520
- "frames": len(manifest["frames"]),
521
- "wall_min": round(wall / 60, 1),
522
- "published": PUBLISH})
523
- rec["visits"] = n
524
- rec["last"] = time.time()
525
- # bound disk: only the just-published visit's dirs remain
526
- for stale in os.listdir(rec["root"]):
527
- if stale not in (f"forcing_{n}", f"out_{n}"):
528
- shutil.rmtree(os.path.join(rec["root"], stale),
529
- ignore_errors=True)
530
- finally:
531
- W["done"].add((ev, queued))
532
- W.update(job=None, runner=None, workdir=None, phase="idle")
533
-
534
-
535
- def _is_quota_error(exc) -> bool:
536
- m = f"{type(exc).__name__}: {exc}".lower()
537
- return "gpu limit" in m or "quota" in m
538
-
539
-
540
- def _fail_job(exc):
541
- job = W["job"]
542
- ev, queued = job["id"], W["queued"]
543
- tb = "\n".join(traceback.format_exc().strip().splitlines()[-8:])
544
- if _is_quota_error(exc):
545
- # not this job's fault: hand it back untouched and back off
546
- W["quota_until"] = time.time() + QUOTA_BACKOFF_S
547
- log(f"{ev}: ZeroGPU quota exhausted β€” releasing the claim, no "
548
- f"claims for {QUOTA_BACKOFF_S / 60:.0f} min ({exc})")
549
- if PUBLISH:
550
- from huggingface_hub import CommitOperationDelete
551
- _commit([CommitOperationDelete(f"{QPREFIX}/{ev}.claim")],
552
- f"zero-worker quota β€” released {ev}")
553
- rec = W["sessions"].pop(ev, None) # mid-visit state is untrustworthy
554
- if rec:
555
- shutil.rmtree(rec["root"], ignore_errors=True)
556
- W.update(job=None, runner=None, workdir=None, phase="idle")
557
- return # NOT added to done: retryable
558
- log(f"{ev}: FAILED β€” {type(exc).__name__}: {exc}")
559
- if PUBLISH: # shadow held no claim
560
- from huggingface_hub import CommitOperationAdd, CommitOperationDelete
561
- ops = [CommitOperationDelete(f"{QPREFIX}/{ev}.claim")]
562
- cur = _read_json(f"{QPREFIX}/{ev}.json")
563
- if cur and cur.get("queued", "") == queued:
564
- body = json.dumps({"worker": IDENT,
565
- "failed": datetime.datetime.utcnow().strftime(TS),
566
- "error": f"{type(exc).__name__}: {exc}"[:2000],
567
- "traceback_tail": tb[:4000]})
568
- ops += [CommitOperationAdd(f"{QPREFIX}/{ev}.failed.json",
569
- io.BytesIO(body.encode())),
570
- CommitOperationDelete(f"{QPREFIX}/{ev}.json"),
571
- CommitOperationDelete(f"{QPREFIX}/{ev}.forcing.tar.gz")]
572
- _commit(ops, f"zero-worker failed {ev}")
573
- W["done"].add((ev, queued))
574
- # mid-visit state is not trustworthy β€” drop the session; the next
575
- # bundle rebuilds via init_depth (state-hierarchy level 2)
576
- rec = W["sessions"].pop(ev, None)
577
- if rec:
578
- shutil.rmtree(rec["root"], ignore_errors=True)
579
- W.update(job=None, runner=None, workdir=None, phase="idle")
580
-
581
-
582
- def _evict_sessions():
583
- """Idle and LRU eviction of parked sessions (never the active one)."""
584
- now = time.time()
585
- active = (W["job"] or {}).get("id")
586
- for ev in list(W["sessions"]):
587
- if ev == active:
588
- continue
589
- idle_h = (now - W["sessions"][ev]["last"]) / 3600.0
590
- if idle_h > EVICT_IDLE_H:
591
- log(f"evicting session {ev} (idle {idle_h:.0f} h)")
592
- shutil.rmtree(W["sessions"][ev]["root"], ignore_errors=True)
593
- del W["sessions"][ev]
594
- while len(W["sessions"]) > MAX_SESSIONS:
595
- parked = [e for e in W["sessions"] if e != active]
596
- if not parked:
597
- break
598
- ev = min(parked, key=lambda e: W["sessions"][e]["last"])
599
- log(f"evicting session {ev} (LRU, cap {MAX_SESSIONS})")
600
- shutil.rmtree(W["sessions"][ev]["root"], ignore_errors=True)
601
- del W["sessions"][ev]
602
-
603
-
604
- def view():
605
- """What a browser sees: status + log only. It never advances the state
606
- machine β€” a visitor's request would run @spaces.GPU under the VISITOR's
607
- ZeroGPU identity, and an anonymous visitor's quota error could stall the
608
- only 2-D engine for a day."""
609
- return _status(), "\n".join(LOG[-40:])
610
-
611
-
612
- def tick_api(key: str = ""):
613
- """The /tick endpoint: advances the worker only for the self-drive loop,
614
- which passes this process's secret; anyone else just gets the status.
615
-
616
- Every call is counted (status "tick_api"): from 10-19 to 10-20 the worker
617
- sat idle next to an eligible job with no log line at all, which means
618
- either the key never arrived or tick() found nothing. These counters tell
619
- the two apart from outside."""
620
- c = W["tick_api"]
621
- if not key:
622
- c["empty"] += 1
623
- return view()
624
- if key != _SELF_KEY:
625
- c["rejected"] += 1
626
- if time.time() - c["last_reject_log"] > 3600:
627
- c["last_reject_log"] = time.time()
628
- log(f"tick_api: rejected a key (len {len(key)}; "
629
- f"{c['rejected']} so far)")
630
- return view()
631
- c["accepted"] += 1
632
- if c["accepted"] == 1:
633
- log("tick_api: self-drive key accepted β€” worker advancing")
634
- return tick()
635
-
636
-
637
- def tick():
638
- """One state-machine step (single-flight); called via tick_api by the
639
- self-drive loop."""
640
- if not _tick_lock.acquire(blocking=False):
641
- return _status(), "\n".join(LOG[-40:])
642
- try:
643
- now = time.time()
644
- wait = _gpu_wait_reason(now)
645
- if (wait or "") != W["paused_note"]:
646
- W["paused_note"] = wait or ""
647
- log(f"paused: {wait}" if wait else "GPU available β€” resuming")
648
- if W["phase"] == "idle":
649
- # never claim a job we couldn't start computing on right away
650
- if now - W["last_poll"] >= POLL_S and wait is None:
651
- W["last_poll"] = now
652
- _evict_sessions()
653
- job = _scan()
654
- if (job and job["id"] not in W["sessions"]
655
- and _budget_left(now) < MIN_NEW_S):
656
- job = None # a new event waits for a fuller budget
657
- if job:
658
- kind = ("catch-up" if job["id"] in W["sessions"]
659
- else "new")
660
- log(f"claiming {job['id']} (backup rung, {kind}, "
661
- f"queued {job['spec'].get('queued')})")
662
- try:
663
- _start_job(job)
664
- except Exception as e:
665
- W["job"] = job
666
- W["queued"] = job["spec"].get("queued", "")
667
- _fail_job(e)
668
- elif W["phase"] == "solving":
669
- ev = W["job"]["id"]
670
- if PUBLISH and now - W["last_hb"] >= HB_S:
671
- _commit_claim(ev, W["queued"])
672
- W["last_hb"] = now
673
- session = W["runner"]
674
- if wait is not None:
675
- # paused mid-visit: the session keeps its state in RAM and
676
- # the claim stays fresh (heartbeat above); resume later
677
- return _status(), "\n".join(LOG[-40:])
678
- t0 = time.time()
679
- try:
680
- done, state = gpu_chunk(session, W["chunk_sim_s"])
681
- session.set_state(state)
682
- except Exception as e:
683
- # admission failures leave the session untouched (the chunk
684
- # never ran): hold the job and retry the same chunk later β€”
685
- # but only when ZeroGPU names a retry time; a quota error
686
- # without one (the anonymous pool: our identity was not
687
- # recognised) takes the old release-and-back-off path
688
- retry = _quota_retry_s(e) if _is_quota_error(e) else None
689
- if _is_capacity_error(e):
690
- W["capacity_until"] = time.time() + CAPACITY_BACKOFF_S
691
- log(f"{ev}: no ZeroGPU capacity β€” retry in "
692
- f"{CAPACITY_BACKOFF_S:.0f} s ({e})")
693
- elif retry is not None:
694
- W["quota_until"] = time.time() + retry
695
- log(f"{ev}: ZeroGPU quota exhausted β€” holding the job "
696
- f"until {_hhmm(W['quota_until'])} ({e})")
697
- else:
698
- if _auth_suspect(e):
699
- W["auth_suspect"] = True
700
- log(f"{ev}: anonymous-pool quota error β€” GPU calls are "
701
- f"NOT attributed to the account (ZERO_AUTH failed?)")
702
- elif not _is_quota_error(e):
703
- _charge(time.time() - t0, t0) # it ran, then failed
704
- _fail_job(e)
705
- return _status(), "\n".join(LOG[-40:])
706
- _charge(time.time() - t0, t0)
707
- wall = max(time.time() - t0, 1e-3)
708
- rate = W["chunk_sim_s"] / wall # sim-s per wall-s
709
- W["chunk_sim_s"] = float(min(max(GPU_TARGET_S * rate, 60.0), 14400.0))
710
- total_h = session.visit["t_end_rel"] / 3600.0
711
- log(f"{ev}: sim t={session.t / 3600:.2f}/{total_h:.0f} h "
712
- f"(chunk {wall:.0f} s wall, next {W['chunk_sim_s'] / 60:.0f} sim-min)")
713
- if done:
714
- try:
715
- _finish_job()
716
- except Exception as e:
717
- _fail_job(e)
718
- return _status(), "\n".join(LOG[-40:])
719
- finally:
720
- _tick_lock.release()
721
-
722
-
723
- def _status():
724
- s = W["runner"]
725
- return {"ident": IDENT, "phase": W["phase"], "publish": PUBLISH,
726
- "job": (W["job"] or {}).get("id"),
727
- "sim_h": round(s.t / 3600, 2) if s else None,
728
- "total_h": (round(s.visit["t_end_rel"] / 3600, 1)
729
- if s and s.visit else None),
730
- "chunk_sim_min": round(W["chunk_sim_s"] / 60, 1),
731
- # seconds since the server-side loop last completed a tick β€” the
732
- # liveness signal: if this grows without bound the worker is idle
733
- # even though the container reads RUNNING (the 2026-08-24 failure)
734
- "tick_age_s": (round(time.time() - W["last_tick"], 1)
735
- if W["last_tick"] else None),
736
- "sessions": {ev: {"visits": r["visits"],
737
- "held_at": (r["session"].state_time()
738
- .strftime("%m-%d %H:%MZ")
739
- if r["session"] else None)}
740
- for ev, r in W["sessions"].items()},
741
- "solved": W["solved"][-5:], "zero_gpu": HAS_ZERO,
742
- "tick_api": {k: W["tick_api"][k]
743
- for k in ("accepted", "rejected", "empty")},
744
- "scan": W["scan"],
745
- "fork_sha": _FORK_SHA,
746
- "gpu_budget": {"used_min_24h": round(_used_s(time.time()) / 60, 1),
747
- "budget_min": round(DAILY_BUDGET_S / 60, 1),
748
- "authenticated": bool(ZERO_AUTH and _tok),
749
- "auth_suspect": W["auth_suspect"],
750
- "paused": W["paused_note"] or None}}
751
-
752
-
753
- def _keep_awake():
754
- import requests
755
- host = os.environ.get("SPACE_HOST", "")
756
- while host:
757
- time.sleep(900)
758
- try:
759
- requests.get(f"https://{host}", timeout=30)
760
- except Exception:
761
- pass
762
-
763
-
764
- def _self_drive():
765
- """Server-side tick loop β€” the worker's actual engine.
766
-
767
- gr.Timer is a FRONTEND timer: it only fires while a browser session has
768
- the Space open, and demo.load only fires on a page load. _keep_awake's
769
- plain GET keeps the container off sleep but opens no Gradio session, so
770
- with nobody watching this worker sat RUNNING and completely idle β€” found
771
- 2026-08-24: 59 min of silence between boot and the first externally
772
- triggered /tick, and every past ZeroGPU solve coincided with someone
773
- polling it. The fleet/ungauged workers self-drive with a plain thread;
774
- this brings the GPU worker in line.
775
-
776
- The tick is driven through the Space's own /tick API over localhost
777
- rather than by calling tick() inline, so the state machine advances on
778
- exactly the request path it was built for (queue single-flight, and the
779
- @spaces.GPU allocation happens inside a normal Gradio request)."""
780
- from gradio_client import Client
781
- base = f"http://127.0.0.1:{os.environ.get('GRADIO_SERVER_PORT', '7860')}"
782
- space = os.environ.get("SPACE_ID", "")
783
- if ZERO_AUTH and _tok and space:
784
- # through HF's proxy with the account token: the proxy attaches the
785
- # caller identity (X-IP-Token), so @spaces.GPU draws on the account's
786
- # ZeroGPU quota instead of the anonymous per-Space pool
787
- src, kw, how = space, {"hf_token": _tok}, "authenticated (account quota)"
788
- else:
789
- src, kw, how = base, {}, "localhost, unauthenticated (anonymous pool)"
790
- client = None
791
- while True:
792
- try:
793
- if client is None:
794
- # token positionally: gradio_client 1.x names it hf_token,
795
- # 2.x names it token
796
- client = (Client(src, kw["hf_token"], verbose=False) if kw
797
- else Client(src, verbose=False))
798
- log(f"self-drive: tick loop connected β€” {how}, "
799
- f"{SELF_TICK_S:.0f} s cadence")
800
- # bounded wait: a half-open proxy connection must not hang the
801
- # only engine forever (a healthy tick is <= ~3 min)
802
- client.submit(_SELF_KEY, api_name="/tick").result(timeout=900)
803
- W["last_tick"] = time.time()
804
- except Exception as e: # server still booting, etc.
805
- try:
806
- if client is not None:
807
- client.close()
808
- except Exception:
809
- pass
810
- client = None
811
- log(f"self-drive: {type(e).__name__}: {str(e)[:200]} β€” retrying")
812
- time.sleep(15)
813
- continue
814
- time.sleep(SELF_TICK_S)
815
-
816
-
817
- threading.Thread(target=_keep_awake, daemon=True).start()
818
- threading.Thread(target=_self_drive, daemon=True).start()
819
- log(f"boot: {IDENT} publish={PUBLISH} zero={HAS_ZERO} "
820
- f"delay={ZERO_DELAY_S:.0f}s budget={MAX_CELLS / 1e6:.1f}M cells "
821
- f"gpu/day={DAILY_BUDGET_S / 60:.0f} min auth={bool(ZERO_AUTH and _tok)} "
822
- f"crestimap={_FORK_SHA}")
823
-
824
- with gr.Blocks(title="CREST GPU worker (ZeroGPU)") as demo:
825
- gr.Markdown("## CREST event worker β€” ZeroGPU\n"
826
- "The 2-D inundation engine since the HPC worker was released "
827
- "(2026-10-16). Works within a daily GPU budget; shadow mode "
828
- "unless PUBLISH=1.")
829
- st = gr.JSON(label="worker state")
830
- lg = gr.Textbox(label="log", lines=18, max_lines=18)
831
- # the page only DISPLAYS (view); the worker advances through /tick,
832
- # which needs the in-process key the self-drive loop sends
833
- t = gr.Timer(5)
834
- t.tick(view, outputs=[st, lg], api_name=False)
835
- demo.load(view, outputs=[st, lg], api_name=False)
836
- key_in = gr.Textbox(visible=False)
837
- tick_btn = gr.Button(visible=False)
838
- tick_btn.click(tick_api, inputs=[key_in], outputs=[st, lg],
839
- api_name="tick", concurrency_limit=1)
840
-
841
- demo.queue().launch()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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()