src/job_runner.py

"""Pluggable job execution.

`HFJobRunner` submits real HF Jobs; `LocalPodmanRunner` runs the identical command
in a local container against fixtures so the whole flow can be verified without HF
or AWS access. Both expose the same small interface used by the Gradio app.
"""
from __future__ import annotations

import os
import subprocess
import tempfile
import time
import uuid
from abc import ABC, abstractmethod
from typing import List, Optional, Tuple

from . import config
from .jobs import JobSpec, RepackageRequest, build_estimate_spec, build_fetch_spec

# Terminal stages reported by HF Jobs.
TERMINAL_STAGES = {"COMPLETED", "ERROR", "CANCELED", "DELETED"}


class JobRunner(ABC):
    """Build + submit the two jobs and report their status/logs."""

    @abstractmethod
    def estimate(self, req: RepackageRequest, token: Optional[str]) -> str: ...

    @abstractmethod
    def fetch(self, req: RepackageRequest, token: Optional[str], n_records: int = 0) -> str: ...

    @abstractmethod
    def status(self, job_id: str, token: Optional[str]) -> str: ...

    @abstractmethod
    def logs(self, job_id: str, token: Optional[str]) -> List[str]: ...

    @abstractmethod
    def cancel(self, job_id: str, token: Optional[str]) -> None: ...

    @abstractmethod
    def report(self, job_id: str, token: Optional[str]) -> Optional[Tuple[int, str]]:
        """(billed running seconds, flavor) for a finished job, or None if unknown."""

    @abstractmethod
    def job_url(self, job_id: str) -> Optional[str]: ...


# --------------------------------------------------------------------------- HF


class HFJobRunner(JobRunner):
    def __init__(self):
        # job_id -> namespaced job URL (e.g. https://huggingface.co/jobs/<ns>/<id>),
        # captured from JobInfo at submit time since the id alone lacks the namespace.
        self._urls: dict[str, str] = {}

    def _submit(self, spec: JobSpec, token: Optional[str]) -> str:
        from huggingface_hub import Volume, run_job

        volumes = [
            Volume(type="bucket", source=v.source, mount_path=v.mount_path, read_only=v.read_only)
            for v in spec.volumes
        ]
        job = run_job(
            image=spec.image,
            command=spec.command,
            flavor=spec.flavor,
            timeout=spec.timeout,
            volumes=volumes,
            # The job needs the user's token to read/write the buckets at runtime.
            secrets={"HF_TOKEN": token} if token else None,
            token=token,
        )
        if getattr(job, "url", None):
            self._urls[job.id] = job.url
        return job.id

    def estimate(self, req: RepackageRequest, token: Optional[str]) -> str:
        return self._submit(build_estimate_spec(req, index_mode="cdn", skip_pip=False), token)

    def fetch(self, req: RepackageRequest, token: Optional[str], n_records: int = 0) -> str:
        return self._submit(
            build_fetch_spec(
                req,
                warc_download_prefix=config.CC_WARC_CDN_PREFIX,
                hf_reader="cdn",
                processes=config.compute_processes(n_records, req.flavor),
                skip_pip=False,
            ),
            token,
        )

    def status(self, job_id: str, token: Optional[str]) -> str:
        from huggingface_hub import inspect_job

        info = inspect_job(job_id=job_id, token=token)
        return (info.status.stage if info and info.status else "UNKNOWN") or "UNKNOWN"

    def logs(self, job_id: str, token: Optional[str]) -> List[str]:
        from huggingface_hub import fetch_job_logs

        try:
            return list(fetch_job_logs(job_id=job_id, token=token))
        except Exception as e:  # noqa: BLE001
            return [f"(could not fetch logs: {type(e).__name__}: {e})"]

    def cancel(self, job_id: str, token: Optional[str]) -> None:
        from huggingface_hub import cancel_job

        cancel_job(job_id=job_id, token=token)

    def report(self, job_id: str, token: Optional[str]) -> Optional[Tuple[int, str]]:
        from huggingface_hub import inspect_job

        try:
            info = inspect_job(job_id=job_id, token=token)
        except Exception:  # noqa: BLE001
            return None
        dur = getattr(info, "durations", None)
        secs = getattr(dur, "running_secs", None) if dur else None
        flavor = getattr(info, "flavor", None)
        if secs is None or not flavor:
            return None
        return int(secs), flavor

    def job_url(self, job_id: str) -> Optional[str]:
        # Prefer the namespaced URL captured at submit; fall back to inspect_job.
        if job_id in self._urls:
            return self._urls[job_id]
        try:
            from huggingface_hub import inspect_job

            url = getattr(inspect_job(job_id=job_id), "url", None)
            if url:
                self._urls[job_id] = url
                return url
        except Exception:  # noqa: BLE001
            pass
        return None


