Repository: dask/distributed

Issue:
Dask distributed client doesn't keep the connection alive
We connect client machine to a distributed dask scheduler, submit a long running task graph and eventually call progress() to wait for completion.
In this case dask client doesn't seem to keep the connection alive.
If the connection goes through ELB or any other element that has an idle timeout, the connection is disconnected and the job is dropped (in case of ELB max timeout is 4000s).


Discussion:
By ELB do you mean Amazon's elastic load balancer?  

My guess is that you're suggesting a regular heartbeat from clients to the scheduler?
Right, we're using the AWS ELB at the moment.
Something to keep the client connection alive would help.
So far we were able to mitigate it only by sending a small request to scheduler once in a while.

Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -456,6 +456,8 @@ class resembles executors in ``concurrent.futures`` but also allows
     name: string (optional)
         Gives the client a name that will be included in logs generated on
         the scheduler for matters relating to this client
+    heartbeat_interval: int
+        Time in milliseconds between heartbeats to scheduler
 
     Examples
     --------
@@ -484,7 +486,7 @@ class resembles executors in ``concurrent.futures`` but also allows
     def __init__(self, address=None, loop=None, timeout=5,
                  set_as_default=True, scheduler_file=None,
                  security=None, asynchronous=False,
-                 name=None, **kwargs):
+                 name=None, heartbeat_interval=None, **kwargs):
 
         self.futures = dict()
         self.refcount = defaultdict(lambda: 0)
@@ -512,9 +514,13 @@ def __init__(self, address=None, loop=None, timeout=5,
         self._loop_runner = LoopRunner(loop=loop, asynchronous=asynchronous)
         self.loop = self._loop_runner.loop
 
-        self._periodic_callbacks = []
-        pc = PeriodicCallback(self._update_scheduler_info, 2000, io_loop=self.loop)
-        self._periodic_callbacks.append(pc)
+        self._periodic_callbacks = dict()
+        self._periodic_callbacks['scheduler-info'] = PeriodicCallback(
+                self._update_scheduler_info, 2000, io_loop=self.loop)
+        self._periodic_callbacks['heartbeat'] = PeriodicCallback(
+                self._heartbeat,
+                heartbeat_interval or config.get('client-heartbeat-interval', 5000),
+                io_loop=self.loop)
 
         if address is None and 'scheduler-address' in config:
             address = config['scheduler-address']
@@ -738,7 +744,7 @@ def _start(self, timeout=5, **kwargs):
 
         yield self._ensure_connected(timeout=timeout)
 
-        for pc in self._periodic_callbacks:
+        for pc in self._periodic_callbacks.values():
             pc.start()
 
         self.coroutines.append(self._handle_report())
@@ -801,6 +807,9 @@ def _ensure_connected(self, timeout=None):
     def _update_scheduler_info(self):
         self._scheduler_identity = yield self.scheduler.identity()
 
+    def _heartbeat(self):
+        self.scheduler_comm.send({'op': 'heartbeat'})
+
     def __enter__(self):
         if not self._loop_runner.is_started():
             self.start()
@@ -935,7 +944,7 @@ def _close(self, fast=False):
         self.status = 'closing'
 
         with log_errors():
-            for pc in self._periodic_callbacks:
+            for pc in self._periodic_callbacks.values():
                 pc.stop()
             self._scheduler_identity = {}
             with ignoring(AttributeError):
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -98,11 +98,13 @@ class ClientState(object):
     __slots__ = (
         'client_key',
         'wants_what',
+        'last_seen',
         )
 
     def __init__(self, client):
         self.client_key = client
         self.wants_what = set()
+        self.last_seen = time()
 
     def __repr__(self):
         return "<Client %r>" % (self.client_key,)
@@ -851,6 +853,7 @@ def __init__(
                                 'update-data': self.update_data,
                                 'report-key': self.report_on_key,
                                 'client-releases-keys': self.client_releases_keys,
+                                'heartbeat': self.client_heartbeat,
                                 'restart': self.restart}
 
         self.handlers = {'register-client': self.add_client,
@@ -1639,6 +1642,10 @@ def client_releases_keys(self, keys=None, client=None):
 
         self.transitions(recommendations)
 
+    def client_heartbeat(self, client=None):
+        """ Handle heartbeats from Client """
+        self.clients[client].last_seen = time()
+
     ###################
     # Task Validation #
     ###################


Test patch:
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
--- a/distributed/tests/test_client.py
+++ b/distributed/tests/test_client.py
@@ -5123,5 +5123,18 @@ def test_warn_when_submitting_large_values(c, s, a, b):
     assert len(record) < 2
 
 
+@gen_cluster()
+def test_scatter_direct(s, a, b):
+    c = yield Client(s.address, asynchronous=True, heartbeat_interval=10)
+
+    last = s.clients[c.id].last_seen
+
+    start = time()
+    while s.clients[c.id].last_seen == last:
+        yield gen.sleep(0.10)
+        assert time() < start + 5
+
+    yield c._close()
+
 if sys.version_info >= (3, 5):
     from distributed.tests.py3_test_client import *  # flake8: noqa