Repository: dask/dask

Issue:
`dask.config.set(num_workers=n)` context doesn't seem to work anymore
I don't know how to make a self contained example, but suffice to say that if I try something like

```python
with dask.config.set(num_workers=2):
    # code the does something
```

All cores start firing, not just 2.

```
%watermark -v -m -p dask
CPython 3.6.6
IPython 7.0.1

dask 0.19.3

compiler   : MSC v.1900 64 bit (AMD64)
system     : Windows
release    : 10
machine    : AMD64
processor  : Intel64 Family 6 Model 63 Stepping 2, GenuineIntel
CPU cores  : 48
interpreter: 64bit
```


Discussion:
mrocklin: Seems fine to me.  I wonder if it would be easy to test this with a diagnostic, similar to what produced the image you made in the MRE.

Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index 07aff0ac2..00d09a00c 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -2387,6 +2387,8 @@ def from_delayed(value, shape, dtype, name=None):
     name = name or 'from-value-' + tokenize(value, shape, dtype)
     dsk = {(name,) + (0,) * len(shape): value.key}
     chunks = tuple((d,) for d in shape)
+    # TODO: value._key may not be the name of the layer in value.dask
+    # This should be fixed after we build full expression graphs
     return Array(sharedict.merge(value.dask, (name, dsk),
                                  dependencies={name: {value._key}}),
                  name, chunks, dtype)
diff --git a/dask/array/linalg.py b/dask/array/linalg.py
index a12c06b89..cd9929b58 100644
--- a/dask/array/linalg.py
+++ b/dask/array/linalg.py
@@ -128,21 +128,23 @@ def tsqr(data, compute_svd=False, _max_vchunk_size=None):
     name_qr_st1 = 'qr' + token
     dsk_qr_st1 = top(_wrapped_qr, name_qr_st1, 'ij', data.name, 'ij',
                      numblocks={data.name: numblocks})
-    dsk.update_with_key(dsk_qr_st1, key=name_qr_st1)
+    dsk.update_with_key(dsk_qr_st1, key=name_qr_st1,
+                        dependencies={name_qr_st1, data.name})
 
     # Block qr[0]
     name_q_st1 = 'getitem' + token + '-q1'
     dsk_q_st1 = dict(((name_q_st1, i, 0),
                       (operator.getitem, (name_qr_st1, i, 0), 0))
                      for i in range(numblocks[0]))
-    dsk.update_with_key(dsk_q_st1, key=name_q_st1)
+    dsk.update_with_key(dsk_q_st1, key=name_q_st1,
+                        dependencies={name_qr_st1})
 
     # Block qr[1]
     name_r_st1 = 'getitem' + token + '-r1'
     dsk_r_st1 = dict(((name_r_st1, i, 0),
                       (operator.getitem, (name_qr_st1, i, 0), 1))
                      for i in range(numblocks[0]))
-    dsk.update_with_key(dsk_r_st1, key=name_r_st1)
+    dsk.update_with_key(dsk_r_st1, key=name_r_st1, dependencies={name_qr_st1})
 
     # Next step is to obtain a QR decomposition for the stacked R factors, so either:
     # - gather R factors into a single core and do a QR decomposition
@@ -183,10 +185,12 @@ def tsqr(data, compute_svd=False, _max_vchunk_size=None):
                                            [(name_r_st1, idx, 0)
                                             for idx, _ in sub_block_info])))
                              for i, sub_block_info in enumerate(all_blocks))
-        dsk.update_with_key(dsk_r_stacked, key=name_r_stacked)
+        dsk.update_with_key(dsk_r_stacked, key=name_r_stacked,
+                            dependencies={name_r_st1})
 
         # retrieve R_stacked for recursion with tsqr
         vchunks_rstacked = tuple([sum(map(lambda x: x[1], sub_block_info)) for sub_block_info in all_blocks])
+        dsk.dependencies[name_r_stacked] = {data.name}
         r_stacked = Array(dsk, name_r_stacked,
                           shape=(sum(vchunks_rstacked), n), chunks=(vchunks_rstacked, (n)), dtype=rr.dtype)
 
@@ -204,19 +208,22 @@ def tsqr(data, compute_svd=False, _max_vchunk_size=None):
                          for i, sub_block_info in enumerate(all_blocks)
                          for j, e in zip([x[0] for x in sub_block_info],
                                          _cumsum_blocks([x[1] for x in sub_block_info])))
-        dsk.update_with_key(dsk_q_st2, key=name_q_st2)
+        dsk.update_with_key(dsk_q_st2, key=name_q_st2,
+                            dependencies={q_inner.name})
 
         # R: R_inner
         name_r_st2 = 'r-inner-' + token
         dsk_r_st2 = {(name_r_st2, 0, 0): (r_inner.name, 0, 0)}
