SlurmRunner proof-of-concept

[lgarrison]

This is a proof-of-concept of a SlurmRunner, following the idea of MPIRunner. Slurm tasks identify their "rank" with the SLURM_PROCID environment variable, and the scheduler address is handled by the scheduler file. Addresses #1.

In particular, because workers/clients/schedulers already know how to deal with scheduler files, I think(?) there's no reason the SlurmRunner needs to open the file or set a scheduler_address. This implementation just passes the scheduler file through. I could be wrong though, happy to revisit this. The API could be improved, too; right now scheduler_file="scheduler.json" has to be specified in a few places.

One possible problem with the currently implementation is that it doesn't wait for the scheduler to start before starting the client and workers (no MPI barrier). I'm not sure if this is allowed in dask. Do clients and workers have a retry mechanism, or will they fail if the scheduler hasn't started yet? I may have seen hangs because of this.

Also, shutdown isn't graceful yet; the scheduler/workers stay alive even after the client disconnects. I didn't look to me like shutdown was graceful in MPIRunner yet either (it just issues an MPI_ABORT), so I figured this was WIP.

Despite these problems, the basic idea seems to work.

[jacobtomlinson]

Thanks so much for jumping in here @lgarrison!

I put together a similar implementation a few months ago but didn't get around to upstreaming it here yet. It looks like we came to a lot of the same conclusions. I'm happy to merge this and keep iterating on it. I'm still not 100% sold on the API and where things things should live though.

I would love to hear your input on dask/community#346 where a lot of the high-level conversations around this topic are being had.

[lgarrison]

Oops, I didn't see your slurm branch when I was working on this! Glad we came to similar conclusions.

The hangs due to the client/workers starting before the scheduler were getting annoying, so I did add some scheduler-file awareness. Specifically, the client and workers now wait for the scheduler file to appear before starting. Additionally, to prevent a situation where the client/workers read an old scheduler.json file before the new scheduler starts and writes its file, I made the default filename scheduler-{}.json, where {} is formatted with the Slurm job ID. This may be overkill, though, since Dask seems reasonably good about cleaning up old scheduler files.

Additionally, to improve compatibility with Client, the SlurmRunner now reads the scheduler address from the scheduler file and stores it in self.scheduler_address. This means that this now works:

with SlurmRunner() as runner:
    with Client(runner):
        ...

It does lessen the "purity" of SlurmRunner, because it now parses the scheduler file, but I think the improved API ergonomics are worth it.

Thanks for letting me know about the community thread! Will hop in there for sure.

[jacobtomlinson]

Thanks for being patient with things here. My attention has been elsewhere but I'm coming back around to thinking about this again. I'm going to merge this in and we can move forward from there!