Cross-system Dask

[fnattino]

Can we run Dask Jobqueue outside the SLURM system (e.g. on SRC) and have workers submitted to SLURM? Dask jobqueue uses sbatch/scancel to manage jobs, can one can provide custom commands that involve connecting to the remote system to submit/delete jobs? In the worst case, one can create aliases for sbatch/scancel that do the job?

Questions:

• Where should the Dask scheduler run? Locally (on SRC) or remotely (on SLURM system)?
• The default of Dask Jobqueue is that the Dask scheduler starts when the instance of SLURMCluster is created, can one get it start remotely? (check out existing issues on remote scheduler)
• If we decide for a remote scheduler, and there is no way of having it started using the SLURMCluster, can we get the setup working manually (using dask-scheduler/dask-worker CL tools)?

[fnattino]

Ideas to setup SSH connections to SLURM system from SRC:

• Upon creating of the SRC workspace, generate a SSH key pair. The public key is uploaded on SURFcua for the SLURM system of choice. The private key is passed on to the SRC, and saved on the host (i.e. the SRC VM).
• The private key is added to the SSH agent running on the host (the agent must be running): ssh-add /path/to/the/key
• When starting the container, make sure the env variable SSH_AUTH_SOCK is passed on to the container:
  • in singularity (see reference gist/comment) with export APPTAINER_SSH_AUTH_SOCK=$SSH_AUTH_SOCK. This is enough to be able to SSH to the SLURM system without pointing to the SSH key (tested).
  • in Docker (see reference gist), by mounting the directory of the socket and by sharing the path to the socket as an environment variable, e.g. as in:

    sudo docker run --rm -t -i -v $(dirname $SSH_AUTH_SOCK):/ssh-agent -e SSH_AUTH_SOCK=/ssh-agent/agent.<PID>  ubuntu /bin/bash

    This is enough to be able to SSH into the SLURM system without pointing to the SSH key (tested) if the container is running as root user. For non-root users (see comment to gist above), we need to expose the socket to the user. For the ubuntu user, running socat UNIX-LISTEN:/home/ubuntu/.ssh/socket,fork,user=ubuntu UNIX-CONNECT:$SSH_AUTH_SOCK & as root user, and SSH_AUTH_SOCK=/home/ubuntu/.ssh/socket as ubuntu user, makes the trick - the ubuntu user can also SSH into the SLURM system without pointing to the SSH key (tested)

[fnattino]

Two issues in the "standard" setup with Dask Jobqueue running on SRC (with a local scheduler) and workers running on the SLURM system:

• First, the batch job scripts are saved locally in a temporary file, and the file is then submitted. However, the temporary file is not visible on the SLURM system, thus sbatch fails. This is different from e.g. https://gist.github.com/willirath/2176a9fa792577b269cb393995f43dda, where everything runs on the same system.
• Second, workers need to connect back to the scheduler, thus they need to be able to connect to SRC. Maybe the accessibility of SRC turns out to be an issue also for other setups, including with dask gateway.

[fnattino]

Possible solution to the first issue above: in the sbatch alias, copy the temporary file to the SLURM system and submit it there. Seems to work (jobs get submitted).

[fnattino]

Possible solution to the second issue above (tested):

• start scheduler on SRC
• on SRC, run remote port forwarding to make the scheduler reachable on SLURM system:

  ssh -i /path/to/the/key -N -R 8786:localhost:8786 spider-username@spider.surf.nl

• workers can start on SLURM system (right now only on login node):

  apptainer exec oras://ghcr.io/fnattino/test-jupyterdask-image-apptainer:latest python -m distributed.cli.dask_worker tcp://localhost:8786 --nthreads 1 --memory-limit 8.00GiB --name SLURMCluster-0 --nanny --death-timeout 600 

• get the port the worker is listening to, and, on SRC, run local port forwarding for the scheduler to reach the worker:

  ssh -i  /path/to/the/key -N -L <WORKERPORT>:ui-01:<WORKERPORT> spider-username@spider.surf.nl

[fnattino]

Update on above, now also working on compute nodes, but hacky and likely over-complicated:

• start scheduler on SRC
• on SRC, run remote port forwarding to make the scheduler reachable on SLURM system. However, since GatewayPorts=no for the SSH daemon on the server, the remote forwarding only allows the login node localhost as remote address. Thus we need to bind the forward port on all interfaces on the server side with ssh -g -L .... We can run the two things in a singe command:

  ssh -A -t -R 0.0.0.0:8786:localhost:8786 -J <USERNAME>@spider.surf.nl <USERNAME>@ui-01  "ssh -g -N -L 8787:localhost:8786 localhost"

  Note the following:

  • -t to stop the execution on the server side when running Ctrl+C.
  • -A to forward the SSH agent in order to allow login to localhost (!).
  • -J to make sure we know which of the login nodes we actually forward to.
• Workers can now be started on the SLURM system, submitting a job like the following:

  #!/usr/bin/env bash
  #SBATCH -J dask-worker
  #SBATCH -p normal
  #SBATCH -n 1 
  #SBATCH --cpus-per-task=1
  #SBATCH --mem=8G
  #SBATCH -t 1:00:00
  APPTAINER_TMPDIR=${TMPDIR}
  apptainer exec oras://ghcr.io/fnattino/test-jupyterdask-image-apptainer:latest python -m distributed.cli.dask_worker tcp://ui-01:8787 --listen-address  tcp://0.0.0.0:6000 --contact-address tcp://localhost:6000 --nthreads 1 --memory-limit 8.00GiB --name SLURMCluster-0 --nanny --death-timeout 600 --local-directory $TMPDIR

  Note the following:

  • We need to specify the port where the worker is running
• once we know the node the worker is running on, we can run the following on the SRC for the scheduler to reach the worker:

  ssh -N -L 6000:<WORKER_HOSTNAME>:6000 <USERNAME>@spider.surf.nl

[fnattino]

Useful resource for SSH tunnels: https://iximiuz.com/en/posts/ssh-tunnels/

[fnattino]

Apart of the SSH forwarding-madness, the following configure most of the elements for Dask Jobqueue:

from dask_jobqueue import SLURMCluster
cluster = SLURMCluster(
    scheduler_options={
        'port': 8786, 
        'host': 'localhost',
        'contact_address': 'tcp://ui-01:8787', 
    },
    cores=1,
    memory='8GiB',
    queue='normal',
    processes=1,
    death_timeout=600,
    local_directory='$TMPDIR',
    walltime='1:00:00',
    job_script_prologue=[
        'APPTAINER_TMPDIR=${TMPDIR}',
    ],
    python='apptainer exec oras://ghcr.io/fnattino/test-jupyterdask-image-apptainer:latest python',
)

with bin/sbatch:

#!/bin/bash 

content=`cat $@`

ssh -t <USERNAME>@spider.surf.nl "
TMPFILE=`mktemp`
cat << 'EOF' > \$TMPFILE
${content}
EOF
sbatch \$TMPFILE
"

bin/scancel:

#!/bin/bash 

ssh -q -t <USERNAME>@spider.surf.nl "scancel $@"

with bin/squeue:

#!/bin/bash 

ssh -q -t <USERNAME>@spider.surf.nl "squeue $@"

Fields that I cannot seem to set directly via the Dask JobQueue interface (but could be set via worker_extra_args) are:

• listen address (it requires the worker's port)
• contact address (same as above)