Memory leak in gather()

EDIT: We are able to avoid the `StreamClosedError` by upscaling the scheduler to 500GB memory, and chunking the dataset to four smaller runs that are aggregated seperately. However, memory usage on `gather(futures, direct=True)` still causes a memory spike several times larger than any individual worker has. The error message when this occurs is:

```
distributed.core - INFO - Event loop was unresponsive in Scheduler for 7.96s. This is often caused by long-running GIL-holding functions or moving large chunks of data. This can cause timeouts and instability.
```

Not chunking the dataset for the download runs fine on the worker nodes... just can't gather that much data in one pass without errors, which will often trigger the futures to recompute faster than the data is able to transfer =/

**What happened**:

Relatively (conceptually) simple big data task. Data is temporally indexed as orbital files-- they look something like this over Antarctica (each overlapping line is 1/12 of an orbit segment, there are ~1300):


Want them to be spatially sharded, here's what ~60 of the shards might look like:


Task is three parts:

1. Read in ~130GB data from ~1300 files in parallel
2. Aggregate data
3. Spatially re-index, write to parquet (hive) in S3

Task 1 works fine and is what most of the example code does.

Task 2 is failing. Task 3 isn't included in the minimal example.

**What you expected to happen**:

It would be nice if `gather()` worked consistently. Sometimes it works; if I use a ~60GB dataset, it **can** complete. If I split the futures list into 5 or 10 smaller lists and call `gather()` on those subsets, it has more of a chance of working. If I increase the worker size, it also has a better chance of working.

The example given is designed to fail-- consistently. The provided example is failing a 100% of the time for me. This is a scaling issue, so the example has been scaled to a size that will bring out the problem.

I can make things work by not using `gather` and doing something like:

```python
agglist = []
for future in reslist0:
agglist.append(future.result())
```

...but it's awful. Running `gather(futures, direct=True)` on a smaller 60GB dataset will complete in about 2 minutes. Doing in a loop like above takes over 20 minutes. The orbital pattern is on a 90 day repeat pattern, so I have like dozens of these things to process and would like to do it the faster way. Using `gather()` without the direct flag overwhelms the scheduler node. Originally this was returning vaex dataframes, which were unmanaged in the worker memory, but changing it to pandas dataframes gives the exact same error message.

**Minimal Complete Verifiable Example**:

