.github/container/Dockerfile.jax

# syntax=docker/dockerfile:1-labs
ARG BASE_IMAGE=ghcr.io/nvidia/jax:base
ARG BUILD_PATH_JAXLIB=/opt/jaxlibs
# Extra targets to build and copy outputs of, can be used for HLO tools. For example
# @xla//xla/hlo/tools:hlo-opt,...
ARG EXTRA_BAZEL_TARGETS="@xla//xla/tools/multihost_hlo_runner:hlo_runner_main,@xla//xla/tools/multihost_hlo_runner:py_hlo_multihost_runner.so"

# Extra arguments to append to build-jax.sh
ARG EXTRA_BUILD_JAX_ARGS=""
ARG URLREF_JAX=https://github.com/google/jax.git#main
ARG URLREF_XLA=https://github.com/openxla/xla.git#main
ARG URLREF_FLAX=https://github.com/google/flax.git#main
ARG URLREF_TRANSFORMER_ENGINE=https://github.com/NVIDIA/TransformerEngine.git#main
ARG SRC_PATH_JAX=/opt/jax
ARG SRC_PATH_XLA=/opt/xla
ARG SRC_PATH_FLAX=/opt/flax
ARG SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine
ARG GIT_USER_NAME="JAX Toolbox"
ARG GIT_USER_EMAIL=jax@nvidia.com

ARG BAZEL_CACHE=/tmp
ARG BUILD_DATE

###############################################################################
## Build JAX
###############################################################################

FROM ${BASE_IMAGE} AS builder
ARG URLREF_JAX
ARG URLREF_TRANSFORMER_ENGINE
ARG URLREF_XLA
ARG SRC_PATH_JAX
ARG SRC_PATH_TRANSFORMER_ENGINE
ARG SRC_PATH_XLA
ARG BAZEL_CACHE
ARG BUILD_PATH_JAXLIB
ARG EXTRA_BAZEL_TARGETS
ARG EXTRA_BUILD_JAX_ARGS
ARG GIT_USER_NAME
ARG GIT_USER_EMAIL

RUN --mount=type=ssh \
    --mount=type=secret,id=SSH_KNOWN_HOSTS,target=/root/.ssh/known_hosts \
<<"EOF" bash -ex
    git-clone.sh ${URLREF_JAX} ${SRC_PATH_JAX}
    git-clone.sh ${URLREF_XLA} ${SRC_PATH_XLA}
EOF

ADD build-jax.sh build-te.sh local_cuda_arch /usr/local/bin/
# Install bazelisk
RUN ARCH="$(dpkg --print-architecture)" && \
    wget -O /usr/local/bin/bazel https://github.com/bazelbuild/bazelisk/releases/latest/download/bazelisk-linux-${ARCH} && \
    chmod +x /usr/local/bin/bazel
# Populate ${BUILD_PATH_JAXLIB} with editable wheels; --no-install because
# (a) this is the builder stage, and (b) pip-finalize.sh does the install
RUN mkdir -p /builder/extra-targets && \
    build-jax.sh \
    --bazel-cache ${BAZEL_CACHE} \
    --build-path-jaxlib ${BUILD_PATH_JAXLIB} \
    --extra-targets "${EXTRA_BAZEL_TARGETS}" \
    --extra-target-dest /builder/extra-targets \
    --no-install \
    --src-path-jax ${SRC_PATH_JAX} \
    --src-path-xla ${SRC_PATH_XLA} \
    --sm all \
    --clean \
    ${EXTRA_BUILD_JAX_ARGS}

## Transformer engine: check out source and build wheel
RUN --mount=type=ssh \
    --mount=type=secret,id=SSH_KNOWN_HOSTS,target=/root/.ssh/known_hosts \
    git-clone.sh ${URLREF_TRANSFORMER_ENGINE} ${SRC_PATH_TRANSFORMER_ENGINE}
# Populate ${SRC_PATH_TRANSFORMER_ENGINE}/dist with [a] .whl file(s); --no-install
# because (a) this is the builder stage, and (b) pip-finalize.sh does the install
RUN build-te.sh \
    --clean \
    --no-install \
    --sm all \
    --src-path-te ${SRC_PATH_TRANSFORMER_ENGINE} \
    --src-path-xla ${SRC_PATH_XLA}

###############################################################################
## Pack jaxlib wheel and various source dirs into a pre-installation image
###############################################################################

ARG BASE_IMAGE
FROM ${BASE_IMAGE} AS mealkit
ARG URLREF_FLAX
ARG SRC_PATH_JAX
ARG SRC_PATH_XLA
ARG SRC_PATH_FLAX
ARG SRC_PATH_TRANSFORMER_ENGINE
ARG BUILD_DATE
ARG BUILD_PATH_JAXLIB

ENV BUILD_DATE=${BUILD_DATE}
# The following environment variables tune performance
ENV XLA_FLAGS=""
ENV XLA_FLAGS="${XLA_FLAGS} --xla_gpu_enable_latency_hiding_scheduler=true"

COPY --from=builder ${BUILD_PATH_JAXLIB} ${BUILD_PATH_JAXLIB}
COPY --from=builder ${SRC_PATH_JAX} ${SRC_PATH_JAX}
COPY --from=builder ${SRC_PATH_XLA} ${SRC_PATH_XLA}
COPY --from=builder /usr/local/bin/bazel /usr/local/bin/bazel
COPY --from=builder /builder/extra-targets/bin/* /usr/local/bin
COPY --from=builder /builder/extra-targets/python/*.so /usr/lib/python3/dist-packages/

# Preserve the versions of jax and xla
COPY --from=builder /opt/manifest.d/git-clone.yaml /opt/manifest.d/git-clone.yaml
ADD build-jax.sh build-te.sh local_cuda_arch pytest-xdist.sh test-jax.sh /usr/local/bin/

RUN mkdir -p /opt/pip-tools.d

## Editable installations of jax and jaxlib
# Note that jax now is an independent wheel, extra [k8s] needs to be from build path also
RUN <<"EOF" bash -ex
for component in $(ls ${BUILD_PATH_JAXLIB}); do
    echo "-e file://${BUILD_PATH_JAXLIB}/${component}" >> /opt/pip-tools.d/requirements-jax.in;
done
echo "-e file://${SRC_PATH_JAX}[k8s]" >> /opt/pip-tools.d/requirements-jax.in
EOF

## Flax
RUN <<"EOF" bash -ex
git-clone.sh ${URLREF_FLAX} ${SRC_PATH_FLAX}
echo "-e file://${SRC_PATH_FLAX}" >> /opt/pip-tools.d/requirements-flax.in
EOF

# Copy TransformerEngine wheel from the builder stage
ENV SRC_PATH_TRANSFORMER_ENGINE=${SRC_PATH_TRANSFORMER_ENGINE}
COPY --from=builder ${SRC_PATH_TRANSFORMER_ENGINE} ${SRC_PATH_TRANSFORMER_ENGINE}
RUN <<"EOF" bash -ex
ls ${SRC_PATH_TRANSFORMER_ENGINE}/dist/*.whl
echo "transformer-engine @ file://$(ls ${SRC_PATH_TRANSFORMER_ENGINE}/dist/*.whl)" > /opt/pip-tools.d/requirements-te.in
EOF

###############################################################################
## Install primary packages and transitive dependencies
###############################################################################

FROM mealkit AS final
RUN pip-finalize.sh