Paralelització amb arrays

Introducció

A vegades necessitem fer un submit de molts jobs amb paràmetres diferents. En aquest cas, una opció és fer servir arrays al job de SLURM. Aquesta opció és molt útil quan tenim molts paràmetres i no volem fer un submit per cada combinació de paràmetres.

Codi d’exemple

Suposem que volem executar el nostre script anomenat process_data.py, que rep 1 valors per paràmetres d’entrada. En el nostre cas els imprimirà per pantalla. Un exemple de comanda seria:

python3 process_data.py 3

I el codi seria:

import pandas as pd
import sys

# Obtenim l'ID de la tasca des dels arguments del sistema
try:
   task_id = int(sys.argv[1])
except IndexError:
   task_id = 0

# Creem un DataFrame d'exemple amb usuaris i saldos
data = {
   'usuari': ['Alice', 'Bob', 'Charlie', 'David'],
   'saldo': [100.50, 250.00, 50.75, 500.20]
}

df = pd.DataFrame(data)

# Simulem un càlcul que depèn de l'ID de la tasca. En aquest cas, simplement multipliquem el saldo per l'ID de la tasca + 1.
# Això és només un exemple, en un cas real podríem fer càlculs més complexos que depenguin de l'ID de la tasca.
df['saldo'] = df['saldo'] * (task_id + 1)

print(f"--- Resultats per a la Tasca ID: {task_id} ---")
print(df[['usuari', 'saldo']])

Crearem un arxiu per exemple anomenat launch_python_array.sh en el què inclourem l’array de valors que volem llançar. Fixem-nos que en aquest cas només tenim un paràmetre d’entrada, però podríem tenir-ne més i fer combinacions de paràmetres. A més, s’activa un venv per a garantir que tenim un entorn aïllat de python i s’usen les comandes de SLURM per a configurar el job.

#!/bin/bash

#SBATCH --job-name=pandas_venv
#SBATCH --output=resultat_%A_%a.out
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=1
#SBATCH --mem=2G
#SBATCH --time=00:10:00
#SBATCH --array=0-2

# 1. Ruta al venv (ajusta aquesta ruta al teu cas)
VENV_PATH="$HOME/tests/arrays_example/venv/bin/activate"

# 2. Activa l'entorn virtual
if [ -f "$VENV_PATH" ]; then
   echo "Activant venv..."
   source "$VENV_PATH"
else
   echo "Error, no s'ha trobat el venv a: $VENV_PATH"
   exit 1
fi

# 3. Executar el script de Python amb l'ID de la tasca com a argument
# En aquest exemple, el valor de $SLURM_ARRAY_TASK_ID serà 0, 1 o 2, depenent de la tasca que s'estigui executant.
echo "Executant Python per a la tasca $SLURM_ARRAY_TASK_ID..."
python3 process_data.py $SLURM_ARRAY_TASK_ID

# Afegim un sleep per a permetre la visualització dels resultats abans de desactivar el venv
echo "Sleep durant una estona:"
sleep 60

# 4. Desactiva l'entorn virtual
echo "Desactivant entorn..."

deactivate

echo "Procés completat."

Execució

sbatch launch_python_array.sh

Això llançarà en aquest exemple 3 tasques diferents, cada una amb un valor diferent per a l’ID de la tasca (0, 1 i 2). El nostre script de Python farà càlculs diferents en funció d’aquest ID de la tasca, i els resultats es guardaran en fitxers de sortida diferents (resultat_%A_%a.out) on %A és l’ID del job i %a és l’ID de la tasca. El resultat del nostre programa un cop finalitzada l’execució serà visible al mateix directori a on hem fet l’execució.