ARG BASE_IMAGE=nvcr.io/nvidia/nemo:26.06
FROM ${BASE_IMAGE}

ARG INSTALL_FLASH_ATTN4=0
ARG PIP_INDEX_URL=https://mirrors.aliyun.com/pypi/simple/
ARG PIP_EXTRA_INDEX_URL=https://pypi.org/simple

ENV PIP_INDEX_URL=${PIP_INDEX_URL}
ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
ENV PIP_DISABLE_PIP_VERSION_CHECK=1

RUN if [ "${INSTALL_FLASH_ATTN4}" = "1" ]; then \
      python3 -m pip uninstall -y flash-attn flash-attn-3 flash-attn-4 || true; \
      MAX_JOBS="${MAX_JOBS:-16}" python3 -m pip install --no-cache-dir --no-build-isolation \
        "flash-attn-4" "nvidia-cutlass-dsl"; \
    fi

RUN set -eux; \
    python3 - <<'PY'
from pathlib import Path
import shutil

site = Path("/usr/local/lib/python3.12/dist-packages")
src = site / "vllm/third_party/triton_kernels"
dst = site / "triton_kernels"
if not (src / "matmul_ogs.py").exists() or not (src / "tensor.py").exists():
    raise SystemExit(f"missing vLLM triton_kernels source: {src}")
if not dst.exists():
    raise SystemExit(f"missing top-level triton_kernels package: {dst}")
backup = site / "triton_kernels_nv26_backup"
if backup.exists():
    shutil.rmtree(backup)
shutil.copytree(dst, backup, ignore=shutil.ignore_patterns("__pycache__"))
shutil.rmtree(dst)
shutil.copytree(src, dst, ignore=shutil.ignore_patterns("__pycache__"))
PY

RUN python3 - <<'PY'
import importlib.metadata as metadata
import importlib.util
required = ["torch", "transformer-engine", "megatron-core", "megatron-bridge"]
for package in required:
    print(package, metadata.version(package))
try:
    print("flash-attn-4", metadata.version("flash-attn-4"))
except metadata.PackageNotFoundError:
    print("flash-attn-4 MISSING")
try:
    print("nvidia-cutlass-dsl", metadata.version("nvidia-cutlass-dsl"))
except metadata.PackageNotFoundError:
    print("nvidia-cutlass-dsl MISSING")
spec = importlib.util.find_spec("triton_kernels.matmul_ogs")
if spec is None:
    raise SystemExit("triton_kernels.matmul_ogs MISSING")
print("triton_kernels.matmul_ogs", spec.origin)
PY
