CropSuite does not respect allocated CPUs by SLURM

[LenLon]

I start CropSuite using SLURM with the following SBATCH arguments:

#!/bin/bash
#SBATCH --job-name=crop_suite
#SBATCH --qos=short
#SBATCH --partition=standard
#SBATCH --cpus-per-task=24
#SBATCH --account=agrica
#SBATCH --output=slurm/%x_%j.log      
#SBATCH --error=slurm/%x_%j.log    

While squeue shows the job with the correctly allocated 24 CPUs per the #SBATCH --cpus-per-task=24 argument, CropSuite's internal multiprocessing.cpu_count() method seems to just give back the number of CPUs on the current node (128), as evidenced by these outputs in the SLURM job logs:

-> Downscaling RRPCF data for alfalfa under rf conditions
 -> Downscaling RRPCF data for avocado under rf conditions
 -> Downscaling RRPCF data for banana under rf conditions
 -> Downscaling RRPCF data for beans under rf conditions
Using 127 workers

-> Data loaded
Limiting to 127 cores
 -> Processing day #0                                      
 -> Processing day #2                                      
 -> Processing day #1                                      
 -> Processing day #4        

This gridlocks the whole HPC node :-(
Our IT demands that an environment variable in the Python code has to be set, namely PYTHON_CPU_COUNT . This limits the used CPUs if used in tandem with os.process_cpu_count(), though from Python 3.13 onwards this environment variable also works with multiprocessing.cpu_count().

Given that I have no idea about Python, where should I set this environment variable. Just in CropSuite.py?

[LenLon]

Hello again, I have fixed this by now. CropSuite does not respect PYTHON_CPU_COUNT, nor NUMBA_NUM_THREADS or OMP_NUM_THREADS or MKL_NUM_THREADS. What I ended up doing was defining my own environment variable CROPSUITE_MAX_WORKERS and adding to any code in climate_suitability_main.py and downscaling.py that defined max_proc this snippet:

env_cap = os.environ.get("CROPSUITE_MAX_WORKERS")
if env_cap:
    try:
        max_proc = max(1, min(max_proc, int(env_cap)))
    except ValueError:
        pass

This might also be fixed by changing the no_threads, av_ram = dt.get_cpu_ram() line to be a bit more restrictive?

Incidentally, this also fixed any OOM errors I had. Possible reason given by Claude:

Your SLURM job allocates 48 CPUs and 62.5 GB RAM:

    NumCPUs=48
    ReqTRES=cpu=48,mem=62.50G

But CropSuite internally decides it can use 127 workers:

    Using 127 workers (downscaling)
    Limiting to 127 cores (climsuit)

That means it is spawning far more parallel workers/threads than Slurm gave you. On many systems this causes:

    much higher temporary memory usage (each worker has buffers/arrays),
    worse cache pressure,
    and faster hitting of the cgroup memory limit → kernel OOM kill.