Crea immagini Docker Slurm personalizzate

Questo documento spiega come creare immagini Docker personalizzate per i cluster Slurm su Google Kubernetes Engine (GKE). Puoi estendere le immagini Slurm di base fornite da GKE per includere strumenti, librerie o configurazioni aggiuntivi necessari per i tuoi carichi di lavoro di computing ad alte prestazioni (HPC).

Prima di leggere questo documento, assicurati di conoscere il componente aggiuntivo Slurm Operator per GKE.

Prima di iniziare

Prima di iniziare, assicurati di aver eseguito le seguenti operazioni:

  • Attiva l'API Google Kubernetes Engine.
    • Attiva l'API Kubernetes Engine
  • Se vuoi utilizzare Google Cloud CLI per questa attività, installala e poi inizializza gcloud CLI. Se hai già installato gcloud CLI, scarica l'ultima versione eseguendo il comando gcloud components update. Le versioni precedenti di gcloud CLI potrebbero non supportare l'esecuzione dei comandi in questo documento.

Prerequisiti

Questo documento presuppone che tu abbia già un cluster Slurm in esecuzione su GKE con il componente aggiuntivo Slurm Operator per GKE installato. Completa le procedure nelle pagine seguenti:

  1. Completa la guida rapida: esegui il deployment di un cluster Slurm su GKE.
  2. Configura un repository Artifact Registry nel tuo progetto per archiviare le immagini personalizzate.

Immagini di base Slurm

GKE fornisce immagini Slurm di base nel repository Artifact Registry gcr.io/gke-release/. GKE aggiorna spesso queste immagini per motivi di sicurezza e prestazioni. Queste immagini sono disponibili in varianti che includono le versioni più recenti di Slurm e due distribuzioni Linux, Ubuntu e Rocky Linux.

Puoi personalizzare le seguenti immagini di base:

  • gcr.io/gke-release/slinky/slurmd: utilizzato per i nodi di computing Slurm.
  • gcr.io/gke-release/slinky/login: utilizzato per i nodi di accesso.

Crea un'immagine personalizzata

L'esempio seguente mostra come creare un'immagine di calcolo Slurm personalizzata che include un ambiente virtuale Python con JAX installato. Crei anche un'immagine di accesso corrispondente che rispecchia la variabile di ambiente PATH dell'immagine di calcolo senza installare effettivamente le librerie JAX.

Seleziona la versione dell'immagine

Quando selezioni un'immagine di base, assicurati che soddisfi le seguenti condizioni:

  • La versione corrisponde a quella di Slurm utilizzata da altri componenti del cluster Slurm.
  • Per una versione specifica di Slurm, scegli il tag dell'immagine più recente disponibile, che include gli aggiornamenti della sicurezza e le correzioni di bug più recenti.

Ad esempio, se la versione predefinita di Slurm nel cluster è 25.11, devi scegliere un tag che inizi con 25.11-, ad esempio 25.11-ubuntu24.04-gke.6.

Crea un Dockerfile

  1. Seleziona un tag immagine slurmd basato su Ubuntu:

    1. Nella console Google Cloud , vai alla pagina Repository Artifact Registry che include il pacchetto slinky/slurmd.

      Vai al repository Artifact Registry

    2. Trova un'immagine con un tag che includa ubuntu e corrisponda alla tua versione di Slurm, ad esempio 25.11-ubuntu24.04-gke.6.

    3. Copia il tag. Utilizza questo tag per sostituire il segnaposto VERSION_TAG nel seguente file di configurazione.

  2. Crea un file denominato Dockerfile con i seguenti contenuti:

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

    Sostituisci VERSION_TAG con il tag della versione di Slurm che corrisponde alla versione predefinita di Slurm del cluster.

  3. Crea le immagini utilizzando il 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 .

    Sostituisci quanto segue:

    • AR_PATH: il percorso del repository Artifact Registry, ad esempio gcr.io/my-project.
    • CUSTOM_SLURMD_TAG: un nome tag slurmd-custom a tua scelta.
    • CUSTOM_LOGIN_TAG: un nome tag login-custom a tua scelta.

  4. Esegui il push delle immagini personalizzate nel repository:

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

Utilizzare le immagini personalizzate in GKE

Per utilizzare le immagini personalizzate:

  1. Come mostrato nell'esempio seguente, aggiorna il repository e il tag dell'immagine per il nodeset slurmd e il loginset login modificando il file 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. Esegui l'upgrade del deployment esistente:

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

  3. Testa le nuove funzionalità del nodo di calcolo accedendo al nodo di accesso ed eseguendo il seguente 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}')
"

    L'output è simile al seguente:

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

    Questo output conferma che Slurm esegue lo script sui pod worker che eseguono l'immagine personalizzata e che l'immagine contiene le funzionalità richieste di Python e JAX.

Esegui la pulizia

Per liberare spazio dalle risorse che hai utilizzato in questo tutorial:

  1. Disinstalla il deployment Helm: sh helm uninstall slurm --namespace slurm

    Questo comando rimuove tutte le risorse Kubernetes di cui è stato eseguito il deployment dal grafico Helm.

  2. Elimina lo spazio dei nomi Slurm:

    kubectl delete namespace slurm

  3. Elimina il cluster GKE:

    gcloud container clusters delete CLUSTER_NAME

    Sostituisci CLUSTER_NAME con il nome del cluster.

  4. Elimina le immagini personalizzate da 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. Rimuovi le immagini personalizzate dall'ambiente Docker locale:

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

Passaggi successivi