miles/docker/build.py

#!/usr/bin/env python3
# doc-dev: docs/developer/ci/02-docker-build.md
"""Build and push Miles Docker images.

Usage:
    python docker/build.py --variant cu13 --image-tag dev --push          # multi-arch (amd64+arm64)
    python docker/build.py --variant cu13-x86 --image-tag dev --push      # single arch
    python docker/build.py --variant cu12-x86 --image-tag latest
    python docker/build.py --variant cu13 --image-tag dev --dry-run
"""

import os
import subprocess
from datetime import datetime, timezone
from enum import Enum
from pathlib import Path

import image_inputs
import typer

CACHE_DIR = "/tmp/miles-docker-cache"
REPO_ROOT = Path(__file__).resolve().parent.parent

VARIANTS = {
    "cu13": {
        "image": "radixark/miles",
        "platforms": ["linux/amd64", "linux/arm64"],
        "tag_postfix": "",
        "build_args": {},
    },
    "cu13-x86": {
        "image": "radixark/miles",
        "platforms": ["linux/amd64"],
        "tag_postfix": "",
        "build_args": {},
    },
    "cu13-aarch64": {
        "image": "radixark/miles",
        "platforms": ["linux/arm64"],
        "tag_postfix": "",
        "build_args": {},
    },
    "cu12-x86": {
        "image": "radixark/miles",
        "platforms": ["linux/amd64"],
        "tag_postfix": "-cu12",
        "build_args": {
            "ENABLE_CUDA_13": "0",
            "SGLANG_IMAGE_TAG": "v0.5.19-cu129",
            # Frozen with the base image: sglang publishes no cu12 image from v0.5.20 on,
            # so this variant cannot follow the branch the cu13 images track.
            "SGLANG_BRANCH": "sglang-miles-v0.5.19-final",
            "WHEELS_TAG_X86": "cu129-x86_64",
        },
    },
    "rocm724-mi35x": {
        "image": "rocm/sgl-dev",
        "tag_postfix": "-rocm724-mi35x",
        "tag_prefix": "miles",
        "dockerfile": "docker/Dockerfile.rocm",
        "build_args": {
            "GPU_ARCH": "gfx950",
            "SGLANG_IMAGE_REPO": "lmsysorg/sglang",
            "SGLANG_IMAGE_TAG": "v0.5.20-rocm724-mi35x",
            "WHEELS_TAG_ROCM": "rocm724-gfx950-v0.5.20",
            "APPLY_ROCR_VMMFIX": "1",
            "TE_USE_WHEEL": "1",
        },
    },
    "rocm10-mi35x": {
        "image": "rocm/sgl-dev",
        "tag_postfix": "-rocm10-mi35x",
        "tag_prefix": "miles",
        "dockerfile": "docker/Dockerfile.rocm",
        "build_args": {
            "GPU_ARCH": "gfx950",
            "SGLANG_IMAGE_REPO": "lmsysorg/sglang",
            "SGLANG_IMAGE_TAG": "v0.5.20-rocm10-mi35x",
            "WHEELS_TAG_ROCM": "rocm10-gfx950-v0.5.18",
            "APEX_USE_PREBUILT": "1",
            "TE_USE_WHEEL": "1",
        },
    },
}


def run(cmd: list[str], dry_run: bool) -> None:
    print(f"+ {' '.join(cmd)}", flush=True)
    if dry_run:
        return
    subprocess.run(cmd, check=True)


