Await close_rpc() in client

[patrix58]

Missing await in client.py causes a warning when shutting down the cluster.

Test

Tested on Ubuntu 22.04 with Python 3.10, dask-gateway==2022.6.1 and Slurm backend.

Test script:

from dask_gateway import Gateway
from dask.distributed import Client, as_completed

import time

gateway = Gateway(...)

cluster = gateway.new_cluster()
cluster.scale(3)

try:
    client = Client(cluster)
    def calc(x):
        time.sleep(10 - x)
        return x 

    jobs = client.map(calc, range(10))
    for job in as_completed(jobs):
        print(job.result())
finally:
    cluster.shutdown()

Actual output

The printed warning:

/lib/python3.10/site-packages/dask_gateway/client.py:1014: RuntimeWarning: coroutine 'rpc.close_rpc' was never awaited
  self.scheduler_comm.close_rpc()
RuntimeWarning: Enable tracemalloc to get the object allocation traceback

After

No warning is displayed, rpc close coroutine is awaited.

[consideRatio]

This makes sense, while the close_rpc function was previously a non-async function, it returned a coroutine before by returning the return value of asyncio.gather() before which makes its return value awaitable just like if it were an async function.

Thanks for your investigative work to catch this @patrix58! ❤️ 🎉

Related

Relevant upstream changes in the close_rpc() function, as changed ~3 months ago and then later ~3 years ago. This change makes sense as long as we use a version of distributed less than 3 years old in my mind, so should be fine to merge without further changes.

• https://github.com/dask/distributed/pull/6082/files#diff-581957c552b88dd04319efca1429c0c0827daff509a115542e70b661c4c04914L880
• https://github.com/dask/distributed/pull/3171/files#diff-581957c552b88dd04319efca1429c0c0827daff509a115542e70b661c4c04914R692