[WIP] enable AMD GPU

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

Author: vickytsang

monai/deploy/packager/util.py

@@ -42,7 +42,11 @@ 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
     """
-    valid_prefixes = {"nvcr.io/nvidia/cuda": "ubuntu", "nvcr.io/nvidia/pytorch": "pytorch"}
+    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:
+        valid_prefixes = {"rocm": "ubuntu", "rocm/pytorch": "pytorch"}
 
     for prefix, template in valid_prefixes.items():
         if prefix in base_image:
@@ -89,12 +93,22 @@ def initialize_args(args: Namespace) -> Dict:
     if args.base:
         dockerfile_type = verify_base_image(args.base)
         if not dockerfile_type:
-            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(
-                    args.base
+            import torch
+            if "AMD" not in torch.cuda.get_device_name(0):
+                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(
+                        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(
+                        args.base
+                    )
+                )
+
             sys.exit(1)
 
     processed_args["dockerfile_type"] = dockerfile_type if args.base else DefaultValues.DOCKERFILE_TYPE

monai/deploy/runner/runner.py

@@ -87,7 +87,9 @@ 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:
-        cmd = "nvidia-docker run --rm -a STDERR"
+        import torch
+        if "AMD" not in torch.cuda.get_device_name(0):
+           cmd = "nvidia-docker run --rm -a STDERR"
 
     if not quiet:
         cmd += " -a STDOUT"
@@ -160,6 +162,8 @@ 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)