```python

import h5py
import numpy as np
import pandas as pd
from astropy.time import Time
import os
import vaex
import time
import pickle
import dask

from dask_gateway import Gateway
gateway = Gateway()
options = gateway.cluster_options()
options.worker_specification = '8CPU, 32GB'
cluster = gateway.new_cluster(options)
cluster.scale(16)
client = cluster.get_client(set_as_default=True)

# worker functions

def gps2dyr(time):
"""Converts GPS time to datetime (can also do decimal years)."""
return Time(time, format='gps').datetime

def read_atl06(fname, cycle):
"""Read one ATL06 file and output 6 reduced files.

Extract variables of interest and separate the ATL06 file
into each beam (ground track) and ascending/descending orbits.
"""

# Each beam is a group
group = ['/gt1l', '/gt1r', '/gt2l', '/gt2r', '/gt3l', '/gt3r']

# Loop trough beams
dataframes = []

with h5py.File(fname, 'r') as fi:
# Check which ground tracks are present in this file
gtracks = sorted(['/'+k for k in fi.keys() if k.startswith('gt')])

for k, g in enumerate(gtracks):
# Read in data for a single beam
data = {}
# this is our unique key (per beam)
data['id'] = fi[g+'/land_ice_segments/segment_id'][:]
npts = len(data['id'])
# Load vars into memory (include as many as you want)
data['lat'] = fi[g+'/land_ice_segments/latitude'][:]
data['lon'] = fi[g+'/land_ice_segments/longitude'][:]

data['slope_y'] = fi[g+'/land_ice_segments/fit_statistics/dh_fit_dy'][:]
data['slope_x'] = fi[g+'/land_ice_segments/fit_statistics/dh_fit_dx'][:]
data['slope_x_sigma'] = fi[g+'/land_ice_segments/fit_statistics/dh_fit_dx_sigma'][:]

data['h_li'] = fi[g+'/land_ice_segments/h_li'][:]
data['s_li'] = fi[g+'/land_ice_segments/h_li_sigma'][:]
data['q_flag'] = fi[g+'/land_ice_segments/atl06_quality_summary'][:]
data['s_fg'] = fi[g+'/land_ice_segments/fit_statistics/signal_selection_source'][:]
data['snr'] = fi[g+'/land_ice_segments/fit_statistics/snr_significance'][:]
data['h_rb'] = fi[g+'/land_ice_segments/fit_statistics/h_robust_sprd'][:]
data['bsnow_conf'] = fi[g+'/land_ice_segments/geophysical/bsnow_conf'][:]

data['cloud_flg_asr'] = fi[g+'/land_ice_segments/geophysical/cloud_flg_asr'][:]
data['cloud_flg_atm'] = fi[g+'/land_ice_segments/geophysical/cloud_flg_atm'][:]
data['msw_flag'] = fi[g+'/land_ice_segments/geophysical/msw_flag'][:]
data['fbsnow_h'] = fi[g+'/land_ice_segments/geophysical/bsnow_h'][:]
data['bsnow_od'] = fi[g+'/land_ice_segments/geophysical/bsnow_od'][:]
data['layer_flag'] = fi[g+'/land_ice_segments/geophysical/layer_flag'][:]
data['bckgrd'] = fi[g+'/land_ice_segments/geophysical/bckgrd'][:]
data['e_bckgrd'] = fi[g+'/land_ice_segments/geophysical/e_bckgrd'][:]
data['n_fit_photons'] = fi[g+'/land_ice_segments/fit_statistics/n_fit_photons'][:]
data['w_surface_window_final'] = fi[g+'/land_ice_segments/fit_statistics/w_surface_window_final'][:]

delta_t = fi[g+'/land_ice_segments/delta_time'][:] # for time conversion
t_ref = fi['/ancillary_data/atlas_sdp_gps_epoch'][:] # single value

# Time in GPS seconds (secs since Jan 5, 1980)
t_gps = t_ref + delta_t

# GPS sec to datetime
data['t_year'] = gps2dyr(t_gps)
data['cycle'] = np.ones(npts, dtype=np.int8)*cycle
data['track'] = np.repeat(g[1:], npts)

# Make a dataframe out of our data dict and store it.

dataframes.append(vaex.from_dict(data))
if len(dataframes) > 0:
result = vaex.concat(dataframes).to_pandas_df()
return result

def get_thing(thing, user='', passw='', delayMax=199):
delay = np.random.randint(0, delayMax)
time.sleep(delay*0.01)
preamble = "wget --http-user=" + user + " --http-password=" + passw
middle = ' --load-cookies mycookies.txt --no-check-certificate --auth-no-challenge -r --reject "index.html*" -np -e robots=off --show-progress=off --cut-dirs=6 '
cmmd = preamble + middle + thing + " -P /tmp/"
os.system(cmmd)
lfilepath = '/tmp/n5eil01u.ecs.nsidc.org' + thing[-40:]
res = read_atl06(lfilepath, cycle=int(3))
return res

# run tasks

with open("test.txt", "rb") as fp:
iceFiles11 = pickle.load(fp)

reslist0 = []
for ice0 in iceFiles11[::2]:
reslist0.append(client.submit(get_thing, ice0, retries=1000))

# Fail
stuff1 = list(filter(None, client.gather(reslist0, errors='skip', direct=True)))

```

Here's the traceback:

