Merge branch 'main' into bump-versions-and-release

Author: alvarobartt

.github/workflows/docker-build-action.yaml

@@ -28,7 +28,7 @@ on:
 
 jobs:
   buildx:
-    runs-on: 
+    runs-on:
       group: aws-highmemory-32-plus-priv
     steps:
       - name: Check out
@@ -56,6 +56,6 @@ jobs:
           context: ${{ inputs.context }}
           build-args: ${{ inputs.build_args }}
           target: base
-          file:  ${{ inputs.context }}/${{ inputs.dockerfile }}
+          outputs: type=image,compression=zstd,force-compression=true,push=true
+          file: ${{ inputs.context }}/${{ inputs.dockerfile }}
           tags: ${{ inputs.repository }}/${{ inputs.image }}:sha-${{ env.GITHUB_SHA_SHORT }},${{ inputs.repository }}/${{ inputs.image }}:latest
-    

dockerfiles/pytorch/Dockerfile

@@ -16,13 +16,13 @@ RUN apt-get update && \
     apt-get install -y \
     build-essential \
     bzip2 \
+    cmake \
     curl \
     git \
     git-lfs \
     tar \
     gcc \
     g++ \
-    cmake \
     libprotobuf-dev \
     protobuf-compiler \
     python3.11 \
@@ -44,14 +44,17 @@ RUN curl -O https://bootstrap.pypa.io/get-pip.py && \
     python get-pip.py && \
     rm get-pip.py
 
-# install wheel and setuptools
+# Upgrade pip
+RUN pip install --no-cache-dir --upgrade pip
+
+# Install wheel and setuptools
 RUN pip install --no-cache-dir --upgrade pip ".[torch,st,diffusers]"
 
-# copy application
+# Copy application
 COPY src/huggingface_inference_toolkit huggingface_inference_toolkit
 COPY src/huggingface_inference_toolkit/webservice_starlette.py webservice_starlette.py
 
-# copy entrypoint and change permissions
+# Copy entrypoint and change permissions
 COPY --chmod=0755 scripts/entrypoint.sh entrypoint.sh
 
 ENTRYPOINT ["bash", "-c", "./entrypoint.sh"]

dockerfiles/pytorch/Dockerfile.inf2

@@ -1,115 +1,69 @@
 # Build based on https://github.com/aws/deep-learning-containers/blob/master/huggingface/pytorch/inference/docker/2.1/py3/sdk2.18.0/Dockerfile.neuronx
-FROM ubuntu:20.04 as base
+FROM ubuntu:22.04 AS base
 
 LABEL maintainer="Hugging Face"
 
-ARG PYTHON=python3.10
-ARG PYTHON_VERSION=3.10.12
-ARG MAMBA_VERSION=23.1.0-4
-
-# Neuron SDK components version numbers
-ARG NEURONX_FRAMEWORK_VERSION=2.1.2.2.1.0
-ARG NEURONX_DISTRIBUTED_VERSION=0.7.0
-ARG NEURONX_CC_VERSION=2.13.66.0
-ARG NEURONX_TRANSFORMERS_VERSION=0.10.0.21
-ARG NEURONX_COLLECTIVES_LIB_VERSION=2.20.22.0-c101c322e
-ARG NEURONX_RUNTIME_LIB_VERSION=2.20.22.0-1b3ca6425
-ARG NEURONX_TOOLS_VERSION=2.17.1.0
+ARG NEURONX_COLLECTIVES_LIB_VERSION=2.22.33.0-d2128d1aa
+ARG NEURONX_RUNTIME_LIB_VERSION=2.22.19.0-5856c0b42
+ARG NEURONX_TOOLS_VERSION=2.19.0.0
 
 # HF ARGS
-ARG OPTIMUM_NEURON_VERSION=0.0.23
+ARG OPTIMUM_NEURON_VERSION=0.0.28
 
 # See http://bugs.python.org/issue19846
-ENV LANG C.UTF-8
-ENV LD_LIBRARY_PATH /opt/aws/neuron/lib:/lib/x86_64-linux-gnu:/opt/conda/lib/:$LD_LIBRARY_PATH
-ENV PATH /opt/conda/bin:/opt/aws/neuron/bin:$PATH
+ENV LANG=C.UTF-8
+ENV LD_LIBRARY_PATH=/opt/aws/neuron/lib:/lib/x86_64-linux-gnu:/opt/conda/lib/:$LD_LIBRARY_PATH
+ENV PATH=/opt/aws/neuron/bin:$PATH
 
 RUN apt-get update \
  && apt-get upgrade -y \
  && apt-get install -y --no-install-recommends software-properties-common \
