Background
In this article, I wrote about setting up a Python development environment that utilizes a GPU by connecting via SSH from my local VS Code to a Docker container running on a GPU-equipped remote workstation. Although I initially set it up on my own, I wondered if there might be room for improvement, so I’ll discuss how I used generative AI to enhance it.
Conclusion: Key Points for Improvement
The key points from the generative AI’s response are as follows:
It would be easier to understand if we switched from the current “method of attaching to a running Docker container” to the **“Dev Container” method, where the project itself has Docker, uv, and VS Code configurations**.
Diagram of the Final Configuration
My Mac
└─ VS Code
│
│ Remote SSH
▼
Ubuntu GPU Host
├─ NVIDIA Driver
├─ Docker
├─ NVIDIA Container Toolkit
│
└─ ~/workspace/pinn/
├─ .devcontainer/
│ ├─ Dockerfile
│ └─ devcontainer.json
├─ pyproject.toml
├─ uv.lock
├─ .python-version
├─ .gitignore
├─ .dockerignore
├─ src/
└─ notebooks/
│
│ VS Code: Reopen in Container
▼
Docker Container
├─ CUDA Runtime
├─ uv(Fixed Version)
├─ Python 3.11
└─ /workspace/.venv
Configuration Files
pyproject.toml
Python libraries are consolidated into the dependencies section as follows. They are not specified in the Dockerfile.
[project]
name = "pinn"
version = "0.1.0"
description = "PyTorch GPU development environment for PINN"
requires-python = ">=3.11,<3.12"
dependencies = [
"numpy",
"pandas",
"matplotlib",
"scipy",
"scikit-learn",
"seaborn",
"plotly",
"tqdm",
"jupyter",
"ipykernel",
"torch",
"torchvision",
"torchaudio",
]
[dependency-groups]
dev = [
"ruff",
"pytest",
]
Dockerfile
Do not include Python packages when building the Docker image. The Docker container should only contain CUDA, uv, and Python.
FROM nvidia/cuda:12.4.0-runtime-ubuntu22.04
ARG DEBIAN_FRONTEND=noninteractive
# --------------------------------------------------
# OS packages
# --------------------------------------------------
RUN apt-get update && apt-get install -y \
ca-certificates \
curl \
git \
build-essential \
sudo \
&& rm -rf /var/lib/apt/lists/*
# --------------------------------------------------
# uv
# バージョンを固定
# --------------------------------------------------
COPY --from=ghcr.io/astral-sh/uv:0.11.32 \
/uv /uvx /bin/
# --------------------------------------------------
# Python
# --------------------------------------------------
RUN uv python install 3.11
# --------------------------------------------------
# Workspace
# --------------------------------------------------
WORKDIR /workspace
# uvのキャッシュと.venvが別FSの場合の警告回避にも有効
ENV UV_LINK_MODE=copy
# .venvをPATHの先頭にする
ENV PATH="/workspace/.venv/bin:${PATH}"
CMD ["sleep", "infinity"]
devcontainer.json
{
"name": "PyTorch GPU Development for PINN",
"build": {
"dockerfile": "Dockerfile",
"context": ".."
},
"runArgs": [
"--gpus=all"
],
"workspaceMount": "source=${localWorkspaceFolder},target=/workspace,type=bind",
"workspaceFolder": "/workspace",
"postCreateCommand": "uv sync --locked",
"customizations": {
"vscode": {
"extensions": [
"ms-python.python",
"ms-toolsai.jupyter"
],
"settings": {
"python.defaultInterpreterPath": "/workspace/.venv/bin/python"
}
}
},
"remoteEnv": {
"PATH": "/workspace/.venv/bin:${containerEnv:PATH}"
}
}
The key point here is the following section.
"postCreateCommand": "uv sync --locked",
With the settings above, if you try to open the Dev Container, you’ll see the following message. (This occurs when uv is not installed on the Ubuntu GPU host. uv is installed in the Docker container.)
Running the postCreateCommand from devcontainer.json...
[88364 ms] Start: Run in container: /bin/sh -c uv sync --locked
Using CPython 3.11.15
Creating virtual environment at: .venv
error: Unable to find lockfile at `uv.lock`, but `--locked` was provided. To
create a lockfile, run `uv lock` or `uv sync` without the flag.
[89785 ms] postCreateCommand from devcontainer.json failed with exit code 2. Skipping any further user-provided commands.
任意のキーを押してターミナルを終了します。
If this happens, temporarily change the section ““postCreateCommand”: “uv sync –locked,”” to the following:
"postCreateCommand": "uv sync"
After making the above changes, run “Dev Containers: Rebuild and Reopen in Container” from VS Code. This will execute “uv sync,” and the following will be created inside the container.
/workspace/uv.lock
/workapace/.env/
As a result, the following are created on the Ubuntu host.
~/workspace/pinn/uv.lock
~/workapace/pinn/.env/
Next, revert the relevant section of the devcontainer.json file you modified earlier (by adding “–locked”).
Connection Instructions
Connecting to the Remote (Ubuntu Host)
First, open the Command Palette in VS Code on your Mac by pressing “Command + Shift + p,” then run the following command.
Remote-SSH: Connect to Host...
Select an SSH host and enter the password for the relevant user on the Ubuntu host. This will open a new instance of VS Code and connect it to the Ubuntu host.
For information on “configuring SSH connections,” refer to this section in the article linked at the beginning.
Opening a Project on the Remote (Ubuntu Host)
In VS Code while connected to the Ubuntu host, select “File” → “Open Folder…” to open the following location on the Ubuntu host.
~/workspace/pinn
Start the Dev Container
With the folder on the Ubuntu host open, open the Command Palette (Command+Shift+p) and run the following command.
Dev Containers: Reopen in Container
At this point, it takes a while for the Dockerfile to be built.
During this time, the following actions are performed from VS Code:
1. Read .devcontainer/devcontainer.json
2. Build the Dockerfile
3. Start the container with --gpus=all
4. Mount the Ubuntu project to /workspace
5. Run `uv sync --locked`
6. Create /workspace/.venv
7. Install Python and Jupyter extensions inside the container
8. Connect VS Code to the container
The above explains generative AI and states that the process of manually connecting to a running container, which was previously required, is no longer necessary.
Connection Procedure for Subsequent Sessions
- Launch VS Code on your Mac
- From the Command Palette (Command+Shift+P), select “Remote-SSH: Connect to Host…”
- Select the host and enter your password; a new VS Code window will open
- In VS Code, open the project directory (the directory where
pyproject.tomlis located) - A pop-up will appear in the bottom-right corner of VS Code stating, “The folder contains configuration files for the development container. Reopen the folder to develop in the container.” Click “Reopen in Container.”
- The container will start up, and you’ll be able to use a Python development environment (project) that supports the GPU.
Summary
The improved Python development environment described here is very convenient because simply navigating to the desired development environment (project) in VS Code automatically starts and connects the container.
Additionally, responsibilities are clearly separated, making it easy to understand once you get the hang of it.