Dockerfile

# syntax=docker/dockerfile:1.7
# Generated by main.py. Regenerate it after changing build options.
FROM nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04 AS builder

ARG TF_REF=v2.21.0
ARG TF_COMMIT=a481b10260dfdf833a1b16007eead49c1d7febf3
ARG TF_PYTHON_VERSION=3.12
ARG TF_WHEEL_SUFFIX=+selfbuild
ARG BAZELISK_VERSION=v1.26.0
ARG BAZEL_JOBS=48

ENV DEBIAN_FRONTEND=noninteractive \
    PIP_DISABLE_PIP_VERSION_CHECK=1 \
    PYTHONUNBUFFERED=1

RUN mkdir /workspace
WORKDIR /workspace

# Install build dependencies.
RUN --mount=target=/var/lib/apt/lists,type=cache,id=apt-lists \
    --mount=target=/var/cache/apt,type=cache,id=apt-cache \
    rm -f /etc/apt/apt.conf.d/docker-clean && \
    apt update && \
    apt install -y --no-install-recommends \
      build-essential \
      ca-certificates \
      curl \
      git \
      gnupg \
      libbz2-dev \
      libffi-dev \
      liblzma-dev \
      libreadline-dev \
      libsqlite3-dev \
      libssl-dev \
      lsb-release \
      make \
      patchelf \
      pkg-config \
      python3 \
      python3-dev \
      python3-pip \
      python3-venv \
      software-properties-common \
      unzip \
      wget \
      xxd \
      zip \
      zlib1g-dev

# Install LLVM/Clang 20.
RUN --mount=target=/var/lib/apt/lists,type=cache,id=apt-lists \
    --mount=target=/var/cache/apt,type=cache,id=apt-cache \
    wget -O llvm.sh https://apt.llvm.org/llvm.sh && \
    chmod +x llvm.sh && \
    ./llvm.sh 20 all && \
    update-alternatives --install /usr/bin/clang clang /usr/bin/clang-20 100 && \
    update-alternatives --install /usr/bin/clang++ clang++ /usr/bin/clang++-20 100 && \
    rm llvm.sh

# Clone the requested TensorFlow ref. Commit verification is skipped when
# TF_COMMIT is empty.
RUN git clone --depth 1 --branch "${TF_REF}" https://github.com/tensorflow/tensorflow.git tensorflow && \
    cd tensorflow && \
    if [ -n "${TF_COMMIT}" ]; then \
      test "$(git rev-parse HEAD)" = "${TF_COMMIT}"; \
    fi
WORKDIR /workspace/tensorflow

# Install Bazelisk as the bazel entry point.
RUN wget "https://github.com/bazelbuild/bazelisk/releases/download/${BAZELISK_VERSION}/bazelisk-linux-amd64" -O /usr/bin/bazel && \
    chmod +x /usr/bin/bazel && \
    bazel version

COPY .tf_configure.bazelrc .

RUN --mount=type=cache,target=/root/.cache/bazel,id=bazel-cache \
    --mount=target=/root/.cache/pip,type=cache,id=pip-global-cache \
    bazel build //tensorflow/tools/pip_package:wheel \
      --repo_env=USE_PYWRAP_RULES=1 \
      --repo_env=WHEEL_NAME=tensorflow \
      --repo_env=TF_PYTHON_VERSION="${TF_PYTHON_VERSION}" \
      --repo_env=HERMETIC_PYTHON_VERSION="${TF_PYTHON_VERSION}" \
      --repo_env=ML_WHEEL_TYPE=release \
      --repo_env=ML_WHEEL_VERSION_SUFFIX="${TF_WHEEL_SUFFIX}" \
      --config=cuda_nvcc \
      --config=cuda_wheel \
      --config=opt \
      -c opt \
      --repo_env=HERMETIC_CUDA_VERSION="12.8.1" \
      --repo_env=HERMETIC_CUDNN_VERSION="9.8.0" \
      --repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_120,compute_120" \
      --jobs="${BAZEL_JOBS}" \
      '--local_ram_resources=HOST_RAM*.80' \
      --verbose_failures && \
    mkdir -p /workspace/wheelhouse && \
    cp bazel-bin/tensorflow/tools/pip_package/wheel_house/*.whl /workspace/wheelhouse/

