Make remote tracebacks easier to read and distinguish

[gjoseph92]

When an exception occurs on the cluster, the traceback that shows up is long, and includes both a local component and a remote component. I suspect this is a stumbling block for beginners because:

1. When tracebacks are too long, people glaze over and don't read them
2. For those that do read them, separating the local client boilerplate from the remote part requires practice

Specifically, I always skim for something like raise exception.with_traceback(traceback) in distributed/client.py, ignore everything above it, and just look at the remote part of the traceback. Could we format our error messages differently so users don't have to learn this unintuitive skill?

Goals:

1. Users can easily tell whether an exception happened locally or on the cluster
2. Minimal internal distributed code is shown in tracebacks when we know the error wasn't an internal distributed error.

We already do a good job with 2 on the worker side thanks to get_traceback, which removes irrelevant frames. So this might be as simple as raise ... from None on the client when re-raising a remote exception. Plus somehow adding a prefix like:

------------------------------------------------------------------------------------------------
ValueError (remote)                               Traceback from cluster (most recent call last)

While calling Client.compute, this error occurred on worker 'worker-abcde'
while executing task ('map-blocks-12345', 0, 0):

so that remote exceptions are easily distinguishable from local ones.

[mrocklin]

There are some cases where I or others want to see the full traceback, but I agree that this probably shouldn't be the default. I'd be comfortable cutting off the traceback at a certain point. Ideally this would be controllable by a configuration value.

[gjoseph92]

I think the key is "when we know the error wasn't an internal distributed error". Specifically, if the client is re-raising an exception sent from the cluster, I think we know that the preceding client frames aren't that relevant (besides maybe knowing the top-level client function, hence the While calling Client.compute).

Ideally this would be controllable by a configuration value.

Also agreed.

[j-bennet]

Hi @gjoseph92, as far as I can see, distributed Client already drops irrelevant frames, by calling raise exc.with_traceback(tb). But maybe I'm looking for the problem in the wrong place. Could you help with a reproducible example, to make sure I understand what is needed?

[j-bennet]

@jrbourbeau @fjetter

TL;DR:

• raise... from None won't solve the problem
• but manipulating traceback before raising the exception might
• there are several different ways to do it, we need to pick one.

I experimented a bit on a branch here. Here are some findings.

1. raise ... from None is not helpful in this case, since in distributed.Client, we are not raising a new Exception, but re-raising the original exception, already with specific traceback (raise exc.with_traceback(tb)).
2. before re-raising the exception, we can remove more of the irrelevant traceback frames from traceback, however, this will only help partially. Some irrelevant frames will be pushed into the trace stack after Client raised. Here is an example unit test in my branch:

def throws_map_blocks(x, block_id=None):
    if block_id is not None and block_id[0] % 2 == 0:
        raise ValueError(f"Bad block id {block_id}")
    return x / 2


@pytest.mark.parametrize("clean_traceback,expected_levels", [(True, 9), (False, 16)])
def test_map_blocks_error(client, clean_traceback, expected_levels):
    da = pytest.importorskip("dask.array")
    arr = da.arange(30, chunks=3)
    op = arr.map_blocks(throws_map_blocks)
    try:
        with dask.config.set(
            {"distributed.exceptions.clean_traceback": clean_traceback}
        ):
            op.compute()
    except Exception:
        _, _, tb = sys.exc_info()
        traceback.print_tb(tb)

The traceback before cleanup has 16 stack frames:

