Zenn
🧚

Rust製パッケージマネージャpixiをABCIで使うには？

2025/02/19に公開
GPU
abci
pixish
tech

はじめに

こんにちは、manuriya（@lTlanual）と申します。今回はRust製パッケージマネージャのpixiを使ってABCI上で実行する環境を構築する方法について解説します。

背景

衛星画像の深層学習モデルの開発を行なうに辺り、学習や推論効率の向上のため最新のGPUを大量に使用する場合があります。一部企業を除きそれだけの計算資源を確保/管理するのは非常に難しいです。そのため計算資源としてABCIやTSUBAMEのようなHPC(High Performance Computer)、いわゆるスパコンを使用する場合があります。

問題点

ABCI上でpixiを使用して環境構築をする際に問題になってくる点は、インタラクティブノード（いわゆるログインノード）と計算ノードが分かれている点です。
ABCIではこのように高性能なCPU/GPUが搭載されている計算ノード、計算ノードにバッチジョブを投げるためのインタラクティブノードに分かれていることが多いです。インタラクティブノードにはGPU等の高性能なハードウェアは搭載されておりません。環境構築を行なう際は基本的にインタラクティブノード上で行ない、作成された環境を計算ノードでも使用するといった具合です。

実行例

実際にインタラクティブノードにて環境構築をしてみます。例としてpytorchのCPU版とGPU版をそれぞれインストールし、モデル学習コードを実行してみます。使用したモデル学習コードはpytorchのサンプルコードをベースにしたものを使用しています。

モデル学習コード
train.py
import random

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from rich import print, progress

torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


def seed_everything(seed: int) -> None:
    """
    Set seed for reproducibility

    Args:
        seed (int): seed value
    """
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)


def get_dataloaders(batch_size: int) -> tuple[torch.utils.data.DataLoader]:
    """
    Get CIFAR10 dataloaders

    Args:
        batch_size (int): batch size

    Returns:
        (tuple[torch.utils.data.DataLoader]): trainloader, testloader
    """
    transform = transforms.Compose(
        [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
    )

    trainset = torchvision.datasets.CIFAR10(
        root="./data", train=True, download=True, transform=transform
    )
    trainloader = torch.utils.data.DataLoader(
        trainset, batch_size=batch_size, shuffle=True, pin_memory=True, num_workers=2
    )

    testset = torchvision.datasets.CIFAR10(
        root="./data", train=False, download=True, transform=transform
    )
    testloader = torch.utils.data.DataLoader(
        testset, batch_size=batch_size, shuffle=False, pin_memory=True, num_workers=2
    )
    return trainloader, testloader


class Net(nn.Module):
    """
    Sample neural network model
    """

    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = torch.flatten(x, 1)  # flatten all dimensions except batch
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


def prepare_training(lr: float) -> tuple[nn.Module, nn.Module, optim.Optimizer]:
    """
    Prepare training components

    Args:
        lr (float): learning rate

    Returns:
        (tuple[nn.Module, nn.Module, optim.Optimizer]): net, criterion, optimizer
    """
    net = Net()
    net.to(DEVICE)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.9)
    return net, criterion, optimizer


def training(
    epoch: int,
    net: nn.Module,
    dataloader: torch.utils.data.DataLoader,
    criterion: nn.Module,
    optimizer: optim.Optimizer,
) -> nn.Module:
    """
    Training loop

    Args:
        epoch (int): epoch number
        net (nn.Module): sample neural network model
        dataloader (torch.utils.data.DataLoader): dataloader
        criterion (nn.Module): loss function
        optimizer (optim.Optimizer): optimizer

    Returns:
        (nn.Module): updated model
    """
    running_loss = 0.0
    for data in progress.track(
        dataloader, description="Training...", total=len(dataloader)
    ):
        # get the inputs; data is a list of [inputs, labels]
        inputs, labels = data
        inputs, labels = (
            inputs.to(DEVICE, non_blocking=True),
            labels.to(DEVICE, non_blocking=True),
        )

        # zero the parameter gradients
        optimizer.zero_grad(set_to_none=True)

        # forward + backward + optimize
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    print(f"[{epoch + 1}] loss: {running_loss / len(dataloader):.4f}")
    return net


