ARG INPUT_IMAGE_NAME

from ${INPUT_IMAGE_NAME} AS llm_custom

ARG EXP_DIR
WORKDIR $EXP_DIR

# RUN DEBIAN_FRONTEND=noninteractive apt-get install \
#     --no-install-recommends -y --reinstall \
#     libaio-dev
# # Install the deepseed ops, see
# # https://github.com/microsoft/DeepSpeed/blob/master/docs/_tutorials/advanced-install.md
# https://www.deepspeed.ai/tutorials/advanced-install/
# RUN DS_BUILD_OPS=1 \
#     TORCH_CUDA_ARCH_LIST="12.0" \
#     pip install deepspeed \
#     --global-option="build_ext" \
#     --global-option="-j8"
#     # --global-option="bdist_wheel"

# 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_dl/pyproject.toml ./src/lib_dl/poetry.lock $EXP_DIR/src/lib_dl/
COPY ./src/lib_llm/pyproject.toml ./src/lib_llm/poetry.lock $EXP_DIR/src/lib_llm/
# We create these files here to make poetry and pip happy
# and later copy them into the target image later
RUN touch $EXP_DIR/README.md \
    && touch $EXP_DIR/src/__init__.py \
    && touch $EXP_DIR/src/lib_dl/README.md \
    && mkdir -p $EXP_DIR/src/lib_dl/lib_dl/ \
    && touch $EXP_DIR/src/lib_dl/lib_dl/__init__.py \
    && touch $EXP_DIR/src/lib_llm/README.md \
    && mkdir -p $EXP_DIR/src/lib_llm/lib_llm/ \
    && touch $EXP_DIR/src/lib_llm/lib_llm/__init__.py

# Create some content inside the lib_dl directory to make poetry happy
# RUN mkdir -p ./src/lib_dl/lib_dl/ ./src/lib_llm/lib_llm/ \
#     && touch ./README.md \
#         ./src/lib_dl/lib_dl/__init__.py ./src/lib_dl/README.md \
#         ./src/lib_llm/lib_llm/__init__.py ./src/lib_llm/README.md \
#     && poetry install --no-root --no-cache --without dev --with docker

# TODO: using deepspeed stage 3 requires precompiling 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