# Add relative RUNPATHs so TensorFlow can find CUDA libraries from
# the nvidia-* pip packages without requiring LD_LIBRARY_PATH.
RUN python3 - <<'PY'
import base64
import csv
import hashlib
import os
import shutil
import subprocess
import tempfile
import zipfile
from pathlib import Path

wheelhouse = Path("/workspace/wheelhouse")
nvidia_lib_dirs = [
    Path("nvidia/cublas/lib"),
    Path("nvidia/cuda_cupti/lib"),
    Path("nvidia/cuda_nvrtc/lib"),
    Path("nvidia/cuda_runtime/lib"),
    Path("nvidia/cudnn/lib"),
    Path("nvidia/cufft/lib"),
    Path("nvidia/curand/lib"),
    Path("nvidia/cusolver/lib"),
    Path("nvidia/cusparse/lib"),
    Path("nvidia/nccl/lib"),
    Path("nvidia/nvjitlink/lib"),
]


def wheel_hash(path: Path) -> str:
    digest = hashlib.sha256(path.read_bytes()).digest()
    encoded = base64.urlsafe_b64encode(digest).rstrip(b"=").decode("ascii")
    return f"sha256={encoded}"


def refresh_record(root: Path) -> None:
    record = next(root.glob("*.dist-info/RECORD"))
    rows = []
    for path in sorted(p for p in root.rglob("*") if p.is_file()):
        rel = path.relative_to(root).as_posix()
        if path == record:
            rows.append([rel, "", ""])
        else:
            rows.append([rel, wheel_hash(path), str(path.stat().st_size)])
    with record.open("w", newline="", encoding="utf-8") as handle:
        csv.writer(handle).writerows(rows)


def patch_rpaths(root: Path) -> None:
    tf_root = root / "tensorflow"
    so_files = sorted(tf_root.rglob("*.so")) + sorted(tf_root.rglob("*.so.*"))
    for so_file in so_files:
        entries = [
            "$ORIGIN",
            "$ORIGIN/..",
            "$ORIGIN/" + os.path.relpath(tf_root, so_file.parent),
        ]
        for lib_dir in nvidia_lib_dirs:
            entries.append("$ORIGIN/" + os.path.relpath(root / lib_dir, so_file.parent))
        entries = list(dict.fromkeys(entries))
        subprocess.run(
            ["patchelf", "--set-rpath", ":".join(entries), str(so_file)],
            check=True,
        )


def repack(root: Path, wheel: Path) -> None:
    with zipfile.ZipFile(wheel, "w", compression=zipfile.ZIP_DEFLATED) as archive:
        for path in sorted(p for p in root.rglob("*") if p.is_file()):
            rel = path.relative_to(root).as_posix()
            info = zipfile.ZipInfo.from_file(path, arcname=rel)
            info.compress_type = zipfile.ZIP_DEFLATED
            archive.writestr(info, path.read_bytes())


for wheel in wheelhouse.glob("*.whl"):
    with tempfile.TemporaryDirectory() as temp_dir:
        root = Path(temp_dir) / "wheel"
        root.mkdir()
        with zipfile.ZipFile(wheel) as archive:
            archive.extractall(root)
        patch_rpaths(root)
        refresh_record(root)
        repaired = wheel.with_suffix(".repaired.whl")
        repack(root, repaired)
        shutil.move(repaired, wheel)
        print(f"Repaired CUDA RUNPATHs in {wheel.name}")
PY


FROM scratch AS wheel-export
COPY --from=builder /workspace/wheelhouse/ /