Merge branch 'main' into ayushmishra/batch-score-oss

Azure · Sep 18, 2024 · c3b6ebf · c3b6ebf
2 parents 8c29298 + 2bcd175
commit c3b6ebf
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/assets/training/automl/environments/ai-ml-automl-dnn-forecasting-gpu/context/Dockerfile b/assets/training/automl/environments/ai-ml-automl-dnn-forecasting-gpu/context/Dockerfile
@@ -29,8 +29,7 @@ RUN conda create -p $AZUREML_CONDA_ENVIRONMENT_PATH \
     # Install pytorch separately to speed up image build
     -c conda-forge -c pytorch -c anaconda && \
     conda install -p $AZUREML_CONDA_ENVIRONMENT_PATH \
-    pytorch=1.13.1 \
-    pytorch-cuda=11.6 \
+    pytorch=2.2.2 \
     -c pytorch -c nvidia -y && \
     # end conda dependencies
     conda run -p $AZUREML_CONDA_ENVIRONMENT_PATH && \
@@ -54,10 +53,9 @@ RUN pip install \
                 azureml-train-automl-runtime=={{latest-pypi-version}} \
                 azureml-dataset-runtime=={{latest-pypi-version}} \
                 azureml-train-automl=={{latest-pypi-version}} \
-                azureml-contrib-automl-dnn-forecasting==1.57.0 \
+                azureml-contrib-automl-dnn-forecasting=={{latest-pypi-version}} \
                 'azure-identity>=1.16.1' \
                 'inference-schema' \
-                'horovod==0.28.1' \
                 'xgboost==1.5.2' \
                 'cryptography>=42.0.5' \
                 'requests>=2.31.0' \
@@ -67,4 +65,6 @@ RUN pip install \
                 'https://aka.ms/automl-resources/packages/en_core_web_sm-3.7.1.tar.gz' \
                 'py-cpuinfo==5.0.0'
                 # end pypi dependencies
+
+RUN HOROVOD_WITH_PYTORCH=1 pip install --no-cache-dir git+https://github.com/horovod/horovod@3a31d933a13c7c885b8a673f4172b17914ad334d
 # end pip install
diff --git a/assets/training/automl/environments/ai-ml-automl-dnn-gpu/context/Dockerfile b/assets/training/automl/environments/ai-ml-automl-dnn-gpu/context/Dockerfile
@@ -57,7 +57,7 @@ SHELL ["conda", "run", "-p", "$AZUREML_CONDA_ENVIRONMENT_PATH", "/bin/bash", "-c
 # end conda create
 
 # Conda installs from extra channels. Separate these to avoid solver OOMs.
-RUN conda install pytorch=1.13.1 pytorch-cuda=11.6 -c pytorch -c nvidia && conda clean -a -y
+RUN conda install pytorch=2.2.2 -c pytorch -c nvidia && conda clean -a -y
 # end conda create
 
 # begin pip install
@@ -79,7 +79,6 @@ RUN pip install \
                 'azureml-model-management-sdk==1.0.1b6.post1' \
                 'azure-identity>=1.16.1' \
                 'inference-schema' \
-                'horovod==0.28.1' \
                 'prophet==1.1.4' \
                 'pytorch-transformers==1.0.0' \
                 'spacy==3.7.4' \
@@ -91,6 +90,8 @@ RUN pip install \
                 'cryptography>=42.0.5'
                 # end pypi dependencies
 
+RUN HOROVOD_WITH_PYTORCH=1 pip install --no-cache-dir git+https://github.com/horovod/horovod@3a31d933a13c7c885b8a673f4172b17914ad334d
+
 # end pip install
 
 ENV LD_LIBRARY_PATH $AZUREML_CONDA_ENVIRONMENT_PATH/lib:$LD_LIBRARY_PATH