Run fits on a computing cluster#

When you have many participants, or a model that is slow to fit, you can spread the fits over the nodes of a computing cluster. cpm fits participants independently of each other, so the simplest approach is to split the participants into chunks, and fit each chunk in a separate job.

This guide uses a SLURM job array, which runs the same script many times, once per chunk. The same idea works with other schedulers.

1. A script that fits one chunk#

Each task of a job array gets its own number in the environment variable SLURM_ARRAY_TASK_ID, and the number of tasks in SLURM_ARRAY_TASK_COUNT. The script splits the participants into as many chunks as there are tasks, fits the chunk of its own task, and saves the results in a file named after the task. Within a task, it fits participants in parallel on the cores that SLURM gave it (SLURM_CPUS_PER_TASK), see Parallelise model fitting.

fit_chunk.py#
"""Fit the participants of one chunk of the data, as one task of a SLURM job array."""

import os
from pathlib import Path

import numpy as np

from cpm.applications.reinforcement_learning import RLRW
from cpm.datasets import load_bandit_data
from cpm.optimisation import FminBound, minimise


def main():
    # which task of the job array this is, and how many tasks there are
    task = int(os.environ.get("SLURM_ARRAY_TASK_ID", 0))
    tasks = int(os.environ.get("SLURM_ARRAY_TASK_COUNT", 1))
    cores = int(os.environ.get("SLURM_CPUS_PER_TASK", 1))

    data = load_bandit_data()
    data["observed"] = data["response"]

    # split the participants into one chunk per task, and keep this task's chunk
    chunks = np.array_split(np.sort(data.ppt.unique()), tasks)
    chunk = data[data.ppt.isin(chunks[task])]
    print(f"task {task + 1} of {tasks}: {chunk.ppt.nunique()} participants on {cores} cores")

    model = RLRW(data=chunk[chunk.ppt == chunk.ppt.iloc[0]], dimensions=4, parameters_settings=[[0.5, 0, 1], [5, 0, 10]])
    fit = FminBound(
        model=model,
        data=chunk.groupby("ppt"),
        minimisation=minimise.LogLikelihood.bernoulli,
        prior=True,
        number_of_starts=5,
        ppt_identifier="ppt",
        parallel=cores > 1,
        cl=cores,
        display=False,
        approx_grad=True,
    )
    fit.optimise()

    results = Path("results")
    results.mkdir(exist_ok=True)
    fit.export().to_csv(results / f"fit_{task:03d}.csv", index=False)


# the guard is required for parallel fitting, see the how-to guide on parallel fitting
if __name__ == "__main__":
    main()

Replace the data and the model with your own. You can test the script on your own computer: without the SLURM variables, it runs as a single task on one core.

2. A job script#

The job script asks for resources for each task, and runs the fitting script. --array=0-9 creates 10 tasks, numbered 0 to 9.

job.sbatch#
#!/bin/bash
#SBATCH --job-name=cpm-fit
#SBATCH --array=0-9               # 10 tasks, each fitting a tenth of the participants
#SBATCH --cpus-per-task=16        # cores per task, used to fit participants in parallel
#SBATCH --mem-per-cpu=2G
#SBATCH --time=02:00:00           # hours:minutes:seconds, per task
#SBATCH --output=logs/%x_%A_%a.out

# load Python and activate the environment in which cpm is installed;
# the module name and the path depend on your cluster
module load python
source ~/venvs/cpm/bin/activate

python fit_chunk.py

How you load Python and cpm depends on your cluster: ask your administrators, or see its documentation. Create the logs folder before you submit the job, then submit it with:

mkdir -p logs
sbatch job.sbatch

squeue --me shows the state of your tasks.

3. Combine the results#

When all tasks have finished, the results folder holds one file per task. Combine them into one file:

combine.py#
"""Combine the results of all tasks of the job array into one file."""

from pathlib import Path

import pandas as pd

files = sorted(Path("results").glob("fit_*.csv"))
fits = pd.concat([pd.read_csv(file) for file in files], ignore_index=True)
fits.to_csv("fits.csv", index=False)
print(f"combined {len(files)} files: {fits.ppt.nunique()} participants")

Choosing the size of the chunks#

  • The number of tasks times the number of cores per task is the number of participants fitted at the same time. There is no benefit in more tasks than participants.

  • Each task starts Python and loads the data, which takes a few seconds; make the chunks large enough that fitting takes much longer than that.

  • If a task fails, only its chunk has to be fitted again: resubmit that task alone with sbatch --array=3 job.sbatch.

  • For hierarchical estimation (cpm.hierarchical), all participants must be fitted together in every iteration, so do not split them into chunks. Run the whole analysis as a single task, with many cores.