В dask_jobqueue, как добавить дополнительные строки в job_script?

Question

Я настраиваю dask для запуска в кластере HPC. Я настраиваю клиента следующим образом:

Сначала измените ~/.config/dask/*.yaml, затем запустите такой код:

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())

Вот как выглядит результирующий сценарий job_script:

#!/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

Таким образом, скрипт немедленно запускает python3, но мне нужно заставить его сделать что-то вроде активации среды conda или виртуального окружения python перед запуском python. Как добавить некоторые предварительные команды в сценарий job_script?

Answer 1

Я понял, прочитав исходный код dask_jobqueue/core.py, что, к счастью, было очень легко.

В ~/.config/dask/jobqueue.yaml, редактировать env-extra. Каждая строка в списке добавляется в скрипт как команда. Например, при использовании

env-extra: ['cd foo', 'mkdir bar', 'cd bar', 'conda foo']

Job_script выглядит так:

#!/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