def test(net: nn.Module, dataloader: torch.utils.data.DataLoader) -> None:
    """
    Test the model

    Args:
        net (nn.Module): sample neural network model
        dataloader (torch.utils.data.DataLoader): dataloader
    """
    correct = 0
    total = 0
    # since we're not training, we don't need to calculate the gradients for our outputs
    with torch.no_grad():
        for data in progress.track(
            dataloader, description="Test...", total=len(dataloader)
        ):
            inputs, labels = data
            inputs, labels = (
                inputs.to(DEVICE, non_blocking=True),
                labels.to(DEVICE, non_blocking=True),
            )
            # calculate outputs by running images through the network
            outputs = net(inputs)
            # the class with the highest energy is what we choose as prediction
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print(
        f"Accuracy of the network on the 10000 test images: {100 * correct // total} %"
    )


print(f"Use CUDA: {torch.cuda.is_available()} | Use training device: {DEVICE}")
seed_everything(3407)
trainloader, testloader = get_dataloaders(4)
net, criterion, optimizer = prepare_training(0.001)

for epoch in range(2):
    net = training(epoch, net, trainloader, criterion, optimizer)
    test(net, testloader)
print("Finished Training")

まずはCPU版で環境構築を行ないます。

> pixi add python pytorch-cpu torchvision rich
✔ Added python >=3.13.1,<3.14
✔ Added pytorch-cpu >=2.5.1,<3
✔ Added torchvision >=0.20.1,<0.21
✔ Added rich >=13.9.4,<14

それぞれCPUのみの環境とGPU環境でモデル学習コードを実行してみます。PyTorchがCPU実行版のため、CUDAの使用ができず、学習時にCPUが使用されていることが確認できます。

# インタラクティブノードで実行
> pixi run python train.py
Use CUDA: False | Use training device: cpu
Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to ./data/cifar-10-python.tar.gz
100.0%
Extracting ./data/cifar-10-python.tar.gz to ./data
Files already downloaded and verified
Training... ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:47
[1] loss: 1.7006
Test... ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:05
Accuracy of the network on the 10000 test images: 48 %
Training... ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:47
[2] loss: 1.3407
Test... ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:05
Accuracy of the network on the 10000 test images: 56 %
Accuracy of the network on the 10000 test images: 56 %
Finished Training

# 計算ノードにログインして実行
> pixi run python train.py
Use CUDA: False | Use training device: cpu
（省略）...

次にGPU版で環境構築します。この時、pixi.tomlにCUDAのバージョンを記載してからインストールを実施します。

pixi.toml
[project]
channels = ["conda-forge"]
description = "Add a short description here"
name = "workspace"
platforms = ["linux-64"]
version = "0.1.0"

[tasks]

[dependencies]

[system-requirements]
cuda = "12.4"

> pixi add python pytorch-gpu torchvision rich
✔ Added python >=3.13.1,<3.14
✔ Added pytorch-gpu >=2.5.1,<3
✔ Added torchvision >=0.20.1,<0.21
✔ Added rich >=13.9.4,<14

環境が構築出来たらモデル学習コードを実行します。そうすると、GPUがない環境で実行しようとしているため、CUDAバージョンが確認できないというエラーが発生します。

# インタラクティブノードで実行
> pixi run python train.py
Error:
  × The platform you are running on should at least have the virtual package __cuda on version 12.4, build_string: 0

# 計算ノードにログインして実行
> pixi run python train.py
Use CUDA: True | Use training device: cuda:0
Files already downloaded and verified
Files already downloaded and verified
Training... ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:14
[1] loss: 1.7010
Test... ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:01
Accuracy of the network on the 10000 test images: 48 %
Training... ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:12
[2] loss: 1.3486
Test... ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:01
Accuracy of the network on the 10000 test images: 55 %
Finished Training

GPU版の環境構築を実施し、インタラクティブノードで実行できないということはデバッグ作業をする場合にも計算ノード上で行う必要が出てくるということです。計算ノードの使用には料金が発生するため、低負荷な処理や依存関係の確認などはなるべくインタラクティブノードで済ませたいです。しかし、GPU版ではCUDAが認識できないリソースでエラーが発生し、CPU版では本来GPUで実行したい処理がCPUでの実行となり処理時間がかかってしまいます。
この問題を解決するため、1つのプロジェクト内でCPU/GPUの両環境を作成します。

CPU/GPU実行環境の構築

pixi.shのMulti Environmentに記載されている方法を用いてCPU（インタラクティブノード）とGPU（計算ノード）の実行環境を分けて作成します。
https://prefix-dev.github.io/pixi/v0.40.3/features/multi_environment/

環境構築

まずCPU/GPU両環境にて共通で使用するライブラリを追加します。

pixi add python rich
✔ Added python >=3.13.2,<3.14
✔ Added rich >=13.9.4,<14

次にCPU環境のみで使用するライブラリを追加します。この時、--feature cpuというオプションを記載します。