-        dsk.update_with_key(dsk_r_st2, key=name_r_st2)
+        dsk.update_with_key(dsk_r_st2, key=name_r_st2,
+                            dependencies={r_inner.name})
 
         # Q: Block qr[0] (*) Q_inner
         name_q_st3 = 'dot-' + token + '-q3'
         dsk_q_st3 = top(np.dot, name_q_st3, 'ij', name_q_st1, 'ij',
                         name_q_st2, 'ij', numblocks={name_q_st1: numblocks,
                                                      name_q_st2: numblocks})
-        dsk.update_with_key(dsk_q_st3, key=name_q_st3)
+        dsk.update_with_key(dsk_q_st3, key=name_q_st3,
+                            dependencies={name_q_st1, name_q_st2, name_q_st3})
     else:
         # Do single core computation
 
@@ -225,19 +232,22 @@ def tsqr(data, compute_svd=False, _max_vchunk_size=None):
         name_r_st1_stacked = 'stack' + token + '-r1'
         dsk_r_st1_stacked = {(name_r_st1_stacked, 0, 0): (np.vstack,
                                                           (tuple, to_stack))}
-        dsk.update_with_key(dsk_r_st1_stacked, key=name_r_st1_stacked)
+        dsk.update_with_key(dsk_r_st1_stacked, key=name_r_st1_stacked,
+                            dependencies={name_r_st1})
 
         # In-core QR computation
         name_qr_st2 = 'qr' + token + '-qr2'
         dsk_qr_st2 = top(np.linalg.qr, name_qr_st2, 'ij', name_r_st1_stacked, 'ij',
                          numblocks={name_r_st1_stacked: (1, 1)})
-        dsk.update_with_key(dsk_qr_st2, key=name_qr_st2)
+        dsk.update_with_key(dsk_qr_st2, key=name_qr_st2,
+                            dependencies={name_r_st1_stacked})
 
         # In-core qr[0]
         name_q_st2_aux = 'getitem' + token + '-q2-aux'
         dsk_q_st2_aux = {(name_q_st2_aux, 0, 0): (operator.getitem,
                                                   (name_qr_st2, 0, 0), 0)}
-        dsk.update_with_key(dsk_q_st2_aux, key=name_q_st2_aux)
+        dsk.update_with_key(dsk_q_st2_aux, key=name_q_st2_aux,
+                            dependencies={name_qr_st2})
 
         if not any(np.isnan(c) for cs in data.chunks for c in cs):
             # when chunks are all known...
@@ -246,6 +256,7 @@ def tsqr(data, compute_svd=False, _max_vchunk_size=None):
             block_slices = [(slice(e[0], e[1]), slice(0, n))
                             for e in _cumsum_blocks(q2_block_sizes)]
             dsk_q_blockslices = {}
+            dependencies = set()
         else:
             # when chunks are not already known...
 
@@ -276,28 +287,33 @@ def tsqr(data, compute_svd=False, _max_vchunk_size=None):
                                             dsk_q2_cumsum,
                                             dsk_block_slices)
 
+            dependencies = {data.name, name_q2bs, name_q2cs}
             block_slices = [(name_blockslice, i) for i in range(numblocks[0])]
 
-        dsk.update_with_key(dsk_q_blockslices, key='q-blocksizes' + token)
+        dsk.update_with_key(dsk_q_blockslices, key='q-blocksizes' + token,
+                            dependencies=dependencies)
 
         # In-core qr[0] unstacking
         name_q_st2 = 'getitem' + token + '-q2'
         dsk_q_st2 = dict(((name_q_st2, i, 0),
                           (operator.getitem, (name_q_st2_aux, 0, 0), b))
                          for i, b in enumerate(block_slices))
-        dsk.update_with_key(dsk_q_st2, key=name_q_st2)
+        dsk.update_with_key(dsk_q_st2, key=name_q_st2,
+                            dependencies={name_q_st2_aux})
 
         # Q: Block qr[0] (*) In-core qr[0]
         name_q_st3 = 'dot' + token + '-q3'
         dsk_q_st3 = top(np.dot, name_q_st3, 'ij', name_q_st1, 'ij',
                         name_q_st2, 'ij', numblocks={name_q_st1: numblocks,
                                                      name_q_st2: numblocks})
