ARG INPUT_IMAGE_NAME

from ${INPUT_IMAGE_NAME} AS llm_custom

ARG EXP_DIR
WORKDIR $EXP_DIR

# RUN DEBIAN_FRONTEND=noninteractive \
#     apt-get update \
#     && apt-get install \
#     --no-install-recommends -y --reinstall \
#     libaio-dev

# Install the deepseed ops, see: https://www.deepspeed.ai/tutorials/advanced-install/
# See here for compute capabilities: https://developer.nvidia.com/cuda-gpus
# RUN TORCH_CUDA_ARCH_LIST="7.0;8.0;8.6;9.0" \
#     DS_BUILD_OPS=1 \
#     DS_BUILD_EVOFORMER_ATTN=0 \
#     DS_BUILD_CUTLASS_OPS=0 \
#     DS_BUILD_RAGGED_DEVICE_OPS=0 \
#     DS_BUILD_SPARSE_ATTN=0 \
#     pip wheel deepspeed==0.13.1 \
#         --wheel-dir=./dist
        # --global-option="build_ext" \
        # --global-option="-j8"
        # --config-settings="--build-ext -j8"
    # DS_BUILD_CPU_ADAM=1 \


# Copy and install dependencies seperately first to avoid triggering
# a dependency reinstallation on every change in the code
# Copy using poetry.lock* in case it doesn't exist yet
COPY ./pyproject.toml ./poetry.lock* $EXP_DIR/
COPY ./src/lib_llm/pyproject.toml ./src/lib_llm/poetry.lock $EXP_DIR/src/lib_llm/
COPY ./src/lib_project/pyproject.toml ./src/lib_project/poetry.lock $EXP_DIR/src/lib_project/
COPY ./src/lib_dl_base/pyproject.toml ./src/lib_dl_base/poetry.lock $EXP_DIR/src/lib_dl_base/

# We create these files here to make poetry and pip happy
# and later copy them into the target image
RUN touch $EXP_DIR/README.md \
    && mkdir -p $EXP_DIR/src/ \
    && touch $EXP_DIR/src/__init__.py
ARG LIBRARIES="lib_llm lib_project lib_dl_base"
RUN for lib in $LIBRARIES; do \
        touch $EXP_DIR/src/$lib/README.md && \
        mkdir -p $EXP_DIR/src/$lib/$lib/ && \
        touch $EXP_DIR/src/$lib/$lib/__init__.py; \
    done

# TODO: using deepspeed stage 3 requires precoANONYMOUSling plugins.
# We should do this here, because the final image does not contain
# devel packages.
RUN poetry build --format wheel \
    && poetry export --without dev --with docker \
        --without-hashes --format constraints.txt --output constraints.txt
