Criar imagens personalizadas do Docker do Slurm

Este documento explica como criar imagens personalizadas do Docker para clusters do Slurm no Google Kubernetes Engine (GKE). É possível estender as imagens base do Slurm fornecidas pelo GKE para incluir outras ferramentas, bibliotecas ou configurações necessárias para cargas de trabalho de computação de alto desempenho (HPC).

Antes de ler este documento, você precisa conhecer o complemento do operador do Slurm para GKE.

Antes de começar

Antes de começar, verifique se você realizou as tarefas a seguir:

  • Ativar a API Google Kubernetes Engine.
    • Ativar a API Google Kubernetes Engine
  • Se você quiser usar a Google Cloud CLI para essa tarefa, instale e, em seguida, inicialize a CLI gcloud. Se você instalou a CLI gcloud anteriormente, instale a versão mais recente executando o comando gcloud components update. Talvez as versões anteriores da CLI gcloud não sejam compatíveis com a execução dos comandos neste documento.

Pré-requisitos

Este documento pressupõe que você já tenha um cluster do Slurm em execução no GKE com o complemento do operador do Slurm para GKE instalado. Conclua os procedimentos nas páginas a seguir:

  1. Conclua o Guia de início rápido: implantar um cluster do Slurm no GKE.
  2. Configure um repositório do Artifact Registry no seu projeto para armazenar as imagens personalizadas.

Imagens de base do Slurm

O GKE fornece imagens de base do Slurm no gcr.io/gke-release/ repositório do Artifact Registry. O GKE atualiza essas imagens com frequência para segurança e performance. Essas imagens vêm em variantes que incluem as versões mais recentes do Slurm e duas distribuições do Linux, Ubuntu e Rocky Linux.

É possível personalizar as seguintes imagens de base:

  • gcr.io/gke-release/slinky/slurmd: usado para nós de computação do Slurm.
  • gcr.io/gke-release/slinky/login: usado para nós de login.

Criar uma imagem personalizada

O exemplo a seguir demonstra como criar uma imagem de computação personalizada do Slurm que inclui um ambiente virtual do Python com o JAX instalado. Você também cria uma imagem de login correspondente que espelha a variável de ambiente PATH da imagem de computação sem instalar as bibliotecas JAX.

Selecionar a versão de imagem

Ao selecionar uma imagem de base, verifique se ela atende às seguintes condições:

  • A versão corresponde à versão do Slurm usada por outros componentes no cluster do Slurm.
  • Para uma versão específica do Slurm, escolha a tag da imagem mais recente disponível, que inclui as atualizações de segurança e correções de bugs mais recentes.

Por exemplo, se a versão padrão do Slurm no cluster for 25.11, você deve escolher uma tag que comece com 25.11-, por exemplo 25.11-ubuntu24.04-gke.6.

Criar um DockerFile

  1. Selecione uma tag de imagem slurmd baseada no Ubuntu:

    1. No Google Cloud console, acesse a página do repositório do Artifact Registry que inclui o pacote slinky/slurmd.

      Acessar o repositório do Artifact Registry

    2. Encontre uma imagem com uma tag que inclua ubuntu e corresponda à versão do Slurm , por exemplo, 25.11-ubuntu24.04-gke.6.

    3. Copie a tag. Use essa tag para substituir o marcador VERSION_TAG no arquivo de configuração a seguir.

  2. Crie um arquivo chamado Dockerfile com o conteúdo a seguir:

    # --- 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

    Substitua VERSION_TAG pela tag de versão do Slurm que corresponde à versão padrão do Slurm do cluster.

  3. Crie as imagens usando o comando 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 .

    Substitua:

    • AR_PATH: o caminho para o repositório do Artifact Registry, por exemplo, gcr.io/my-project.
    • CUSTOM_SLURMD_TAG: um nome de tag slurmd-custom de sua escolha.
    • CUSTOM_LOGIN_TAG: um nome de tag login-custom de sua escolha.

  4. Envie as imagens personalizadas para o repositório:

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

Usar as imagens personalizadas no GKE

Para usar as imagens personalizadas, siga estas etapas:

  1. Conforme mostrado no exemplo a seguir, atualize o repositório e a tag da imagem para o conjunto de nós slurmd e o conjunto de logins login modificando o arquivo values.yaml:

    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. Faça upgrade da implantação atual:

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

  3. Teste os novos recursos do nó de computação fazendo login no nó de login e executando o seguinte comando 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}')
"

    O resultado será assim:

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

    Essa saída confirma que o Slurm executa o script em pods de worker que executam a imagem personalizada, e a imagem contém os recursos necessários do Python e do JAX.

Limpar

Para limpar os recursos usados neste tutorial, faça o seguinte:

  1. Desinstale a implantação do Helm: sh helm uninstall slurm --namespace slurm

    Esse comando remove todos os recursos do Kubernetes implantados pelo gráfico do Helm.

  2. Exclua o namespace do Slurm:

    kubectl delete namespace slurm

  3. Exclua o cluster do GKE:

    gcloud container clusters delete CLUSTER_NAME

    Substitua CLUSTER_NAME pelo nome do seu cluster.

  4. Exclua as imagens personalizadas do 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. Remova as imagens personalizadas do ambiente local do Docker:

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

A seguir