update inferentia2 image

dockerfiles/pytorch/Dockerfile.inf2

@@ -1,116 +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.12.313
 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.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
@@ -119,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/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 $@