# coach — build & publish the baked training images.
#
# Each model is one self-contained image with its venv baked (see
# images/<model>/Containerfile). Build, push to the registry, then pin by digest.
#
#   just build-sam3        # build the cu128 image (self-contained)
#   just push-sam3         # push to the registry
#   just digest-sam3       # the repo@sha256:… to pin in your app
#   just fetch-base        # pull the base checkpoint (to bake into the image)
#
# cu128 wheels cover Ampere/Hopper/Blackwell, so one image serves every GPU as
# long as the host driver is >= R570. (To add another backend later, build with
# --build-arg TORCH_BACKEND=cu126 + a matching CUDA_IMAGE.)

set shell := ["bash", "-uc"]

# Registry (coach is podman-only).
registry := env_var_or_default("COACH_REGISTRY", "git.deepvis.ai/deepvis")

# CUDA runtime base (torch wheels bring their own CUDA libs; this base supplies
# cuDNN + the userspace CUDA the toolkit injects the driver into).
cuda_image := "docker.io/nvidia/cuda:12.8.0-cudnn-runtime-ubuntu24.04"

# Optional base-checkpoint bake. Set COACH_SAM3_BASE to a directory containing
# the base checkpoint named `base.pt` to bake it into the image (then runs can
# use `checkpoint: None`). Defaults to an empty dir → nothing baked. The SAM3
# base is license-gated: only bake into images bound for a private registry.
base_ctx := env_var_or_default("COACH_SAM3_BASE", "images/sam3/no-base")

# The base checkpoint's home: the Forgejo generic package registry. `fetch-base`
# downloads it (auth via COACH_FORGEJO_TOKEN) into `.base/` — OUTSIDE the build
# context, so a plain `build-sam3` never tars up the 3.3GB file.
base_url := "https://git.deepvis.ai/api/packages/deepvis/generic/sam3-base/3.1/sam3.1_multiplex.pt"
base_dir := ".base"

default:
    @just --list

# ── sam3 ────────────────────────────────────────────────────────────────────

# Build the cu128 image (self-contained context: sam3-python + train_coco.py).
build-sam3:
    podman build \
      -f images/sam3/Containerfile \
      --build-arg TORCH_BACKEND=cu128 \
      --build-arg CUDA_IMAGE={{cuda_image}} \
      --build-context base={{base_ctx}} \
      -t {{registry}}/coach-sam3:cu128 \
      images/sam3

# Push the built image to the registry.
push-sam3:
    podman push {{registry}}/coach-sam3:cu128

# Print the registry digest to pin a run (run after push).
digest-sam3:
    @podman image inspect --format '{{{{index .RepoDigests 0}}}}' \
      {{registry}}/coach-sam3:cu128

# Fetch the base checkpoint from the Forgejo package into .base/ (needs COACH_FORGEJO_TOKEN).
fetch-base:
    @test -n "${COACH_FORGEJO_TOKEN:-}" || { echo "set COACH_FORGEJO_TOKEN (package read scope)"; exit 1; }
    mkdir -p {{base_dir}}
    curl -sS --fail-with-body -H "Authorization: token ${COACH_FORGEJO_TOKEN}" \
      -o {{base_dir}}/base.pt {{base_url}}
    @echo "fetched -> {{base_dir}}/base.pt  (bake: COACH_SAM3_BASE={{base_dir}} just build-sam3)"
