Add support for --gpus options in Rocker (shm merged separately)

setup.py

@@ -52,6 +52,8 @@
             'expose = rocker.extensions:Expose',
             'git = rocker.git_extension:Git',
             'group_add = rocker.extensions:GroupAdd',
+            'shm_size = rocker.extensions:ShmSize',
+            'gpus = rocker.extensions:Gpus',
             'home = rocker.extensions:HomeDir',
             'hostname = rocker.extensions:Hostname',
             'ipc = rocker.extensions:Ipc',

src/rocker/extensions.py

@@ -468,3 +468,52 @@ def register_arguments(parser, defaults):
             default=defaults.get(GroupAdd.get_name(), None),
             action='append',
             help="Add additional groups to join.")
+
+class ShmSize(RockerExtension):
+    @staticmethod
+    def get_name():
+        return 'shm_size'
+
+    def __init__(self):
+        self.name = ShmSize.get_name()
+
+    def get_preamble(self, cliargs):
+        return ''
+
+    def get_docker_args(self, cliargs):
+        args = ''
+        shm_size = cliargs.get('shm_size', None)
+        if shm_size:
+            args += f' --shm-size={shm_size} '
+        return args
+
+    @staticmethod
+    def register_arguments(parser, defaults={}):
+        parser.add_argument('--shm-size',
+                            default=defaults.get('shm_size', None),
+                            help="Set the size of the shared memory for the container (e.g., 512m, 1g).")
+
+
+class Gpus(RockerExtension):
+    @staticmethod
+    def get_name():
+        return 'gpus'
+
+    def __init__(self):
+        self.name = Gpus.get_name()
+
+    def get_preamble(self, cliargs):
+        return ''
+
+    def get_docker_args(self, cliargs):
+        args = ''
+        shm_size = cliargs.get('gpus', None)
+        if shm_size:
+            args += f' --gpus {shm_size} '
+        return args
+
+    @staticmethod
+    def register_arguments(parser, defaults={}):
+        parser.add_argument('--gpus',
+                            default=defaults.get('gpus', None),
+                            help="Set the indices of GPUs to use")