inf2: docker image update

- + fix broken dep install (kenlm)

Signed-off-by: Raphael Glon <[email protected]>

Author: oOraph

dockerfiles/pytorch/Dockerfile.inf2

@@ -1,12 +1,8 @@
 # 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:22.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
-
 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
@@ -15,9 +11,9 @@ ARG NEURONX_TOOLS_VERSION=2.19.0.0
 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 \
@@ -34,62 +30,40 @@ RUN apt-get update \
     wget \
     unzip \
     zlib1g-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
 
-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 \
-    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