Add maisi bundle (#612)

Fixes # .

### Description
A few sentences describing the changes proposed in this pull request.

### Status
**Ready**

### Please ensure all the checkboxes:
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Codeformat tests passed locally by running `./runtests.sh
--codeformat`.
- [ ] In-line docstrings updated.
- [ ] Update `version` and `changelog` in `metadata.json` if changing an
existing bundle.
- [ ] Please ensure the naming rules in config files meet our
requirements (please refer to: `CONTRIBUTING.md`).
- [ ] Ensure versions of packages such as `monai`, `pytorch` and `numpy`
are correct in `metadata.json`.
- [ ] Descriptions should be consistent with the content, such as
`eval_metrics` of the provided weights and TorchScript modules.
- [ ] Files larger than 25MB are excluded and replaced by providing
download links in `large_file.yml`.
- [ ] Avoid using path that contains personal information within config
files (such as use `/home/your_name/` for `"bundle_root"`).

---------

Signed-off-by: Yiheng Wang <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: yiheng-wang-nv

.github/workflows/code-format-check.yml

@@ -17,10 +17,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v2
-    - name: Set up Python 3.9
+    - name: Set up Python 3.10
       uses: actions/setup-python@v2
       with:
-        python-version: 3.9
+        python-version: 3.10.14
     - name: cache weekly timestamp
       id: pip-cache
       run: |

.github/workflows/premerge-cpu.yml

@@ -17,10 +17,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v2
-    - name: Set up Python 3.9
+    - name: Set up Python 3.10
       uses: actions/setup-python@v2
       with:
-        python-version: 3.9
+        python-version: 3.10.14
     - name: cache weekly timestamp
       id: pip-cache
       run: |

ci/bundle_custom_data.py

@@ -19,11 +19,12 @@
     "pathology_nuclei_segmentation_classification",
     "brats_mri_generative_diffusion",
     "brats_mri_axial_slices_generative_diffusion",
+    "maisi_ct_generative",
 ]
 
 # This list is used for our CI tests to determine whether a bundle contains the preferred files.
 # If a bundle does not have any of the preferred files, please add the bundle name into the list.
-exclude_verify_preferred_files_list = []
+exclude_verify_preferred_files_list = ["maisi_ct_generative"]
 
 # This list is used for our CI tests to determine whether a bundle needs to be tested with
 # the `verify_export_torchscript` function in `verify_bundle.py`.
@@ -37,12 +38,13 @@
     "mednist_reg",
     "brats_mri_axial_slices_generative_diffusion",
     "vista3d",
+    "maisi_ct_generative",
 ]
 
 # This dict is used for our CI tests to install required dependencies that cannot be installed by `pip install` directly.
 # If a bundle has this kind of dependencies, please add the bundle name (key), and the path of the install script (value)
 # into the dict.
-install_dependency_dict = {}
+install_dependency_dict = {"maisi_ct_generative": "ci/install_scripts/install_maisi_ct_generative_dependency.sh"}
 
 # This list is used for our CI tests to determine whether a bundle supports TensorRT export. Related
 # test will be employed for bundles in the dict.

ci/get_bundle_requirements.py

@@ -19,6 +19,8 @@
 
 ALLOW_MONAI_RC = os.environ.get("ALLOW_MONAI_RC", "false").lower() in ("true", "1", "t", "y", "yes")
 
