gpu/single-gpu/home_benjamin_hdbasin/code/scripts/a100_worker.py

#!/usr/bin/env python3
"""A100 worker for the HD-Basin public demo queue.

The worker uses Supabase REST/RPC with a service-role key. Keep it on the
private GPU box; do not expose it to the public internet.
"""

from __future__ import annotations

import argparse
import json
import os
import socket
import subprocess
import sys
import tempfile
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Any


def env(name: str, required: bool = True, default: str | None = None) -> str:
    value = os.environ.get(name, default)
    if required and not value:
        raise SystemExit(f"Missing required environment variable: {name}")
    return value or ""


class SupabaseClient:
    def __init__(self) -> None:
        self.url = env("SUPABASE_URL").rstrip("/")
        self.key = env("SUPABASE_SERVICE_ROLE_KEY")

    def _request(self, method: str, path: str, body: Any | None = None, headers: dict[str, str] | None = None) -> Any:
        data = None if body is None else json.dumps(body).encode("utf-8")
        request = urllib.request.Request(
            f"{self.url}{path}",
            method=method,
            data=data,
            headers={
                "apikey": self.key,
                "authorization": f"Bearer {self.key}",
                "content-type": "application/json",
                **(headers or {}),
            },
        )
        try:
            with urllib.request.urlopen(request, timeout=60) as response:
                payload = response.read().decode("utf-8")
                return json.loads(payload) if payload else None
        except urllib.error.HTTPError as exc:
            detail = exc.read().decode("utf-8", errors="replace")
            raise RuntimeError(f"Supabase {method} {path} failed: {exc.code} {detail}") from exc

    def rpc(self, name: str, body: dict[str, Any]) -> Any:
        return self._request("POST", f"/rest/v1/rpc/{name}", body)

    def claim_next_job(self, worker_id: str) -> dict[str, Any] | None:
        job = self.rpc("claim_next_hdbasin_job", {"p_worker_id": worker_id})
        return job if job and job.get("id") else None

    def heartbeat(self, job_id: str, worker_id: str) -> None:
        self.rpc("heartbeat_hdbasin_job", {"p_job_id": job_id, "p_worker_id": worker_id})

    def complete(
        self,
        job_id: str,
        worker_id: str,
        report_uri: str,
        result_uri: str,
        best_config_uri: str,
        best_validation_loss: float | None,
    ) -> None:
        self.rpc(
            "complete_hdbasin_job",
            {
                "p_job_id": job_id,
                "p_worker_id": worker_id,
                "p_report_uri": report_uri,
                "p_result_uri": result_uri,
                "p_best_config_uri": best_config_uri,
                "p_best_validation_loss": best_validation_loss,
            },
        )

    def fail(self, job_id: str, worker_id: str, reason: str) -> None:
        self.rpc("fail_hdbasin_job", {"p_job_id": job_id, "p_worker_id": worker_id, "p_failure_reason": reason[:1000]})

    def upload_result(self, local_path: Path, object_path: str, content_type: str) -> str:
        data = local_path.read_bytes()
        encoded_path = "/".join(urllib.parse.quote(part) for part in object_path.split("/"))
        request = urllib.request.Request(
            f"{self.url}/storage/v1/object/hdbasin-results/{encoded_path}",
            method="PUT",
            data=data,
            headers={
                "apikey": self.key,
                "authorization": f"Bearer {self.key}",
                "content-type": content_type,
                "x-upsert": "true",
            },
        )
        try:
            with urllib.request.urlopen(request, timeout=120) as response:
                response.read()
        except urllib.error.HTTPError as exc:
            detail = exc.read().decode("utf-8", errors="replace")
            raise RuntimeError(f"Upload failed: {exc.code} {detail}") from exc
        return f"supabase://hdbasin-results/{object_path}"


