check for AMD GPU device and rocm installation with rocminfo

Signed-off-by: Vicky Tsang <[email protected]>

Author: vickytsang

monai/deploy/packager/util.py

@@ -23,6 +23,7 @@
 from monai.deploy.exceptions import WrongValueError
 from monai.deploy.packager.constants import DefaultValues
 from monai.deploy.packager.templates import Template
+from monai.deploy.utils.deviceutil import has_rocm
 from monai.deploy.utils.fileutil import checksum
 from monai.deploy.utils.importutil import dist_module_path, dist_requires, get_application
 from monai.deploy.utils.spinner import ProgressSpinner
@@ -42,11 +43,10 @@ def verify_base_image(base_image: str) -> str:
         str: returns string identifier of the dockerfile template to build MAP
         if valid base image provided, returns empty string otherwise
     """
-    import torch
-    if "AMD" not in torch.cuda.get_device_name(0):
-        valid_prefixes = {"nvcr.io/nvidia/cuda": "ubuntu", "nvcr.io/nvidia/pytorch": "pytorch"}
-    else:
+    if has_rocm():
         valid_prefixes = {"rocm": "ubuntu", "rocm/pytorch": "pytorch"}
+    else:
+        valid_prefixes = {"nvcr.io/nvidia/cuda": "ubuntu", "nvcr.io/nvidia/pytorch": "pytorch"}
 
     for prefix, template in valid_prefixes.items():
         if prefix in base_image:
@@ -93,18 +93,17 @@ def initialize_args(args: Namespace) -> Dict:
     if args.base:
         dockerfile_type = verify_base_image(args.base)
         if not dockerfile_type:
-            import torch
-            if "AMD" not in torch.cuda.get_device_name(0):
+            if has_rocm():
                 logger.error(
                     "Provided base image '{}' is not supported \n \
-                            Please provide a Cuda or Pytorch image from https://ngc.nvidia.com/ (nvcr.io/nvidia)".format(
+                            Please provide a ROCm or Pytorch image from https://hub.docker.com/r/rocm/pytorch".format(
                         args.base
                     )
                 )
             else:
                 logger.error(
                     "Provided base image '{}' is not supported \n \
-                            Please provide a ROCm or Pytorch image from https://hub.docker.com/r/rocm/pytorch".format(
+                            Please provide a Cuda or Pytorch image from https://ngc.nvidia.com/ (nvcr.io/nvidia)".format(
                         args.base
                     )
                 )

monai/deploy/runner/runner.py

@@ -20,6 +20,7 @@
 from typing import Tuple
 
 from monai.deploy.runner.utils import get_requested_gpus, run_cmd, verify_image
+from monai.deploy.utils.deviceutil import has_rocm
 
 logger = logging.getLogger("app_runner")
 
@@ -87,9 +88,8 @@ def run_app(map_name: str, input_path: Path, output_path: Path, app_info: dict,
     # Use nvidia-docker if GPU resources are requested
     requested_gpus = get_requested_gpus(pkg_info)
     if requested_gpus > 0:
-        import torch
-        if "AMD" not in torch.cuda.get_device_name(0):
-           cmd = "nvidia-docker run --rm -a STDERR"
+        if not has_rocm():
+            cmd = "nvidia-docker run --rm -a STDERR"
 
     if not quiet:
         cmd += " -a STDOUT"
@@ -162,14 +162,13 @@ def pkg_specific_dependency_verification(pkg_info: dict) -> bool:
     """
     requested_gpus = get_requested_gpus(pkg_info)
     if requested_gpus > 0:
-        import torch
-        if "AMD" not in torch.cuda.get_device_name(0):
-        # check for nvidia-docker
-        prog = "nvidia-docker"
-        logger.info('--> Verifying if "%s" is installed...\n', prog)
-        if not shutil.which(prog):
-            logger.error('ERROR: "%s" not installed, please install nvidia-docker.', prog)
-            return False
+        if not has_rocm():
+            # check for nvidia-docker
+            prog = "nvidia-docker"
+            logger.info('--> Verifying if "%s" is installed...\n', prog)
+            if not shutil.which(prog):
+                logger.error('ERROR: "%s" not installed, please install nvidia-docker.', prog)
+                return False
 
     return True
 

monai/deploy/utils/deviceutil.py

@@ -0,0 +1,33 @@
+# Copyright 2023 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import subprocess
+
+
+def has_rocm():
+    """Return True if ROCm is installed and GPU device is detected.
+
+    Args:
+
+    Returns:
+        True if ROCm is installed and GPU device is detected, otherwise False.
+    """
+    cmd = "rocminfo"
+    process = subprocess.run([cmd], stdout=subprocess.PIPE)
+    for line_in in process.stdout.decode().splitlines():
+        if "Device Type" in line_in and "GPU" in line_in:
+            return True
+
+    return False
+
+
+if __name__ == "__main__":
+    print(has_rocm())