-        dsk.update_with_key(dsk_q_st3, key=name_q_st3)
+        dsk.update_with_key(dsk_q_st3, key=name_q_st3,
+                            dependencies={name_q_st1, name_q_st2})
 
         # R: In-core qr[1]
         name_r_st2 = 'getitem' + token + '-r2'
         dsk_r_st2 = {(name_r_st2, 0, 0): (operator.getitem, (name_qr_st2, 0, 0), 1)}
-        dsk.update_with_key(dsk_r_st2, key=name_r_st2)
+        dsk.update_with_key(dsk_r_st2, key=name_r_st2,
+                            dependencies={name_qr_st2})
 
     if not compute_svd:
         is_unknown_m = np.isnan(data.shape[0]) or any(np.isnan(c) for c in data.chunks[0])
@@ -327,6 +343,8 @@ def tsqr(data, compute_svd=False, _max_vchunk_size=None):
             r_shape = (n, n) if data.shape[0] >= data.shape[1] else data.shape
             r_chunks = r_shape
 
+        dsk.dependencies[name_q_st3] = {data.name}
+        dsk.dependencies[name_r_st2] = {data.name}
         q = Array(dsk, name_q_st3,
                   shape=q_shape, chunks=q_chunks, dtype=qq.dtype)
         r = Array(dsk, name_r_st2,
@@ -355,11 +373,11 @@ def tsqr(data, compute_svd=False, _max_vchunk_size=None):
                         name_u_st2, 'kj', numblocks={name_q_st3: numblocks,
                                                      name_u_st2: (1, 1)})
 
-        dsk.update_with_key(dsk_svd_st2, key=name_svd_st2)
-        dsk.update_with_key(dsk_u_st2, key=name_u_st2)
-        dsk.update_with_key(dsk_u_st4, key=name_u_st4)
-        dsk.update_with_key(dsk_s_st2, key=name_s_st2)
-        dsk.update_with_key(dsk_v_st2, key=name_v_st2)
+        dsk.update_with_key(dsk_svd_st2, key=name_svd_st2, dependencies={name_r_st2})
+        dsk.update_with_key(dsk_u_st2, key=name_u_st2, dependencies={name_svd_st2})
+        dsk.update_with_key(dsk_u_st4, key=name_u_st4, dependencies={name_q_st3, name_u_st2})
+        dsk.update_with_key(dsk_s_st2, key=name_s_st2, dependencies={name_svd_st2})
+        dsk.update_with_key(dsk_v_st2, key=name_v_st2, dependencies={name_svd_st2})
 
         uu, ss, vvh = np.linalg.svd(np.ones(shape=(1, 1), dtype=data.dtype))
 
@@ -434,11 +452,11 @@ def sfqr(data, name=None):
     name_A_rest = prefix + 'A_rest'
     dsk.update_with_key({
         (name_A_1, 0, 0): (data.name, 0, 0)
-    }, key=name_A_1)
+    }, key=name_A_1, dependencies={data.name})
     dsk.update_with_key({
         (name_A_rest, 0, idx): (data.name, 0, 1 + idx)
         for idx in range(nc - 1)
-    }, key=name_A_rest)
+    }, key=name_A_rest, dependencies={data.name})
 
     # Q R_1 = A_1
     name_Q_R1 = prefix + 'Q_R_1'
@@ -446,13 +464,13 @@ def sfqr(data, name=None):
     name_R_1 = prefix + 'R_1'
     dsk.update_with_key({
         (name_Q_R1, 0, 0): (np.linalg.qr, (name_A_1, 0, 0))
-    }, key=name_Q_R1)
+    }, key=name_Q_R1, dependencies={name_A_1})
     dsk.update_with_key({
         (name_Q, 0, 0): (operator.getitem, (name_Q_R1, 0, 0), 0)
-    }, key=name_Q)
+    }, key=name_Q, dependencies={name_Q_R1})
     dsk.update_with_key({
         (name_R_1, 0, 0): (operator.getitem, (name_Q_R1, 0, 0), 1),
-    }, key=name_R_1)
+    }, key=name_R_1, dependencies={name_Q_R1})
 
     Q = Array(dsk, name_Q,
               shape=(m, min(m, n)), chunks=(m, min(m, n)), dtype=qq.dtype)
diff --git a/dask/array/top.py b/dask/array/top.py
index 7db97a7d2..266ba37bb 100644
--- a/dask/array/top.py
+++ b/dask/array/top.py
@@ -69,7 +69,7 @@ def _top(func, output, output_indices, *arrind_pairs, **kwargs):
                           for a in arrind_pairs[1::2]]
     arrind_pairs[1::2] = [index_subs(a, sub)
                           for a in arrind_pairs[1::2]]
