Créer des images Docker Slurm personnalisées

Ce document explique comment créer des images Docker personnalisées pour vos clusters Slurm sur Google Kubernetes Engine (GKE). Vous pouvez étendre les images Slurm de base fournies par GKE pour inclure des outils, des bibliothèques ou des configurations supplémentaires requis pour vos charges de travail de calcul hautes performances (HPC).

Avant de lire ce document, assurez-vous de connaître le module complémentaire Slurm Operator pour GKE.

Avant de commencer

Avant de commencer, effectuez les tâches suivantes :

  • Activez l'API Google Kubernetes Engine.
    • Activer l'API Google Kubernetes Engine
  • Si vous souhaitez utiliser la Google Cloud CLI pour cette tâche, installez et initialisez la gcloud CLI. Si vous avez déjà installé la gcloud CLI, obtenez la dernière version en exécutant la commande gcloud components update. Il est possible que les versions antérieures de la gcloud CLI ne permettent pas d'exécuter les commandes de ce document.

Prérequis

Ce document suppose que vous disposez déjà d'un cluster Slurm exécuté sur GKE avec le module complémentaire Slurm Operator pour GKE installé. Suivez les procédures décrites sur les pages suivantes :

  1. Suivez le guide de démarrage rapide : déployer un cluster Slurm sur GKE.
  2. Configurez un dépôt Artifact Registry dans votre projet pour stocker vos images personnalisées.

Images de base Slurm

GKE fournit des images Slurm de base dans le dépôt Artifact Registry gcr.io/gke-release/. GKE met fréquemment à jour ces images pour des raisons de sécurité et de performances. Ces images sont disponibles dans des variantes qui incluent les dernières versions de Slurm et deux distributions Linux, Ubuntu et Rocky Linux.

Vous pouvez personnaliser les images de base suivantes :

  • gcr.io/gke-release/slinky/slurmd : utilisé pour les nœuds de calcul Slurm.
  • gcr.io/gke-release/slinky/login : utilisé pour les nœuds de connexion.

Créer une image personnalisée

L'exemple suivant montre comment créer une image de calcul Slurm personnalisée qui inclut un environnement virtuel Python avec JAX installé. Vous créez également une image de connexion correspondante qui reflète la variable d'environnement PATH de l'image de calcul sans installer réellement les bibliothèques JAX.

Sélectionnez la version de l'image.

Lorsque vous sélectionnez une image de base, assurez-vous qu'elle remplit les conditions suivantes :

  • La version correspond à la version de Slurm utilisée par les autres composants de votre cluster Slurm.
  • Pour une version spécifique de Slurm, choisissez le tag de l'image la plus récente disponible, qui inclut les dernières mises à jour de sécurité et corrections de bugs.

Par exemple, si la version Slurm par défaut de votre cluster est 25.11, vous devez choisir un tag qui commence par 25.11-, par exemple 25.11-ubuntu24.04-gke.6.

Créer un fichier Dockerfile

  1. Sélectionnez un tag d'image slurmd basé sur Ubuntu :

    1. Dans la console Google Cloud , accédez à la page du dépôt Artifact Registry qui inclut le package slinky/slurmd.

      Accéder au dépôt Artifact Registry

    2. Recherchez une image avec un tag qui inclut ubuntu et correspond à votre version de Slurm, par exemple 25.11-ubuntu24.04-gke.6.

    3. Copiez le tag. Vous utilisez cette balise pour remplacer l'espace réservé VERSION_TAG dans le fichier de configuration suivant.

  2. Créez un fichier nommé Dockerfile avec le contenu suivant :

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

    Remplacez VERSION_TAG par le tag de version Slurm qui correspond à la version Slurm par défaut de votre cluster.

  3. Créez les images à l'aide de la commande 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 .

    Remplacez les éléments suivants :

    • AR_PATH : chemin d'accès à votre dépôt Artifact Registry, par exemple gcr.io/my-project.
    • CUSTOM_SLURMD_TAG : nom de tag slurmd-custom de votre choix.
    • CUSTOM_LOGIN_TAG : nom de tag login-custom de votre choix.

  4. Transférez les images personnalisées vers votre dépôt :

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

Utiliser les images personnalisées dans GKE

Pour utiliser vos images personnalisées, procédez comme suit :

  1. Comme indiqué dans l'exemple suivant, mettez à jour le dépôt et le tag d'image pour le nodeset slurmd et le loginset login en modifiant le fichier 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. Mettez à niveau le déploiement existant :

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

  3. Testez les nouvelles fonctionnalités de votre nœud de calcul en vous connectant au nœud de connexion et en exécutant la commande srun suivante :

    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}')
"

    Le résultat ressemble à ce qui suit :

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

    Ce résultat confirme que Slurm exécute le script sur les pods de nœuds de calcul exécutant votre image personnalisée, et que l'image contient les fonctionnalités Python et JAX requises.

Effectuer un nettoyage

Pour nettoyer les ressources que vous avez utilisées dans ce tutoriel, procédez comme suit :

  1. Désinstallez le déploiement Helm : sh helm uninstall slurm --namespace slurm

    Cette commande supprime toutes les ressources Kubernetes déployées par le chart Helm.

  2. Supprimez l'espace de noms Slurm :

    kubectl delete namespace slurm

  3. Supprimez le cluster GKE :

    gcloud container clusters delete CLUSTER_NAME

    Remplacez CLUSTER_NAME par le nom de votre cluster.

  4. Supprimez les images personnalisées d'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. Supprimez les images personnalisées de votre environnement Docker local :

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

Étapes suivantes