- && add-apt-repository ppa:openjdk-r/ppa \
  && apt-get update \
  && apt-get install -y --no-install-recommends \
     build-essential \
     apt-transport-https \
     ca-certificates \
     cmake \
     curl \
-    emacs \
     git \
     jq \
-    libgl1-mesa-glx \
-    libsm6 \
-    libxext6 \
-    libxrender-dev \
-    openjdk-11-jdk \
-    vim \
     wget \
     unzip \
     zlib1g-dev \
-    libcap-dev \
-    gpg-agent \
- && rm -rf /var/lib/apt/lists/* \
- && rm -rf /tmp/tmp* \
- && apt-get clean
+    gpg-agent
 
-RUN echo "deb https://apt.repos.neuron.amazonaws.com focal main" > /etc/apt/sources.list.d/neuron.list
+RUN echo "deb https://apt.repos.neuron.amazonaws.com jammy main" > /etc/apt/sources.list.d/neuron.list
 RUN wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | apt-key add -
 
 # Install Neuronx tools
 RUN apt-get update \
  && apt-get install -y \
     aws-neuronx-tools=$NEURONX_TOOLS_VERSION \
     aws-neuronx-collectives=$NEURONX_COLLECTIVES_LIB_VERSION \
-    aws-neuronx-runtime-lib=$NEURONX_RUNTIME_LIB_VERSION \
- && rm -rf /var/lib/apt/lists/* \
+    aws-neuronx-runtime-lib=$NEURONX_RUNTIME_LIB_VERSION
+
+RUN apt-get install -y \
+    python3 \
+    python3-pip \
+    python-is-python3
+
+RUN rm -rf /var/lib/apt/lists/* \
  && rm -rf /tmp/tmp* \
  && apt-get clean
 
-# https://github.com/docker-library/openjdk/issues/261 https://github.com/docker-library/openjdk/pull/263/files
-RUN keytool -importkeystore -srckeystore /etc/ssl/certs/java/cacerts -destkeystore /etc/ssl/certs/java/cacerts.jks -deststoretype JKS -srcstorepass changeit -deststorepass changeit -noprompt; \
-    mv /etc/ssl/certs/java/cacerts.jks /etc/ssl/certs/java/cacerts; \
-    /var/lib/dpkg/info/ca-certificates-java.postinst configure;
-
-RUN curl -L -o ~/mambaforge.sh https://github.com/conda-forge/miniforge/releases/download/${MAMBA_VERSION}/Mambaforge-${MAMBA_VERSION}-Linux-x86_64.sh \
- && chmod +x ~/mambaforge.sh \
- && ~/mambaforge.sh -b -p /opt/conda \
- && rm ~/mambaforge.sh \
- && /opt/conda/bin/conda update -y conda \
- && /opt/conda/bin/conda install -c conda-forge -y \
-    python=$PYTHON_VERSION \
-    pyopenssl \
-    cython \
-    mkl-include \
-    mkl \
-    botocore \
-    parso \
-    scipy \
-    typing \
-    # Below 2 are included in miniconda base, but not mamba so need to install
-    conda-content-trust \
-    charset-normalizer \
- && /opt/conda/bin/conda update -y conda \
- && /opt/conda/bin/conda clean -ya
-
-RUN conda install -c conda-forge \
-    scikit-learn \
-    h5py \
-    requests \
- && conda clean -ya \
- && pip install --upgrade pip --trusted-host pypi.org --trusted-host files.pythonhosted.org \
- && ln -s /opt/conda/bin/pip /usr/local/bin/pip3 \
- && pip install --no-cache-dir "protobuf>=3.18.3,<4" setuptools==69.5.1 packaging
-
+RUN pip install --no-cache-dir "protobuf>=3.18.3,<4" setuptools==69.5.1 packaging
+ 
 WORKDIR /
 
 # install Hugging Face libraries and its dependencies
-RUN pip install --extra-index-url https://pip.repos.neuron.amazonaws.com --no-cache-dir optimum-neuron[neuronx]==${OPTIMUM_NEURON_VERSION}  \
+RUN pip install --extra-index-url https://pip.repos.neuron.amazonaws.com --no-cache-dir optimum-neuron[neuronx]==${OPTIMUM_NEURON_VERSION} \
  && pip install --no-deps --no-cache-dir -U torchvision==0.16.*
 
+# FIXME
+RUN pip install --extra-index-url https://pip.repos.neuron.amazonaws.com git+https://github.com/huggingface/optimum-neuron.git@5237fb0ada643ba471f60ed3a5d2eef3b66e8e59
 
 COPY . .
-# install wheel and setuptools
+
 RUN pip install --no-cache-dir -U pip ".[st]"
 
 # copy application
@@ -118,5 +72,7 @@ COPY src/huggingface_inference_toolkit/webservice_starlette.py webservice_starle
 
 # copy entrypoint and change permissions
 COPY --chmod=0755  scripts/entrypoint.sh entrypoint.sh
+COPY --chmod=0755  scripts/inf2_env.py inf2_env.py
+COPY --chmod=0755  scripts/inf2_entrypoint.sh inf2_entrypoint.sh
 
-ENTRYPOINT ["bash", "-c", "./entrypoint.sh"]
+ENTRYPOINT ["bash", "-c", "./inf2_entrypoint.sh"]

scripts/entrypoint.sh

@@ -1,14 +1,55 @@
 #!/bin/bash
 
-# Define the default port
+# Set the default port
 PORT=5000
 
 # Check if AIP_MODE is set and adjust the port for Vertex AI
 if [[ ! -z "${AIP_MODE}" ]]; then
     PORT=${AIP_HTTP_PORT}
 fi
 
-# Check if HF_MODEL_DIR is set and if not skip installing custom dependencies
+# Check that only one of HF_MODEL_ID or HF_MODEL_DIR is provided
+if [[ ! -z "${HF_MODEL_ID}" && ! -z "${HF_MODEL_DIR}" ]]; then
+    echo "Error: Both HF_MODEL_ID and HF_MODEL_DIR are set. Please provide only one."
+    exit 1
+elif [[ -z "${HF_MODEL_ID}" && -z "${HF_MODEL_DIR}" ]]; then
+    echo "Error: Neither HF_MODEL_ID nor HF_MODEL_DIR is set. Please provide one of them."
+    exit 1
+fi
+
+# If HF_MODEL_ID is provided, download handler.py and requirements.txt if available
+if [[ ! -z "${HF_MODEL_ID}" ]]; then
+    filename=${HF_DEFAULT_PIPELINE_NAME:-handler.py}
+    revision=${HF_REVISION:-main}
+
+    echo "Downloading $filename for model ${HF_MODEL_ID}"
+    huggingface-cli download ${HF_MODEL_ID} "$filename" --revision "$revision" --local-dir /tmp
+
+    # Check if handler.py was downloaded successfully
+    if [ -f "/tmp/$filename" ]; then
+        echo "$filename downloaded successfully, checking if there's a requirements.txt file..."
+        rm /tmp/$filename
+
+        # Attempt to download requirements.txt
+        echo "Downloading requirements.txt for model ${HF_MODEL_ID}"
+        huggingface-cli download "${HF_MODEL_ID}" requirements.txt --revision "$revision" --local-dir /tmp
+
+        # Check if requirements.txt was downloaded successfully
+        if [ -f "/tmp/requirements.txt" ]; then
+            echo "requirements.txt downloaded successfully, now installing the dependencies..."
+            
+            # Install dependencies
+            pip install -r /tmp/requirements.txt --no-cache-dir
+            rm /tmp/requirements.txt
+        else
+            echo "${HF_MODEL_ID} with revision $revision contains a custom handler at $filename but doesn't contain a requirements.txt file, so skipping downloading and installing extra requirements from it."
+        fi
+    else
+        echo "${HF_MODEL_ID} with revision $revision doesn't contain a $filename file, so skipping download."
+    fi
+fi
+
+# If HF_MODEL_DIR is provided, check for requirements.txt and install dependencies if available
 if [[ ! -z "${HF_MODEL_DIR}" ]]; then
     # Check if requirements.txt exists and if so install dependencies
     if [ -f "${HF_MODEL_DIR}/requirements.txt" ]; then

scripts/inf2_entrypoint.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+set -e -o pipefail -u
+
+export ENV_FILEPATH=$(mktemp)
+
+trap "rm -f ${ENV_FILEPATH}" EXIT
+
+touch $ENV_FILEPATH
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+
+${SCRIPT_DIR}/inf2_env.py $@
+
+source $ENV_FILEPATH
+
+rm -f $ENV_FILEPATH
+
+exec ${SCRIPT_DIR}/entrypoint.sh $@