Click me!

  File "/Users/jbennet/src/distributed/distributed/tests/test_as_completed.py", line 196, in test_map_blocks_error
    op.compute()
  File "/Users/jbennet/mambaforge/envs/distributed/lib/python3.11/site-packages/dask/base.py", line 314, in compute
    (result,) = compute(self, traverse=False, **kwargs)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jbennet/mambaforge/envs/distributed/lib/python3.11/site-packages/dask/base.py", line 599, in compute
    results = schedule(dsk, keys, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jbennet/src/distributed/distributed/client.py", line 3229, in get
    raise exc_value.with_traceback(exc_traceback)
  File "/Users/jbennet/src/distributed/distributed/client.py", line 3225, in get
    results = self.gather(packed, asynchronous=asynchronous, direct=direct)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jbennet/src/distributed/distributed/client.py", line 2358, in gather
    return self.sync(
           ^^^^^^^^^^
  File "/Users/jbennet/src/distributed/distributed/utils.py", line 362, in sync
    return sync(
           ^^^^^
  File "/Users/jbennet/src/distributed/distributed/utils.py", line 430, in sync
    raise exc.with_traceback(tb)
  File "/Users/jbennet/src/distributed/distributed/utils.py", line 402, in f
    result = yield future
             ^^^^^^^^^^^^
  File "/Users/jbennet/mambaforge/envs/distributed/lib/python3.11/site-packages/tornado/gen.py", line 767, in run
    value = future.result()
            ^^^^^^^^^^^^^^^
  File "/Users/jbennet/src/distributed/distributed/client.py", line 2221, in _gather
    raise exception.with_traceback(traceback)
  File "/Users/jbennet/mambaforge/envs/distributed/lib/python3.11/site-packages/dask/optimization.py", line 990, in __call__
    return core.get(self.dsk, self.outkey, dict(zip(self.inkeys, args)))
    ^^^^^^^^^^^^^^^^^
  File "/Users/jbennet/mambaforge/envs/distributed/lib/python3.11/site-packages/dask/core.py", line 149, in get
    result = _execute_task(task, cache)
    ^^^^^^^^^^^^^^^^^
  File "/Users/jbennet/mambaforge/envs/distributed/lib/python3.11/site-packages/dask/core.py", line 119, in _execute_task
    return func(*(_execute_task(a, cache) for a in args))
    ^^^^^^^^^^^^^^^^^
  File "/Users/jbennet/mambaforge/envs/distributed/lib/python3.11/site-packages/dask/array/core.py", line 523, in _pass_extra_kwargs
    return func(*args[len(keys) :], **kwargs)
      ^^^^^^^^^^^^^^^^^
  File "/Users/jbennet/src/distributed/distributed/tests/test_as_completed.py", line 185, in throws_map_blocks
    raise ValueError(f"Bad block id {block_id}")
    ^^^^^^^^^^^^^^^^^]

The traceback after cleanup has 6 stack frames:

Click me!

  File "/Users/jbennet/src/distributed/distributed/tests/test_as_completed.py", line 199, in test_map_blocks_error
    op.compute()
  File "/Users/jbennet/mambaforge/envs/distributed/lib/python3.11/site-packages/dask/base.py", line 314, in compute
    (result,) = compute(self, traverse=False, **kwargs)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jbennet/mambaforge/envs/distributed/lib/python3.11/site-packages/dask/base.py", line 599, in compute
    results = schedule(dsk, keys, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jbennet/src/distributed/distributed/client.py", line 3228, in get
    raise exc_value.with_traceback(exc_traceback)
  File "/Users/jbennet/mambaforge/envs/distributed/lib/python3.11/site-packages/dask/array/core.py", line 523, in _pass_extra_kwargs
    return func(*args[len(keys) :], **kwargs)
      ^^^^^^^^^^^^^^^^^
  File "/Users/jbennet/src/distributed/distributed/tests/test_as_completed.py", line 186, in throws_map_blocks
    raise ValueError(f"Bad block id {block_id}")
    ^^^^^^^^^^^^^^^^^

This is helpful, right? Except, you may notice that I cheated. I included some dask code into my list of "internal code":

This brings me to my next finding.

3. There are a number of Python libraries that have to fudge their traceback. For example, ipython, jinja2, IDLE.

They are using several different methods to mark code internal.

• Keep a list of internal code paths (optionally including function names), like I did here.
• Decorate internal function with something that will push it into a singleton list.
• Set a magic local value inside this frame, ipython sets __tracebackhide__.

2nd and 3rd are more explicit and feel more like "the right way" to do this. With the first, I can remove traceback frames that don't belong to distributed code, which is useful, on one hand, but on the other hand, I'm messing with the code that's not part of distributed, and that's a dangerous path to take. Another advantage of the 1st method is that I could put "internal" code patterns inside dask config file, and then it would be extremely easy to add or change them.

[fjetter]

I just ran the above example

import dask.array as da
from distributed import Client
client = Client(silence_logs="ERROR")

def throws_map_blocks(x, block_id=None):
    if block_id is not None and block_id[0] % 2 == 0:
        raise ValueError(f"Bad block id {block_id}")
    return x / 2

arr = da.arange(30, chunks=3)
op = arr.map_blocks(throws_map_blocks)
op.compute()
    

and was surprised at how verbose the output is. What I get is

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[1], line 12
     10 arr = da.arange(30, chunks=3)
     11 op = arr.map_blocks(throws_map_blocks)
---> 12 op.compute(scheduler=client)

File ~/miniforge3/envs/test-environment/lib/python3.10/site-packages/dask/base.py:314, in DaskMethodsMixin.compute(self, **kwargs)
    290 def compute(self, **kwargs):
    291     """Compute this dask collection
    292 
    293     This turns a lazy Dask collection into its in-memory equivalent.
   (...)
    312     dask.compute
    313     """
--> 314     (result,) = compute(self, traverse=False, **kwargs)
    315     return result

File ~/miniforge3/envs/test-environment/lib/python3.10/site-packages/dask/base.py:599, in compute(traverse, optimize_graph, scheduler, get, *args, **kwargs)
    596     keys.append(x.__dask_keys__())
    597     postcomputes.append(x.__dask_postcompute__())
--> 599 results = schedule(dsk, keys, **kwargs)
    600 return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])

File ~/workspace/distributed/distributed/client.py:3225, in Client.get(self, dsk, keys, workers, allow_other_workers, resources, sync, asynchronous, direct, retries, priority, fifo_timeout, actors, **kwargs)
   3223         should_rejoin = False
   3224 try:
-> 3225     results = self.gather(packed, asynchronous=asynchronous, direct=direct)
   3226 finally:
   3227     for f in futures.values():

File ~/workspace/distributed/distributed/client.py:2360, in Client.gather(self, futures, errors, direct, asynchronous)
   2358 else:
   2359     local_worker = None
-> 2360 return self.sync(
   2361     self._gather,
   2362     futures,
   2363     errors=errors,
   2364     direct=direct,
   2365     local_worker=local_worker,
   2366     asynchronous=asynchronous,
   2367 )

File ~/workspace/distributed/distributed/utils.py:351, in SyncMethodMixin.sync(self, func, asynchronous, callback_timeout, *args, **kwargs)
    349     return future
    350 else:
--> 351     return sync(
    352         self.loop, func, *args, callback_timeout=callback_timeout, **kwargs
    353     )

File ~/workspace/distributed/distributed/utils.py:418, in sync(loop, func, callback_timeout, *args, **kwargs)
    416 if error:
    417     typ, exc, tb = error
--> 418     raise exc.with_traceback(tb)
    419 else:
    420     return result

File ~/workspace/distributed/distributed/utils.py:391, in sync.<locals>.f()
    389         future = wait_for(future, callback_timeout)
    390     future = asyncio.ensure_future(future)
--> 391     result = yield future
    392 except Exception:
    393     error = sys.exc_info()

File ~/miniforge3/envs/test-environment/lib/python3.10/site-packages/tornado/gen.py:767, in Runner.run(self)
    765 try:
    766     try:
--> 767         value = future.result()
    768     except Exception as e:
    769         # Save the exception for later. It's important that
    770         # gen.throw() not be called inside this try/except block
    771         # because that makes sys.exc_info behave unexpectedly.
    772         exc: Optional[Exception] = e

File ~/workspace/distributed/distributed/client.py:2223, in Client._gather(self, futures, errors, direct, local_worker)
   2221         exc = CancelledError(key)
   2222     else:
-> 2223         raise exception.with_traceback(traceback)
   2224     raise exc
   2225 if errors == "skip":

File ~/miniforge3/envs/test-environment/lib/python3.10/site-packages/dask/optimization.py:990, in __call__()
    988 if not len(args) == len(self.inkeys):
    989     raise ValueError("Expected %d args, got %d" % (len(self.inkeys), len(args)))
--> 990 return core.get(self.dsk, self.outkey, dict(zip(self.inkeys, args)))

File ~/miniforge3/envs/test-environment/lib/python3.10/site-packages/dask/core.py:149, in get()
    147 for key in toposort(dsk):
    148     task = dsk[key]
--> 149     result = _execute_task(task, cache)
    150     cache[key] = result
    151 result = _execute_task(out, cache)

File ~/miniforge3/envs/test-environment/lib/python3.10/site-packages/dask/core.py:119, in _execute_task()
    115     func, args = arg[0], arg[1:]
    116     # Note: Don't assign the subtask results to a variable. numpy detects
    117     # temporaries by their reference count and can execute certain
    118     # operations in-place.
--> 119     return func(*(_execute_task(a, cache) for a in args))
    120 elif not ishashable(arg):
    121     return arg

File ~/miniforge3/envs/test-environment/lib/python3.10/site-packages/dask/array/core.py:523, in _pass_extra_kwargs()
    516 """Helper for :func:`dask.array.map_blocks` to pass `block_info` or `block_id`.
    517 
    518 For each element of `keys`, a corresponding element of args is changed
    519 to a keyword argument with that key, before all arguments re passed on
    520 to `func`.
    521 """
    522 kwargs.update(zip(keys, args))
--> 523 return func(*args[len(keys) :], **kwargs)

Cell In[1], line 7, in throws_map_blocks()
      5 def throws_map_blocks(x, block_id=None):
      6     if block_id is not None and block_id[0] % 2 == 0:
----> 7         raise ValueError(f"Bad block id {block_id}")
      8     return x / 2

ValueError: Bad block id (0,)

As Gabe already suggested, a user friendly way to raise this would be

------------------------------------------------------------------------------------------------
ValueError (remote)                               Traceback from cluster (most recent call last)

While calling Client.compute, this error occurred on worker 'worker-abcde'
while executing task ('map-blocks-12345', 0, 0):

Cell In[1], line 7, in throws_map_blocks()
      5 def throws_map_blocks(x, block_id=None):
      6     if block_id is not None and block_id[0] % 2 == 0:
----> 7         raise ValueError(f"Bad block id {block_id}")
      8     return x / 2

ValueError: Bad block id (0,)

I don't think the other part of this traceback is actually relevant. The OP mentions that get_traceback was doing a decent job on the worker side but it doesn't look like it since the output is very verbose even on worker side. This may be a regression that slipped in at some point.

[j-bennet]

In my branch here, the same traceback would be quite a bit shorter.

import dask
import dask.array as da
from distributed import Client

def throws_map_blocks(x, block_id=None):
    if block_id is not None and block_id[0] % 2 == 0:
        raise ValueError(f"Bad block id {block_id}")
    return x / 2

dask.config.set({'distributed.exceptions.clean_traceback': True})
client = Client(silence_logs="ERROR")
arr = da.arange(30, chunks=3)
op = arr.map_blocks(throws_map_blocks)

op.compute()

The output is:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[1], line 15
     12 arr = da.arange(30, chunks=3)
     13 op = arr.map_blocks(throws_map_blocks)
---> 15 op.compute()

File ~/mambaforge/envs/distributed/lib/python3.11/site-packages/dask/base.py:314, in DaskMethodsMixin.compute(self, **kwargs)
    290 def compute(self, **kwargs):
    291     """Compute this dask collection
    292
    293     This turns a lazy Dask collection into its in-memory equivalent.
   (...)
    312     dask.compute
    313     """
--> 314     (result,) = compute(self, traverse=False, **kwargs)
    315     return result

File ~/mambaforge/envs/distributed/lib/python3.11/site-packages/dask/base.py:599, in compute(traverse, optimize_graph, scheduler, get, *args, **kwargs)
    596     keys.append(x.__dask_keys__())
    597     postcomputes.append(x.__dask_postcompute__())
--> 599 results = schedule(dsk, keys, **kwargs)
    600 return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])

File ~/src/distributed/distributed/client.py:3237, in Client.get(self, dsk, keys, workers, allow_other_workers, resources, sync, asynchronous, direct, retries, priority, fifo_timeout, actors, **kwargs)
   3235     exc_type, exc_value, exc_traceback = sys.exc_info()
   3236     exc_traceback = drop_internal_traceback(exc_traceback)
-> 3237     raise exc_value.with_traceback(exc_traceback)
   3238 finally:
   3239     for f in futures.values():

Cell In[1], line 7, in throws_map_blocks()
      5 def throws_map_blocks(x, block_id=None):
      6     if block_id is not None and block_id[0] % 2 == 0:
----> 7         raise ValueError(f"Bad block id {block_id}")
      8     return x / 2

ValueError: Bad block id (2,)

So it's better, but it doesn't quite do what you want. This is because we can only do so much in distributed, before the exception bubbles up to dask/base. Dask frames will be added to the stack after the stack was filtered in distributed.client.

A better way might be overriding sys.excepthook and in IPython, shostraceback. I'll experiment with that one.