def write_smoke_report(job: dict[str, Any], outdir: Path) -> tuple[Path, Path, Path, float | None]:
    outdir.mkdir(parents=True, exist_ok=True)
    report = outdir / "report.md"
    result = outdir / "result.json"
    best_config = outdir / "best_config.json"
    best_validation_loss = 0.0

    report.write_text(
        "\n".join(
            [
                "# HD-Basin Demo Report",
                "",
                "## Result",
                "",
                "Demo pipeline passed.",
                "",
                "Your dataset sample was accepted, saved as a queued job, picked up by the A100 worker,",
                "and returned through the website result system.",
                "",
                "## What this proves",
                "",
                "- The website upload and validation flow works.",
                "- The Supabase job queue works.",
                "- The A100 worker can claim and finish jobs.",
                "- Result storage and website display work.",
                "",
                "## What this does not prove yet",
                "",
                "This free demo did not train a real classifier or compare validation loss against a baseline yet.",
                "It is a system test showing that the full product pipeline is connected.",
                "",
                "## What the real report will include",
                "",
                "- Best hyperparameter setting found.",
                "- Validation loss and whether HD-Basin improved it.",
                "- Baseline comparison at the same trial budget.",
                "- GPU time used.",
                "- Downloadable result files.",
                "",
                "## Job details",
                "",
                f"Job: `{job['id']}`",
                f"File: `{job.get('file_name', 'unknown')}`",
                f"Workload: `{job.get('workload', 'unknown')}`",
                f"Trial budget: `{job.get('trial_budget', 'unknown')}`",
                "",
            ]
        ),
        encoding="utf-8",
    )
    result.write_text(
        json.dumps(
            {
                "job": job["id"],
                "status": "demo_pipeline_passed",
                "summary": "Upload, queue, A100 worker, result storage, and website display all passed.",
                "real_training_completed": False,
                "next_step": "Connect HDBASIN_RUN_COMMAND to the real HD-Basin training/report runner.",
            },
            indent=2,
        ),
        encoding="utf-8",
    )
    best_config.write_text(json.dumps({"mode": "smoke", "optimizer": "hdbasinflow"}, indent=2), encoding="utf-8")
    return report, result, best_config, best_validation_loss


def run_configured_command(job: dict[str, Any], outdir: Path) -> tuple[Path, Path, Path, float | None]:
    command_template = os.environ.get("HDBASIN_RUN_COMMAND")
    if not command_template:
        return write_smoke_report(job, outdir)

    outdir.mkdir(parents=True, exist_ok=True)
    command = command_template.format(
        job_id=job["id"],
        outdir=str(outdir),
        text_column=job.get("text_column", "text"),
        label_column=job.get("label_column", "label"),
        trial_budget=job.get("trial_budget", 12),
        dataset_uri=job.get("dataset_uri") or "",
        workload=job.get("workload", "unknown"),
        file_name=job.get("file_name", "unknown"),
    )
    subprocess.run(command, shell=True, check=True, cwd=Path(__file__).resolve().parents[1])

    report = outdir / os.environ.get("HDBASIN_REPORT_PATH", "report.md")
    result = outdir / os.environ.get("HDBASIN_RESULT_PATH", "result.json")
    best_config = outdir / os.environ.get("HDBASIN_BEST_CONFIG_PATH", "best_config.json")
    for path in (report, result, best_config):
        if not path.exists():
            raise RuntimeError(f"Expected worker output was not created: {path}")
    best_validation_loss = None
    try:
        result_data = json.loads(result.read_text(encoding="utf-8"))
        best_validation_loss = result_data.get("hdbasin_best_validation_loss")
    except (json.JSONDecodeError, OSError):
        best_validation_loss = None
    return report, result, best_config, best_validation_loss


def process_one_job(client: SupabaseClient, worker_id: str, workdir: Path) -> bool:
    job = client.claim_next_job(worker_id)
    if not job:
        return False

    job_id = job["id"]
    outdir = workdir / job_id
    print(f"claimed job {job_id}", flush=True)
    try:
        client.heartbeat(job_id, worker_id)
        report, result, best_config, best_validation_loss = run_configured_command(job, outdir)
        prefix = f"{job_id}"
        report_uri = client.upload_result(report, f"{prefix}/{report.name}", "text/markdown")
        result_uri = client.upload_result(result, f"{prefix}/{result.name}", "application/json")
        config_uri = client.upload_result(best_config, f"{prefix}/{best_config.name}", "application/json")
        client.complete(job_id, worker_id, report_uri, result_uri, config_uri, best_validation_loss)
        print(f"completed job {job_id}", flush=True)
    except Exception as exc:  # noqa: BLE001 - this is the worker failure boundary.
        client.fail(job_id, worker_id, str(exc))
        print(f"failed job {job_id}: {exc}", file=sys.stderr, flush=True)
    return True


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--once", action="store_true", help="Process at most one queued job.")
    parser.add_argument("--poll-seconds", type=int, default=15)
    parser.add_argument("--worker-id", default=f"{socket.gethostname()}-{os.getpid()}")
    parser.add_argument("--workdir", default=os.environ.get("HDBASIN_WORKER_DIR", "runs/public_demo_worker"))
    args = parser.parse_args()

    client = SupabaseClient()
    workdir = Path(args.workdir)
    workdir.mkdir(parents=True, exist_ok=True)

    while True:
        processed = process_one_job(client, args.worker_id, workdir)
        if args.once:
            return 0 if processed else 2
        if not processed:
            time.sleep(args.poll_seconds)


if __name__ == "__main__":
    raise SystemExit(main())