Membangun image Docker Slurm kustom

Dokumen ini menjelaskan cara membuat image Docker kustom untuk cluster Slurm di Google Kubernetes Engine (GKE). Anda dapat memperluas image dasar Slurm yang disediakan oleh GKE untuk menyertakan alat, library, atau konfigurasi tambahan yang diperlukan untuk workload komputasi berperforma tinggi (HPC).

Sebelum membaca dokumen ini, pastikan Anda sudah memahami add-on Slurm Operator untuk GKE.

Sebelum memulai

Sebelum memulai, pastikan Anda telah melakukan tugas berikut:

  • Aktifkan Google Kubernetes Engine API.
    • Aktifkan Google Kubernetes Engine API
  • Jika ingin menggunakan Google Cloud CLI untuk tugas ini, instal lalu lakukan inisialisasi gcloud CLI. Jika sebelumnya Anda telah menginstal gcloud CLI, dapatkan versi terbaru dengan menjalankan perintah gcloud components update. gcloud CLI versi sebelumnya mungkin tidak mendukung menjalankan perintah dalam dokumen ini.

Prasyarat

Dokumen ini mengasumsikan bahwa Anda telah menjalankan cluster Slurm di GKE dengan add-on Slurm Operator untuk GKE yang diinstal. Selesaikan prosedur di halaman berikut:

  1. Selesaikan Panduan memulai: Men-deploy cluster Slurm di GKE.
  2. Konfigurasi repositori Artifact Registry di project Anda untuk menyimpan image kustom Anda.

Image dasar Slurm

GKE menyediakan image dasar Slurm di repositori Artifact Registry gcr.io/gke-release/. GKE sering memperbarui image ini untuk keamanan dan performa. Image ini hadir dalam varian yang mencakup versi Slurm terbaru dan dua distribusi Linux, Ubuntu dan Rocky Linux.

Anda dapat menyesuaikan image dasar berikut:

  • gcr.io/gke-release/slinky/slurmd: digunakan untuk node komputasi Slurm.
  • gcr.io/gke-release/slinky/login: digunakan untuk node login.

Membangun image kustom

Contoh berikut menunjukkan cara membuat image komputasi Slurm kustom yang menyertakan lingkungan virtual Python dengan JAX yang diinstal. Anda juga membuat image login yang sesuai yang mencerminkan variabel lingkungan PATH image komputasi tanpa benar-benar menginstal library JAX.

Pilih versi gambar

Saat Anda memilih gambar dasar, pastikan gambar tersebut memenuhi kondisi berikut:

  • Versi ini cocok dengan versi Slurm yang digunakan oleh komponen lain di cluster Slurm Anda.
  • Untuk versi Slurm tertentu, pilih tag image terbaru yang tersedia, yang mencakup update keamanan dan perbaikan bug terbaru.

Misalnya, jika versi Slurm default di cluster Anda adalah 25.11, Anda harus memilih tag yang dimulai dengan 25.11-, misalnya 25.11-ubuntu24.04-gke.6.

Membuat Dockerfile

  1. Pilih tag image slurmd berbasis Ubuntu:

    1. Di konsol Google Cloud , buka halaman repositori Artifact Registry yang menyertakan paket slinky/slurmd.

      Buka repositori Artifact Registry

    2. Temukan gambar dengan tag yang menyertakan ubuntu dan cocok dengan versi Slurm Anda, misalnya 25.11-ubuntu24.04-gke.6.

    3. Salin tag. Anda menggunakan tag ini untuk mengganti placeholder VERSION_TAG dalam file konfigurasi berikut.

  2. Buat file bernama Dockerfile dengan konten berikut:

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

    Ganti VERSION_TAG dengan tag versi Slurm yang cocok dengan versi Slurm default cluster Anda.

  3. Bangun image menggunakan perintah 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 .

    Ganti kode berikut:

    • AR_PATH: jalur ke repositori Artifact Registry Anda, misalnya gcr.io/my-project.
    • CUSTOM_SLURMD_TAG: nama tag slurmd-custom pilihan Anda.
    • CUSTOM_LOGIN_TAG: nama tag login-custom pilihan Anda.

  4. Kirim gambar kustom ke repositori Anda:

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

Menggunakan image kustom di GKE

Untuk menggunakan image kustom, selesaikan langkah-langkah berikut:

  1. Seperti yang ditunjukkan dalam contoh berikut, perbarui repositori dan tag image untuk nodeset slurmd dan loginset login dengan mengubah 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. Upgrade deployment yang ada:

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

  3. Uji kemampuan baru node komputasi Anda dengan login ke node login dan menjalankan perintah srun berikut:

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

    Outputnya mirip dengan hal berikut ini:

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

    Output ini mengonfirmasi bahwa Slurm menjalankan skrip di Pod pekerja yang menjalankan image kustom Anda, dan image tersebut berisi kemampuan Python dan JAX yang diperlukan.

Pembersihan

Untuk membersihkan resource yang Anda gunakan dalam tutorial ini, lakukan hal berikut:

  1. Uninstal deployment Helm: sh helm uninstall slurm --namespace slurm

    Perintah ini akan menghapus semua resource Kubernetes yang di-deploy oleh Helm chart.

  2. Hapus namespace Slurm:

    kubectl delete namespace slurm

  3. Hapus cluster GKE:

    gcloud container clusters delete CLUSTER_NAME

    Ganti CLUSTER_NAME dengan nama cluster Anda.

  4. Hapus image kustom dari 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. Hapus image kustom dari lingkungan Docker lokal Anda:

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

Langkah berikutnya