pixi add --feature cpu pytorch-cpu torchvision
✔ Added pytorch-cpu
✔ Added torchvision
Added these only for feature: cpu

すると、pixi.tomlの記述に下記のブロックが追加されます。

[feature.cpu.dependencies]
pytorch-cpu = "*"
torchvision = "*"

次にGPU環境でのみ使用するライブラリを追加します。CPUの時と同様に--feature gpuというオプションを記載します。この時、警告が出ますが気にせず進めます。

pixi add --feature gpu pytorch-gpu torchvision
 WARN The feature 'cpu' is defined but not used in any environment
✔ Added pytorch-gpu
✔ Added torchvision
Added these only for feature: gpu

CPUの時と同様にpixi.tomlにfeature.gpu.dependenciesブロックが追加されます。
実行環境への紐づけのため、pixi project environment add XXX --feature xxxを利用して両環境をそれぞれgpuとdefaultというようにします。実行環境の紐づけが出来るとpixi.tomlにenvironmentsブロックが記載されます。

pixi project environment add gpu --feature gpu
✔ Added environment gpu

# defaultへの紐づけ時は--forceオプションをつけること
pixi project environment add default --force --feature cpu
✔ Updated environment default

最後にGPU環境におけるCUDAのバージョン情報をpixi.tomlのfeature.gpu.system-requirementsブロックに記載します。ここまでの作業を実施した際のpixi.tomlを下記に示します。

[project]
channels = ["conda-forge"]
description = "Sample project"
name = "workspace"
platforms = ["linux-64"]
version = "0.1.0"

[tasks]

[dependencies]
python = ">=3.13.2,<3.14"
rich = ">=13.9.4,<14"

[feature.gpu.system-requirements]
cuda = "12.4"

[feature.cpu.dependencies]
pytorch-cpu = "*"
torchvision = "*"

[feature.gpu.dependencies]
pytorch-gpu = "*"
torchvision = "*"

[environments]
gpu = ["gpu"]
default = ["cpu"]

実行方法

まずはCPU環境の方を試します。実行方法はpixi shellとpixi runの2種類あります。pixi shellはcondaでいうconda activateに該当し、pixi runはconda runに該当します。ここでは後の比較のためpixi runを用いた実行方法を記載します。

インタラクティブノード上でCPU環境によるモデル学習スクリプトの実行を行ないます。初回実行時のみライブラリのインストールが開始されます。

pixi run python train.py
Use CUDA: False | Use training device: cpu
Training... ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:45
[1] loss: 1.7006
Test... ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:03
Accuracy of the network on the 10000 test images: 48 %
Training... ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:42
[2] loss: 1.3407
Test... ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:04
Accuracy of the network on the 10000 test images: 56 %
Finished Training

次に計算ノードにログインしてGPU環境によるモデル学習スクリプトの実行を行ないます。CPU環境と違う点はオプションに-e gpuという実行環境指定オプションが含まれているところです。

pixi run -e gpu python train.py 
Use CUDA: True | Use training device: cuda:0
Training... ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:19
[1] loss: 1.7010
Test... ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:01
Accuracy of the network on the 10000 test images: 48 %
Training... ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:13
[2] loss: 1.3486
Test... ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100% 0:00:01
Accuracy of the network on the 10000 test images: 55 %
Finished Training

task機能による実行コマンド管理

これで両環境の構築ができましたが、毎回pixi run -e gpuと打つのは面倒です。あらかじめ処理によってCPUまたはGPUを使うのが分かっている場合、task機能を使ってコマンドを登録すると便利です。

# CPU環境で実行するタスク
pixi task add --feature cpu training-cpu "python train.py"
✔ Added task `training-cpu`: python train.py

# GPU環境で実行するタスク
> pixi task add --feature gpu training-gpu "python train.py"
✔ Added task `training-gpu`: python train.py

実行時はpixi run [登録したタスク]で実行することが出来ます。

# CPU環境で実行
pixi run training-cpu
✨ Pixi task (training-cpu in default): python train.py
Use CUDA: False | Use training device: cpu
（省略）...

# GPU環境で実行
pixi run training-gpu
✨ Pixi task (training-gpu in gpu): python train.py
Use CUDA: True | Use training device: cuda:0
（省略）...

さらに細かい指定をしたい場合は公式ページを参照してください。

おわりに

今回はABCI上でpixiの実行環境を作成する方法についてまとめました。ここで記載した方法はABCIに限らずログインノード/計算ノードで分かれているHPCであれば同じように環境構築が可能です（要：インターネット環境）。今までログインノード/計算ノードでの実行環境構築が分からなかった方は、ぜひ記載の方法をお試しください。

Discussion