From 3fa158d774b279d5c7f4d0ac1535a4c1eba2ec43 Mon Sep 17 00:00:00 2001 From: indirected <46082152+indirected@users.noreply.github.com> Date: Thu, 29 Dec 2022 18:49:59 +0330 Subject: [PATCH 1/4] added easy gpu sharing to DockerSpawner --- dockerspawner/dockerspawner.py | 42 ++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/dockerspawner/dockerspawner.py b/dockerspawner/dockerspawner.py index 41ab5789..5f759b4f 100644 --- a/dockerspawner/dockerspawner.py +++ b/dockerspawner/dockerspawner.py @@ -1113,6 +1113,42 @@ def _cast_cpu_limit(self, proposal): """cast cpu_limit to a float if it's callable""" return self._eval_if_callable(proposal.value) + gpu_ids = Union( + [Callable(), Unicode(allow_none=True)], + help=""" + GPU IDs to share with containers + + Default is None which means no GPUs are shared + + Acceptable values are None or a string containing list of + GPU integer IDs or uuids or the literal 'all' + + Examples: + c.DockerSpawner.gpu_ids = None # Will not share GPUs with containers + c.DockerSpawner.gpu_ids = 'all' # Shares all the GPUs with containers + c.DockerSpawner.gpu_ids = '0' # Shares GPU of ID 0 with containers + c.DockerSpawner.gpu_ids = '0,1,2' # Shares GPUs of IDs 0, 1 and 2 with containers + + Alternatively, you can pass a callable that takes the spawner as + the only argument and returns one of the above acceptable values: + + def per_user_gpu_ids(spawner): + username = spawner.user.name + gpu_assign = {'alice': '0', 'bob': '1,2'} + return gpu_assign.get(username, None) + c.DockerSpawner.gpu_ids = per_user_gpu_ids + + Note that before using this config option, you have to: + 1- Install the Nvidia Container Toolkit and make sure your docker is able to run containers with gpu + 2- Use an image with a CUDA version supported by your hosts GPU driver + """, + ).tag(config=True) + + @validate('gpu_ids') + def _cast_gpu_ids(self, proposal): + """cast gpu_ids to a string if it's callable""" + return self._eval_if_callable(proposal.value) + async def create_object(self): """Create the container/service object""" @@ -1135,6 +1171,7 @@ async def create_object(self): binds=self.volume_binds, links=self.links, mounts=self.mount_binds, + device_requests = [], ) if getattr(self, "mem_limit", None) is not None: @@ -1149,6 +1186,11 @@ async def create_object(self): ) host_config["cpu_quota"] = int(self.cpu_limit * cpu_period) + if getattr(self, "gpu_ids", None) is not None: + host_config['device_requests'].append( + docker.types.DeviceRequest(device_ids=[f"{self.gpu_ids}"], capabilities=[['gpu']]) + ) + if not self.use_internal_ip: host_config["port_bindings"] = {self.port: (self.host_ip,)} host_config.update(self._render_templates(self.extra_host_config)) From 0ebfc6d50181a223367202afa30d178debe4b65f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 29 Dec 2022 15:25:28 +0000 Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- dockerspawner/dockerspawner.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/dockerspawner/dockerspawner.py b/dockerspawner/dockerspawner.py index 5f759b4f..cc4a4fef 100644 --- a/dockerspawner/dockerspawner.py +++ b/dockerspawner/dockerspawner.py @@ -1171,7 +1171,7 @@ async def create_object(self): binds=self.volume_binds, links=self.links, mounts=self.mount_binds, - device_requests = [], + device_requests=[], ) if getattr(self, "mem_limit", None) is not None: @@ -1188,7 +1188,9 @@ async def create_object(self): if getattr(self, "gpu_ids", None) is not None: host_config['device_requests'].append( - docker.types.DeviceRequest(device_ids=[f"{self.gpu_ids}"], capabilities=[['gpu']]) + docker.types.DeviceRequest( + device_ids=[f"{self.gpu_ids}"], capabilities=[['gpu']] + ) ) if not self.use_internal_ip: From 661ee8c98362e7e3a7f21e454385113e66f43243 Mon Sep 17 00:00:00 2001 From: indirected <46082152+indirected@users.noreply.github.com> Date: Thu, 29 Dec 2022 19:41:26 +0330 Subject: [PATCH 3/4] updated readme --- README.md | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/README.md b/README.md index aedcf505..6abb1f73 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,33 @@ # DockerSpawner +## Fork README +This fork adds a config option to easily share specific GPUs with containers spawned by the DockerSpawner. I have created a pull request but until it is merged, you can install the package by: +```bash +pip install git+https://github.com/indirected/dockerspawner.git +``` +### Here is how you use the config option: +```python +c.DockerSpawner.gpu_ids = None # Will not share GPUs with containers +c.DockerSpawner.gpu_ids = 'all' # Shares all the GPUs with containers +c.DockerSpawner.gpu_ids = '0' # Shares GPU of ID 0 with containers +c.DockerSpawner.gpu_ids = '0,1,2' # Shares GPUs of IDs 0, 1 and 2 with containers +``` +Alternatively, you can pass a callable that takes the spawner as +the only argument and returns one of the above acceptable values: +```python +def per_user_gpu_ids(spawner): + username = spawner.user.name + gpu_assign = {'alice': '0', 'bob': '1,2'} + return gpu_assign.get(username, None) +c.DockerSpawner.gpu_ids = per_user_gpu_ids +``` +Note that before using this config option, you have to: + 1. Install the Nvidia Container Toolkit and make sure your docker is able to run containers with gpu + 2. Use an image with a CUDA version supported by your host's GPU driver + +--- + +## Original README [![GitHub Workflow Status - Test](https://img.shields.io/github/workflow/status/jupyterhub/dockerspawner/Tests?logo=github&label=tests)](https://github.com/jupyterhub/dockerspawner/actions) [![Latest PyPI version](https://img.shields.io/pypi/v/dockerspawner?logo=pypi)](https://pypi.org/project/dockerspawner/) [![Documentation build status](https://img.shields.io/readthedocs/jupyterhub?logo=read-the-docs)](https://jupyterhub-dockerspawner.readthedocs.org/en/latest/) From 934deb8bb3e1e26ab4e8e4ff186ae3e2db524ef1 Mon Sep 17 00:00:00 2001 From: indirected <46082152+indirected@users.noreply.github.com> Date: Thu, 29 Dec 2022 19:46:45 +0330 Subject: [PATCH 4/4] Revert "updated readme" This reverts commit 661ee8c98362e7e3a7f21e454385113e66f43243. --- README.md | 28 ---------------------------- 1 file changed, 28 deletions(-) diff --git a/README.md b/README.md index 6abb1f73..aedcf505 100644 --- a/README.md +++ b/README.md @@ -1,33 +1,5 @@ # DockerSpawner -## Fork README -This fork adds a config option to easily share specific GPUs with containers spawned by the DockerSpawner. I have created a pull request but until it is merged, you can install the package by: -```bash -pip install git+https://github.com/indirected/dockerspawner.git -``` -### Here is how you use the config option: -```python -c.DockerSpawner.gpu_ids = None # Will not share GPUs with containers -c.DockerSpawner.gpu_ids = 'all' # Shares all the GPUs with containers -c.DockerSpawner.gpu_ids = '0' # Shares GPU of ID 0 with containers -c.DockerSpawner.gpu_ids = '0,1,2' # Shares GPUs of IDs 0, 1 and 2 with containers -``` -Alternatively, you can pass a callable that takes the spawner as -the only argument and returns one of the above acceptable values: -```python -def per_user_gpu_ids(spawner): - username = spawner.user.name - gpu_assign = {'alice': '0', 'bob': '1,2'} - return gpu_assign.get(username, None) -c.DockerSpawner.gpu_ids = per_user_gpu_ids -``` -Note that before using this config option, you have to: - 1. Install the Nvidia Container Toolkit and make sure your docker is able to run containers with gpu - 2. Use an image with a CUDA version supported by your host's GPU driver - ---- - -## Original README [![GitHub Workflow Status - Test](https://img.shields.io/github/workflow/status/jupyterhub/dockerspawner/Tests?logo=github&label=tests)](https://github.com/jupyterhub/dockerspawner/actions) [![Latest PyPI version](https://img.shields.io/pypi/v/dockerspawner?logo=pypi)](https://pypi.org/project/dockerspawner/) [![Documentation build status](https://img.shields.io/readthedocs/jupyterhub?logo=read-the-docs)](https://jupyterhub-dockerspawner.readthedocs.org/en/latest/)