Handoff: fixed supervisor + full findings/state so work can resume

pipeline.py

@@ -0,0 +1,214 @@
+# SecureCoder v2 pipeline supervisor (stdlib only).
+#
+# Waits for the v2 training job to leave the active set, then chains the merge
+# job and the GGUF quantise job so the pipeline finishes unattended.
+#
+#   HF_TOKEN  - fine-grained token with job.write + repo.write
+#   TRAIN_JOB - id of the training job to wait on (required unless SKIP_MERGE=1)
+#   QUANT     - GGUF quantisation to produce (default Q4_K_M)
+#
+# Changes vs the original supervisor:
+#   * BUG FIX: the SKIP_MERGE=1 path no longer dies with UnboundLocalError when
+#     logging the final summary. Every step's outcome is seeded in RESULT up
+#     front, so the summary can never reference an unbound name.
+#   * TRAIN_JOB is validated before polling instead of hitting a malformed URL.
+#   * api()/repo_ok() retry transient network errors, 5xx and 429 with backoff.
+#   * stage_of() treats a not-yet-visible job as NOT_FOUND instead of raising,
+#     which matters right after launch().
+#   * Failures are reported as clean log lines + exit codes, not tracebacks.
+
+import json
+import os
+import time
+import urllib.error
+import urllib.request
+
+HUB = "https://huggingface.co"
+NS = "Taimwe"
+TOKEN = os.environ.get("HF_TOKEN", "")
+MERGE_REPO = os.environ.get("MERGE_REPO", "Taimwe/securecoder-30b-pro-v2-merged")
+TRAIN_JOB = os.environ.get("TRAIN_JOB", "")
+ADAPTER = os.environ.get("ADAPTER_REPO", "Taimwe/securecoder-30b-pro-v2")
+MERGED = MERGE_REPO
+GGUF = os.environ.get("GGUF_REPO", "Taimwe/securecoder-30b-pro-v2-GGUF")
+IMAGE = "ghcr.io/astral-sh/uv:python3.12-bookworm"
+ACTIVE = {"CREATED", "SCHEDULING", "RUNNING", "QUEUED", "PENDING"}
+OK = {"COMPLETED", "SUCCEEDED"}
+SKIP_MERGE = os.environ.get("SKIP_MERGE", "0") == "1"
+QUANT = os.environ.get("QUANT", "Q4_K_M")
+RETRY_ATTEMPTS = int(os.environ.get("RETRY_ATTEMPTS", "4"))
+
+# Seeded before any work starts - this is what makes the summary crash-proof.
+RESULT = {"train": "NOT_RUN", "merge": "NOT_RUN", "gguf": "NOT_RUN",
+          "adapter": ADAPTER, "merged": MERGED, "gguf_repo": GGUF, "quant": QUANT}
+
+
+def log(msg):
+    print("[supervisor %s] %s" % (time.strftime("%H:%M:%S"), msg), flush=True)
+
+
+def _backoff(attempt):
+    return min(60, 5 * 2 ** (attempt - 1))
+
+
+def api(path, payload=None):
+    """Call the HF API, retrying transient failures with exponential backoff."""
+    data = json.dumps(payload).encode() if payload is not None else None
+    last = "unknown"
+    for attempt in range(1, RETRY_ATTEMPTS + 1):
+        headers = {"Authorization": "Bearer " + TOKEN}
+        if data is not None:
+            headers["Content-Type"] = "application/json"
+        req = urllib.request.Request(
+            HUB + path, data=data, headers=headers,
+            method="POST" if data is not None else "GET")
+        try:
+            with urllib.request.urlopen(req, timeout=120) as resp:
+                return json.loads(resp.read().decode() or "null")
+        except urllib.error.HTTPError as exc:
+            body = exc.read().decode("utf-8", "replace")[:300]
+            last = "HTTP %s %s %s" % (exc.code, exc.reason, body)
+            if exc.code < 500 and exc.code != 429:
+                raise RuntimeError("%s -> %s" % (path, last))
+        except (urllib.error.URLError, TimeoutError, OSError) as exc:
+            last = "%s: %s" % (type(exc).__name__, exc)
+        if attempt < RETRY_ATTEMPTS:
+            wait = _backoff(attempt)
+            log("retry %d/%d in %ds (%s)" % (attempt, RETRY_ATTEMPTS, wait, last))
+            time.sleep(wait)
+    raise RuntimeError("%s failed after %d attempts -> %s" % (path, RETRY_ATTEMPTS, last))
+def stage_of(job_id):
+    """Current stage of a job. A job that is not visible yet is NOT_FOUND."""
+    try:
+        return api("/api/jobs/%s/%s" % (NS, job_id))["status"]["stage"]
+    except RuntimeError as exc:
+        if "HTTP 404" in str(exc):
+            return "NOT_FOUND"
+        raise
+
+
+def repo_ok(repo):
+    """True if the model repo exists. 404 -> False; transient errors -> retried."""
+    for attempt in range(1, RETRY_ATTEMPTS + 1):
+        try:
+            req = urllib.request.Request(
+                HUB + "/api/models/" + repo,
+                headers={"Authorization": "Bearer " + TOKEN})
+            with urllib.request.urlopen(req, timeout=60) as resp:
+                return resp.status == 200
+        except urllib.error.HTTPError as exc:
+            if exc.code == 404:
+                return False
+            if (exc.code >= 500 or exc.code == 429) and attempt < RETRY_ATTEMPTS:
+                time.sleep(_backoff(attempt))
+                continue
+            raise RuntimeError("repo_ok(%s) -> HTTP %s" % (repo, exc.code))
+        except (urllib.error.URLError, TimeoutError, OSError):
+            if attempt == RETRY_ATTEMPTS:
+                raise
+            time.sleep(_backoff(attempt))
+    return False
+
+
+def wait_for(job_id, label, budget_s):
+    """Poll until the job leaves the active set, logging only on change."""
+    deadline = time.time() + budget_s
+    seen = None
+    while True:
+        stage = stage_of(job_id)
+        if stage != seen:
+            log("%s %s stage=%s" % (label, job_id, stage))
+            seen = stage
+        if stage not in ACTIVE:
+            return stage
+        if time.time() > deadline:
+            log("%s %s exceeded its %ds budget" % (label, job_id, budget_s))
+            return "TIMEOUT"
+        time.sleep(60)
+
+
+def launch(command, flavor, timeout_s, name):
+    # The secret must carry the real token value: an empty string is injected
+    # literally and every push/upload then fails with 401.
+    job = api("/api/jobs/" + NS, {
+        "flavor": flavor,
+        "dockerImage": IMAGE,
+        "command": command,
+        "arguments": [],
+        "environment": {},
+        "secrets": {"HF_TOKEN": TOKEN},
+        "timeoutSeconds": timeout_s,
+        "labels": {"name": name},
+    })
+    log("launched %s %s (%s)" % (name, job["id"], flavor))
+    return job["id"]
+def main():
+    if not TOKEN.strip():
+        log("HF_TOKEN is empty - set it to a token with job.write + repo.write")
+        return 2
+
+    if SKIP_MERGE:
+        log("merge already done upstream; going straight to GGUF")
+        if not repo_ok(MERGED):
+            log("merged repo %s missing; stopping" % MERGED)
+            RESULT["merge"] = "MISSING_REPO"
+            log("SUMMARY " + json.dumps(RESULT))
+            return 1
+        log("merged repo %s is live" % MERGED)
+        RESULT["merge"] = "SKIPPED"
+    else:
+        if not TRAIN_JOB.strip():
+            log("TRAIN_JOB is required unless SKIP_MERGE=1")
+            return 2
+        log("watching training job %s" % TRAIN_JOB)
+        final = wait_for(TRAIN_JOB, "train", 6 * 3600)
+        RESULT["train"] = final
+        log("training finished stage=%s" % final)
+        if final not in OK:
+            log("training did not complete cleanly; not chaining")
+            log("SUMMARY " + json.dumps(RESULT))
+            return 1
+        if not repo_ok(ADAPTER):
+            log("adapter %s missing after training; stopping" % ADAPTER)
+            log("SUMMARY " + json.dumps(RESULT))
+            return 1
+        log("adapter %s is live" % ADAPTER)
+
+        merge_job = launch(
+            ["uv", "run", "--no-project",
+             HUB + "/" + NS + "/securecoder-scripts/resolve/main/merge_securecoder.py",
+             "--adapter", ADAPTER, "--output-repo", MERGED],
+            "a100-large", 3 * 3600, "securecoder-v2-merge")
+        merge_final = wait_for(merge_job, "merge", 4 * 3600)
+        RESULT["merge"] = merge_final
+        log("merge finished stage=%s" % merge_final)
+        if merge_final not in OK or not repo_ok(MERGED):
+            log("merge failed or repo missing; skipping GGUF")
+            log("SUMMARY " + json.dumps(RESULT))
+            return 1
+        log("merged repo %s is live" % MERGED)
+
+    gguf_job = launch(
+        ["uv", "run", "--no-project",
+         HUB + "/" + NS + "/securecoder-scripts/resolve/main/quantise_securecoder.py",
+         "--merged-repo", MERGED, "--gguf-repo", GGUF, "--quant", QUANT],
+        "cpu-performance", 6 * 3600, "securecoder-v2-gguf")
+    gguf_final = wait_for(gguf_job, "gguf", 7 * 3600)
+    RESULT["gguf"] = gguf_final
+
+    if repo_ok(GGUF):
+        log("PIPELINE COMPLETE gguf=%s merged=%s quant=%s" % (GGUF, MERGED, QUANT))
+    else:
+        log("PIPELINE ENDED but %s not found" % GGUF)
+
+    log("SUMMARY " + json.dumps(RESULT))
+    return 0
+
+
+if __name__ == "__main__":
+    try:
+        raise SystemExit(main())
+    except RuntimeError as exc:
+        log("FATAL %s" % exc)
+        log("SUMMARY " + json.dumps(RESULT))
+        raise SystemExit(1)