update HPCCM recipe and build script to receive repo branch as arg an…

…d also remove unecessary pip installations

Signed-off-by: João Seródio <[email protected]>

build/build_container.sh

@@ -10,6 +10,7 @@ RAPIDS_VERSION="23.08"
 CUDA_VERSION="12.0"
 UBUNTU_VERSION="20.04"
 PYTHON_VERSION="3.10"
+REPO_BRANCH="main"
 DOCKERFILE_DIR=docker/
 RECIPE_ONLY=0
 
@@ -29,6 +30,7 @@ function print_help() {
     echo "    --cuda-version CUDA_VERSION      Defines which version of CUDA will be used (default='$CUDA_VERSION')."
     echo "    --os-version OS_VERSION          Defines which version of the container will be used (default='$UBUNTU_VERSION')."
     echo "    --python-version PYTHON_VERSION  Defines which version of the python interpreter will be used (default='$PYTHON_VERSION')."
+    echo "    --repo-branch REPO_BRANCH        Defines which branch from the DASF repositary to install (default='$REPO_BRANCH')."
     echo "    --format FORMAT                  Select the container backend for this build."
     echo "                                     Use 'docker' for Docker images."
     echo "                                     Use 'singularity' for SIF images (default='$FORMAT')."
@@ -70,6 +72,11 @@ while [[ $# -gt 0 ]]; do
       shift
       shift
       ;;
+    --repo-branch)
+      REPO_BRANCH="$2"
+      shift
+      shift
+      ;;
     --format)
       FORMAT="$2"
       shift
@@ -151,6 +158,7 @@ hpccm --recipe hpccm/build_docker.py \
                 cuda-version=$CUDA_VERSION \
                 ubuntu-version=$UBUNTU_VERSION \
                 python-version=$PYTHON_VERSION \
+                repo-branch=$REPO_BRANCH \
       --format $FORMAT > $DOCKERFILE_DIR/$OUTPUT_FILE
 
 if [ $RECIPE_ONLY -eq 1 ]; then

build/hpccm/build_docker.py

@@ -19,6 +19,7 @@ def str2bool(string):
 rapidsai_version = USERARG.get('rapids-version', '23.08')
 ubuntu_version = USERARG.get('ubuntu-version', '20.04')
 python_version = USERARG.get('python-version', '3.10')
+repo_branch = USERARG.get('repo-branch', 'main')
 
 if python_version:
     python_version = f"-py{python_version}"
@@ -54,7 +55,7 @@ def str2bool(string):
     # Install NVIDIA NSight packages
     package_list += ["nsight-compute"]
 
-pip_package_install = "pip3 install --extra-index-url https://test.pypi.org/simple/ XPySom-dask git+https://github.com/discovery-unicamp/dasf-core.git"
+pip_package_install = f"pip3 install -U git+https://github.com/discovery-unicamp/dasf-core.git@{repo_branch}"
 
 if device_target.lower() == "cpu":
     packages_list.extend(["python3-pip"])
@@ -70,22 +71,18 @@ def str2bool(string):
         Stage0 += apt_get(keys=apt_keys, ospackages=packages_list)
         Stage0 += apt_get(ospackages=packages_list)
 
-    pip_package_install = ("%s cupy_xarray" % pip_package_install)
+    Stage0 += shell(commands=["pip install --no-dependencies cupy_xarray==0.1.3"]) # this avoids CuPy being installed twice and taking too long (installation process doesn't find CuPy because its named cupy_cuda12x)
 
     if is_devel:
         pip_package_install = ("%s %s" % (pip_package_install, "git+https://github.com/cupy/cupy.git"))
     else:
-        pip_package_install = ("%s %s" % (pip_package_install, "cupy==13.0.0b1"))
-        Stage0 += shell(commands=["rm -r /usr/local/lib/python3.10/dist-packages/cupy_cuda12x-12.0.0b3.dist-info"]) # not the best solution but it works
+        pip_package_install = ("%s %s" % (pip_package_install, "cupy-cuda12x==13.2.0"))
 
 
 Stage0 += shell(commands=["pip3 install pip --upgrade"])
 
 Stage0 += shell(commands=[pip_package_install])
 
-# TODO: fix numpy issue with version 1.24 and other fixed reqs
-Stage0 += shell(commands=["pip install \"numpy<1.24\" \"pandas==1.5.2\" bokeh==2.4.3 \"protobuf<=3.20.1\" \"charset-normalizer<3.0\" \"tornado<6.2\""])
-
 if is_devel:
     Stage0 += shell(commands=["wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && bash Miniconda3-latest-Linux-x86_64.sh -b && cp /root/miniconda3/bin/conda /usr/bin/conda"])