Please find this functionality in https://github.com/dask/knit/blob/master/knit/dask_yarn.py

Example

>>> from dask_yarn import YARNCluster
>>> cluster = YARNCluster()

>>> from dask.distributed import Client
>>> client = Client(cluster)
>>> cluster.start(2, cpus=1, memory=500)

>>> future = client.submit(lambda x: x + 1, 10)
>>> future.result()
11