Compila imágenes personalizadas de Docker de Slurm

En este documento, se explica cómo compilar imágenes de Docker personalizadas para tus clústeres de Slurm en Google Kubernetes Engine (GKE). Puedes extender las imágenes base de Slurm que proporciona GKE para incluir herramientas, bibliotecas o configuraciones adicionales que se requieran para tus cargas de trabajo de computación de alto rendimiento (HPC).

Antes de leer este documento, asegúrate de conocer el complemento de Slurm Operator para GKE.

Antes de comenzar

Antes de comenzar, asegúrate de haber realizado las siguientes tareas:

  • Habilita la API de Google Kubernetes Engine.
    • Habilitar la API de Google Kubernetes Engine
  • Si deseas usar Google Cloud CLI para esta tarea, instala y, luego, inicializa gcloud CLI. Si ya instalaste la gcloud CLI, ejecuta el comando gcloud components update para obtener la versión más reciente. Es posible que las versiones anteriores de la gcloud CLI no admitan la ejecución de los comandos que se describen en este documento.

Requisitos previos

En este documento, se supone que ya tienes un clúster de Slurm en ejecución en GKE con el complemento de Slurm Operator para GKE instalado. Completa los procedimientos de las siguientes páginas:

  1. Completa la guía de inicio rápido para implementar un clúster de Slurm en GKE.
  2. Configura un repositorio de Artifact Registry en tu proyecto para almacenar tus imágenes personalizadas.

Imágenes base de Slurm

GKE proporciona imágenes base de Slurm en el repositorio de Artifact Registry gcr.io/gke-release/. GKE actualiza estas imágenes con frecuencia para mejorar la seguridad y el rendimiento. Estas imágenes vienen en variantes que incluyen las versiones más recientes de Slurm y dos distribuciones de Linux, Ubuntu y Rocky Linux.

Puedes personalizar las siguientes imágenes base:

  • gcr.io/gke-release/slinky/slurmd: Se usa para los nodos de procesamiento de Slurm.
  • gcr.io/gke-release/slinky/login: Se usa para los nodos de acceso.

Compila una imagen personalizada

En el siguiente ejemplo, se muestra cómo compilar una imagen de procesamiento de Slurm personalizada que incluye un entorno virtual de Python con JAX instalado. También compilas una imagen de acceso correspondiente que refleja la variable de entorno PATH de la imagen de procesamiento sin instalar realmente las bibliotecas de JAX.

Selecciona la versión de la imagen

Cuando selecciones una imagen base, asegúrate de que cumpla con las siguientes condiciones:

  • La versión coincide con la versión de Slurm que usan otros componentes de tu clúster de Slurm.
  • Para una versión específica de Slurm, elige la etiqueta de la imagen disponible más reciente, que incluye las actualizaciones de seguridad y las correcciones de errores más recientes.

Por ejemplo, si la versión predeterminada de Slurm en tu clúster es 25.11, debes elegir una etiqueta que comience con 25.11-, por ejemplo, 25.11-ubuntu24.04-gke.6.

Crea un Dockerfile

  1. Selecciona una etiqueta de imagen slurmd basada en Ubuntu:

    1. En la consola de Google Cloud , ve a la página del repositorio de Artifact Registry que incluye el paquete slinky/slurmd.

      Ir al repositorio de Artifact Registry

    2. Busca una imagen con una etiqueta que incluya ubuntu y coincida con tu versión de Slurm, por ejemplo, 25.11-ubuntu24.04-gke.6.

    3. Copia la etiqueta. Usas esta etiqueta para reemplazar el marcador de posición VERSION_TAG en el siguiente archivo de configuración.

  2. Crea un archivo llamado Dockerfile con el siguiente contenido:

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

    Reemplaza VERSION_TAG por la etiqueta de versión de Slurm que coincida con la versión predeterminada de Slurm de tu clúster.

  3. Compila las imágenes con el 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 .

    Reemplaza lo siguiente:

    • AR_PATH: Es la ruta de acceso a tu repositorio de Artifact Registry, por ejemplo, gcr.io/my-project.
    • CUSTOM_SLURMD_TAG: Es el nombre de la etiqueta slurmd-custom que elijas.
    • CUSTOM_LOGIN_TAG: Es el nombre de la etiqueta login-custom que elijas.

  4. Envía las imágenes personalizadas a tu repositorio:

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

Usa las imágenes personalizadas en GKE

Para usar tus imágenes personalizadas, completa los siguientes pasos:

  1. Como se muestra en el siguiente ejemplo, actualiza el repositorio y la etiqueta de la imagen para los conjuntos de nodos slurmd y de acceso login modificando el archivo 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. Actualiza la implementación existente:

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

  3. Para probar las nuevas capacidades de tu nodo de procesamiento, accede al nodo de acceso y ejecuta el siguiente 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}')
"

    El resultado es similar a lo siguiente:

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

    Este resultado confirma que Slurm ejecuta el script en los Pods de trabajador que ejecutan tu imagen personalizada y que la imagen contiene las capacidades requeridas de Python y JAX.

Realiza una limpieza

Para limpiar los recursos que usaste en este instructivo, haz lo siguiente:

  1. Desinstala la implementación de Helm: sh helm uninstall slurm --namespace slurm

    Este comando quita todos los recursos de Kubernetes implementados por el gráfico de Helm.

  2. Borra el espacio de nombres de Slurm:

    kubectl delete namespace slurm

  3. Borra el clúster de GKE:

    gcloud container clusters delete CLUSTER_NAME

    Reemplaza CLUSTER_NAME por el nombre del clúster.

  4. Borra las imágenes personalizadas de 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. Quita las imágenes personalizadas de tu entorno local de Docker:

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

¿Qué sigue?