-    new_axes = {index_subs(k, sub)[0]: v for k, v in new_axes.items()}
+    new_axes = {index_subs((k,), sub)[0]: v for k, v in new_axes.items()}
 
     # Unpack dask values in non-array arguments
     argpairs = list(toolz.partition(2, arrind_pairs))
@@ -572,7 +572,7 @@ def optimize_atop(full_graph, keys=()):
     keep = {k[0] if type(k) is tuple else k for k in keys}
     layers = full_graph.dicts
     dependents = core.reverse_dict(full_graph.dependencies)
-    roots = {k for k, v in full_graph.dicts.items()
+    roots = {k for k in full_graph.dicts
              if not dependents.get(k)}
     stack = list(roots)
 
@@ -582,7 +582,7 @@ def optimize_atop(full_graph, keys=()):
 
     while stack:
         layer = stack.pop()
-        if layer in seen:
+        if layer in seen or layer not in layers:
             continue
         seen.add(layer)
 
@@ -695,11 +695,16 @@ def rewrite_atop(inputs):
             # Bump new inputs up in list
             sub = {}
             for i, index in enumerate(new_indices):
-                if index not in indices:  # use old inputs if available
+                try:
+                    contains = index in indices
+                except (ValueError, TypeError):
+                    contains = False
+
+                if contains:  # use old inputs if available
+                    sub[atop_token(i)] = atop_token(indices.index(index))
+                else:
                     sub[atop_token(i)] = atop_token(len(indices))
                     indices.append(index)
-                else:
-                    sub[atop_token(i)] = atop_token(indices.index(index))
             new_dsk = subs(inputs[dep].dsk, sub)
 
             # indices.extend(new_indices)
@@ -725,9 +730,10 @@ def rewrite_atop(inputs):
     sub = {atop_token(k): atop_token(v) for k, v in sub.items()}
     dsk = {k: subs(v, sub) for k, v in dsk.items()}
 
+    indices_check = {k for k, v in indices if v is not None}
     numblocks = toolz.merge([inp.numblocks for inp in inputs.values()])
     numblocks = {k: v for k, v in numblocks.items()
-                 if k in toolz.pluck(0, indices)}
+                 if v is None or k in indices_check}
 
     out = TOP(root, inputs[root].output_indices, dsk, new_indices,
               numblocks=numblocks, new_axes=new_axes, concatenate=concatenate)
diff --git a/dask/delayed.py b/dask/delayed.py
index a072c8706..ff3cbe358 100644
--- a/dask/delayed.py
+++ b/dask/delayed.py
@@ -502,7 +502,8 @@ class DelayedLeaf(Delayed):
 
     @property
     def dask(self):
