カスタム Slurm Docker イメージをビルドする

このドキュメントでは、Google Kubernetes Engine（GKE）上の Slurm クラスタ用にカスタム Docker イメージをビルドする方法について説明します。GKE が提供するベース Slurm イメージを拡張して、ハイ パフォーマンス コンピューティング（HPC）ワークロードに必要な追加のツール、ライブラリ、構成を含めることができます。

このドキュメントを読む前に、GKE 向け Slurm Operator アドオンについて理解しておいてください。

始める前に

作業を始める前に、次のタスクが完了していることを確認してください。

  • Google Kubernetes Engine API を有効にする。
    • Google Kubernetes Engine API の有効化
  • このタスクに Google Cloud CLI を使用する場合は、gcloud CLI をインストールして初期化する。gcloud CLI をインストール済みの場合は、gcloud components update コマンドを実行して最新のバージョンを取得します。以前のバージョンの gcloud CLI では、このドキュメントのコマンドを実行できない場合があります。

前提条件

このドキュメントでは、GKE 向け Slurm Operator アドオンがインストールされた GKE で Slurm クラスタがすでに実行されていることを前提としています。次のページの操作を完了してください。

  1. クイックスタート: GKE に Slurm クラスタをデプロイするを完了します。
  2. カスタム イメージを保存する Artifact Registry リポジトリをプロジェクトに構成します。

Slurm ベースイメージ

GKE は、 gcr.io/gke-release/ Artifact Registry リポジトリにベース Slurm イメージを提供します。GKE は、セキュリティとパフォーマンスのためにこれらのイメージを頻繁に更新します。これらのイメージには、最新の Slurm バージョンと 2 つの Linux ディストリビューション（Ubuntu と Rocky Linux）を含むバリアントがあります。

次のベースイメージをカスタマイズできます。

  • gcr.io/gke-release/slinky/slurmd: Slurm コンピューティング ノードに使用されます。
  • gcr.io/gke-release/slinky/login: ログインノードに使用されます。

カスタム イメージをビルドする

次の例は、JAX がインストールされた Python 仮想環境を含むカスタム Slurm コンピューティング イメージをビルドする方法を示しています。また、JAX ライブラリを実際にインストールせずに、コンピューティング イメージ PATH 環境変数をミラーリングする対応するログイン イメージもビルドします。

画像モードを選択する

ベースイメージを選択する場合は、次の条件を満たしていることを確認してください。

  • バージョンが、Slurm クラスタ内の他のコンポーネントで使用されている Slurm バージョンと一致している。
  • 特定の Slurm バージョンについては、最新のセキュリティ アップデートとバグ修正を含む、利用可能な最新のイメージのタグを選択します。

たとえば、クラスタのデフォルトの Slurm バージョンが 25.11 の場合は、 25.11- で始まるタグ（25.11-ubuntu24.04-gke.6 など）を選択する必要があります。

Dockerfile の作成

  1. Ubuntu ベースの slurmd イメージタグを選択します。

    1. コンソールで、パッケージを含む [Artifact Registry リポジトリ] ページに移動します。 Google Cloud slinky/slurmd

      [Artifact Registry リポジトリ] に移動

    2. ubuntu を含み、Slurm バージョンと一致するタグ（25.11-ubuntu24.04-gke.6 など）を持つイメージを見つけます。

    3. タグをコピーします。このタグを使用して、次の構成ファイルの VERSION_TAG プレースホルダを置き換えます。

  2. 次の内容で Dockerfile という名前のファイルを作成します。

    # --- Target 1: The Worker Node (slurmd) ---
FROM gcr.io/gke-release/slinky/slurmd:VERSION_TAG AS slurmd-custom
USER root

