Repository: dask/distributed

Issue:
diagnostics.progress doesn't support TLS authenticated connections
**Version: 0.17.4**

I'm trying to display a ProgressBar within a Juypter notebook, I'm using a remote cluster which is authenticated with TLS.  Calling the progress() method with a future throws the exception `TypeError: TLS expects a `ssl_context` argument of type ssl.SSLContext (perhaps check your TLS configuration?)`

The below snippet can recreate the issue, note that execution without the progress() function works fine.

```
from time import sleep
from distributed import Client
from distributed.diagnostics import progress
from distributed.security import Security

sec = Security(tls_ca_file='rootCA.crt',
               tls_client_cert='client.crt',
               tls_client_key='client.key',
               require_encryption=True)

client = Client('tls://<IP_ADDR>:8786', security=sec)

futures = client.map(sleep, [1] * 100)
progress(futures, notebook=True)
```

Full stack trace:
```
tornado.application - ERROR - Exception in callback functools.partial(<function wrap.<locals>.null_wrapper at 0x7f944cbf4bf8>, <tornado.concurrent.Future object at 0x7f944d46eac8>)
Traceback (most recent call last):
  File "/opt/anaconda/envs/Python3/lib/python3.6/site-packages/tornado/ioloop.py", line 605, in _run_callback
    ret = callback()
  File "/opt/anaconda/envs/Python3/lib/python3.6/site-packages/tornado/stack_context.py", line 277, in null_wrapper
    return fn(*args, **kwargs)
  File "/opt/anaconda/envs/Python3/lib/python3.6/site-packages/tornado/ioloop.py", line 626, in _discard_future_result
    future.result()
  File "/opt/anaconda/envs/Python3/lib/python3.6/site-packages/tornado/concurrent.py", line 238, in result
    raise_exc_info(self._exc_info)
  File "<string>", line 4, in raise_exc_info
  File "/opt/anaconda/envs/Python3/lib/python3.6/site-packages/tornado/gen.py", line 1063, in run
    yielded = self.gen.throw(*exc_info)
  File "/opt/anaconda/envs/Python3/lib/python3.6/site-packages/distributed/diagnostics/progressbar.py", line 188, in listen
    self.comm = yield connect(self.scheduler)
  File "/opt/anaconda/envs/Python3/lib/python3.6/site-packages/tornado/gen.py", line 1055, in run
    value = future.result()
  File "/opt/anaconda/envs/Python3/lib/python3.6/site-packages/tornado/concurrent.py", line 238, in result
    raise_exc_info(self._exc_info)
  File "<string>", line 4, in raise_exc_info
  File "/opt/anaconda/envs/Python3/lib/python3.6/site-packages/tornado/gen.py", line 1063, in run
    yielded = self.gen.throw(*exc_info)
  File "/opt/anaconda/envs/Python3/lib/python3.6/site-packages/distributed/comm/core.py", line 185, in connect
    quiet_exceptions=EnvironmentError)
  File "/opt/anaconda/envs/Python3/lib/python3.6/site-packages/tornado/gen.py", line 1055, in run
    value = future.result()
  File "/opt/anaconda/envs/Python3/lib/python3.6/site-packages/tornado/concurrent.py", line 238, in result
    raise_exc_info(self._exc_info)
  File "<string>", line 4, in raise_exc_info
  File "/opt/anaconda/envs/Python3/lib/python3.6/site-packages/tornado/gen.py", line 307, in wrapper
    yielded = next(result)
  File "/opt/anaconda/envs/Python3/lib/python3.6/site-packages/distributed/comm/tcp.py", line 303, in connect
    kwargs = self._get_connect_args(**connection_args)
  File "/opt/anaconda/envs/Python3/lib/python3.6/site-packages/distributed/comm/tcp.py", line 336, in _get_connect_args
    ctx = _expect_tls_context(connection_args)
  File "/opt/anaconda/envs/Python3/lib/python3.6/site-packages/distributed/comm/tcp.py", line 282, in _expect_tls_context
    raise TypeError("TLS expects a `ssl_context` argument of type "
TypeError: TLS expects a `ssl_context` argument of type ssl.SSLContext (perhaps check your TLS configuration?)
```


Discussion:

Resolving patch:
diff --git a/distributed/diagnostics/progressbar.py b/distributed/diagnostics/progressbar.py
--- a/distributed/diagnostics/progressbar.py
+++ b/distributed/diagnostics/progressbar.py
@@ -3,6 +3,7 @@
 import logging
 from timeit import default_timer
 import sys
