Skip to content

Commit

Permalink
Add host_config_options to allow configuration of DockerTask for GPU …
Browse files Browse the repository at this point in the history
…resources (#3234)

* feat: expose host config options to DockerTask

* docs: add docstrings with links to docker api
  • Loading branch information
acmiyaguchi authored Apr 21, 2023
1 parent 18bf067 commit 154a4f7
Showing 1 changed file with 18 additions and 1 deletion.
19 changes: 18 additions & 1 deletion luigi/contrib/docker_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit 154a4f7

Please sign in to comment.