Build GPU Variants of Current Images

[sdwalker62]

Hello everyone! First, I just want to say that at work we almost exclusively use Jupyter as our data science platform and appreciate everything that the community and the Jupyter team has done to create such a wonderful tool!

I have searched through the current list of issues and pull requests open in this repository and have not found and that match this request so I am starting a new one. Apologies if I have missed any. Our work necessitates the use of cuda-enabled containers as we primarily train deep learning models. We have made a slight modification to the makefile that allows us to re-use most of the infrastructure already in place to build the various levels of notebooks.

---

I am aware of the fantastic work being done at https://github.com/iot-salzburg/gpu-jupyter/ and suggest anyone who needs gpu-enabled containers to check it out as a first stop. This solution will work for many, and it contains some additional features which are nice to have, but it is slightly different than what we prepose.

The current docker-stacks starts with the base-notebook which is built from an Ubuntu 20.04 layer. Our preposal is to build a separate set of notebooks based on the nvidia/cuda:11.3.0-cudnn8-runtime-ubuntu20.04 image (we chose this for compatibility with Pytorch). This is similar to what the iot-salzburg repository does in their images but maintains the minimalism found in some of the base notebooks for those who are looking for a solid foundation for their own custom images. With the change to the makefiles it would be easy to build two sets of images: one for the standard notebooks as they exist now, and another for the cuda-enabled variants.

To address the selection criteria found in the documention for new features:

• The majority of the data science community working on neural networks uses GPUs to train their models and being able to pull down mirrors of the docker-stacks containers with cuda backed in would be tremendously useful.
• One of the most beautiful things about the docker-stacks project is the ability to base custom images on any of the notebook images depending on the requirements of the user. Building two sets of containers would be inline with the current design philosophy of the project.
• The complexity of this change is low. We have a simply python script that replaces lines in the base-notebook and a new section in the Makefile that is minimal and could be hooked into the current build process.
• The build times are identical to the current build times as only the base layers changes. None of the current images will change, only new tags added.
• Sustainment of the change should not be difficult at all as the nvidia/cuda containers are regularly updated by Nvidia. Any test that need to be written we can do to support the pull request.

I would be happy to answer any questions about this proposal!

[mathbunnyru]

@sdwalker62 first of all, thank you for your suggestion!
I'm sorry it took such a long time to answer you.

I see a growing interest in the GPU-related docker stacks because people ask more and more about them.

But I think right now it's not easy to add a new image. And in this case a whole set of new images.
The main reason is our complicated and slow build system.

The build times are identical to the current build times as only the base layers changes.

Right now, if we add new images, this will add time to our workflow.
Please take a look at our proposal and what our build system might look like in the future.
#1407

The complexity of this change is low. We have a simply python script that replaces lines in the base-notebook and a new section in the Makefile that is minimal and could be hooked into the current build process

I think you can actually get rid of this patching because we have a way to pass custom arguments to docker build like this:
DOCKER_BUILD_ARGS='--build-arg ROOT_CONTAINER="nvidia/cuda:12.3.1-devel-ubuntu22.04"' make build/docker-stacks-foundation.
I think it builds the image you wanted it to build.

Please, tell me, if this helps.

Update 2024.01.17: fixed the command to use more recent image and docker-stacks-foundation image (which is the root image now).

[mathbunnyru]

This works because we have ARG ROOT_CONTAINER in our docker-stacks-foundation image.
https://github.com/jupyter/docker-stacks/blob/main/images/docker-stacks-foundation/Dockerfile#L6

[mathbunnyru]

@sdwalker62 if you're still interested in this issue - we've changed our build system completely, so it should probably be easier now to add cuda based images.

[markusschmaus]

It would be awesome if such cuda based images were available.

[mathbunnyru]

I think it's actually currently possible to have cuda based images and not increase build time at all.
If someone wants to do it, please, take a look here, this is where you will need to add new steps:
https://github.com/jupyter/docker-stacks/blob/main/.github/workflows/docker.yml

[benz0li]

@sdwalker62 You may be interested in my/b-data's CUDA-enabled JupyterLab Python docker stack:

[...]

Similar projects

• iot-salzburg/gpu-jupyter
• prp/jupyter-stack

What makes this project different:

1. Multi-arch: linux/amd64, linux/arm64/v8
2. Derived from nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04
   • including development libraries and headers
3. TensortRT and TensorRT plugin libraries
   • including development libraries and headers
4. IDE: code-server next to JupyterLab
5. Just Python – no Conda / Mamba

See Notes for tweaks, settings, etc.

P.S.: I have whitelisted your GitHub account at https://demo.cuda.jupyter.b-data.ch so you may quickly test online.

[mathbunnyru]

I updated some advices in my comments above, so they work with current set of images.

If someone comes across this issue and feels they need GPU images, please, upvote this issue - this will make it more clear for me and other maintainers, that community is really interested in such images.
Also, I would highly appreciate the information on how to choose the relevant nvidia/cuda base image - there are many of them and I'm not sure which version community needs and wants the most.