gpu support (#203)

* gpu support * update version Co-authored-by: Sherry Wang <[email protected]>
carsdotcom · Dec 8, 2021 · ca7a8b5 · ca7a8b5
1 parent eb23ab9
commit ca7a8b5
Show file tree

Hide file tree

Showing 7 changed files with 42 additions and 3 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,9 +1,14 @@
 # Skelebot - Changelog
 Documenting All Changes to the Skelebot Project
 
+## v1.25.0
+#### Changed
+- **Dependencies** | Added the ability to run docker with gpu
 ---
 
 ## v1.24.0
+#### Merged: 2021-11-17
+#### Released: 2021-11-17
 #### Changed
 - **Dependencies** | Added the ability for python dependencies to be specified and installed via a txt file such as `requirements.txt`
 

diff --git a/VERSION b/VERSION
@@ -1 +1 @@
-1.24.0
+1.25.0
diff --git a/skelebot/objects/config.py b/skelebot/objects/config.py
@@ -36,7 +36,8 @@ class Config(SkeleYaml):
         Optional('components'): And(dict, error='\'components\' must be a Dictionary'),
         Optional('params'): And(list, error='\'params\' must be a List'),
         Optional('commands'): And(list, error='\'commands\' must be a List'),
-        Optional('pythonVersion'): And(str, Or(*PYTHON_VERSIONS), error='\'pythonVersion\' must be one of:' + ', '.join(PYTHON_VERSIONS))
+        Optional('pythonVersion'): And(str, Or(*PYTHON_VERSIONS), error='\'pythonVersion\' must be one of:' + ', '.join(PYTHON_VERSIONS)),
+        Optional('gpu'): And(bool, error='\'gpu\' must be a Boolean')
     }, ignore_extra_keys=True)
 
     name = None
@@ -60,11 +61,13 @@ class Config(SkeleYaml):
     params = None
     commands = None
     pythonVersion = '3.6'
+    gpu = False
 
     def __init__(self, name=None, env=None, description=None, version=None, maintainer=None,
                  contact=None, host=None, language=None, baseImage=None, timezone=None,
                  primaryJob=None, primaryExe=None, ephemeral=None, dependencies=None, ignores=None,
-                 jobs=None, ports=None, components=None, params=None, commands=None, pythonVersion = '3.6'):
+                 jobs=None, ports=None, components=None, params=None, commands=None, pythonVersion = '3.6',
+                 gpu = False):
         """Initialize the config object with all provided optional attributes"""
 
         self.name = name
@@ -88,6 +91,7 @@ def __init__(self, name=None, env=None, description=None, version=None, maintain
         self.params = params if params is not None else []
         self.commands = commands if commands is not None else []
         self.pythonVersion = pythonVersion
+        self.gpu = gpu
 
     def toDict(self):
         """

diff --git a/skelebot/systems/execution/docker.py b/skelebot/systems/execution/docker.py
@@ -79,6 +79,9 @@ def run(config, command, mode, ports, mappings, task, host=None, verbose=False):
     run_name = "{image}-{job}".format(image=image, job=task)
     runCMD = DockerCommandBuilder(host=host).run(image).set_name(run_name).set_rm().set_mode(mode)
 
+    if config.gpu:
+        runCMD = runCMD.set_gpu()
+
     # Construct the port mappings
     if (ports):
         for port in ports:

diff --git a/skelebot/systems/execution/dockerCommand.py b/skelebot/systems/execution/dockerCommand.py
@@ -55,6 +55,10 @@ def set_entrypoint(self, parameters):
         self.entrypoint = True
         self.cmd += " --entrypoint"
         return self
+
+    def set_gpu(self):
+        self.cmd += " --gpus all --ipc=host"
+        return self
 
     def set_mode(self, mode):
         self.cmd += " -{}".format(mode)

diff --git a/test/test_objects_config_validate.py b/test/test_objects_config_validate.py
@@ -26,6 +26,7 @@ class TestConfigValidate(unittest.TestCase):
         'params': [1, 2],
         'commands': [],
         'pythonVersion': '3.8',
+        'gpu': True
     }
 
     def validate_error(self, attr, reset, expected):
@@ -75,6 +76,7 @@ def test_invalid(self):
         self.validate_error('commands', 123, 'a List')
         self.validate_error('host', 123, 'a String')
         self.validate_error('pythonVersion', 123, 'one of:' + ', '.join(sb.common.PYTHON_VERSIONS))
+        self.validate_error('gpu', 123, 'a Boolean')
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/test_systems_execution_docker.py b/test/test_systems_execution_docker.py
@@ -344,6 +344,27 @@ def test_run_docker_params_entrypoint(self, mock_getcwd, mock_call, mock_expandu
         sb.systems.execution.docker.run(config, command, job.mode, config.ports, job.mappings, job.name, host=None)
         mock_call.assert_called_once_with(expected, shell=True)
 
+    @mock.patch('os.path.expanduser')
+    @mock.patch('skelebot.systems.execution.docker.call')
+    @mock.patch('os.getcwd')
+    def test_run_gpu(self, mock_getcwd, mock_call, mock_expanduser):
+        folderPath = "{path}/test/files".format(path=self.path)
+        args = Namespace(version='0.1')
+
+        homePath = "{path}/test/plugins".format(path=self.path)
+        mock_expanduser.return_value = homePath
+        mock_getcwd.return_value = folderPath
+        mock_call.return_value = 0
+
+        config = sb.systems.generators.yaml.loadConfig()
+        config.gpu = True
+        job = sb.objects.job.Job(name='some_command', help='Dummy', source='echo some_command')
+        command = sb.systems.execution.commandBuilder.build(config, job, args)
+
+        expected = "docker run --name test-some_command --rm -i --gpus all --ipc=host test /bin/bash -c \"echo some_command\""
+        sb.systems.execution.docker.run(config, command, job.mode, config.ports, job.mappings, job.name, host=None)
+        mock_call.assert_called_once_with(expected, shell=True)
+
     @mock.patch('skelebot.systems.execution.docker.call')
     def test_save(self, mock_call):
         mock_call.return_value = 0