From 154a4f75fc592dcfe1caa65db834b0d66c266313 Mon Sep 17 00:00:00 2001 From: Anthony Miyaguchi Date: Thu, 20 Apr 2023 20:05:43 -0700 Subject: [PATCH] Add host_config_options to allow configuration of DockerTask for GPU resources (#3234) * feat: expose host config options to DockerTask * docs: add docstrings with links to docker api --- luigi/contrib/docker_runner.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/luigi/contrib/docker_runner.py b/luigi/contrib/docker_runner.py index 9655190760..7fde4d47a1 100644 --- a/luigi/contrib/docker_runner.py +++ b/luigi/contrib/docker_runner.py @@ -77,8 +77,24 @@ def command(self): def name(self): return None + @property + def host_config_options(self): + ''' + Override this to specify host_config options like gpu requests or shm + size e.g. `{"device_requests": [docker.types.DeviceRequest(count=1, capabilities=[["gpu"]])]}` + + See https://docker-py.readthedocs.io/en/stable/api.html#docker.api.container.ContainerApiMixin.create_host_config + ''' + return {} + @property def container_options(self): + ''' + Override this to specify container options like user or ports e.g. + `{"user": f"{os.getuid()}:{os.getgid()}"}` + + See https://docker-py.readthedocs.io/en/stable/api.html#docker.api.container.ContainerApiMixin.create_container + ''' return {} @property @@ -192,7 +208,8 @@ def run(self): % (self._image, self.command, self._binds)) host_config = self._client.create_host_config(binds=self._binds, - network_mode=self.network_mode) + network_mode=self.network_mode, + **self.host_config_options) container = self._client.create_container(self._image, command=self.command,