+SPECIAL_LIB_LIST = ["xformers"]
+
 
 def increment_version(version):
     """
@@ -75,10 +77,16 @@ def get_requirements(bundle, models_path):
         if "numpy_version" in metadata.keys():
             numpy_version = metadata["numpy_version"]
             libs.append(f"numpy=={numpy_version}")
-        if "optional_packages_version" in metadata.keys():
-            optional_dict = metadata["optional_packages_version"]
-            for name, version in optional_dict.items():
-                libs.append(f"{name}=={version}")
+        for package_key in ["optional_packages_version", "required_packages_version"]:
+            if package_key in metadata.keys():
+                optional_dict = metadata[package_key]
+                for name, version in optional_dict.items():
+                    if name not in SPECIAL_LIB_LIST:
+                        libs.append(f"{name}=={version}")
+                    else:
+                        if "pytorch_version" in metadata.keys():
+                            # remove torch from libs
+                            libs = [lib for lib in libs if "torch" not in lib]
 
         if len(libs) > 0:
             requirements_file_name = f"requirements_{bundle}.txt"

ci/install_scripts/install_maisi_ct_generative_dependency.sh

@@ -0,0 +1 @@
+pip install --extra-index-url https://urm.nvidia.com/artifactory/api/pypi/sw-dlmed-pypi-local/simple xformers==0.0.26+622595c.d20240617

ci/run_premerge_cpu.sh

@@ -30,6 +30,17 @@ elif [[ $# -gt 1 ]]; then
     exit 1
 fi
 
+# Usually, CPU test is required, but for some bundles that are too large to run in Github Actions, we can exclude them.
+exclude_test_list=("maisi_ct_generative")
+is_excluded() {
+    for item in "${exclude_list[@]}"; do
+        if [ "$1" == "$item" ]; then
+            return 0 # Return true (0) if excluded
+        fi
+    done
+    return 1 # Return false (1) if not excluded
+}
+
 verify_bundle() {
     for dir in /opt/hostedtoolcache/*; do
         if [[ $dir != "/opt/hostedtoolcache/Python" ]]; then
@@ -52,21 +63,25 @@ verify_bundle() {
             echo $bundle_list
             for bundle in $bundle_list;
             do
-                pip install -r requirements-dev.txt
-                # get required libraries according to the bundle's metadata file
-                requirements=$(python $(pwd)/ci/get_bundle_requirements.py --b "$bundle")
-                # check if ALLOW_MONAI_RC is set to 1, if so, append --pre to the pip install command
-                if [ $ALLOW_MONAI_RC = true ]; then
-                    include_pre_release="--pre"
+                if is_excluded "$bundle"; then
+                    echo "skip '$bundle' cpu premerge tests."
                 else
-                    include_pre_release=""
-                fi
-                if [ ! -z "$requirements" ]; then
-                    echo "install required libraries for bundle: $bundle"
-                    pip install $include_pre_release -r "$requirements"
+                    pip install -r requirements-dev.txt
+                    # get required libraries according to the bundle's metadata file
+                    requirements=$(python $(pwd)/ci/get_bundle_requirements.py --b "$bundle")
+                    # check if ALLOW_MONAI_RC is set to 1, if so, append --pre to the pip install command
+                    if [ $ALLOW_MONAI_RC = true ]; then
+                        include_pre_release="--pre"
+                    else
+                        include_pre_release=""
+                    fi
+                    if [ ! -z "$requirements" ]; then
+                        echo "install required libraries for bundle: $bundle"
+                        pip install $include_pre_release -r "$requirements"
+                    fi
+                    # verify bundle
+                    python $(pwd)/ci/verify_bundle.py -b "$bundle" -m "min"  # min tests on cpu
                 fi
-                # verify bundle
-                python $(pwd)/ci/verify_bundle.py -b "$bundle" -m "min"  # min tests on cpu
             done
         else
             echo "this pull request does not change any bundles, skip verify."
@@ -81,6 +96,7 @@ case $BUILD_TYPE in
 
     all)
         echo "Run all tests..."
+        verify_bundle
         ;;
     changed)
         echo "Run changed tests..."

ci/run_premerge_gpu.sh

@@ -16,95 +16,122 @@
 #
 
 # Argument(s):
-#   BUILD_TYPE:   all/specific_test_name, tests to execute
+#   $1 - Dist flag (True/False)
+
+dist_flag=$1
 
 set -ex
-BUILD_TYPE=all
-export ALLOW_MONAI_RC=true
 
-if [[ $# -eq 1 ]]; then
-    BUILD_TYPE=$1
+export ALLOW_MONAI_RC=true
 
-elif [[ $# -gt 1 ]]; then
+if [[ $# -gt 1 ]]; then
     echo "ERROR: too many parameters are provided"
     exit 1
 fi
 
-init_pipenv() {
-    echo "initializing pip environment: $1"
-    pipenv install update pip wheel
-    pipenv install --python=3.9 -r $1
-    export PYTHONPATH=$PWD
+init_venv() {
+    if [ ! -d "model_zoo_venv" ]; then  # Check if the venv directory does not exist
+        echo "initializing pip environment: $1"
+        python -m venv model_zoo_venv
+        source model_zoo_venv/bin/activate
+        pip install --upgrade pip wheel
+        pip install -r $1
+        export PYTHONPATH=$PWD
+    else
+        echo "Virtual environment model_zoo_venv already exists. Activating..."
+        source model_zoo_venv/bin/activate
+    fi
+}
+
+remove_venv() {
+    if [ -d "model_zoo_venv" ]; then  # Check if the venv directory exists
+        echo "Removing virtual environment..."
+        deactivate 2>/dev/null || true  # Deactivate venv, ignore errors if not activated
+        rm -rf model_zoo_venv  # Remove the venv directory
+    else
+        echo "Virtual environment not found. Skipping removal."
+    fi
 }
 
-remove_pipenv() {
-    echo "removing pip environment"
-    pipenv --rm
-    rm Pipfile Pipfile.lock
+set_local_env() {
+    echo "set local pip environment: $1"
+    pip install --upgrade pip wheel
+    pip install -r $1
+    export PYTHONPATH=$PWD
 }
 
 verify_bundle() {
     echo 'Run verify bundle...'
-    init_pipenv requirements-dev.txt
     head_ref=$(git rev-parse HEAD)
     git fetch origin dev $head_ref
     # achieve all changed files in 'models'
     changes=$(git diff --name-only $head_ref origin/dev -- models)
     if [ ! -z "$changes" ]
     then
         # get all changed bundles
-        bundle_list=$(pipenv run python $(pwd)/ci/get_changed_bundle.py --f "$changes")
+        bundle_list=$(python $(pwd)/ci/get_changed_bundle.py --f "$changes")
         if [ ! -z "$bundle_list" ]
         then
-            pipenv run python $(pwd)/ci/prepare_schema.py --l "$bundle_list"
-            for bundle in $bundle_list;
-            do
-                init_pipenv requirements-dev.txt
-                # get required libraries according to the bundle's metadata file
-                requirements=$(pipenv run python $(pwd)/ci/get_bundle_requirements.py --b "$bundle")
-                # check if ALLOW_MONAI_RC is set to 1, if so, append --pre to the pip install command
-                if [ $ALLOW_MONAI_RC = true ]; then
-                    include_pre_release="--pre"
+            python $(pwd)/ci/prepare_schema.py --l "$bundle_list"
+        for bundle in $bundle_list;
+        do
+            # Check if the bundle is "maisi_ct_generative", if so, set local environment (venv cannot work with xformers)
+            if [ "$bundle" == "maisi_ct_generative" ]; then
+                echo "Special handling for maisi_ct_generative bundle"
+                set_local_env requirements-dev.txt
+            else
+                init_venv requirements-dev.txt
+            fi
+            # get required libraries according to the bundle's metadata file
+            requirements=$(python $(pwd)/ci/get_bundle_requirements.py --b "$bundle")
+            # check if ALLOW_MONAI_RC is set to 1, if so, append --pre to the pip install command
+            if [ $ALLOW_MONAI_RC = true ]; then
+                include_pre_release="--pre"
+            else
+                include_pre_release=""
+            fi
+            if [ ! -z "$requirements" ]; then
+                echo "install required libraries for bundle: $bundle"
+                pip install $include_pre_release -r "$requirements"
+            fi
+            # get extra install script if exists
+            extra_script=$(python $(pwd)/ci/get_bundle_requirements.py --b "$bundle" --get_script True)
+            if [ ! -z "$extra_script" ]; then
+                echo "install extra libraries with script: $extra_script"
+                bash $extra_script
+            fi
+            # verify bundle
+            python $(pwd)/ci/verify_bundle.py --b "$bundle"
+            # unzip data and do unit tests
+            DATA_DIR="$(pwd)/models/maisi_ct_generative/datasets"
+            ZIP_FILE="$DATA_DIR/all_masks_flexible_size_and_spacing_3000.zip"
+            UNZIP_DIR="$DATA_DIR/all_masks_flexible_size_and_spacing_3000"
+            if [ -f "$ZIP_FILE" ]; then
+                if [ ! -d "$UNZIP_DIR" ]; then
+                    echo "Unzipping files for MAISI Bundle..."
+                    unzip $ZIP_FILE -d $DATA_DIR
+                    echo "Unzipping complete."
                 else
-                    include_pre_release=""
-                fi
-                if [ ! -z "$requirements" ]; then
-                    echo "install required libraries for bundle: $bundle"
-                    pipenv install $include_pre_release -r "$requirements"
+                    echo "Unzipped content already exists, continuing..."
                 fi
-                # get extra install script if exists
-                extra_script=$(pipenv run python $(pwd)/ci/get_bundle_requirements.py --b "$bundle" --get_script True)
-                if [ ! -z "$extra_script" ]; then
-                    echo "install extra libraries with script: $extra_script"
-                    bash $extra_script
-                fi
-                # verify bundle
-                pipenv run python $(pwd)/ci/verify_bundle.py --b "$bundle"
-                # do unit tests
-                pipenv run python $(pwd)/ci/unit_tests/runner.py --b "$bundle"
-                remove_pipenv
-            done
+            fi
+            test_cmd="python $(pwd)/ci/unit_tests/runner.py --b \"$bundle\""
+            if [ "$dist_flag" = "True" ]; then
+                test_cmd="$test_cmd --dist True"
+            fi
+            eval $test_cmd
+            # if not maisi_ct_generative, remove venv
+            if [ "$bundle" != "maisi_ct_generative" ]; then
+                remove_venv
+            fi
+        done
         else
             echo "this pull request does not change any bundles, skip verify."
         fi
     else
         echo "this pull request does not change any files in 'models', skip verify."
-        remove_pipenv
+        remove_venv
     fi
 }
 
-case $BUILD_TYPE in
-
-    all)
-        echo "Run all tests..."
-        verify_bundle
-        ;;
-
-    verify_bundle)
-        verify_bundle
-        ;;
-
-    *)
-        echo "ERROR: unknown parameter: $BUILD_TYPE"
-        ;;
-esac
+verify_bundle