FROM python:3.11-slim AS cache_warmer

ENV APP_HOME=/app
RUN mkdir -p ${APP_HOME}/model_cache
WORKDIR ${APP_HOME}

{% if hf_access_token %}
ENV HUGGING_FACE_HUB_TOKEN="{{hf_access_token}}"
{% endif %}

RUN apt-get -y update; apt-get -y install curl; curl -s https://baseten-public.s3.us-west-2.amazonaws.com/bin/b10cp-5fe8dc7da-linux-amd64 -o /app/b10cp; chmod +x /app/b10cp
ENV B10CP_PATH_TRUSS="${APP_HOME}/b10cp"
COPY --chown={{ default_owner }} ./cache_requirements.txt ${APP_HOME}/cache_requirements.txt
RUN pip install -r ${APP_HOME}/cache_requirements.txt --no-cache-dir && rm -rf /root/.cache/pip
COPY --chown={{ default_owner }} ./cache_warmer.py /cache_warmer.py

{% for credential in credentials_to_cache  %}
COPY ./{{credential}} ${APP_HOME}/{{credential}}
{% endfor %}

{% for repo, hf_dir in models.items() %}
        {% for file in hf_dir.files %}
{{ "RUN --mount=type=secret,id=" + hf_access_token_file_name + ",dst=/etc/secrets/" + hf_access_token_file_name if use_hf_secret else "RUN" }} python3 /cache_warmer.py {{file}} {{repo}} {% if hf_dir.revision != None %}{{hf_dir.revision}}{% endif %}
        {% endfor %}
{% endfor %}