def build_and_push(
    variant: str,
    image_tag: str,
    dry_run: bool,
    dockerfile: str,
    push: bool = False,
    custom_tag: str = "",
    extra_build_args: list[str] | None = None,
    context: Path = REPO_ROOT,
    output: str = "",
) -> None:
    extra_build_args = extra_build_args or []
    config = VARIANTS[variant]
    # A variant may pin its own Dockerfile (e.g. ROCm); otherwise use the CLI default.
    dockerfile = config.get("dockerfile", dockerfile)
    image = config["image"]
    postfix = config.get("tag_postfix", "")
    platforms = config.get("platforms")

    if image_tag == "latest":
        tags = [f"{image}:latest{postfix}"]
    elif image_tag == "dev":
        prefix = config.get("tag_prefix", "dev")
        timestamp = datetime.now(timezone.utc).strftime("%Y%m%d%H%M")
        tags = [f"{image}:{prefix}{postfix}", f"{image}:{prefix}{postfix}-{timestamp}"]
    elif image_tag == "custom":
        if not custom_tag:
            raise typer.BadParameter("--custom-tag is required when --image-tag is custom")
        tags = [f"{image}:{custom_tag}{postfix}"]
    else:
        raise typer.BadParameter(f"Unknown image tag: {image_tag}")

    cmd = [
        "docker",
        "buildx",
        "build",
        "-f",
        str(context / dockerfile),
    ]

    if platforms:
        cmd += ["--platform", ",".join(platforms)]

    if push:
        cmd += ["--push"]
    if output:
        cmd += ["--output", output]

    # Proxy args (pass through if set in environment, check both cases)
    for arg_name in ["HTTP_PROXY", "HTTPS_PROXY"]:
        value = os.environ.get(arg_name.lower()) or os.environ.get(arg_name)
        if value:
            cmd += ["--build-arg", f"{arg_name}={value}"]

    cmd += ["--build-arg", "NO_PROXY=localhost,127.0.0.1"]

    # Variant-specific build args
    for key, value in config.get("build_args", {}).items():
        cmd += ["--build-arg", f"{key}={value}"]

    # Caller overrides (e.g. release builds pinning SGLANG_COMMIT / MILES_COMMIT
    # from release-lock.json) come last so they win over variant defaults.
    for spec in extra_build_args:
        assert "=" in spec, f"--build-arg expects KEY=VALUE, got {spec!r}"
        cmd += ["--build-arg", spec]

    # CI reads this back off the published tag to skip rebuilds whose inputs are unchanged.
    cmd += ["--label", f"{image_inputs.LABEL_KEY}={image_inputs.compute(root=context)}"]

    for tag in tags:
        cmd += ["-t", tag]

    # Context is repo root
    cmd += [str(context)]

    print(f"\n=== Building {' '.join(tags)} ===", flush=True)
    run(cmd, dry_run)


class Variant(str, Enum):
    cu13 = "cu13"
    cu13_x86 = "cu13-x86"
    cu13_aarch64 = "cu13-aarch64"
    cu12_x86 = "cu12-x86"
    rocm724_mi35x = "rocm724-mi35x"
    rocm10_mi35x = "rocm10-mi35x"


class ImageTag(str, Enum):
    latest = "latest"
    dev = "dev"
    custom = "custom"


def main(
    variant: Variant = typer.Option(..., help="Build variant to use."),  # noqa: B008
    image_tag: ImageTag = typer.Option(..., help="Tag mode: latest, dev, or custom."),  # noqa: B008
    dockerfile: str = typer.Option("docker/Dockerfile", help="Path to the Dockerfile."),  # noqa: B008
    dry_run: bool = typer.Option(False, help="Print commands without executing them."),  # noqa: B008
    push: bool = typer.Option(False, help="Push images to registry after building."),  # noqa: B008
    custom_tag: str = typer.Option("", help="Custom tag name (required when --image-tag is custom)."),  # noqa: B008
    build_arg: list[str] = typer.Option([], help="Extra KEY=VALUE build-arg (repeatable)."),  # noqa: B008
    context: Path = typer.Option(REPO_ROOT, help="Repository build context, separate from this driver."),  # noqa: B008
    output: str = typer.Option("", help="Buildx exporter, e.g. type=oci,dest=/tmp/image,tar=false."),  # noqa: B008
) -> None:
    build_and_push(
        variant.value,
        image_tag.value,
        dry_run,
        dockerfile,
        push=push,
        custom_tag=custom_tag,
        extra_build_args=build_arg,
        context=context.resolve(),
        output=output,
    )


if __name__ == "__main__":
    typer.run(main)