From 3c87248bacc0a986818a2f2e3bde01e39b3f0103 Mon Sep 17 00:00:00 2001 From: Samrat Thapa Date: Thu, 5 Dec 2024 22:01:23 +0900 Subject: [PATCH 1/5] shm size and gpus --- setup.py | 2 ++ src/rocker/extensions.py | 49 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/setup.py b/setup.py index 2cc1d9d..804aeb1 100644 --- a/setup.py +++ b/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', diff --git a/src/rocker/extensions.py b/src/rocker/extensions.py index 3e22b71..ecc1d39 100644 --- a/src/rocker/extensions.py +++ b/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") \ No newline at end of file From 3de50e0bbdde4a72dc879b4bbe228e45bd75ce5d Mon Sep 17 00:00:00 2001 From: Samrat Thapa Date: Sat, 7 Dec 2024 15:17:10 +0900 Subject: [PATCH 2/5] typo fix --- setup.py | 4 ++-- src/rocker/extensions.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/setup.py b/setup.py index 804aeb1..a4d6479 100644 --- a/setup.py +++ b/setup.py @@ -51,9 +51,8 @@ 'env = rocker.extensions:Environment', 'expose = rocker.extensions:Expose', 'git = rocker.git_extension:Git', - 'group_add = rocker.extensions:GroupAdd', - 'shm_size = rocker.extensions:ShmSize', 'gpus = rocker.extensions:Gpus', + 'group_add = rocker.extensions:GroupAdd', 'home = rocker.extensions:HomeDir', 'hostname = rocker.extensions:Hostname', 'ipc = rocker.extensions:Ipc', @@ -64,6 +63,7 @@ 'privileged = rocker.extensions:Privileged', 'pulse = rocker.extensions:PulseAudio', 'rmw = rocker.rmw_extension:RMW', + 'shm_size = rocker.extensions:ShmSize', 'ssh = rocker.ssh_extension:Ssh', 'ulimit = rocker.ulimit_extension:Ulimit', 'user = rocker.extensions:User', diff --git a/src/rocker/extensions.py b/src/rocker/extensions.py index ecc1d39..c6de230 100644 --- a/src/rocker/extensions.py +++ b/src/rocker/extensions.py @@ -507,9 +507,9 @@ def get_preamble(self, cliargs): def get_docker_args(self, cliargs): args = '' - shm_size = cliargs.get('gpus', None) - if shm_size: - args += f' --gpus {shm_size} ' + gpus = cliargs.get('gpus', None) + if gpus: + args += f' --gpus {gpus} ' return args @staticmethod From 06dbabbf16e7901b97a4a1e049a3cc28e8531c0a Mon Sep 17 00:00:00 2001 From: Samrat Thapa Date: Sat, 7 Dec 2024 16:15:07 +0900 Subject: [PATCH 3/5] added tests --- src/rocker/extensions.py | 2 +- test/test_extension.py | 70 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 71 insertions(+), 1 deletion(-) diff --git a/src/rocker/extensions.py b/src/rocker/extensions.py index c6de230..ff8dddf 100644 --- a/src/rocker/extensions.py +++ b/src/rocker/extensions.py @@ -484,7 +484,7 @@ def get_docker_args(self, cliargs): args = '' shm_size = cliargs.get('shm_size', None) if shm_size: - args += f' --shm-size={shm_size} ' + args += f' --shm-size {shm_size} ' return args @staticmethod diff --git a/test/test_extension.py b/test/test_extension.py index d264c35..26558ab 100644 --- a/test/test_extension.py +++ b/test/test_extension.py @@ -617,3 +617,73 @@ def test_group_add_extension(self): args = p.get_docker_args(mock_cliargs) self.assertIn('--group-add sudo', args) self.assertIn('--group-add docker', args) + +class ShmSizeExtensionTest(unittest.TestCase): + + def setUp(self): + # Work around interference between empy Interpreter + # stdout proxy and test runner. empy installs a proxy on stdout + # to be able to capture the information. + # And the test runner creates a new stdout object for each test. + # This breaks empy as it assumes that the proxy has persistent + # between instances of the Interpreter class + # empy will error with the exception + # "em.Error: interpreter stdout proxy lost" + em.Interpreter._wasProxyInstalled = False + + @pytest.mark.docker + def test_shm_size_extension(self): + plugins = list_plugins() + shm_size_plugin = plugins['shm_size'] + self.assertEqual(shm_size_plugin.get_name(), 'shm_size') + + p = shm_size_plugin() + self.assertTrue(plugin_load_parser_correctly(shm_size_plugin)) + + mock_cliargs = {} + self.assertEqual(p.get_snippet(mock_cliargs), '') + self.assertEqual(p.get_preamble(mock_cliargs), '') + args = p.get_docker_args(mock_cliargs) + self.assertNotIn('--shm-size', args) + + mock_cliargs = {'shm_size': '12g'} + args = p.get_docker_args(mock_cliargs) + self.assertIn('--shm-size 12g', args) + +class GpusExtensionTest(unittest.TestCase): + + def setUp(self): + # Work around interference between empy Interpreter + # stdout proxy and test runner. empy installs a proxy on stdout + # to be able to capture the information. + # And the test runner creates a new stdout object for each test. + # This breaks empy as it assumes that the proxy has persistent + # between instances of the Interpreter class + # empy will error with the exception + # "em.Error: interpreter stdout proxy lost" + em.Interpreter._wasProxyInstalled = False + + @pytest.mark.docker + def test_gpus_extension(self): + plugins = list_plugins() + gpus_plugin = plugins['gpus'] + self.assertEqual(gpus_plugin.get_name(), 'gpus') + + p = gpus_plugin() + self.assertTrue(plugin_load_parser_correctly(gpus_plugin)) + + # Test when no GPUs are specified + mock_cliargs = {} + self.assertEqual(p.get_snippet(mock_cliargs), '') + self.assertEqual(p.get_preamble(mock_cliargs), '') + args = p.get_docker_args(mock_cliargs) + self.assertNotIn('--gpus', args) + + # Test when GPUs are specified + mock_cliargs = {'gpus': 'all'} + args = p.get_docker_args(mock_cliargs) + self.assertIn('--gpus all', args) + + mock_cliargs = {'gpus': '0,1'} + args = p.get_docker_args(mock_cliargs) + self.assertIn('--gpus 0,1', args) From bd7fcb5c6704f2897b8ccc93ee41d5f42f8c95fe Mon Sep 17 00:00:00 2001 From: Samrat Thapa Date: Sun, 8 Dec 2024 22:32:04 +0900 Subject: [PATCH 4/5] add nvidia compatibility --- src/rocker/extensions.py | 3 +++ src/rocker/nvidia_extension.py | 7 +++++-- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/src/rocker/extensions.py b/src/rocker/extensions.py index ff8dddf..3a0474c 100644 --- 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 index 3644c22..deedeb9 100644 --- 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 From 5c2eba8c7ffb4998266f9eaf47f538cbdf4adaf9 Mon Sep 17 00:00:00 2001 From: Samrat Thapa Date: Thu, 12 Dec 2024 15:05:19 +0900 Subject: [PATCH 5/5] move gpus extension --- setup.py | 2 +- src/rocker/extensions.py | 30 +------------------------ src/rocker/nvidia_extension.py | 27 +++++++++++++++++++++++ test/test_extension.py | 40 +--------------------------------- test/test_nvidia.py | 38 ++++++++++++++++++++++++++++++++ 5 files changed, 68 insertions(+), 69 deletions(-) diff --git a/setup.py b/setup.py index a4d6479..41d904f 100644 --- a/setup.py +++ b/setup.py @@ -51,7 +51,7 @@ 'env = rocker.extensions:Environment', 'expose = rocker.extensions:Expose', 'git = rocker.git_extension:Git', - 'gpus = rocker.extensions:Gpus', + 'gpus = rocker.nvidia_extension:Gpus', 'group_add = rocker.extensions:GroupAdd', 'home = rocker.extensions:HomeDir', 'hostname = rocker.extensions:Hostname', diff --git a/src/rocker/extensions.py b/src/rocker/extensions.py index 3a0474c..bee9363 100644 --- a/src/rocker/extensions.py +++ b/src/rocker/extensions.py @@ -491,32 +491,4 @@ def get_docker_args(self, cliargs): 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): - # 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: - args += f' --gpus {gpus} ' - 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") \ No newline at end of file + help="Set the size of the shared memory for the container (e.g., 512m, 1g).") \ No newline at end of file diff --git a/src/rocker/nvidia_extension.py b/src/rocker/nvidia_extension.py index deedeb9..e4e47e3 100644 --- a/src/rocker/nvidia_extension.py +++ b/src/rocker/nvidia_extension.py @@ -239,3 +239,30 @@ def register_arguments(parser, defaults): action='store_true', default=defaults.get('cuda', None), help="Install cuda and nvidia-cuda-dev into the container") + +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): + # 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: + args += f' --gpus {gpus} ' + 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") \ No newline at end of file diff --git a/test/test_extension.py b/test/test_extension.py index 26558ab..1ad88d8 100644 --- a/test/test_extension.py +++ b/test/test_extension.py @@ -648,42 +648,4 @@ def test_shm_size_extension(self): mock_cliargs = {'shm_size': '12g'} args = p.get_docker_args(mock_cliargs) - self.assertIn('--shm-size 12g', args) - -class GpusExtensionTest(unittest.TestCase): - - def setUp(self): - # Work around interference between empy Interpreter - # stdout proxy and test runner. empy installs a proxy on stdout - # to be able to capture the information. - # And the test runner creates a new stdout object for each test. - # This breaks empy as it assumes that the proxy has persistent - # between instances of the Interpreter class - # empy will error with the exception - # "em.Error: interpreter stdout proxy lost" - em.Interpreter._wasProxyInstalled = False - - @pytest.mark.docker - def test_gpus_extension(self): - plugins = list_plugins() - gpus_plugin = plugins['gpus'] - self.assertEqual(gpus_plugin.get_name(), 'gpus') - - p = gpus_plugin() - self.assertTrue(plugin_load_parser_correctly(gpus_plugin)) - - # Test when no GPUs are specified - mock_cliargs = {} - self.assertEqual(p.get_snippet(mock_cliargs), '') - self.assertEqual(p.get_preamble(mock_cliargs), '') - args = p.get_docker_args(mock_cliargs) - self.assertNotIn('--gpus', args) - - # Test when GPUs are specified - mock_cliargs = {'gpus': 'all'} - args = p.get_docker_args(mock_cliargs) - self.assertIn('--gpus all', args) - - mock_cliargs = {'gpus': '0,1'} - args = p.get_docker_args(mock_cliargs) - self.assertIn('--gpus 0,1', args) + self.assertIn('--shm-size 12g', args) \ No newline at end of file diff --git a/test/test_nvidia.py b/test/test_nvidia.py index 6fb7e63..9cfd2f5 100644 --- a/test/test_nvidia.py +++ b/test/test_nvidia.py @@ -328,3 +328,41 @@ def test_cuda_env_subs(self): with self.assertRaises(SystemExit) as cm: p.get_environment_subs(mock_cliargs) self.assertEqual(cm.exception.code, 1) + +class GpusExtensionTest(unittest.TestCase): + + def setUp(self): + # Work around interference between empy Interpreter + # stdout proxy and test runner. empy installs a proxy on stdout + # to be able to capture the information. + # And the test runner creates a new stdout object for each test. + # This breaks empy as it assumes that the proxy has persistent + # between instances of the Interpreter class + # empy will error with the exception + # "em.Error: interpreter stdout proxy lost" + em.Interpreter._wasProxyInstalled = False + + @pytest.mark.docker + def test_gpus_extension(self): + plugins = list_plugins() + gpus_plugin = plugins['gpus'] + self.assertEqual(gpus_plugin.get_name(), 'gpus') + + p = gpus_plugin() + self.assertTrue(plugin_load_parser_correctly(gpus_plugin)) + + # Test when no GPUs are specified + mock_cliargs = {} + self.assertEqual(p.get_snippet(mock_cliargs), '') + self.assertEqual(p.get_preamble(mock_cliargs), '') + args = p.get_docker_args(mock_cliargs) + self.assertNotIn('--gpus', args) + + # Test when GPUs are specified + mock_cliargs = {'gpus': 'all'} + args = p.get_docker_args(mock_cliargs) + self.assertIn('--gpus all', args) + + mock_cliargs = {'gpus': '0,1'} + args = p.get_docker_args(mock_cliargs) + self.assertIn('--gpus 0,1', args)