add nvidia compatibility

osrf · Dec 8, 2024 · bd7fcb5 · bd7fcb5
1 parent 06dbabb
commit bd7fcb5
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/src/rocker/extensions.py b/src/rocker/extensions.py
@@ -506,6 +506,9 @@ def get_preamble(self, cliargs):
         return ''
 
     def get_docker_args(self, cliargs):
+        # The gpu ids will be set in the nvidia extension, if the nvidia argument is passed.
+        if cliargs.get('nvidia', None):
+            return ''
         args = ''
         gpus = cliargs.get('gpus', None)
         if gpus:

diff --git a/src/rocker/nvidia_extension.py b/src/rocker/nvidia_extension.py
@@ -152,12 +152,15 @@ def get_snippet(self, cliargs):
 
     def get_docker_args(self, cliargs):
         force_flag = cliargs.get('nvidia', None)
+        gpus_ids_flag = cliargs.get('gpus', None)
+        if gpus_ids_flag is None:
+            gpus_ids_flag = 'all'
         if force_flag == 'runtime':
             return "  --runtime=nvidia"
         if force_flag == 'gpus':
-            return "  --gpus all"
+            return f"  --gpus {gpus_ids_flag}"
         if get_docker_version() >= Version("19.03"):
-            return "  --gpus all"
+            return f"  --gpus {gpus_ids_flag}"
         return "  --runtime=nvidia"
 
     @staticmethod