Launching dask-mpi clusters

[rcthomas]

Is this supported? We have the ability (through Slurm) to srun ... dask-mpi... which will launch the cluster on a queue set aside for interactive work. From the command line it's a single command. Other facilities may not have this capability (it may require a batch script).

We're willing to take a look into how to make this work if it's needed. cc @shreddd

[mrocklin]

Moving this to dask/dask-mpi .

This isn't currently supported, but it could be. In principle we need to create a class that meets some of the interface in https://github.com/dask/distributed/blob/master/distributed/deploy/cluster.py

If you wanted to get a start on this, I would recommend subclassing from Cluster, and implementing _start and _close. I suspect that that will get you most of the way there.

[mrocklin]

I think also that @andersy005 and @kmpaul may have had something like this in the past. I'm not seeing it now though, so I'm curiuos if maybe I removed it when we redid things for 2.0. (although the old version probably wouldn't have met the current-day constraints of the dask.distributed.Cluster interface)

[kmpaul]

I feel like some context was lost in the transfer of this issue to dask-mpi...? Not sure what exactly is being asked.

@rcthomas I hate to ask you to repeat yourself, but can you describe what kind of functionality (exactly) you are asking about?

[kmpaul]

Is this related to #9 in any way?

[mrocklin]

I think that the ask here is for a Python class that follows the Cluster interface so that it can be integrated into the JupyterLab "new cluster" button, as is depicted here: https://youtu.be/EX_voquHdk0

[kmpaul]

Ah! Yes. This would be very cool. ...Let's talk about what this would look like, then, because I'm actually wondering if this is more of a fit for Dask-Jobqueue.

[rcthomas]

Thanks for moving this to the right place. We're looking to enable both dask-jobqueue and dask-mpi style interactive HPC. We have a configuration that allows interactive use via either salloc or srun. It looks like dask-jobqueue sets up a job and submits it, here we'd want to do the same but there's no batch script involved. I'm sure this can be accommodated, it's just a question of where...

[kmpaul]

@rcthomas Is there a difference between the two approaches, other than that one uses a batch script and the other doesn't? In other words, would a user notice a difference if one of these approaches or the other was implemented behind the "new cluster" button?

[rcthomas]

That is one difference and it is not clear to me that it's something that needs to be included or whether it can be handled via abstraction. Another difference is vanilla dask-jobqueue has a scheduler running somewhere already and workers are added or come and go. With vanilla dask-mpi the scheduler and workers all start up at once. This latter difference may not matter but I'm not that familiar with the dask-labextension to know what matters.

[mrocklin]

A minimal implementation will probably look like this ...

class MPICluster(dask.distributed.Cluster):
    async def _start(self):
        self.proc = subprocess.popen("dask-mpi", ...)
        await scheduler_has_started_up()

    async def _close(self):
        self.proc.close()

And then you'll add a config file like the following:

labextension:
  factory:
    module: 'dask_mpi'
    class: 'MPICluster'
    args: []
    kwargs: {}

I'm sure that it will be more complex than that, but in theory that should work.

[rcthomas]

I've been studying the code a bit and I have a few questions and comments. To make things a bit more concrete I'm looking to encapsulate something like the following. Many examples will be less complicated but we want to make a more user friendly way to accomplish something like the following real (working!) example:

salloc \
    -N 40 \
    -n 1280 \
    -c 2 \
    -t 60 \
    -C haswell \
    -q interactive \
    -A m3384 \
    --image=registry.services.nersc.gov/rthomas/nersc-dask-example:v0.1.0 \
    srun -u shifter \
        /opt/conda/bin/python -u \
        /opt/conda/bin/dask-mpi \
            --scheduler-file=scheduler.json \
            --dashboard-address=0 \
            --nthreads=1 \
            --memory-limit=0 \
            --no-nanny \
            --local-directory=/tmp

So @kmpaul if dask-jobqueue can be made to do that then OK (there is just no concept of a batch script here).

The next thing I wondered was whether @mrocklin you really meant to say SpecCluster as the base class and not Cluster?

Taking a step back though all I want to do is subprocess something like that command line that starts a cluster, but it could be started with straight-up srun if say, Jupyter is running on a compute node already. That (starting a cluster via subprocess) is probably happening in quite a few different Cluster implementations but I wonder if that ought to just be a base class itself? My plan was to just make something that does that first, and then move on to getting the MPI and job launcher stuff right.

[mrocklin]

I meant Cluster. SpecCluster is based around running many individual worker objects, which you won't have. You have one big job, rather than a configuration of many jobs.

[kmpaul]

@rcthomas Yes. I think you are correct that this doesn't fit into the Dask-Jobqueue paradigm. This cluster object should probably be defined in Dask-MPI using the subprocess algorithm like the one @mrocklin suggests.

[kmpaul]

@rcthomas Ok. A question for you about what you actually need.

The command that you quoted above. Is this what you want to run when you click on the Dask Labextension button?

salloc \
    -N 40 \
    -n 1280 \
    -c 2 \
    -t 60 \
    -C haswell \
    -q interactive \
    -A m3384 \
    --image=registry.services.nersc.gov/rthomas/nersc-dask-example:v0.1.0 \
    srun -u shifter \
        /opt/conda/bin/python -u \
        /opt/conda/bin/dask-mpi \
            --scheduler-file=scheduler.json \
            --dashboard-address=0 \
            --nthreads=1 \
            --memory-limit=0 \
            --no-nanny \
            --local-directory=/tmp

Or is this what you want to run when you click the Dask Labextension button?

    srun -u shifter \
        /opt/conda/bin/python -u \
        /opt/conda/bin/dask-mpi \
            --scheduler-file=scheduler.json \
            --dashboard-address=0 \
            --nthreads=1 \
            --memory-limit=0 \
            --no-nanny \
            --local-directory=/tmp

If you want to run the former command, then it falls within the purview of Dask-Jobqueue.
And if you really need to avoid creating a batch script and submitting it via SLURM (i.e., you really want to run the single command without a batch script), then I am sure a modification to Dask-Jobqueue's SLURMCluster can be made to accommodate this.

However, if you want to run the latter command, then it falls within the purview of Dask-MPI, and all we would need to do is create the MPICluster class as @mrocklin recommends.

What makes the most sense to you?