Visual Studio Code Dev Containers でシンプルに Python w/ CUDA を動かす

はじめに

なるべく少ない設定でDev Container上でPythonを動かす方法を検討してみました。

2024年11月末現在、Python 環境のセットアップには uv を使うのがよさそうな雰囲気です。そこで、Dev Container 側では uv をセットアップするにとどめ、Python やそのパッケージのインストールは uv に任せる方針にします。

Dev Container の設定

ベースコンテナは Microsoft が公開している devcontainers-base を使いました。以下のように .devcontainer/devcontainer.json を作成します:

// .devcontainer/devcontainer.json`

{
	"name": "CUDA + uv",
	"image": "mcr.microsoft.com/devcontainers/base:bookworm",
	"onCreateCommand": "wget -qO- https://astral.sh/uv/install.sh | sh",
	"containerEnv": {
		"UV_LINK_MODE": "copy"
	},
	"runArgs": [
		"--device",
		"nvidia.com/gpu=all"
	]
}

Dev Container の設定としてはこれで全てです。

uv がコンテナ内でファイルシステムを跨いだハードリンクを作成しようとして警告が出るので、 UV_LINK_MODE を指定してあります。とはいえ、現状キャッシュはDev Containerの再構築で捨てられてしまうので、Caching の項にあるようにうまいこと残るようにすると良さそうです。

Dev Container を使ってみる

それでは、試しにこの環境に PyTorch をセットアップしてみましょう。

VS Code で開き、Dev Container を起動します。

初回は uv のプロジェクトをセットアップする必要があります。

vscode ➜ /workspaces/demo $ uv init
Initialized project `demo`

numpy と torch パッケージをインストールしてみましょう:

vscode ➜ /workspaces/demo (main) $ uv add numpy torch
Using CPython 3.13.0
Creating virtual environment at: .venv
(略)

ちゃんと動くでしょうか。試しに以下のようなコードを用意します。

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

x = torch.tensor([1.0, 2.0, 3.0], device=device)
y = torch.tensor([4.0, 5.0, 6.0], device=device)

z = x + y
print(f"Result on GPU: {z}")

vscode ➜ /workspaces/demo (main) $ uv run main.py 
Using device: cuda
Result on GPU: tensor([5., 7., 9.], device='cuda:0')

いいですね、無事GPU上で計算できているようです。

まだ「動いた」というレベルで、使い込んではいないので、これから色々試してみたいところです。

uv sync

ところで、 onCreateCommand で uv をセットアップしたあと、 uv sync するようにしてもよいかもしれません。

"onCreateCommand": "wget -qO- https://astral.sh/uv/install.sh | sh && uv sync",

pyproject.toml が無いと（つまり uv init する前）uv sync が失敗するので、以下のようにしてもよいかもしれません。

"onCreateCommand": "wget -qO- https://astral.sh/uv/install.sh | sh && [ -f pyproject.toml ] && uv sync",

検証環境

私は NixOS 24.05 上で試しました。

❯ docker --version
Docker version 25.0.6, build v25.0.6
❯ code --version
1.92.0
b1c0a14de1414fcdaa400695b4db1c0799bc3124
x64

おわりに

なるべく少ない設定で CUDA と uv をセットアップするための Dev Container 設定と、その上でPyTorch を利用する方法を紹介しました。