# ----------------------------------------------------------------------- podman


class _LocalJob:
    def __init__(self, proc: subprocess.Popen, log_path: str, flavor: str):
        self.proc = proc
        self.log_path = log_path
        self.flavor = flavor
        self.started = time.monotonic()
        self.ended: Optional[float] = None


class LocalPodmanRunner(JobRunner):
    """Runs the JobSpec command in `podman run` against local fixture directories.

    Bucket mounts are resolved to local dirs via config.PODMAN_CC_DIR / PODMAN_OUT_DIR
    keyed by mount_path. Jobs run in the background; status/logs poll the process.
    """

    def __init__(self):
        self._jobs: dict[str, _LocalJob] = {}

    def _local_dir_for(self, mount_path: str) -> str:
        mapping = {
            config.CC_MOUNT: config.PODMAN_CC_DIR,
            config.OUT_MOUNT: config.PODMAN_OUT_DIR,
        }
        local = mapping.get(mount_path, "")
        if not local:
            raise RuntimeError(
                f"No local directory configured for mount {mount_path} "
                "(set CC_PODMAN_CC_DIR / CC_PODMAN_OUT_DIR)."
            )
        os.makedirs(local, exist_ok=True)
        return os.path.abspath(local)

    def _submit(self, spec: JobSpec, token: Optional[str]) -> str:
        argv = ["podman", "run", "--rm"]
        for v in spec.volumes:
            local = self._local_dir_for(v.mount_path)
            suffix = ":ro" if v.read_only else ""
            argv += ["-v", f"{local}:{v.mount_path}{suffix}"]
        argv += [config.PODMAN_IMAGE, *spec.command]

        job_id = f"local-{spec.kind}-{uuid.uuid4().hex[:8]}"
        log_path = os.path.join(tempfile.gettempdir(), f"{job_id}.log")
        log_fh = open(log_path, "w")
        log_fh.write("$ " + " ".join(argv) + "\n\n")
        log_fh.flush()
        proc = subprocess.Popen(argv, stdout=log_fh, stderr=subprocess.STDOUT)
        self._jobs[job_id] = _LocalJob(proc, log_path, spec.flavor)
        return job_id

    def estimate(self, req: RepackageRequest, token: Optional[str]) -> str:
        return self._submit(
            build_estimate_spec(req, index_mode="local", skip_pip=config.PODMAN_SKIP_PIP), token
        )

    def fetch(self, req: RepackageRequest, token: Optional[str], n_records: int = 0) -> str:
        return self._submit(
            build_fetch_spec(
                req,
                warc_download_prefix=config.PODMAN_WARC_HTTP,
                hf_reader=None,
                processes=config.compute_processes(n_records, req.flavor),
                skip_pip=config.PODMAN_SKIP_PIP,
            ),
            token,
        )

    def status(self, job_id: str, token: Optional[str]) -> str:
        job = self._jobs.get(job_id)
        if job is None:
            return "UNKNOWN"
        rc = job.proc.poll()
        if rc is None:
            return "RUNNING"
        if job.ended is None:
            job.ended = time.monotonic()
        return "COMPLETED" if rc == 0 else "ERROR"

    def logs(self, job_id: str, token: Optional[str]) -> List[str]:
        job = self._jobs.get(job_id)
        if job is None:
            return []
        try:
            with open(job.log_path) as f:
                return f.read().splitlines()
        except FileNotFoundError:
            return []

    def cancel(self, job_id: str, token: Optional[str]) -> None:
        job = self._jobs.get(job_id)
        if job is not None and job.proc.poll() is None:
            job.proc.terminate()

    def report(self, job_id: str, token: Optional[str]) -> Optional[Tuple[int, str]]:
        job = self._jobs.get(job_id)
        if job is None:
            return None
        end = job.ended if job.ended is not None else time.monotonic()
        return int(end - job.started), job.flavor

    def job_url(self, job_id: str) -> Optional[str]:
        return None


def make_runner(executor: Optional[str] = None) -> JobRunner:
    executor = executor or config.EXECUTOR
    if executor == "podman":
        return LocalPodmanRunner()
    return HFJobRunner()