-        return sharedict.ShareDict({self._key: {self._key: self._obj}}, {})
+        return sharedict.ShareDict({self._key: {self._key: self._obj}},
+                                   dependencies={self._key: set()})
 
     def __call__(self, *args, **kwargs):
         return call_function(self._obj, self._key, args, kwargs,
diff --git a/dask/multiprocessing.py b/dask/multiprocessing.py
index aa825da6b..3a377b54c 100644
--- a/dask/multiprocessing.py
+++ b/dask/multiprocessing.py
@@ -160,6 +160,7 @@ def get(dsk, keys, num_workers=None, func_loads=None, func_dumps=None,
         If True [default], `fuse` is applied to the graph before computation.
     """
     pool = config.get('pool', None)
+    num_workers = num_workers or config.get('num_workers', None)
     if pool is None:
         context = get_context()
         pool = context.Pool(num_workers,
diff --git a/dask/sharedict.py b/dask/sharedict.py
index ce239e7bd..438d7ea44 100644
--- a/dask/sharedict.py
+++ b/dask/sharedict.py
@@ -53,7 +53,9 @@ class ShareDict(Mapping):
         self.dicts = dicts or dict()
         self.dependencies = dependencies or dict()
 
-    def update_with_key(self, arg, key=None):
+        assert set(self.dependencies) == set(self.dicts)
+
+    def update_with_key(self, arg, key=None, dependencies=None):
         if type(arg) is ShareDict:
             assert key is None or key in arg.dicts
             self.dicts.update(arg.dicts)
@@ -63,6 +65,10 @@ class ShareDict(Mapping):
         if key is None:
             key = id(arg)
 
+        if dependencies:
+            assert isinstance(dependencies, (tuple, list, set))
+            self.dependencies[key] = set(dependencies)
+
         assert isinstance(arg, Mapping)
         if arg:
             self.dicts[key] = arg
diff --git a/dask/threaded.py b/dask/threaded.py
index fd336db63..84a382e9c 100644
--- a/dask/threaded.py
+++ b/dask/threaded.py
@@ -56,6 +56,7 @@ def get(dsk, result, cache=None, num_workers=None, **kwargs):
     """
     global default_pool
     pool = config.get('pool', None)
+    num_workers = num_workers or config.get('num_workers', None)
     thread = current_thread()
 
     with pools_lock:


Test patch:
diff --git a/dask/array/tests/test_atop.py b/dask/array/tests/test_atop.py
index 6ae073d81..7d0541767 100644
--- a/dask/array/tests/test_atop.py
+++ b/dask/array/tests/test_atop.py
@@ -231,6 +231,18 @@ def test_atop_new_axes():
     assert_eq(y, np.ones((4, 5)) * 6)
 
 
+def test_atop_new_axes_2():
+    x = da.ones((2, 2), chunks=(1, 1))
+
+    def func(x):
+        return np.stack([x, -x], axis=-1)
+
+    y = atop(func, ('x', 'y', 'sign'), x, ('x', 'y'), dtype=x.dtype,
+             concatenate=True, new_axes={'sign': 2})
+
+    assert_eq(y, y)
+
+
 @pytest.mark.parametrize('concatenate', [True, False])
 def test_atop_stacked_new_axes(concatenate):
     def f(x):
@@ -319,3 +331,33 @@ def test_atop_raises_on_incorrect_indices():
 
     assert 'ii' in str(info.value)
     assert '1' in str(info.value)
+
+
+def test_atop_numpy_arg():
+    x = da.arange(10, chunks=(5,))
+    y = np.arange(1000)
+
+    x = x.map_blocks(lambda x, y: x, 1.0)
+    x = x.map_blocks(lambda x, y: x, 'abc')
+    x = x.map_blocks(lambda x, y: x, y)
+    x = x.map_blocks(lambda x, y: x, 'abc')
+    x = x.map_blocks(lambda x, y: x, 1.0)
+    x = x.map_blocks(lambda x, y, z: x, 'abc', np.array(['a', 'b'], dtype=object))
+    assert_eq(x, np.arange(10))
+
+
+def test_bag_array_conversion():
+    import dask.bag as db
+    b = db.range(10, npartitions=1)
+    x, = b.map_partitions(np.asarray).to_delayed()
+    x, = [da.from_delayed(a, shape=(10,), dtype=int) for a in [x]]
+    z = da.concatenate([x])
+    assert_eq(z, np.arange(10), check_graph=False)
+
+
+def test_svd():
+    x = da.ones((1, 1), chunks=(1, 1))
+    y = x * 2
+    u, s, v = da.linalg.svd(y)
+    z = y + u
+    assert_eq(z, z)
diff --git a/dask/dataframe/tests/test_rolling.py b/dask/dataframe/tests/test_rolling.py
index 04be39fa1..a37b2dbc7 100644
--- a/dask/dataframe/tests/test_rolling.py
+++ b/dask/dataframe/tests/test_rolling.py
@@ -100,8 +100,8 @@ rolling_method_args_check_less_precise = [
     ('median', (), False),
     ('min', (), False),
     ('max', (), False),
-    ('std', (), False),
-    ('var', (), False),
+    ('std', (), True),
+    ('var', (), True),
     ('skew', (), True),  # here and elsewhere, results for kurt and skew are
     ('kurt', (), True),  # checked with check_less_precise=True so that we are
     # only looking at 3ish decimal places for the equality check
diff --git a/dask/tests/test_base.py b/dask/tests/test_base.py
index 6f1c0a01b..24cef408b 100644
--- a/dask/tests/test_base.py
+++ b/dask/tests/test_base.py
@@ -19,6 +19,7 @@ from dask.delayed import Delayed
 from dask.utils import tmpdir, tmpfile, ignoring
 from dask.utils_test import inc, dec
 from dask.compatibility import long, unicode, PY2
+from dask.diagnostics import Profiler
 
 
 def import_or_none(path):
@@ -868,3 +869,18 @@ def test_callable_scheduler():
 
     assert delayed(lambda: 1)().compute(scheduler=get) == 1
     assert called[0]
+
+
+@pytest.mark.skipif('not da')
+@pytest.mark.parametrize('num_workers', range(1, 4))
+@pytest.mark.parametrize('scheduler', ['threads', 'processes'])
+def test_num_workers_config(num_workers, scheduler):
+    # Regression test for issue #4082
+    a = da.random.randint(0, 10, size=400, chunks=2)
+
+    with dask.config.set(num_workers=num_workers), Profiler() as prof:
+        a.compute(scheduler=scheduler)
+
+    workers = {i.worker_id for i in prof.results}
+
+    assert len(workers) == num_workers