Add a Docker image for Puzzletron v2 workers

.dockerignore

@@ -1,6 +1,7 @@
 docker
 examples/**/.git
 examples/**/saved_models*
+examples/puzzletron/reports
 **/experimental
 
 ##### Copied from .gitignore #####

.github/workflows/unit_tests.yml

@@ -7,6 +7,7 @@ on:
     branches: [main, release/*, feature/*]
     paths:
       - ".github/workflows/unit_tests.yml"
+      - "examples/__init__.py"
       - "examples/puzzletron/**"
       - "modelopt/**"
       - "noxfile.py"
@@ -83,6 +84,8 @@ jobs:
         with:
           files: |
             .github/workflows/unit_tests.yml
+            examples/__init__.py
+            examples/puzzletron/Dockerfile
             examples/puzzletron/**/*.html
             examples/puzzletron/**/*.json
             examples/puzzletron/**/*.md

examples/puzzletron/Dockerfile

@@ -0,0 +1,151 @@
+FROM nvidia/cuda:12.9.2-cudnn-devel-ubuntu24.04@sha256:b4db213759eb86d55a7271909bdad891fab300ec5700fb4f4656463b2f51980f
+
+SHELL ["/bin/bash", "-o", "pipefail", "-c"]
+
+ARG DEBIAN_FRONTEND=noninteractive
+ARG TARGETPLATFORM
+
+ENV VIRTUAL_ENV=/venv
+ENV PATH=/venv/bin/:$PATH
+ENV PIP_NO_CACHE_DIR=1
+ENV NLTK_DATA=/opt/puzzletron/nltk_data
+ENV PUZZLETRON_CI_ENVIRONMENT=/opt/puzzletron/ci_environment.json
+ENV PUZZLETRON_REQUIREMENTS=/opt/puzzletron/requirements.txt
+ENV PUZZLETRON_ROOT=/opt/puzzletron/src/modelopt
+ENV PUZZLETRON_VENV=/venv
+ENV PUZZLETRON_VLLM_ANYMODEL=1
+ENV PYTHONPATH=/opt/puzzletron/src/modelopt
+ENV PYTHONUNBUFFERED=1
+
+COPY examples/puzzletron/ci_environment.json /opt/puzzletron/ci_environment.json
+COPY examples/puzzletron/requirements.txt /opt/puzzletron/requirements.txt
+COPY examples/__init__.py /opt/puzzletron/src/modelopt/examples/__init__.py
+COPY examples/puzzletron/patches /opt/puzzletron/patches
+COPY pyproject.toml LICENSE_HEADER /opt/modelopt-dependencies/
+
+# Install base build tools before compiling CUDA extensions.
+RUN test "${TARGETPLATFORM}" = "linux/amd64" || \
+        (echo "Puzzletron requires linux/amd64 because its CUDA stack and video decoder are not validated for Linux ARM" >&2; exit 1) && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends \
+        build-essential ca-certificates cmake curl git jq ninja-build \
+        python3 python3-dev python3-pip python3-venv unzip && \
+    rm -rf /var/lib/apt/lists/* && \
+    python3 -m venv "${VIRTUAL_ENV}" && \
+    python -m pip install --upgrade \
+        pip "setuptools>=80,<81" "setuptools-scm>=8,<10" setuptools-rust wheel \
+        "packaging>=24.2" "cmake>=3.26.1" ninja jinja2
+
+# `pin` keeps ci_environment.json as the single source of dependency versions
+# while leaving the install commands readable.
+RUN pin() { jq -er ".${1}" "${PUZZLETRON_CI_ENVIRONMENT}"; } && \
+    python -m pip install \
+        "torch==$(pin torch)" \
+        "torchvision==$(pin torchvision)" \
+        "torchaudio==$(pin torch)" \
+        --index-url https://download.pytorch.org/whl/cu129
+
+# Evaluation and worker dependencies. Install Transformers last so the pinned
+# worker version wins over transitive requirements.
+RUN pin() { jq -er ".${1}" "${PUZZLETRON_CI_ENVIRONMENT}"; } && \
+    python -m pip install \
+        -r "${PUZZLETRON_REQUIREMENTS}" \
+        "nemo-automodel @ git+$(pin nemo_automodel.repository)@$(pin nemo_automodel.commit)" \
+        "aiperf==$(pin gpu_image.aiperf)" \
+        "langdetect==$(pin gpu_image.langdetect)" \
+        "nltk==$(pin gpu_image.nltk)" \
+        "nox==$(pin gpu_image.nox)" && \
+    mkdir -p "${NLTK_DATA}/tokenizers" && \
+    while read -r nltk_resource nltk_resource_sha256; do \
+        nltk_archive="/tmp/${nltk_resource}.zip" && \
+        curl --fail --location --silent --show-error \
+            --output "${nltk_archive}" \
+            "https://raw.githubusercontent.com/nltk/nltk_data/$(pin gpu_image.nltk_data_commit)/packages/tokenizers/${nltk_resource}.zip" && \
+        echo "${nltk_resource_sha256}  ${nltk_archive}" | sha256sum --check --strict && \
+        unzip -q "${nltk_archive}" -d "${NLTK_DATA}/tokenizers" && \
+        rm "${nltk_archive}"; \
+    done < <(jq -r \
+        '.gpu_image.nltk_resources[] as $name | "\($name) \(.gpu_image.nltk_resource_sha256[$name])"' \
+        "${PUZZLETRON_CI_ENVIRONMENT}") && \
+    python -m pip install "transformers==$(pin transformers)"
+
+# Compile the pinned AnyModel vLLM fork for every supported worker GPU.
+RUN pin() { jq -er ".${1}" "${PUZZLETRON_CI_ENVIRONMENT}"; } && \
+    export FORCE_CUDA=1 && \
+    export TORCH_CUDA_ARCH_LIST="$(pin runtime_image.torch_cuda_arch_list)" && \
+    python -m pip install --no-build-isolation \
+        "vllm @ git+$(pin vllm.repository)@$(pin vllm.commit)"
+
+# Compile the remaining CUDA extensions. Mamba is built from its pinned source
+# because the compatibility patch is not present in the upstream release.
+RUN pin() { jq -er ".${1}" "${PUZZLETRON_CI_ENVIRONMENT}"; } && \
+    mamba_patch="$(pin runtime_image.mamba_ssm.compatibility_patch)" && \
+    export FORCE_CUDA=1 && \
+    export TORCH_CUDA_ARCH_LIST="$(pin runtime_image.torch_cuda_arch_list)" && \
+    python -m pip install --no-build-isolation \
+        "causal-conv1d==$(pin runtime_image.causal_conv1d)" && \
+    echo "$(pin runtime_image.mamba_ssm.compatibility_patch_sha256)  /opt/puzzletron/patches/${mamba_patch}" | \
+        sha256sum --check --strict && \
+    git clone --filter=blob:none --no-checkout \
+        "$(pin runtime_image.mamba_ssm.repository)" /tmp/mamba-ssm && \
+    git -C /tmp/mamba-ssm checkout --detach "$(pin runtime_image.mamba_ssm.commit)" && \
+    test "$(git -C /tmp/mamba-ssm rev-parse HEAD)" = \
+        "$(pin runtime_image.mamba_ssm.commit)" && \
+    git -C /tmp/mamba-ssm apply "/opt/puzzletron/patches/${mamba_patch}" && \
+    MAMBA_FORCE_BUILD=TRUE python -m pip install --no-build-isolation \
+        /tmp/mamba-ssm && \
+    rm -rf /tmp/mamba-ssm && \
+    python -m pip install \
+        "flash-linear-attention[cuda]==$(pin runtime_image.flash_linear_attention)" && \
+    export TORCH_CUDA_ARCH_LIST="$(pin runtime_image.grouped_gemm_cuda_arch_list)" && \
+    python -m pip install --no-build-isolation \
+        "$(pin runtime_image.grouped_gemm.distribution) @ git+$(pin runtime_image.grouped_gemm.repository)@$(pin runtime_image.grouped_gemm.commit)"
+
+# Cache ModelOpt's ordinary Python dependencies separately from its source.
+RUN mkdir -p /opt/modelopt-dependencies/modelopt && \
+    touch /opt/modelopt-dependencies/modelopt/__init__.py && \
+    python -m pip install "/opt/modelopt-dependencies[hf,puzzletron,dev-test]" && \
+    python -m pip uninstall -y nvidia-modelopt
+
+# Keep the immutable source revision below dependency compilation so source-only
+# rebuilds reuse the pinned CUDA dependency layers.
+ARG MODELOPT_REVISION
+RUN [[ "${MODELOPT_REVISION}" =~ ^[0-9a-f]{40}$ ]] && \
+    printf '%s\n' "${MODELOPT_REVISION}" > /opt/puzzletron/modelopt_revision
+
+COPY pyproject.toml LICENSE_HEADER README.md /opt/puzzletron/src/modelopt/
+COPY modelopt /opt/puzzletron/src/modelopt/modelopt
+COPY modelopt_recipes /opt/puzzletron/src/modelopt/modelopt_recipes
+COPY puzzletron_orchestrator /opt/puzzletron/src/modelopt/puzzletron_orchestrator
+COPY puzzletron_setup /opt/puzzletron/src/modelopt/puzzletron_setup
+COPY examples/puzzletron /opt/puzzletron/src/modelopt/examples/puzzletron
+
+# Smoke-test the installed environment. Pins and source revisions are checked
+# by the repository tests and by the install commands above.
+RUN python -m pip install --no-build-isolation --no-deps -e \
+        "/opt/puzzletron/src/modelopt[hf,puzzletron,dev-test]" && \
+    pin() { jq -er ".${1}" "${PUZZLETRON_CI_ENVIRONMENT}"; } && \
+    for module in \
+        aiperf causal_conv1d decord fla grouped_gemm langdetect lmms_eval \
+        mamba_ssm modelopt nemo_automodel nltk puzzletron_orchestrator \
+        puzzletron_setup tilelang torch transformers vllm; do \
+        python -c "import ${module}"; \
+    done && \
+    test "$(python -c 'import torch; print(torch.version.cuda)')" = \
+        "$(pin gpu_image.torch_cuda)" && \
+    while read -r nltk_resource; do \
+        python -c "import nltk; nltk.data.find('tokenizers/${nltk_resource}')"; \
+    done < <(jq -r '.gpu_image.nltk_resources[]' "${PUZZLETRON_CI_ENVIRONMENT}") && \
+    lmms_root="$(python -c 'import lmms_eval; print(next(iter(lmms_eval.__path__)))')" && \
+    while read -r task_config; do \
+        test -f "${lmms_root}/${task_config}"; \
+    done < <(jq -r '.lmms_eval.task_configs[]' "${PUZZLETRON_CI_ENVIRONMENT}")
+
+LABEL org.opencontainers.image.source="https://github.com/NVIDIA/Model-Optimizer" \
+    org.opencontainers.image.revision="${MODELOPT_REVISION}" \
+    com.nvidia.modelopt.puzzletron.environment="examples/puzzletron/ci_environment.json" \
+    com.nvidia.modelopt.puzzletron.environment-recipe="examples/puzzletron/Dockerfile"
+
+WORKDIR /opt/puzzletron/src/modelopt
+
+CMD ["bash"]

examples/puzzletron/README.md

@@ -36,6 +36,12 @@ conversion, training, evaluation, and benchmarking run in the worker
 environment or container selected during setup. Prepare the
 [worker environment](docs/environment_setup.md) before launching a campaign.
 
+### Worker image
+
+The repository includes a Dockerfile and build command for Puzzletron workers.
+Follow the [image guide](docs/worker_image.md) to build, check, or export an
+Enroot/Pyxis SquashFS image.
+
 ### 2. Generate a campaign
 
 Start the guided setup with the repository defaults:

examples/puzzletron/build_worker_image.py

@@ -0,0 +1,193 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Build and optionally export the Linux amd64 Puzzletron worker image."""
+
+from __future__ import annotations
+
+import argparse
+import os
+import platform
+import re
+import shutil
+import subprocess
+import tempfile
+from pathlib import Path
+
+__all__ = [
+    "build_parser",
+    "main",
+    "sqsh_name",
+]
+
+_PLATFORM = "linux/amd64"
+_REVISION_PATTERN = re.compile(r"[0-9a-f]{40}")
+
+
+def sqsh_name(revision: str) -> str:
+    """Return the SquashFS filename for an image built from ``revision``."""
+
+    if not _REVISION_PATTERN.fullmatch(revision):
+        raise ValueError("Puzzletron image revision must be a full lowercase Git commit")
+    return f"modelopt-puzzletron-linux-amd64-git-{revision[:12]}.sqsh"
+
+
+def _run(command: list[str], **kwargs) -> subprocess.CompletedProcess:
+    return subprocess.run(command, check=True, **kwargs)
+
+
+def _source_revision(repository_root: Path) -> str:
+    status = subprocess.check_output(
+        ["git", "status", "--porcelain", "--untracked-files=all"],
+        cwd=repository_root,
+        text=True,
+    )
+    if status:
+        raise RuntimeError("Build the Puzzletron worker image from a clean Git checkout")
+    revision = subprocess.check_output(
+        ["git", "rev-parse", "HEAD"], cwd=repository_root, text=True
+    ).strip()
+    if not _REVISION_PATTERN.fullmatch(revision):
+        raise RuntimeError(f"Git returned an invalid Puzzletron image revision: {revision!r}")
+    return revision
+
+
+def _require_tools(*tools: str) -> None:
+    missing = [tool for tool in tools if shutil.which(tool) is None]
+    if missing:
+        raise RuntimeError(f"Missing required command(s): {', '.join(missing)}")
+
+
+def _require_linux_amd64() -> None:
+    if platform.system() != "Linux" or platform.machine() not in {"amd64", "x86_64"}:
+        raise RuntimeError("Build the Puzzletron worker image on a Linux amd64 host")
+
+
+def _output_path(output_dir: Path, name: str) -> Path:
+    path = output_dir / name
+    if path.exists():
+        raise FileExistsError(f"Refusing to overwrite existing artifact: {path}")
+    return path
+
+
+def _export_sqsh(image: str, output: Path) -> None:
+    partial = output.with_name(f".{output.stem}.partial.sqsh")
+    if partial.exists():
+        raise FileExistsError(f"Refusing to overwrite incomplete artifact: {partial}")
+
+    with tempfile.TemporaryDirectory(prefix="puzzletron-enroot-") as enroot_root:
+        environment = {
+            **os.environ,
+            "ENROOT_CACHE_PATH": f"{enroot_root}/cache",
+            "ENROOT_DATA_PATH": f"{enroot_root}/data",
+            "ENROOT_RUNTIME_PATH": f"{enroot_root}/runtime",
+        }
+        try:
+            _run(
+                ["enroot", "import", "--output", str(partial), f"dockerd://{image}"],
+                env=environment,
+            )
+        except (OSError, subprocess.CalledProcessError):
+            partial.unlink(missing_ok=True)
+            raise
+    partial.replace(output)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """Build the worker-image command-line parser."""
+
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        help="directory for optional exported artifacts",
+    )
+    parser.add_argument(
+        "--sqsh",
+        action="store_true",
+        help="export an Enroot/Pyxis SquashFS image",
+    )
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    """Build, verify, and optionally export the Puzzletron worker image."""
+
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    if args.sqsh and args.output_dir is None:
+        parser.error("--output-dir is required with --sqsh")
+    if args.output_dir is not None and not args.sqsh:
+        parser.error("--output-dir requires --sqsh")
+
+    repository_root = Path(__file__).resolve().parents[2]
+    revision = _source_revision(repository_root)
+    _require_linux_amd64()
+    image = f"modelopt-puzzletron:linux-amd64-git-{revision[:12]}"
+
+    sqsh = None
+    if args.output_dir is not None:
+        args.output_dir = args.output_dir.expanduser().resolve()
+        if args.output_dir.is_relative_to(repository_root):
+            parser.error("--output-dir must be outside the repository")
+        sqsh = _output_path(args.output_dir, sqsh_name(revision))
+
+    required_tools = ["docker"]
+    if args.sqsh:
+        required_tools.append("enroot")
+    _require_tools(*required_tools)
+    if args.output_dir is not None:
+        args.output_dir.mkdir(parents=True, exist_ok=True)
+
+    _run(
+        [
+            "docker",
+            "build",
+            "--platform",
+            _PLATFORM,
+            "--file",
+            "examples/puzzletron/Dockerfile",
+            "--build-arg",
+            f"MODELOPT_REVISION={revision}",
+            "--tag",
+            image,
+            ".",
+        ],
+        cwd=repository_root,
+    )
+    recorded_revision = subprocess.check_output(
+        [
+            "docker",
+            "image",
+            "inspect",
+            "--format",
+            '{{ index .Config.Labels "org.opencontainers.image.revision" }}',
+            image,
+        ],
+        text=True,
+    ).strip()
+    if recorded_revision != revision:
+        raise RuntimeError(
+            f"Puzzletron image revision mismatch: expected {revision}, found {recorded_revision}"
+        )
+    if sqsh is not None:
+        _export_sqsh(image, sqsh)
+
+    print(f"Docker image: {image}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())