构建自定义 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 版本，请选择包含最新安全更新和 bug 修复的最新可用映像的标记。

例如，如果集群中的默认 Slurm 版本为 25.11，您应选择以 25.11- 开头的标记，例如 25.11-ubuntu24.04-gke.6。

创建 Dockerfile

  1. 选择基于 Ubuntu 的 slurmd 映像标记：

    1. 在 Google Cloud 控制台中，前往包含 slinky/slurmd 软件包的 Artifact Registry 制品库页面。

      前往 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 nodeset 和 login loginset 的映像代码库和标记：

    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

后续步骤