[docker] bump neuron version to 2.20.0 SDK (#2438)

serving/docker/pytorch-inf2.Dockerfile

@@ -14,12 +14,13 @@ ARG djl_version=0.30.0~SNAPSHOT
 ARG torch_version=2.1.2
 ARG torchvision_version=0.16.2
 ARG python_version=3.10
-ARG neuronsdk_version=2.19.1
-ARG torch_neuronx_version=2.1.2.2.2.0
-ARG transformers_neuronx_version=0.11.351
-ARG neuronx_distributed_version=0.8.0
-ARG neuronx_cc_version=2.14.227.0
-ARG protobuf_version=3.19.6
+ARG neuronsdk_version=2.20.0
+ARG torch_neuronx_version=2.1.2.2.3.0
+ARG transformers_neuronx_version=0.12.313
+ARG neuronx_distributed_version=0.9.0
+ARG neuronx_cc_version=2.15.128.0
+ARG neuronx_cc_stubs_version=2.15.128.0
+ARG torch_xla_version=2.1.4
 ARG transformers_version=4.43.2
 ARG accelerate_version=0.29.2
 ARG diffusers_version=0.28.2
@@ -75,7 +76,8 @@ RUN mkdir -p /opt/djl/bin && cp scripts/telemetry.sh /opt/djl/bin && \
     scripts/install_inferentia2.sh && \
     pip install accelerate==${accelerate_version} safetensors torchvision==${torchvision_version} \
     neuronx-cc==${neuronx_cc_version} torch-neuronx==${torch_neuronx_version} transformers-neuronx==${transformers_neuronx_version} \
-    neuronx_distributed==${neuronx_distributed_version} protobuf==${protobuf_version} sentencepiece jinja2 \
+    torch_xla==${torch_xla_version} neuronx-cc-stubs==${neuronx_cc_stubs_version} \
+    neuronx_distributed==${neuronx_distributed_version} protobuf sentencepiece jinja2 \
     diffusers==${diffusers_version} opencv-contrib-python-headless  Pillow --extra-index-url=https://pip.repos.neuron.amazonaws.com \
     pydantic==${pydantic_version} optimum optimum-neuron==${optimum_neuron_version} tiktoken blobfile && \
     pip install transformers==${transformers_version} ${vllm_wheel} && \

serving/docker/scripts/install_inferentia2.sh

@@ -7,6 +7,7 @@ apt-get update -y && apt-get install -y --no-install-recommends \
   git \
   gnupg2 \
   pciutils \
+  libxml2 \
   udev
 
 # Configure Linux for Neuron repository updates
@@ -15,12 +16,12 @@ echo "deb https://apt.repos.neuron.amazonaws.com ${VERSION_CODENAME} main" >/etc
 curl -L https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | apt-key add -
 
 # https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/releasecontent.html#inf2-packages
-apt-get update -y && apt-get install -y aws-neuronx-collectives=2.21.46.0* \
-  aws-neuronx-runtime-lib=2.21.41.0* \
-  aws-neuronx-tools=2.18.3.0
+apt-get update -y && apt-get install -y aws-neuronx-collectives=2.22.26.0* \
+  aws-neuronx-runtime-lib=2.22.14.0* \
+  aws-neuronx-tools=2.19.0.0
 
 # TODO: Remove this hack after aws-neuronx-dkms install no longer throws an error, this bypasses the `set -ex`
 #       exit criteria. The package is installed and functional after running, just throws an error on install.
-apt-get install -y aws-neuronx-dkms=2.17.17.0 || echo "Installed aws-neuronx-dkms with errors"
+apt-get install -y aws-neuronx-dkms=2.18.12.0 || echo "Installed aws-neuronx-dkms with errors"
 
 export PATH=/opt/aws/neuron/bin:$PATH