In dask_jobqueue, how to add extra lines to job_script? - dask

Question

I'm configuring dask to run on an HPC cluster. I setup a client as follows:
First modify ~/.config/dask/*.yaml then run some code like this:
from dask_jobqueue import SLURMCluster
cluster = SLURMCluster()
cluster.scale(100) # Start 100 workers in 100 jobs
from distributed import Client
client = Client(cluster)
print(cluster.job_script())
Here's what the resultant job_script looks like:
#!/bin/bash
#!/usr/bin/env bash
#SBATCH -J dask-worker
#SBATCH -n 1
#SBATCH --cpus-per-task=1
#SBATCH --mem=1GB
#SBATCH -t 00:30:00
JOB_ID=${SLURM_JOB_ID%;*}
/path/to/python3 -m distributed.cli.dask_worker tcp://192.168.*.*:* --nthreads 1 --memory-limit 1000.00MB --name dask-worker--${JOB_ID}-- --death-timeout 60 --local-directory /scratch
So the script launches python3 immediately, but I need to make it do some stuff like activating a conda environment, or python virtual env before launching python. How can I add some pre-commands to the job_script?

Answer 1

I got it, by reading source code of dask_jobqueue/core.py, which was thankfully very easy.
In ~/.config/dask/jobqueue.yaml, edit env-extra. Each string in the list is added to the script as a command. For example, when using
env-extra: ['cd foo', 'mkdir bar', 'cd bar', 'conda foo']
The job_script comes out like this:
#!/bin/bash
#!/usr/bin/env bash
#SBATCH -J dask-worker
#SBATCH -n 1
#SBATCH --cpus-per-task=1
#SBATCH --mem=1GB
#SBATCH -t 00:30:00
JOB_ID=${SLURM_JOB_ID%;*}
cd foo
mkdir bar
cd bar
conda foo
/path/to/python3 -m distributed.cli.dask_worker tcp://192.168.*.*:* --nthreads 1 --memory-limit 1000.00MB --name dask-worker--${JOB_ID}-- --death-timeout 60 --local-directory /scratch