Machine Learning Container with GPU inside Visual Studio Code (Ubuntu)

Now that visual studio code supports dev environment inside a container is it possible to work on multiple platform with the same machine.

And with a few parameter added you can use your GPU thanks to the Nvidia Container Toolkit !

You must have installed:

• nvidia driver installed
• docker installed
• nvidia container
• visual studio code

First install the GPU driver (proprietary, tested)

Then you have to install the container toolkit which allow you to use the GPU inside a container ( you must also install docker as a prerequisite)

Installation Guide - NVIDIA Cloud Native Technologies documentation

Now you need of course Visual Studio Code!

After installing it you also will install some extension

Remote - Containers - Visual Studio Marketplace

Docker - Visual Studio Marketplace

Python - Visual Studio Marketplace

Now you’re all set. To start a project inside a container with GPU:

Create a folder, inside add 2 files : Dockerfile and requirements.txt

I will use here a tensorflow image

put whatever library you must use in the requirements.txt

save the file and then use CTRL + SHIFT + P and type

 reopen in container

You should see the option in the list, click and your container should be created !

official example to open the container

The container is now opened, your dev env is inside the container

Cool … but has the GPU been detected ? let’s try that on a python console

import tensorflow as tf
physical_devices = tf.config.list_physical_devices('GPU')
print("Num GPUs:", len(physical_devices))

no GPU detected :(

and if you type

nvidia-smi

Nothing appears.

We need one more conf to enjoy it

On the previous screen you should notice a new folder .devcontainer

As this name suggest, it’s used by visual studio code for building the container

We are going to modify the file devcontainer.json to enjoy the GPU

"extensions": ["ms-python.python"],
"runArgs": [ "--gpus", "all"]

here we are adding the python extension and the docker running arguments to use the GPU!

the file should look like this

Now you need to rebuild the container, click on the bottom left green area

And choose rebuild container, now if you type on a terminal nvidia-smi

the gpu is now recognized!

That’s it you’re good to go !