Build neuron image

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

Author: oOraph

dockerfiles/pytorch/Dockerfile

@@ -1,6 +1,9 @@
 ARG BASE_IMAGE=nvidia/cuda:12.1.0-devel-ubuntu22.04
 
 FROM $BASE_IMAGE
+
+ARG NEURONX=0
+
 SHELL ["/bin/bash", "-c"]
 
 LABEL maintainer="Hugging Face"
@@ -31,12 +34,12 @@ RUN apt-get update && \
         libsndfile1-dev \
         ffmpeg \
     && apt-get clean autoremove --yes \
-    && rm -rf /var/lib/{apt,dpkg,cache,log}
+    && rm -rf /var/lib/{apt,cache,log}
+
 # Copying only necessary files as filtered by .dockerignore
 COPY . .
 
-# install wheel and setuptools
-RUN pip install --no-cache-dir -U pip ".[torch, st, diffusers]"
+RUN if [[ "$NEURONX" == "1" ]];then /bin/bash -c "./dockerfiles/pytorch/neuronx.sh";else pip install --no-cache-dir -U pip ".[torch, st, diffusers]";fi
 
 # copy application
 COPY src/huggingface_inference_toolkit huggingface_inference_toolkit
@@ -45,4 +48,6 @@ COPY src/huggingface_inference_toolkit/webservice_starlette.py webservice_starle
 # copy entrypoint and change permissions
 COPY --chmod=0755  scripts/entrypoint.sh entrypoint.sh
 
-ENTRYPOINT ["bash", "-c", "./entrypoint.sh"]
+ENTRYPOINT ["bash", "-c", "./entrypoint.sh"]
+
+RUN apt-get update && apt-get install -y vim

dockerfiles/pytorch/neuronx.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+
+set -e
+
+# Install system prerequisites
+apt-get update -y \
+ && apt-get install -y --no-install-recommends \
+    gnupg2 \
+    wget
+
+echo "deb https://apt.repos.neuron.amazonaws.com jammy main" > /etc/apt/sources.list.d/neuron.list
+wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | apt-key add -
+
+apt-get update -y \
+ && apt-get install -y --no-install-recommends \
+    aws-neuronx-dkms=2.* \
+    aws-neuronx-collectives=2.* \
+    aws-neuronx-runtime-lib=2.* \
+    aws-neuronx-tools=2.*
+
+pip install -U pip
+
+pip3 install neuronx-cc==2.12.68.0 \
+    torch-neuronx==1.13.1.1.13.1 \
+    transformers-neuronx==0.9.474 \
+    --extra-index-url=https://pip.repos.neuron.amazonaws.com
+
+pip3 install optimum[neuronx,diffusers]
+
+pip install ".[st,torch1]"

makefile

@@ -26,5 +26,8 @@ inference-pytorch-gpu:
 inference-pytorch-cpu:
 	docker build --build-arg="BASE_IMAGE=ubuntu:22.04" -f dockerfiles/pytorch/Dockerfile -t integration-test-pytorch:cpu .
 
+inference-pytorch-neuron:
+	docker build --build-arg=BASE_IMAGE=ubuntu:22.04 --build-arg=NEURONX=1 -f dockerfiles/pytorch/Dockerfile  -t integration-test-pytorch:neuron .
+
 stop-all:
-	docker stop $$(docker ps -a -q) && docker container prune --force
+	docker stop $$(docker ps -a -q) && docker container prune --force

setup.py

@@ -38,6 +38,8 @@
 extras["st"] = ["sentence_transformers==2.4.0"]
 extras["diffusers"] = ["diffusers==0.26.3", "accelerate==0.27.2"]
 extras["torch"] = ["torch==2.2.0", "torchvision", "torchaudio"]
+# For neuronx
+extras["torch1"] = ["torch==1.13.1", "torchvision", "torchaudio"]
 extras["tensorflow"] = ["tensorflow"]
 extras["test"] = [
     "pytest==7.2.1",