3737"""
3838import argparse
3939import logging
40- import glob
4140
4241from azureml .core .authentication import AzureCliAuthentication
4342from azureml .core import Workspace
@@ -146,7 +145,6 @@ def setup_persistent_compute_target(workspace, cluster_name, vm_size, max_nodes)
146145
147146def create_run_config (
148147 cpu_cluster ,
149- docker_proc_type ,
150148 add_gpu_dependencies ,
151149 add_spark_dependencies ,
152150 conda_pkg_jdk ,
@@ -165,7 +163,6 @@ def create_run_config(
165163 the following:
166164 - Reco_cpu_test
167165 - Reco_gpu_test
168- docker_proc_type (str) : processor type, cpu or gpu
169166 add_gpu_dependencies (bool) : True if gpu packages should be
170167 added to the conda environment, else False
171168 add_spark_dependencies (bool) : True if PySpark packages should be
@@ -179,7 +176,39 @@ def create_run_config(
179176 run_azuremlcompute = RunConfiguration ()
180177 run_azuremlcompute .target = cpu_cluster
181178 run_azuremlcompute .environment .docker .enabled = True
182- run_azuremlcompute .environment .docker .base_image = docker_proc_type
179+ if not add_gpu_dependencies :
180+ # https://github.com/Azure/AzureML-Containers/blob/master/base/cpu/openmpi4.1.0-ubuntu22.04
181+ run_azuremlcompute .environment .docker .base_image = "mcr.microsoft.com/azureml/openmpi4.1.0-ubuntu22.04"
182+ else :
183+ run_azuremlcompute .environment .docker .base_image = None
184+ # Use the latest CUDA
185+ # See
186+ # * https://learn.microsoft.com/en-us/azure/machine-learning/how-to-train-with-custom-image?view=azureml-api-1#use-a-custom-dockerfile-optional
187+ # * https://github.com/Azure/AzureML-Containers/blob/master/base/gpu/openmpi4.1.0-cuda11.8-cudnn8-ubuntu22.04
188+ run_azuremlcompute .environment .docker .base_dockerfile = r"""
189+ FROM nvcr.io/nvidia/cuda:12.3.1-devel-ubuntu22.04
190+ USER root:root
191+ ENV NVIDIA_VISIBLE_DEVICES all
192+ ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
193+ ENV LANG=C.UTF-8 LC_ALL=C.UTF-8
194+ ENV DEBIAN_FRONTEND noninteractive
195+ RUN apt-get update && \
196+ apt-get install -y wget git-all && \
197+ apt-get clean -y && \
198+ rm -rf /var/lib/apt/lists/*
199+ # Conda Environment
200+ ENV MINICONDA_VERSION py38_23.3.1-0
201+ ENV PATH /opt/miniconda/bin:$PATH
202+ ENV CONDA_PACKAGE 23.5.0
203+ RUN wget -qO /tmp/miniconda.sh https://repo.anaconda.com/miniconda/Miniconda3-${MINICONDA_VERSION}-Linux-x86_64.sh && \
204+ bash /tmp/miniconda.sh -bf -p /opt/miniconda && \
205+ conda install conda=${CONDA_PACKAGE} -y && \
206+ conda update --all -c conda-forge -y && \
207+ conda clean -ay && \
208+ rm -rf /opt/miniconda/pkgs && \
209+ rm /tmp/miniconda.sh && \
210+ find / -type d -name __pycache__ | xargs rm -rf
211+ """
183212
184213 # Use conda_dependencies.yml to create a conda environment in
185214 # the Docker image for execution
@@ -195,6 +224,7 @@ def create_run_config(
195224
196225 # install recommenders
197226 reco_extras = "dev"
227+ conda_dep .add_conda_package ("anaconda::git" )
198228 if add_gpu_dependencies and add_spark_dependencies :
199229 conda_dep .add_channel ("conda-forge" )
200230 conda_dep .add_conda_package (conda_pkg_jdk )
@@ -326,13 +356,6 @@ def create_arg_parser():
326356 default = "STANDARD_D3_V2" ,
327357 help = "Set the size of the VM either STANDARD_D3_V2" ,
328358 )
329- # cpu or gpu
330- parser .add_argument (
331- "--dockerproc" ,
332- action = "store" ,
333- default = "cpu" ,
334- help = "Base image used in docker container" ,
335- )
336359 # Azure subscription id, when used in a pipeline, it is stored in keyvault
337360 parser .add_argument (
338361 "--subid" , action = "store" , default = "123456" , help = "Azure Subscription ID"
@@ -421,16 +444,6 @@ def create_arg_parser():
421444
422445 logger = logging .getLogger ("submit_groupwise_azureml_pytest.py" )
423446 args = create_arg_parser ()
424-
425- if args .dockerproc == "cpu" :
426- from azureml .core .runconfig import DEFAULT_CPU_IMAGE
427-
428- docker_proc_type = DEFAULT_CPU_IMAGE
429- else :
430- from azureml .core .runconfig import DEFAULT_GPU_IMAGE
431-
432- docker_proc_type = DEFAULT_GPU_IMAGE
433-
434447 cli_auth = AzureCliAuthentication ()
435448
436449 workspace = setup_workspace (
@@ -450,7 +463,6 @@ def create_arg_parser():
450463
451464 run_config = create_run_config (
452465 cpu_cluster = cpu_cluster ,
453- docker_proc_type = docker_proc_type ,
454466 add_gpu_dependencies = args .add_gpu_dependencies ,
455467 add_spark_dependencies = args .add_spark_dependencies ,
456468 conda_pkg_jdk = args .conda_pkg_jdk ,
0 commit comments