From 2238386ad07c10fd13d3fd64c65857c8cd23b5cd Mon Sep 17 00:00:00 2001
From: Rhett Ying <85214957+Rhett-Ying@users.noreply.github.com>
Date: Mon, 27 Feb 2023 17:46:28 +0800
Subject: [PATCH] [CI] enable more options for conda env creation (#5386)

* [CI] enable more options for conda env creation

* update
---
 script/create_dev_conda_env.sh | 102 ++++++++++++++++++++++++---------
 script/dgl_dev.yml.template    |   2 +-
 2 files changed, 77 insertions(+), 27 deletions(-)

diff --git a/script/create_dev_conda_env.sh b/script/create_dev_conda_env.sh
index c08c402feeec..9d980f7ddc8d 100644
--- a/script/create_dev_conda_env.sh
+++ b/script/create_dev_conda_env.sh
@@ -1,7 +1,8 @@
 #!/bin/bash
 
 readonly CUDA_VERSIONS="10.2,11.3,11.6,11.7"
-readonly TORCH_VERSION="1.12.0"
+readonly TORCH_VERSION="1.12"
+readonly PYTHON_VERSION="3.7"
 
 usage() {
 cat << EOF
@@ -9,14 +10,20 @@ usage: bash $0 OPTIONS
 examples:
   bash $0 -c
   bash $0 -g 11.7
+  bash $0 -g 11.7 -p 3.8
 
 Create a developement environment for DGL developers.
 
 OPTIONS:
   -h           Show this message.
   -c           Create dev environment in CPU mode.
+  -d           Only display environment YAML file instead of creating it.
+  -f           Force creation of environment (removing a previously existing 
+               environment of the same name).
   -g           Create dev environment in GPU mode with specified CUDA version,
                supported: ${CUDA_VERSIONS}.
+  -o           Save environment YAML file to specified path.
+  -p           Create dev environment based on specified python version.
   -s           Run silently which indicates always 'yes' for any confirmation.
 EOF
 }
@@ -41,28 +48,50 @@ confirm() {
 }
 
 # Parse flags.
-while getopts "cg:hs" flag; do
-  if [[ ${flag} == "c" ]]; then
-    cpu=1
-  elif [[ ${flag} == "g" ]]; then
-    gpu=${OPTARG}
-  elif [[ ${flag} == "h" ]]; then
-    usage
-    exit 0
-  elif [[ ${flag} == "s" ]]; then
-    always_yes=1
-  else
-    usage
-    exit 1
-  fi
+while getopts "cdfg:ho:p:s" flag; do
+  case "${flag}" in
+    c)
+      cpu=1
+      ;;
+    d)
+      dry_run=1
+      ;;
+    f)
+      force_create=1
+      ;;
+    g)
+      cuda_version=${OPTARG}
+      ;;
+    h)
+      usage
+      exit 0
+      ;;
+    o)
+      output_path=${OPTARG}
+      ;;
+    p)
+      python_version=${OPTARG}
+      ;;
+    s)
+      always_yes=1
+      ;;
+    :)
+      echo "Error: -${OPTARG} requires an argument."
+      exit 1
+      ;;
+    *)
+      usage
+      exit 1
+      ;;
+  esac
 done
 
-if [[ -n ${gpu} && ${cpu} -eq 1 ]]; then
+if [[ -n ${cuda_version} && ${cpu} -eq 1 ]]; then
   echo "Only one mode can be specified."
   exit 1
 fi
 
-if [[ -z ${gpu} && -z ${cpu} ]]; then
+if [[ -z ${cuda_version} && -z ${cpu} ]]; then
   usage
   exit 1
 fi
@@ -74,8 +103,8 @@ if [[ ${cpu} -eq 1 ]]; then
 fi
 
 # Set up GPU mode.
-if [[ -n ${gpu} ]]; then
-  if ! validate ${CUDA_VERSIONS} ${gpu}; then
+if [[ -n ${cuda_version} ]]; then
+  if ! validate ${CUDA_VERSIONS} ${cuda_version}; then
     echo "Error: Invalid CUDA version."
     usage
     exit 1
@@ -84,10 +113,15 @@ if [[ -n ${gpu} ]]; then
   echo "Confirm the installed CUDA version matches the specified one."
   [[ -n "${always_yes}" ]] || confirm
 
-  torchversion=${TORCH_VERSION}"+cu"${gpu//[-._]/}
+  torchversion=${TORCH_VERSION}"+cu"${cuda_version//[-._]/}
   name="dgl-dev-gpu"
 fi
 
+# Set python version.
+if [[ -z "${python_version}" ]]; then
+  python_version=${PYTHON_VERSION}
+fi
+
 echo "Confirm you are excuting the script from your DGL root directory."
 echo "Current working directory: ${PWD}"
 [[ -n "${always_yes}" ]] || confirm
@@ -95,20 +129,36 @@ echo "Current working directory: ${PWD}"
 # Prepare the conda environment yaml file.
 rand=$(echo "${RANDOM}" | md5sum | head -c 20)
 mkdir -p /tmp/${rand}
-cp script/dgl_dev.yml.template /tmp/${rand}/dgl_dev.yml
-sed -i "s|__NAME__|${name}|g" /tmp/${rand}/dgl_dev.yml
-sed -i "s|__TORCH_VERSION__|${torchversion}|g" /tmp/${rand}/dgl_dev.yml
-sed -i "s|__DGL_HOME__|${PWD}|g" /tmp/${rand}/dgl_dev.yml
+yaml_path="/tmp/${rand}/dgl_dev.yml"
+cp script/dgl_dev.yml.template ${yaml_path}
+sed -i "s|__NAME__|${name}|g" ${yaml_path}
+sed -i "s|__PYTHON_VERSION__|${python_version}|g" ${yaml_path}
+sed -i "s|__TORCH_VERSION__|${torchversion}|g" ${yaml_path}
+sed -i "s|__DGL_HOME__|${PWD}|g" ${yaml_path}
 
 # Ask for final confirmation.
 echo "--------------------------------------------------"
-cat /tmp/${rand}/dgl_dev.yml
+cat ${yaml_path}
 echo "--------------------------------------------------"
 echo "Create a conda enviroment with the config?"
 [[ -n "${always_yes}" ]] || confirm
 
+# Save YAML file to specified path
+if [[ -n "${output_path}" ]]; then
+  cp ${yaml_path} ${output_path}
+  echo "Environment YAML file has been saved to ${output_path}."
+fi
+
 # Create conda environment.
-conda env create -f /tmp/${rand}/dgl_dev.yml
+if [[ -z "${dry_run}" ]]; then
+  conda_args=""
+  if [[ -n "${force_create}" ]]; then
+    conda_args="${conda_args} --force "
+  fi
+  conda env create -f ${yaml_path} ${conda_args}
+else
+  echo "Running in dry mode, so creation of conda environment is skipped."
+fi
 
 # Clean up created tmp conda environment yaml file.
 rm -rf /tmp/${rand}
diff --git a/script/dgl_dev.yml.template b/script/dgl_dev.yml.template
index 3d7178ee3c97..1cfadefa9dfd 100644
--- a/script/dgl_dev.yml.template
+++ b/script/dgl_dev.yml.template
@@ -1,6 +1,6 @@
 name: __NAME__
 dependencies:
-  - python=3.7.0
+  - python=__PYTHON_VERSION__
   - pip
   - pip:
     - --find-links https://download.pytorch.org/whl/torch_stable.html