# Install minimal requirements for venv
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3-pip \
    python3-venv \
    && rm -rf /var/lib/apt/lists/*

# Create and populate the virtual environment
ENV VIRTUAL_ENV=/opt/custom_venv
RUN python3 -m venv ${VIRTUAL_ENV}
ENV PATH="${VIRTUAL_ENV}/bin:$PATH"

# Install JAX (CPU version for general compatibility) and dependencies
RUN pip install --no-cache-dir jax[cpu] numpy

# --- Target 2: The Login Node ---
FROM gcr.io/gke-release/slinky/login:VERSION_TAG AS login-custom
USER root

# Mirror the PATH exactly so that the srun command captures it.
# Note: You don't need to install the JAX libs here,
# but the binary path must exist for the shell to recognize it.
ENV VIRTUAL_ENV=/opt/custom_venv
ENV PATH="${VIRTUAL_ENV}/bin:$PATH"

# Create the directory structure so the PATH is valid on the login node
RUN mkdir -p ${VIRTUAL_ENV}/bin

    VERSION_TAG は、クラスタのデフォルトの Slurm バージョンと一致する Slurm バージョンタグに置き換えます。

  3. docker build コマンドを使用してイメージをビルドします。

    docker build --target=slurmd-custom \
    -t AR_PATH/slinky/slurmd:CUSTOM_SLURMD_TAG \
    -f Dockerfile .
docker build --target=login-custom \
    -t AR_PATH/slinky/login:CUSTOM_LOGIN_TAG \
    -f Dockerfile .

    次のように置き換えます。

    • AR_PATH: Artifact Registry リポジトリのパス（例: gcr.io/my-project）。
    • CUSTOM_SLURMD_TAG: 任意の slurmd-custom タグ名。
    • CUSTOM_LOGIN_TAG: 任意の login-custom タグ名。

  4. カスタム イメージをリポジトリに push します。

    docker push AR_PATH/slinky/slurmd:CUSTOM_SLURMD_TAG
docker push AR_PATH/slinky/login:CUSTOM_LOGIN_TAG

GKE でカスタム イメージを使用する

カスタム イメージを使用するには、次の操作を行います。

  1. 次の例に示すように、values.yaml ファイルを変更して、slurmd ノードセットと login ログインセットのイメージ リポジトリとタグを更新します。

    nodesets:
    slinky:
        replicas: 1
        slurmd:
            image:
                repository: AR_PATH/slinky/slurmd
                tag: CUSTOM_SLURMD_TAG

loginsets:
    slinky:
        enabled: true
        replicas: 1
        login:
            image:
                repository: AR_PATH/slinky/login
                tag: CUSTOM_LOGIN_TAG

  2. 既存のデプロイをアップグレードします。

    helm upgrade slurm oci://ghcr.io/slinkyproject/charts/slurm \
    --namespace slurm \
    --version=1.0.2 \
    -f values.yaml

  3. ログインノードにログインして次の srun コマンドを実行し、コンピューティング ノードの新しい機能をテストします。

    srun python3 -c "
import sys
import jax
import jax.numpy as jnp

print(f'Python Executable: {sys.executable}')
print(f'Using JAX backend: {jax.devices()[0].platform}')

key = jax.random.PRNGKey(42)
x = jax.random.normal(key, (5000, 5000))
result = jnp.dot(x, x)
print(f'Matrix multiplication successful. Shape: {result.shape}')
"

    出力は次のようになります。

    Python Executable: /opt/custom_venv/bin/python3
Using JAX backend: cpu
Matrix multiplication successful. Shape: (5000, 5000)

    この出力は、Slurm がカスタム イメージを実行するワーカー Pod でスクリプトを実行し、イメージに必要な Python 機能と JAX 機能が含まれていることを確認します。

クリーンアップ

このチュートリアルで使用したリソースをクリーンアップするには、次の操作を行います。

  1. Helm デプロイをアンインストールします。 sh helm uninstall slurm --namespace slurm

    このコマンドは、Helm チャートによってデプロイされたすべての Kubernetes リソースを削除します。

  2. Slurm Namespace を削除します。

    kubectl delete namespace slurm

  3. GKE クラスタを削除します。

    gcloud container clusters delete CLUSTER_NAME

    CLUSTER_NAME はクラスタ名で置き換えます。

  4. Artifact Registry からカスタム イメージを削除します。

    gcloud container images delete AR_PATH/slinky/slurmd:CUSTOM_SLURMD_TAG --force-delete-tags
gcloud container images delete AR_PATH/slinky/login:CUSTOM_LOGIN_TAG --force-delete-tags

  5. ローカル Docker 環境からカスタム イメージを削除します。

    docker rmi AR_PATH/slinky/slurmd:CUSTOM_SLURMD_TAG
docker rmi AR_PATH/slinky/login:CUSTOM_LOGIN_TAG

次のステップ