```python-traceback

---------------------------------------------------------------------------
StreamClosedError Traceback (most recent call last)
/srv/conda/envs/notebook/lib/python3.8/site-packages/distributed/comm/tcp.py in read(self, deserializers)
197 try:
--> 198 frames_nbytes = await stream.read_bytes(fmt_size)
199 (frames_nbytes,) = struct.unpack(fmt, frames_nbytes)

StreamClosedError: Stream is closed

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

CommClosedError Traceback (most recent call last)
in

/srv/conda/envs/notebook/lib/python3.8/site-packages/distributed/client.py in gather(self, futures, errors, direct, asynchronous)
1964 else:
1965 local_worker = None
-> 1966 return self.sync(
1967 self._gather,
1968 futures,

/srv/conda/envs/notebook/lib/python3.8/site-packages/distributed/client.py in sync(self, func, asynchronous, callback_timeout, *args, **kwargs)
858 return future
859 else:
--> 860 return sync(
861 self.loop, func, *args, callback_timeout=callback_timeout, **kwargs
862 )

/srv/conda/envs/notebook/lib/python3.8/site-packages/distributed/utils.py in sync(loop, func, callback_timeout, *args, **kwargs)
324 if error[0]:
325 typ, exc, tb = error[0]
--> 326 raise exc.with_traceback(tb)
327 else:
328 return result[0]

/srv/conda/envs/notebook/lib/python3.8/site-packages/distributed/utils.py in f()
307 if callback_timeout is not None:
308 future = asyncio.wait_for(future, callback_timeout)
--> 309 result[0] = yield future
310 except Exception:
311 error[0] = sys.exc_info()

/srv/conda/envs/notebook/lib/python3.8/site-packages/tornado/gen.py in run(self)
760
761 try:
--> 762 value = future.result()
763 except Exception:
764 exc_info = sys.exc_info()

/srv/conda/envs/notebook/lib/python3.8/site-packages/distributed/client.py in _gather(self, futures, errors, direct, local_worker)
1858 else:
1859 self._gather_future = future
-> 1860 response = await future
1861
1862 if response["status"] == "error":

/srv/conda/envs/notebook/lib/python3.8/site-packages/distributed/client.py in _gather_remote(self, direct, local_worker)
1904 if missing_keys:
1905 keys2 = [key for key in keys if key not in data2]
-> 1906 response = await retry_operation(self.scheduler.gather, keys=keys2)
1907 if response["status"] == "OK":
1908 response["data"].update(data2)

/srv/conda/envs/notebook/lib/python3.8/site-packages/distributed/utils_comm.py in retry_operation(coro, operation, *args, **kwargs)
383 dask.config.get("distributed.comm.retry.delay.max"), default="s"
384 )
--> 385 return await retry(
386 partial(coro, *args, **kwargs),
387 count=retry_count,

/srv/conda/envs/notebook/lib/python3.8/site-packages/distributed/utils_comm.py in retry(coro, count, delay_min, delay_max, jitter_fraction, retry_on_exceptions, operation)
368 delay *= 1 + random.random() * jitter_fraction
369 await asyncio.sleep(delay)
--> 370 return await coro()
371
372

/srv/conda/envs/notebook/lib/python3.8/site-packages/distributed/core.py in send_recv_from_rpc(**kwargs)
872 name, comm.name = comm.name, "ConnectionPool." + key
873 try:
--> 874 result = await send_recv(comm=comm, op=key, **kwargs)
875 finally:
876 self.pool.reuse(self.addr, comm)

/srv/conda/envs/notebook/lib/python3.8/site-packages/distributed/core.py in send_recv(comm, reply, serializers, deserializers, **kwargs)
649 await comm.write(msg, serializers=serializers, on_error="raise")
650 if reply:
--> 651 response = await comm.read(deserializers=deserializers)
652 else:
653 response = None

/srv/conda/envs/notebook/lib/python3.8/site-packages/distributed/comm/tcp.py in read(self, deserializers)
212 self._closed = True
213 if not sys.is_finalizing():
--> 214 convert_stream_closed_error(self, e)
215 except Exception:
216 # Some OSError or a another "low-level" exception. We do not really know what

/srv/conda/envs/notebook/lib/python3.8/site-packages/distributed/comm/tcp.py in convert_stream_closed_error(obj, exc)
126 raise CommClosedError(f"in {obj}: {exc.__class__.__name__}: {exc}") from exc
127 else:
--> 128 raise CommClosedError(f"in {obj}: {exc}") from exc
129
130

CommClosedError: in : Stream is closed

```

**Anything else we need to know?**:

The instances that this was run on either had 256GB or 972GB memory, so the error isn't related to where the data is being dumped.

Here's the repo with the docker file to build the workers:

[https://github.com/pangeo-data/jupyter-earth/tree/master/hub.jupytearth.org-image](https://github.com/pangeo-data/jupyter-earth/tree/master/hub.jupytearth.org-image)

Here's the repo with the [test.txt](https://github.com/espg/DaskExampleError/blob/main/test.txt?raw=true) file that's needed to run the example and also a notebook showing the error:
[https://github.com/espg/DaskExampleError](https://github.com/espg/DaskExampleError)

@crusaderky and @ian-r-rose might find this of interest from previous conversations. Related to [https://github.com/pangeo-data/jupyter-earth/issues/89](https://github.com/pangeo-data/jupyter-earth/issues/89)

**Environment**:

- Dask version: 2021.09.1
- Python version: 3.8.12
- Operating System: Linux, 5.4.149-73.259.amzn2.x86_64
- Install method (conda, pip, source): Conda/Mamba/Docker