Update python dependencies for neuronx 2.1.2

huggingface/pytorch/training/docker/2.1/py3/sdk2.20.0/Dockerfile.neuronx

@@ -13,7 +13,6 @@ ARG OPTIMUM_NEURON_VERSION=0.0.25
 ARG TRANSFORMERS_VERSION
 ARG DATASETS_VERSION
 ARG GEVENT_VERSION=24.10.3
-ARG GAUTH_VERSION=1.35.0
 ARG PYTHON=python3
 
  # install Hugging Face libraries and its dependencies
@@ -24,16 +23,23 @@ RUN pip install --no-cache-dir \
 	datasets==${DATASETS_VERSION} \
 	optimum-neuron==${OPTIMUM_NEURON_VERSION} \
     peft \
-	google-auth==${GAUTH_VERSION} \
 	gevent==${GEVENT_VERSION}
 
 # Pin numpy to version required by neuronx-cc
 # Update Pillow and urllib version to fix high and critical vulnerabilities
 RUN pip install -U \
+	"tensorboard>=2.11.0" \
 	"numpy>=1.24.3,<=1.25.2" \
 	"numba==0.58.1" \
 	"Pillow==10.3.0" \
 	"requests<2.32.0" \
+	Jinja2 \
+	mlflow \
+	tornado \
+	"awscli<2" \
+	boto3 \
+	botocore \
+	google-auth \
 	"urllib3>=1.26.17,<1.27"
 
 RUN apt-get update \