what's the correct way to close a client? getting tornado stream is closed errors after a program finishes

[rbavery]

I have a script using dask distributed that calls dask.compute. it completes and saves everything, but returns a traceback at the end of the program

Number of seconds for img tile saving: 259.9357166290283
[###################################     ] | 87% Completed |  4min 51.5sNumber of seconds for coco_output creation: 29.505939245224
[####################################    ] | 91% Completed |  5min  1.7sNumber of seconds for coco_output creation: 19.211583852767944
[######################################  ] | 95% Completed |  5min 35.2sNumber of seconds for coco_output creation: 47.976495027542114
Number of seconds for 1790 images: 337.8815290927887mpleted |  5min 35.3s
Images and COCO JSON have been saved in ../data/partitions/test-no-context/.
2022-05-26 20:47:03,330 - distributed.worker - WARNING - Heartbeat to scheduler failed
Traceback (most recent call last):
  File "/root/miniconda3/envs/fastai2/lib/python3.9/site-packages/distributed/comm/tcp.py", line 226, in read
    frames_nbytes = await stream.read_bytes(fmt_size)
tornado.iostream.StreamClosedError: Stream is closed

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/root/miniconda3/envs/fastai2/lib/python3.9/site-packages/distributed/worker.py", line 1198, in heartbeat
    response = await retry_operation(
  File "/root/miniconda3/envs/fastai2/lib/python3.9/site-packages/distributed/utils_comm.py", line 381, in retry_operation
    return await retry(
  File "/root/miniconda3/envs/fastai2/lib/python3.9/site-packages/distributed/utils_comm.py", line 366, in retry
    return await coro()
  File "/root/miniconda3/envs/fastai2/lib/python3.9/site-packages/distributed/core.py", line 975, in send_recv_from_rpc
    return await send_recv(comm=comm, op=key, **kwargs)
  File "/root/miniconda3/envs/fastai2/lib/python3.9/site-packages/distributed/core.py", line 739, in send_recv
    response = await comm.read(deserializers=deserializers)
  File "/root/miniconda3/envs/fastai2/lib/python3.9/site-packages/distributed/comm/tcp.py", line 242, in read
    convert_stream_closed_error(self, e)
  File "/root/miniconda3/envs/fastai2/lib/python3.9/site-packages/distributed/comm/tcp.py", line 150, in convert_stream_closed_error
    raise CommClosedError(f"in {obj}: {exc}") from exc
distributed.comm.core.CommClosedError: in <TCP (closed) ConnectionPool.heartbeat_worker local=tcp://127.0.0.1:42452 remote=tcp://127.0.0.1:39439>: Stream is closed

this confuses me because I'm using a with statement for the client so I would think that everything gets cleaned up

    start = time.time()
    with Client() as client:  # this needs to be commented out to use single threaded for profiling
        os.makedirs(coco_outdir, exist_ok=True)
        os.makedirs(os.path.join(coco_outdir, "tiled_images_no_context"), exist_ok=True)
        class_foldes_path = Path(class_folder_path)
        class_folders = list(class_foldes_path.glob("*/"))
        coco_output = make_coco_metadata(name=name)
        coco_tiler = data.COCOtiler(
            os.path.join(coco_outdir, "tiled_images_no_context")
        )
        scene_index = 0
        coco_outputs = []
        for class_folder in class_folders:
            for scene_folder in list(class_folder.glob("*GRDH*")):
                assert "S1" in str(scene_folder)
                scene_id = os.path.basename(scene_folder)
                layer_pths = [str(i) for i in list(scene_folder.glob("*png"))]
                scene_data_tuple = dask.delayed(coco_tiler.save_background_img_tiles)(
                    scene_id,
                    layer_pths,
                    aux_datasets=[],
                    aux_resample_ratio=8,
                )
                coco_output = dask.delayed(coco_tiler.create_coco_from_photopea_layers)(
                    scene_index, scene_data_tuple, layer_pths
                )
                coco_outputs.append(coco_output)
                scene_index += 1
        final_coco_output = make_coco_metadata(name=name)
        # when we create a distributed client, dask.compute uses that isntead of thread scheduler by default
        coco_outputs = client.persist(
            coco_outputs
        )  # start computation in the background
        progress(coco_outputs)  # watch progress
        coco_outputs = dask.compute(*coco_outputs)
        for co in coco_outputs:
            final_coco_output["images"].extend(co["images"])
            final_coco_output["annotations"].extend(co["annotations"])
        for i, ann in enumerate(final_coco_output["annotations"]):
            ann["id"] = i
        coco_tiler.save_coco_output(
            final_coco_output,
            os.path.join(coco_outdir, f"./instances_{name.replace(' ', '')}.json"),
        )
        num_images = len(final_coco_output["images"])
        print(f"Number of seconds for {num_images} images: {time.time() - start}")
        print(f"Images and COCO JSON have been saved in {coco_outdir}.")

the final print statements occur and the files are saved which is how I know this completes.