Hi everyone, I am new to prefect but love how it works already. I am trying to submit a Dask SLURMCluster using multi-gpu nodes and I want to let each single prefect task see only one gpu (distributing multiple gpus for many parallel tasks in a flow). Is it possible? I have tried

SLURM_CPU_BIND="cores"

with

--gpus-per-task=1

but no luck. Here is my SLURMCluster input

cluster_kwargs = dict(
    cores=4*nnodes,
    memory="64 GB",
    shebang="#!/bin/bash",
    account="m3828_g",
    walltime="02:00:00",
    job_mem="0",
    job_script_prologue=[
        "source ~/.bashrc",
        "module load python",
        'export SLURM_CPU_BIND="cores"'
        # "CUDA_VISIBLE_DEVICES=$SLURM_PROCID"
    ],
    job_directives_skip=["-n", "--cpus-per-task"],
    job_extra_directives=[
        "-q preempt", "-C gpu",
        f"-N {nnodes}",
        "--ntasks-per-node=4",
        "--cpus-per-task=4",
        "--gpus-per-task=1",
        "--comment=12:00:00",
        "--time-min=01:00:00",
        "--signal=B:USR1@60",
        "--requeue",
        "--open-mode=append"
    ],
    # python="srun python",
    death_timeout=86400,
)

cluster = SLURMCluster(**cluster_kwargs)

👋