[dagster-dask] Support execution on dask_yarn

[natekupp]

Right now we just pass config along to the dask.distributed.Client: https://github.com/dagster-io/dagster/blob/master/python_modules/libraries/dagster-dask/dagster_dask/engine.py#L83

that would need to change to support dask_yarn which needs to pass in a YARN cluster object to the client

[natekupp]

hmm wonder if we should use something like https://jcristharif.com/introducing-skein.html and just support YARN directly

[DavidKatz-il]

I have an idea how to address this issue.
When configuring the dask Client we can submit a specific type of cluster.

We cam amend the execute function in the DaskEngine class to include the following (similar to dask_yarn example)

from dask.distributed import Client, SSHCluster
from dask_yarn import YarnCluster
from dask_jobqueue import PBSCluster
from dask_kubernetes import KubeCluster

client_type = pipeline_context.client_type  # coming from environment_dict
if client_type == 'local':
   cluster = LocalCluster(**dask_config.build_dict(pipeline_name))
elif client_type == 'yarn':
   cluster = YarnCluster(**dask_config.build_dict(pipeline_name))
elif client_type == 'ssh':
   cluster = SSHCluster(**dask_config.build_dict(pipeline_name))
elif client_type == 'pbs':
   cluster = PBSCluster(**dask_config.build_dict(pipeline_name))
elif client_type == 'kube':
   cluster = KubeCluster(**dask_config.build_dict(pipeline_name))
else:
   raise ValueError(f"'client_type' must be one of the following ('local', 'yarn', 'ssh', 'pbs', 'kube') not {client_type}") 

with dask.distributed.Client(cluster) as client:
...

If it make sense i will be happy to submit a PR.

[natekupp]

@DavidKatz-il yes something along those lines could work; one additional concern I should flag here is that to properly support these, we should expose a way to configure the cluster via the executor config schema (currently: https://github.com/dagster-io/dagster/blob/master/python_modules/libraries/dagster-dask/dagster_dask/executor.py#L9-L36)

We could either explicitly create the config schema fields for each cluster configuration parameter, or use Permissive() to allow the end user to to pass through these fields. While the latter would bypass config schema validation, it has the advantage that we won't need to worry about drift between our APIs and the underlying Dask configs, so I think that would be preferable.

Please ping me on the Dagster Slack if you decide to jump in, I'm happy to help you get a PR together!

[natekupp]

Fixed by #2498