Repository: dask/dask

Issue:
Client should verify Scheduler's existence
On startup a Client should perform a handshake with the scheduler to ensure that it exists.  This would presumably be similar to the Worker-Scheduler relationship.



Discussion:

Resolving patch:
diff --git a/dask/distributed/client.py b/dask/distributed/client.py
--- a/dask/distributed/client.py
+++ b/dask/distributed/client.py
@@ -1,9 +1,11 @@
 from __future__ import print_function
 
-import zmq
+import os
 import itertools
 import uuid
 from datetime import datetime
+
+import zmq
 import dill
 from .scheduler import pickle
 from ..compatibility import unicode
@@ -33,6 +35,7 @@ def __init__(self, scheduler, address=None):
         self.socket = context.socket(zmq.DEALER)
         self.socket.setsockopt(zmq.IDENTITY, self.address)
         self.socket.connect(self.address_to_scheduler)
+        self.register_client()
 
     def get(self, dsk, keys):
         header = {'function': 'schedule',
@@ -121,3 +124,8 @@ def close(self, close_scheduler=False):
         if close_scheduler:
             self.close_scheduler()
         self.socket.close(1)
+
+    def register_client(self):
+        self.send_to_scheduler({'function': 'register'}, {'pid': os.getpid()})
+        header, payload = self.recv_from_scheduler()
+        self.registered_workers = payload['workers']
diff --git a/dask/distributed/scheduler.py b/dask/distributed/scheduler.py
--- a/dask/distributed/scheduler.py
+++ b/dask/distributed/scheduler.py
@@ -13,7 +13,6 @@
 from contextlib import contextmanager
 import traceback
 import sys
-from time import sleep
 from ..compatibility import Queue, unicode, Empty
 try:
     import cPickle as pickle
@@ -104,6 +103,9 @@ def __init__(self, port_to_workers=None, port_to_clients=None,
             self.to_clients.bind('tcp://%s:%d' % (bind_to_clients, port_to_clients))
         self.address_to_clients = ('tcp://%s:%d' % (hostname, port_to_clients)).encode()
 
+        # Client state
+        self.clients = dict()
+
         # State about my workers and computed data
         self.workers = dict()
         self.who_has = defaultdict(set)
@@ -133,6 +135,7 @@ def __init__(self, port_to_workers=None, port_to_clients=None,
                                  'setitem-ack': self._setitem_ack,
                                  'getitem-ack': self._getitem_ack}
         self.client_functions = {'status': self._status_to_client,
+                                 'register': self._client_registration,
                                  'schedule': self._schedule_from_client,
                                  'set-collection': self._set_collection,
                                  'get-collection': self._get_collection,
@@ -209,6 +212,15 @@ def block(self):
         self._listen_to_workers_thread.join()
         self._listen_to_clients_thread.join()
 
+    def _client_registration(self, header, payload):
+        """ Client comes in, register it, send back info about the cluster"""
+        payload = pickle.loads(payload)
+        address = header['address']
+        self.clients[address] = payload
+        out_header = {}
+        out_payload = {'workers': self.workers}
+        self.send_to_client(header['address'], out_header, out_payload)
+
     def _worker_registration(self, header, payload):
         """ Worker came in, register them """
         payload = pickle.loads(payload)


Test patch:
diff --git a/dask/distributed/tests/test_client.py b/dask/distributed/tests/test_client.py
--- a/dask/distributed/tests/test_client.py
+++ b/dask/distributed/tests/test_client.py
@@ -1,3 +1,5 @@
+import os
+
 import pytest
 pytest.importorskip('zmq')
 pytest.importorskip('dill')
@@ -111,3 +113,11 @@ def test_register_collections():
         c.close()
         d.close()
 
+
+def test_register_with_scheduler():
+    with scheduler_and_workers() as (s, (a, b)):
+        c = Client(s.address_to_clients)
+        pid = os.getpid()
+        assert s.clients[c.address]['pid'] == os.getpid()
+        assert (c.registered_workers == {a.address: {'pid': pid},
+                                         b.address: {'pid': pid}})