ARG PYVERSION={{ config.python_version }}
ARG HOME
FROM {{ base_image_name_and_tag }} AS truss_server

{%- set python_exec_path = config.base_image.python_executable_path or '$(which python3)' %}
{# NB(nikhil): Our `uv` installation commands require a full path to the python interpreter, but it relies #}
{# on shell expansion to do so if `python_executable_path` is not provided. This expansion doesn't work for certain #}
{# docker commands (ENV, ENTRYPPOINT), so we point that to `python3` as a default. #}
{%- set python_executable = config.base_image.python_executable_path or 'python3' %}
ENV PYTHON_EXECUTABLE="{{ python_executable }}"

{%- set default_owner = "root:root" %}
{%- set ROOT_USER_ID = 0 %}
{%- set NONROOT_DEFAULT_USER_ID = 60000 %}
{#
  run_as_user_id encoding:
  - ROOT_USER_ID = root (non-root disabled)
  - NONROOT_DEFAULT_USER_ID = platform default "app" user (non-root enabled, create user)
  - any other = custom UID (user-specified, don't create user)
#}
{% block user_setup %}
{#
  NB(nikhil): Assume `root` for various setup that requires permissions; we conditionally switch to
  non-root user at the end, if `non_root` behavior is enabled.
#}
USER root
{%- set app_username = "app" %}
{%- set control_server_dir = "/control" %}
{%- if run_as_user_id == NONROOT_DEFAULT_USER_ID %}
{# Platform-managed non-root user's home directory. #}
{# uv will use $HOME to install packages. #}
ENV HOME=${HOME:-/home/{{ app_username }}}
{%- else %}
{# For root or custom run_as_user, use /root during build (we switch user at the end). #}
{# For custom UIDs, we don't know the username and its corresponding $HOME, so /root. We still switch user at the end. #}
ENV HOME=${HOME:-/root}
{%- endif %}
{# Directory containing inference server code - always /app regardless of username #}
ENV APP_HOME=/{{ app_username }}
RUN mkdir -p ${APP_HOME} {{ control_server_dir }}
{%- if run_as_user_id == ROOT_USER_ID or run_as_user_id == NONROOT_DEFAULT_USER_ID %}
{# Create platform-managed non-root user to run model containers. #}
{# For custom UIDs, assume the user already exists and don't create. #}
RUN useradd -u {{ NONROOT_DEFAULT_USER_ID }} -ms /bin/bash {{ app_username }}
{%- endif %}
{% endblock %} {#- endblock user_setup #}

{#- at the very beginning, set non-interactive mode for apt #}
ENV DEBIAN_FRONTEND=noninteractive


{%- set UV_VERSION = "0.8.22" %}
{#
NB(nikhil): We use a semi-complex uv installation command across the board:
- A generous UV_HTTP_TIMEOUT (5m) for packages that take a long time to install.
- We use `unsafe-best-match` since `uv` is stricter about having multiple registries, and we often need to search multiple.
- Use the python executable currently on the path, will respect activated virtual envs.
- `--system --break-system-packages` bypasses checks if we're actually installing into system python, but has no effect otherwise.
#}
{%- set sys_pip_install_command =
    "UV_HTTP_TIMEOUT=${UV_HTTP_TIMEOUT:-300} " +
    "uv pip install --system --break-system-packages " +
    "--index-strategy unsafe-best-match " +
    "--python " + python_exec_path
%}

{% block fail_fast %}
RUN grep -w 'ID=debian\|ID_LIKE=debian' /etc/os-release || { echo "ERROR: Supplied base image is not a debian image"; exit 1; }
RUN {{ python_exec_path }} -c "import sys; \
    sys.exit(0) \
    if sys.version_info.major == {{ supported_python_major_version_in_custom_base_image }} \
    and sys.version_info.minor >= {{ min_supported_python_minor_version_in_custom_base_image }} \
    and sys.version_info.minor <= {{ max_supported_python_minor_version_in_custom_base_image }} \
    else sys.exit(1)" \
    || { echo "ERROR: Supplied base image does not have {{ min_supported_python_version_in_custom_base_image }} <= python <= {{ max_supported_python_version_in_custom_base_image }}"; exit 1; }
{% endblock %}

{# mirrors.ubuntu.com is a redirector service (uses mirrors.txt) that provides the best mirror based on location #}
{# Handle old format (Ubuntu 22.04 and earlier): /etc/apt/sources.list, and new format (Ubuntu 24.04+): /etc/apt/sources.list.d/ubuntu.sources #}
{%- set ubuntu_sources_list_files = ["/etc/apt/sources.list", "/etc/apt/sources.list.d/ubuntu.sources"] %}
{%- for ubuntu_sources_list_file in ubuntu_sources_list_files %}
RUN if [ -f {{ ubuntu_sources_list_file }} ]; then \
    sed -i.bak 's|http://archive.ubuntu.com/ubuntu/|mirror://mirrors.ubuntu.com/mirrors.txt|g' {{ ubuntu_sources_list_file }}; \
fi
{%- endfor %}

{% block install_tools %}
{# Install curl if not already present. Required for uv installation and truss-transfer-cli. #}
RUN command -v curl >/dev/null 2>&1 || (apt update && apt install -y curl)

{# Install uv if not already present in the image. We validate the expected location for `uv` at the very end
   due to limitations with `pipefail` in Docker context. #}
RUN if ! command -v uv >/dev/null 2>&1; then \
    curl -LsSf --retry 5 --retry-delay 5 https://astral.sh/uv/{{ UV_VERSION }}/install.sh | sh && \
    test -x ${HOME}/.local/bin/uv; \
fi
{# Add the user's local bin to the path, used by uv. #}
ENV PATH=${PATH}:${HOME}/.local/bin
{% endblock %}

{% block base_image_patch %}
{% endblock %}

{% if config.model_framework.value == 'huggingface_transformer' %}
    {% if config.resources.use_gpu %}
{# HuggingFace pytorch gpu support needs mkl. #}
RUN {{ sys_pip_install_command }} install mkl
    {% endif %}
{% endif %}

{% block post_base %}
{% endblock %}

{% block install_system_requirements %}
    {%- if should_install_system_requirements %}
COPY --chown={{ default_owner }} ./{{ system_packages_filename }} {{ system_packages_filename }}
RUN apt-get update && apt-get install --yes --no-install-recommends $(cat {{ system_packages_filename }}) \
    && apt-get autoremove -y \
    && apt-get clean -y \
    && rm -rf /var/lib/apt/lists/*
    {%- endif %}
{% endblock %}


{% block install_requirements %}
    {%- if requirements_file_type == "uv_lock" %}
COPY --chown={{ default_owner }} ./{{ pyproject_toml_filename }} {{ pyproject_toml_filename }}
COPY --chown={{ default_owner }} ./{{ uv_lock_filename }} {{ uv_lock_filename }}
{#
  NB(nikhil) By default, `uv sync --frozen` manages its own venv, but we need to install into the provided python
  env. To get around this, we leverage `uv export` to dump the requirements in a format that `uv pip install` can use
#}
RUN uv export --frozen --format requirements-txt > /tmp/uv_requirements.txt && \
    {{ sys_pip_install_command }} -r /tmp/uv_requirements.txt --no-cache-dir && \
    rm -f /tmp/uv_requirements.txt

    {%- elif requirements_file_type == "pyproject" %}
COPY --chown={{ default_owner }} ./{{ pyproject_toml_filename }} {{ pyproject_toml_filename }}
RUN {{ sys_pip_install_command }} -r {{ pyproject_toml_filename }} --no-cache-dir

    {%- elif requirements_file_type == "pip" %}
COPY --chown={{ default_owner }} ./{{ user_supplied_requirements_filename }} {{ user_supplied_requirements_filename }}
RUN {{ sys_pip_install_command }} -r {{ user_supplied_requirements_filename }} --no-cache-dir
    {%- endif %}

    {%- if should_install_requirements %}
COPY --chown={{ default_owner }} ./{{ config_requirements_filename }} {{ config_requirements_filename }}
RUN {{ sys_pip_install_command }} -r {{ config_requirements_filename }} --no-cache-dir
    {%- endif %}
{% endblock %}

{%- if not config.docker_server %}
WORKDIR $APP_HOME
{%- endif %}

{% block app_copy %}
{% endblock %}


{% set packages_dir = "/packages" %}
RUN mkdir -p {{ packages_dir }}
{% block bundled_packages_copy %}
    {%- if bundled_packages_dir_exists %}
COPY --chown={{ default_owner }} ./{{ config.bundled_packages_dir }} {{ packages_dir }}
{%- endif %}
{% endblock %}


{% for env_var_name, env_var_value in passthrough_environment_variables.items() %}
ENV {{ env_var_name }}="{{ env_var_value }}"
{% endfor %}

{% block run %}
{% endblock %}