+import weakref
 
 from toolz import valmap
 from tornado import gen
@@ -31,6 +32,12 @@ class ProgressBar(object):
     def __init__(self, keys, scheduler=None, interval='100ms', complete=True):
         self.scheduler = get_scheduler(scheduler)
 
+        self.client = None
+        for key in keys:
+            if hasattr(key, 'client'):
+                self.client = weakref.ref(key.client)
+                break
+
         self.keys = {k.key if hasattr(k, 'key') else k for k in keys}
         self.interval = parse_timedelta(interval, default='s')
         self.complete = complete
@@ -59,17 +66,21 @@ def function(scheduler, p):
                 result.update(p.extra)
             return result
 
-        self.comm = yield connect(self.scheduler)
+        self.comm = yield connect(self.scheduler,
+                                  connection_args=self.client().connection_args
+                                  if self.client else None)
         logger.debug("Progressbar Connected to scheduler")
 
         yield self.comm.write({'op': 'feed',
                                'setup': dumps(setup),
                                'function': dumps(function),
-                               'interval': self.interval})
+                               'interval': self.interval},
+                              serializers=self.client()._serializers if self.client else None)
 
         while True:
             try:
-                response = yield self.comm.read()
+                response = yield self.comm.read(deserializers=self.client()._deserializers
+                                                if self.client else None)
             except CommClosedError:
                 break
             self._last_response = response
@@ -166,6 +177,12 @@ class MultiProgressBar(object):
     def __init__(self, keys, scheduler=None, func=key_split, interval='100ms', complete=False):
         self.scheduler = get_scheduler(scheduler)
 
+        self.client = None
+        for key in keys:
+            if hasattr(key, 'client'):
+                self.client = weakref.ref(key.client)
+                break
+
         self.keys = {k.key if hasattr(k, 'key') else k for k in keys}
         self.func = func
         self.interval = interval
@@ -196,7 +213,9 @@ def function(scheduler, p):
                 result.update(p.extra)
             return result
 
-        self.comm = yield connect(self.scheduler)
+        self.comm = yield connect(self.scheduler,
+                                  connection_args=self.client().connection_args
+                                  if self.client else None)
         logger.debug("Progressbar Connected to scheduler")
 
         yield self.comm.write({'op': 'feed',
@@ -205,7 +224,8 @@ def function(scheduler, p):
                                'interval': self.interval})
 
         while True:
-            response = yield self.comm.read()
+            response = yield self.comm.read(deserializers=self.client()._deserializers if
+                                            self.client else None)
             self._last_response = response
             self.status = response['status']
             self._draw_bar(**response)


Test patch:
diff --git a/distributed/diagnostics/tests/test_widgets.py b/distributed/diagnostics/tests/test_widgets.py
--- a/distributed/diagnostics/tests/test_widgets.py
+++ b/distributed/diagnostics/tests/test_widgets.py
@@ -78,7 +78,8 @@ def record_display(*args):
 
 from distributed.client import Client, wait
 from distributed.worker import dumps_task
-from distributed.utils_test import (cluster, inc, dec, throws, gen_cluster)
+from distributed.utils_test import (cluster, inc, dec, throws, gen_cluster,
+        gen_tls_cluster)
 from distributed.utils_test import loop  # noqa: F401
 from distributed.utils import sync
 from distributed.diagnostics.progressbar import (ProgressWidget,
@@ -92,13 +93,13 @@ def test_progressbar_widget(c, s, a, b):
     z = c.submit(inc, y)
     yield wait(z)
 
-    progress = ProgressWidget([z.key], scheduler=(s.ip, s.port), complete=True)
+    progress = ProgressWidget([z.key], scheduler=s.address, complete=True)
     yield progress.listen()
 
     assert progress.bar.value == 1.0
     assert '3 / 3' in progress.bar_text.value
 
-    progress = ProgressWidget([z.key], scheduler=(s.ip, s.port))
+    progress = ProgressWidget([z.key], scheduler=s.address)
     yield progress.listen()
 
 
@@ -113,7 +114,7 @@ def test_multi_progressbar_widget(c, s, a, b):
     other = c.submit(inc, 123)
     yield wait([other, e])
 
-    p = MultiProgressWidget([e.key], scheduler=(s.ip, s.port), complete=True)
+    p = MultiProgressWidget([e.key], scheduler=s.address, complete=True)
     yield p.listen()
 
     assert p.bars['inc'].value == 1.0
@@ -155,7 +156,7 @@ def test_multi_progressbar_widget_after_close(s, a, b):
                                  'y-1': {'x-3'}, 'y-2': {'y-1'},
                                  'e': {'y-2'}})
 
