建構自訂 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 叢集，並安裝 GKE 適用的 Slurm Operator 外掛程式。請完成下列頁面的程序：

  1. 完成「快速入門：在 GKE 上部署 Slurm 叢集」。
  2. 在專案中設定 Artifact Registry 存放區，用於儲存自訂映像檔。

Slurm 基本映像檔

GKE 在 gcr.io/gke-release/ Artifact Registry 存放區中提供基本 Slurm 映像檔。GKE 會頻繁更新這些映像檔，確保安全性和效能。這些映像檔包含最新版 Slurm 和兩種 Linux 發行版本 (Ubuntu 和 Rocky Linux)。

您可以自訂下列基礎圖片：

  • gcr.io/gke-release/slinky/slurmd：用於 Slurm 運算節點。
  • gcr.io/gke-release/slinky/login：用於登入節點。

建立自訂映像檔

以下範例說明如何建構自訂 Slurm 計算映像檔，其中包含已安裝 JAX 的 Python 虛擬環境。您也會建構對應的登入映像檔，反映運算映像檔 PATH 環境變數，但不會實際安裝 JAX 程式庫。

選取映像檔版本

選取基礎映像檔時，請確認符合下列條件：

  • 這個版本與 Slurm 叢集中其他元件使用的 Slurm 版本相符。
  • 如要使用特定 Slurm 版本，請選擇最新可用映像檔的標記，其中包含最新的安全性更新和錯誤修正。

舉例來說，如果叢集中的預設 Slurm 版本為 25.11，您應選擇以 25.11- 開頭的標記，例如 25.11-ubuntu24.04-gke.6。

建立 Dockerfile

  1. 選取以 Ubuntu 為基礎的 slurmd 映像檔標記：

    1. 前往 Google Cloud 控制台的「Artifact Registry repository」(構件登錄檔存放區) 頁面，其中包含 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. 將自訂映像檔推送至存放區：

    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 命名空間：

    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

後續步驟