pipeline.py

# 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}


# Explicit dependencies for the child jobs. `uv run --no-project` was the cause of a
# silent 3-second merge failure: --no-project suppresses PEP 723 inline metadata, so torch
# and friends were never installed. Keep in sync with the `# dependencies = [...]` header
# of merge_securecoder.py / quantise_securecoder.py.
DEPS = ["--with", "torch", "--with", "transformers", "--with", "peft",
        "--with", "huggingface_hub", "--with", "accelerate",
        "--with", "safetensors", "--with", "unsloth"]


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"] + DEPS + [
             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"] + DEPS + [
         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)