-    p = MultiProgressWidget(['x-1', 'x-2', 'x-3'], scheduler=(s.ip, s.port))
+    p = MultiProgressWidget(['x-1', 'x-2', 'x-3'], scheduler=s.address)
     yield p.listen()
 
     assert 'x' in p.bars
@@ -235,7 +236,7 @@ def test_multibar_complete(s, a, b):
                                  'y-1': {'x-3'}, 'y-2': {'y-1'},
                                  'e': {'y-2'}})
 
-    p = MultiProgressWidget(['e'], scheduler=(s.ip, s.port), complete=True)
+    p = MultiProgressWidget(['e'], scheduler=s.address, complete=True)
     yield p.listen()
 
     assert p._last_response['all'] == {'x': 3, 'y': 2, 'e': 1}
@@ -253,3 +254,31 @@ def test_fast(loop):
             p = progress(L3, multi=True, complete=True, notebook=True)
             sync(loop, p.listen)
             assert set(p._last_response['all']) == {'inc', 'dec', 'add'}
+
+
+@gen_cluster(client=True, client_kwargs={'serializers': ['msgpack']})
+def test_serializers(c, s, a, b):
+    x = c.submit(inc, 1)
+    y = c.submit(inc, x)
+    z = c.submit(inc, y)
+    yield wait(z)
+
+    progress = ProgressWidget([z], scheduler=s.address, complete=True)
+    yield progress.listen()
+
+    assert progress.bar.value == 1.0
+    assert '3 / 3' in progress.bar_text.value
+
+
+@gen_tls_cluster(client=True)
+def test_tls(c, s, a, b):
+    x = c.submit(inc, 1)
+    y = c.submit(inc, x)
+    z = c.submit(inc, y)
+    yield wait(z)
+
+    progress = ProgressWidget([z], scheduler=s.address, complete=True)
+    yield progress.listen()
+
+    assert progress.bar.value == 1.0
+    assert '3 / 3' in progress.bar_text.value
diff --git a/distributed/tests/test_tls_functional.py b/distributed/tests/test_tls_functional.py
--- a/distributed/tests/test_tls_functional.py
+++ b/distributed/tests/test_tls_functional.py
@@ -10,14 +10,8 @@
 
 from distributed import Nanny, worker_client, Queue
 from distributed.client import wait
-from distributed.utils_test import (gen_cluster, tls_only_security,
-                                    inc, double, slowinc, slowadd)
-
-
-def gen_tls_cluster(**kwargs):
-    kwargs.setdefault('ncores', [('tls://127.0.0.1', 1), ('tls://127.0.0.1', 2)])
-    return gen_cluster(scheduler='tls://127.0.0.1',
-                       security=tls_only_security(), **kwargs)
+from distributed.utils_test import (gen_tls_cluster, inc, double, slowinc,
+        slowadd)
 
 
 @gen_tls_cluster(client=True)
diff --git a/distributed/utils_test.py b/distributed/utils_test.py
--- a/distributed/utils_test.py
+++ b/distributed/utils_test.py
@@ -726,7 +726,8 @@ def iscoroutinefunction(f):
 def gen_cluster(ncores=[('127.0.0.1', 1), ('127.0.0.1', 2)],
                 scheduler='127.0.0.1', timeout=10, security=None,
                 Worker=Worker, client=False, scheduler_kwargs={},
-                worker_kwargs={}, active_rpc_timeout=1, config={}):
+                worker_kwargs={}, client_kwargs={}, active_rpc_timeout=1,
+                config={}):
     from distributed import Client
     """ Coroutine test with small cluster
 
@@ -778,7 +779,7 @@ def coro():
                             args = [s] + workers
                             if client:
                                 c = yield Client(s.address, loop=loop, security=security,
-                                                 asynchronous=True)
+                                                 asynchronous=True, **client_kwargs)
                                 args = [c] + args
                             try:
                                 future = func(*args)
@@ -1265,3 +1266,9 @@ def bump_rlimit(limit, desired):
     except Exception as e:
         pytest.skip("rlimit too low (%s) and can't be increased: %s"
                     % (soft, e))
+
+
+def gen_tls_cluster(**kwargs):
+    kwargs.setdefault('ncores', [('tls://127.0.0.1', 1), ('tls://127.0.0.1', 2)])
+    return gen_cluster(scheduler='tls://127.0.0.1',
+                       security=tls_only_security(), **kwargs)