Upgrade base images for py3.9 through py3.13 (#1704)

Author: nnarayen

.github/actions/setup-python/action.yml

@@ -45,7 +45,7 @@ runs:
       id: cache
       with:
         path: .venv
-        key: venv-${{ runner.os }}-${{ env.version }}-${{ hashFiles('**/poetry.lock') }}
+        key: venv-${{ runner.os }}-${{ inputs.python-version }}-${{ hashFiles('**/poetry.lock') }}
 
     - name: Ensure cache is healthy
       if: steps.cache.outputs.cache-hit == 'true'

docker/base_images/base_image.Dockerfile.jinja

@@ -1,10 +1,10 @@
 {% if use_gpu %}
-FROM nvidia/cuda:12.2.2-base-ubuntu20.04
+FROM nvidia/cuda:12.2.2-base-ubuntu22.04
 ENV CUDNN_VERSION="8.9.5.29"
 ENV CUDA="12.2"
 ENV LD_LIBRARY_PATH="/usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH"
 
-RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub && \
+RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub && \
     apt-get update && apt-get install -y --no-install-recommends \
         ca-certificates \
         cuda-command-line-tools-12-2 \

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truss"
-version = "0.9.102"
+version = "0.9.103rc001"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"

truss/contexts/image_builder/util.py

@@ -6,7 +6,7 @@
 # This needs to be updated whenever we want to update
 # base images on a merge. Updating this version will cause
 # base images to be pushed with this tag.
-TRUSS_BASE_IMAGE_VERSION_TAG = "v0.9.0"
+TRUSS_BASE_IMAGE_VERSION_TAG = "v0.9.103"
 
 
 def file_is_empty(path: Path, ignore_hash_style_comments: bool = True) -> bool:

truss/templates/control_venv.Dockerfile.jinja

@@ -0,0 +1,8 @@
+{%- if requires_live_reload %}
+FROM python:{{ control_python_version }}-slim as control-server-venv
+
+RUN echo "Creating venv for py{{ control_python_version }}";
+RUN python -m pip install --upgrade pip
+RUN python -m venv --copies /control/.env
+RUN echo "Created venv for py{{ control_python_version }}";
+{%- endif %} {#- endif requires_live_reload #}

truss/templates/server.Dockerfile.jinja

@@ -2,8 +2,15 @@
 {%- include "cache.Dockerfile.jinja" %}
 {%- endif %} {#- endif model_cache_v1 #}
 
+{%- set requires_live_reload = config.live_reload and not config.docker_server %}
+
+{%- if requires_live_reload %}
+{% include "control_venv.Dockerfile.jinja" %}
+{%- endif %} {#- endif requires_live_reload #}
+
 {% extends "base.Dockerfile.jinja" %}
 
+
 {% block base_image_patch %}
 {# If user base image is supplied in config, apply build commands from truss base image #}
 {% if config.base_image %}
@@ -27,14 +34,14 @@ COPY ./{{ base_server_requirements_filename }} {{ base_server_requirements_filen
 RUN pip install -r {{ base_server_requirements_filename }} --no-cache-dir && rm -rf /root/.cache/pip
     {%- endif %} {#- endif not config.docker_server #}
 
-    {%- if config.live_reload and not config.docker_server %}
+    {%- if requires_live_reload %}
     {# Create symlink for inference server IF a user base image is supplied and live_reload is enabled. #}
     {# This links the base images primary python executable path to /usr/local/bin/python. #}
     {# This is specific to the user-provided base image scenario. #}
 RUN readlink {{ config.base_image.python_executable_path }} &>/dev/null \
     && echo "WARNING: Overwriting existing link at /usr/local/bin/python"
 RUN ln -sf {{ config.base_image.python_executable_path }} /usr/local/bin/python
-    {%- endif %} {#- endif config.live_reload and not config.docker_server (for symlink) #}
+    {%- endif %} {#- endif requires_live_reload #}
 {% endif %} {#- endif config.base_image #}
 
 {% endblock %} {#- endblock base_image_patch #}
@@ -91,34 +98,24 @@ COPY ./truss /app/truss
 {%- endif %} {#- endif use_local_src #}
 
 COPY ./config.yaml /app/config.yaml
-    {%- if config.live_reload and not config.docker_server %}
+{%- if requires_live_reload %}
 COPY ./control /control
-# Step 1: Ensure a usable python{{ control_python_version }} is available
-RUN if python{{ control_python_version }} -c 'import venv; venv.EnvBuilder(with_pip=True).create("/tmp/__probe_env")' > /dev/null 2>&1; then \
-      echo "Using system python{{ control_python_version }}"; \
-      python{{ control_python_version }} -m pip install --upgrade pip virtualenv && \
-      python{{ control_python_version }} -m virtualenv /control/.env; \
-    else \
-      echo "Installing Miniforge-based python{{ control_python_version }}..."; \
-      curl -fsSL -o miniforge.sh https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh && \
-      bash miniforge.sh -b -p /opt/conda-control && \
-      rm miniforge.sh && \
-      /opt/conda-control/bin/conda create -y -p /opt/conda-control/envs/py{{ control_python_version }} python={{ control_python_version }} && \
-      /opt/conda-control/envs/py{{ control_python_version }}/bin/pip install --upgrade pip virtualenv && \
-      /opt/conda-control/bin/conda clean -afy && \
-      rm -rf /opt/conda-control/pkgs && \
-      /opt/conda-control/envs/py{{ control_python_version }}/bin/python -m virtualenv /control/.env; \
-    fi
-
-# Step 2: Install requirements into the freshly created venv
+
+# NB(nikhil): Copy the virtual env, and system dependencies for the
+# python version specific to the control server. We need to run ldconfig
+# so dynamic linking works.
+COPY --from=control-server-venv /usr/local/lib/ /usr/local/lib/
+COPY --from=control-server-venv /control/.env /control/.env
+RUN ldconfig
+
 RUN /control/.env/bin/pip install -r /control/requirements.txt
-    {%- endif %} {#- endif config.live_reload and not config.docker_server (for control server setup) #}
+{%- endif %} {#- endif requires_live_reload #}
+
 {%- if model_dir_exists %}
 COPY ./{{ config.model_module_dir }} /app/model
 {%- endif %} {#- endif model_dir_exists #}
 {% endblock %} {#- endblock app_copy #}
 
-
 {% block run %}
     {%- if config.docker_server %}
 RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
@@ -136,7 +133,7 @@ COPY supervisord.conf {{ supervisor_config_path }}
 ENV SUPERVISOR_SERVER_URL="{{ supervisor_server_url }}"
 ENV SERVER_START_CMD="supervisord -c {{ supervisor_config_path }}"
 ENTRYPOINT ["supervisord", "-c", "{{ supervisor_config_path }}"]
-    {%- elif config.live_reload %} {#- elif config.live_reload #}
+    {%- elif requires_live_reload %} {#- elif requires_live_reload #}
 ENV HASH_TRUSS="{{ truss_hash }}"
 ENV CONTROL_SERVER_PORT="8080"
 ENV INFERENCE_SERVER_PORT="8090"

truss/templates/server/requirements.txt

@@ -9,7 +9,7 @@ joblib>=1.2.0
 loguru>=0.7.2
 msgpack-numpy==0.4.8
 msgpack==1.1.0  # Numpy/msgpack versions are finniky (1.0.2 breaks), double check when changing.
-numpy>=1.23.5
+numpy>=1.23.5,<2.0 # 2.x versions of numpy are incompatible with many libraries we depend on.
 opentelemetry-api>=1.25.0
 opentelemetry-exporter-otlp>=1.25.0
 opentelemetry-sdk>=1.25.0

truss/tests/test_truss_handle.py

@@ -8,9 +8,11 @@
 from python_on_whales.exceptions import DockerException
 from tenacity import RetryError
 
+from truss.base.constants import SUPPORTED_PYTHON_VERSIONS
 from truss.base.custom_types import Example
 from truss.base.errors import ContainerIsDownError, ContainerNotFoundError
 from truss.base.truss_config import map_local_to_supported_python_version
+from truss.contexts.image_builder.util import TRUSS_BASE_IMAGE_VERSION_TAG
 from truss.local.local_config_handler import LocalConfigHandler
 from truss.templates.control.control.helpers.custom_types import (
     Action,
@@ -100,15 +102,33 @@ def test_build_docker_image(custom_model_truss_dir_with_pre_and_post):
     assert image.repo_tags[0] == tag
 
 
+def _generate_base_image_variations(
+    base_template="baseten/truss-server-base:{}-{}",
+    base_py3_path="/usr/local/bin/python3",
+    gpu_py3_path="/usr/bin/python3",
+    python_versions=SUPPORTED_PYTHON_VERSIONS,
+    version_tag=TRUSS_BASE_IMAGE_VERSION_TAG,
+):
+    variations = []
+    for version in python_versions:
+        variations.append(
+            (base_template.format(version, version_tag), base_py3_path, False)
+        )
+
+        # NB(nikhil): Python 3.8 base images currently don't have support for development models
+        # on GPUs.
+        fail = False if version != "3.8" else True
+        variations.append(
+            (base_template.format(f"{version}-gpu", version_tag), gpu_py3_path, fail)
+        )
+    return variations
+
+
 @pytest.mark.integration
 @pytest.mark.parametrize(
     "base_image, path, expected_fail",
-    [
-        ("baseten/truss-server-base:3.9-v0.9.0", "/usr/local/bin/python3", False),
-        ("baseten/truss-server-base:3.10-v0.9.0", "/usr/local/bin/python3", False),
-        ("baseten/truss-server-base:3.11-v0.9.0", "/usr/local/bin/python3", False),
-        ("baseten/truss-server-base:3.12-v0.9.0", "/usr/local/bin/python3", False),
-        ("baseten/truss-server-base:3.13-v0.9.0", "/usr/local/bin/python3", False),
+    _generate_base_image_variations()
+    + [
         ("python:3.8", "/usr/local/bin/python3", False),
         ("python:3.10", "/usr/local/bin/python3", False),
         ("python:3.11", "/usr/local/bin/python3", False),