Spaces:
Running on Zero
Running on Zero
| # SPDX-License-Identifier: Apache-2.0 | |
| """Immutable, typed stage checkpoints, displayed as native Markdown ports. | |
| The Markdown is a view of a server-owned manifest, not editable configuration. | |
| Only its opaque /artifact/<id> capability is consumed. Browser-supplied paths, | |
| metadata, coordinate claims and HTML are never trusted as artifact contents. | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| import math | |
| import os | |
| import re | |
| import threading | |
| import uuid | |
| from pathlib import Path | |
| from workflow_diagnostics import StageTrace | |
| from workflow_execution import cancellation_check, ensure_gpu_idle, ExecutionBusy | |
| ROOT = Path(os.environ.get("APP_DATA_ROOT", "/tmp/viggle-mine-workflow")).resolve() | |
| PUBLIC = ROOT / "outputs" | |
| STATE = ROOT / "state" | |
| for directory in (PUBLIC, STATE): | |
| directory.mkdir(parents=True, exist_ok=True) | |
| # Also serializes V22's module globals, its progress bridge and CPU model state. | |
| # One process/worker per Space, matching the canonical launch model. | |
| EXECUTION_LOCK = threading.RLock() | |
| POINTS_LOCK = threading.RLock() | |
| ID = re.compile(r"^[0-9a-f]{32}$") | |
| REF = re.compile(r"\[Inspect checkpoint\]\(/artifact/([0-9a-f]{32})\)") | |
| SPACES = { | |
| "prepared": "normalized_source", | |
| "tracked": "normalized_source", | |
| "guarded": "normalized_source", | |
| "camera": "stabilized", | |
| "pass_a": "stabilized_generated", | |
| "recovery": "stabilized", | |
| "pass_b": "normalized_source", | |
| "standard": "normalized_source", | |
| } | |
| def atomic_json(path, value): | |
| path = Path(path) | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| tmp = path.with_name(path.name + "." + uuid.uuid4().hex + ".tmp") | |
| tmp.write_text(json.dumps(value, indent=2, allow_nan=False), encoding="utf-8") | |
| tmp.replace(path) | |
| def digest(value): | |
| return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(",", ":"), | |
| allow_nan=False).encode()).hexdigest() | |
| def file_digest(path): | |
| h = hashlib.sha256() | |
| with open(path, "rb") as f: | |
| for block in iter(lambda: f.read(1024 * 1024), b""): | |
| h.update(block) | |
| return h.hexdigest() | |
| def checkpoint_id(value): | |
| if not isinstance(value, str): | |
| raise ValueError("Use the connected checkpoint produced by the previous stage.") | |
| if ID.fullmatch(value): | |
| return value | |
| matches = REF.findall(value) | |
| if len(matches) != 1: | |
| raise ValueError("Checkpoint missing. Run the previous stage first.") | |
| return matches[0] | |
| def load(value, expected=None): | |
| ident = checkpoint_id(value) | |
| path = STATE / "artifacts" / f"{ident}.json" | |
| if not path.is_file(): | |
| raise ValueError("This checkpoint has expired after a Space restart. Prepare the source again.") | |
| artifact = json.loads(path.read_text()) | |
| if artifact.get("schema") != "viggle.workflow.checkpoint.v1" or artifact.get("id") != ident: | |
| raise ValueError("Invalid checkpoint manifest.") | |
| kind = artifact.get("kind") | |
| if kind not in SPACES or artifact.get("coordinate_space") != SPACES[kind]: | |
| raise ValueError("Checkpoint coordinate-space contract is invalid.") | |
| if expected and kind not in ((expected,) if isinstance(expected, str) else expected): | |
| raise ValueError(f"Expected {expected}, received {kind}. Do not mix coordinate spaces.") | |
| for file in artifact.get("files", []): | |
| p = Path(file).resolve() | |
| if not p.is_relative_to(PUBLIC) or not p.is_file(): | |
| raise ValueError("Checkpoint media has expired. Run its producing stage again.") | |
| return artifact | |
| def create(kind, payload, files=(), parents=()): | |
| ident = uuid.uuid4().hex | |
| value = {"schema": "viggle.workflow.checkpoint.v1", "id": ident, | |
| "kind": kind, "coordinate_space": SPACES[kind], | |
| "parents": list(parents), "files": list(map(str, files)), **payload} | |
| atomic_json(STATE / "artifacts" / f"{ident}.json", value) | |
| return load(ident, kind) | |
| def card(artifact, body): | |
| return f"{body}\n\n[Inspect checkpoint](/artifact/{artifact['id']})" | |
| def cached(stage, inputs, compute, *, slot=None): | |
| """Replay completed work even after canvas refresh/force-run/API traversal. | |
| A generation take is immutable. Reusing its slot with changed inputs fails | |
| instead of silently spending another GPU allocation during Recovery. | |
| A failed GPU attempt is NOT cached; retrying the same take is permitted. | |
| """ | |
| fingerprint = digest([stage, inputs]) | |
| key = digest([stage, slot]) if slot is not None else fingerprint | |
| path = STATE / "cache" / f"{key}.json" | |
| check_cancelled = cancellation_check() | |
| check_cancelled() | |
| # Do not queue indefinitely behind a cancelled or stalled SDK call. V22 | |
| # shares mutable CPU model/progress state across otherwise separate nodes. | |
| if not EXECUTION_LOCK.acquire(blocking=False): | |
| raise ExecutionBusy("Another V22 stage is running. Wait for it to finish before running this stage.") | |
| try: | |
| with StageTrace(stage) as trace: | |
| if path.is_file(): | |
| entry = json.loads(path.read_text()) | |
| if entry["fingerprint"] != fingerprint: | |
| raise ValueError("This take already has different inputs. Increase Take to explicitly generate a new result; Recovery never regenerates an existing Pass A take.") | |
| for ident in entry["checkpoints"]: | |
| load(ident) | |
| trace.set_phase("cache_hit") | |
| return tuple(entry["result"]) | |
| ensure_gpu_idle() | |
| check_cancelled() | |
| trace.set_phase("calling_backend") | |
| result, checkpoints = compute() | |
| check_cancelled() | |
| trace.set_phase("saving_checkpoint") | |
| atomic_json(path, {"fingerprint": fingerprint, "result": result, | |
| "checkpoints": checkpoints}) | |
| return result | |
| finally: | |
| EXECUTION_LOCK.release() | |
| def number(value, name, low, high, integer=False): | |
| x = float(value) | |
| if not math.isfinite(x) or not low <= x <= high or (integer and int(x) != x): | |
| raise ValueError(f"{name} must be {'an integer ' if integer else ''}between {low} and {high}.") | |
| return int(x) if integer else x | |
| def boolean(value, name): | |
| if not isinstance(value, bool): | |
| raise ValueError(f"{name} requires a checkbox boolean, not text.") | |
| return value | |
| def validate_points(value, prepared): | |
| meta = prepared["meta"] | |
| if not isinstance(value, dict) or value.get("coordinate_space") != "normalized_source": | |
| raise ValueError("Picker points must use normalized_source pixel coordinates.") | |
| if value.get("preparation_id") != prepared["id"]: | |
| raise ValueError("Points belong to another prepared source.") | |
| if (value.get("w"), value.get("h")) != (meta["w"], meta["h"]): | |
| raise ValueError("Point dimensions do not match this prepared source.") | |
| anchors = value.get("anchors") | |
| # The data contract is a list. Only this V22 adapter enforces three anchors. | |
| if not isinstance(anchors, list) or len(anchors) != len(meta["frames"]) or len(anchors) != 3: | |
| raise ValueError("V22 currently requires three prepared anchors.") | |
| clean = [] | |
| for i, (a, frame, time) in enumerate(zip(anchors, meta["frames"], meta["times"])): | |
| if not isinstance(a, dict): | |
| raise ValueError(f"Anchor {i + 1} must be an object with frame, time and points.") | |
| actual_time = float(a.get("time", -1)) | |
| if a.get("frame") != frame or not math.isfinite(actual_time) or abs(actual_time - time) > 1e-6: | |
| raise ValueError(f"Anchor {i + 1} timestamp/frame does not match preparation.") | |
| points = a.get("points") | |
| if not isinstance(points, list) or not 3 <= len(points) <= 12: | |
| raise ValueError(f"Anchor {i + 1}: select 3–12 visible points on the same subject.") | |
| pts = [] | |
| for p in points: | |
| if not isinstance(p, (list, tuple)) or len(p) != 2: | |
| raise ValueError("Each point must be [x, y].") | |
| x, y = map(float, p) | |
| if not (math.isfinite(x) and math.isfinite(y) and 0 <= x < meta["w"] and 0 <= y < meta["h"]): | |
| raise ValueError("Point outside the normalized source image.") | |
| pts.append([x, y]) | |
| clean.append({"frame": frame, "time": time, "points": pts}) | |
| return {"coordinate_space": "normalized_source", "preparation_id": prepared["id"], | |
| "w": meta["w"], "h": meta["h"], "anchors": clean} | |
| def save_points(preparation_id, value): | |
| prepared = load(preparation_id, "prepared") | |
| clean = validate_points(value, prepared) | |
| clean["revision"] = digest(clean) | |
| with POINTS_LOCK: | |
| atomic_json(STATE / "points" / f"{prepared['id']}.json", clean) | |
| return clean | |
| def points(prepared): | |
| path = STATE / "points" / f"{prepared['id']}.json" | |
| with POINTS_LOCK: | |
| if not path.is_file(): | |
| raise ValueError("Open Point Picker in the Source checkpoint, select 3–12 points on each anchor, and Save points. Then run Track Subject.") | |
| value = json.loads(path.read_text()) | |
| clean = validate_points(value, prepared) | |
| if value.get("revision") != digest(clean): | |
| raise ValueError("Saved point revision is invalid.") | |
| return value | |
| def assert_current_points(artifact): | |
| current = points(load(artifact["preparation_id"], "prepared")) | |
| if current["revision"] != artifact["points_revision"]: | |
| raise ValueError("Points changed in the picker. Run Track Subject again before building a new camera or Pass A.") | |