Print locally on the worker too

.github/workflows/test_conda.yml

@@ -31,8 +31,7 @@ jobs:
           activate-environment: afar
       - name: Install dependencies
         run: |
-          conda install -y -c conda-forge distributed pytest
-          pip install innerscope
+          conda install -y -c conda-forge distributed pytest innerscope
           pip install -e .
       - name: PyTest
         run: |

.github/workflows/test_pip.yml

@@ -39,7 +39,7 @@ jobs:
         run: |
           pip install black flake8
           flake8 .
-          black afar *.py --check --diff
+          black . --check --diff
       - name: Coverage
         env:
           GITHUB_TOKEN: ${{ secrets.github_token }}

afar/__init__.py

@@ -1,3 +1,18 @@
+"""afar runs code within a context manager or IPython magic on a Dask cluster.
+
+>>> with afar.run, remotely:
+...     import dask_cudf
+...     df = dask_cudf.read_parquet("s3://...")
+...     result = df.sum().compute()
+
+or to use an IPython magic:
+
+>>> %load_ext afar
+>>> %afar z = x + y
+
+Read the documentation at https://github.com/eriknw/afar
+"""
+
 from ._core import get, run  # noqa
 from ._version import get_versions
 from ._where import later, locally, remotely  # noqa

afar/_abra.py

@@ -1,3 +1,8 @@
+"""Perform a magic trick: given lines of code, create a function to run remotely.
+
+This callable object is able to provide the values of the requested argument
+names and return the final expression so it can be displayed.
+"""
 import dis
 from types import FunctionType
 

afar/_core.py

@@ -1,5 +1,7 @@
+"""Define the user-facing `run` object; this is where it all comes together."""
 import dis
 import sys
+import traceback
 from inspect import currentframe
 from uuid import uuid4
 from weakref import WeakKeyDictionary, WeakSet
@@ -344,10 +346,21 @@ def run_afar(magic_func, names, futures, capture_print, channel, unique_key):
             # Hopefully computing the repr is fast.  If it is slow, perhaps it would be
             # better to add the return value to rv and call repr_afar as a separate task.
             # Also, pretty_repr must be msgpack serializable if done via events.
-            # Hence, custom _ipython_display_ probably won't work.
+            # Hence, custom _ipython_display_ probably won't work, and we resort to
+            # trying to use a basic repr (if that fails, we show the first exception).
             pretty_repr = repr_afar(results.return_value, magic_func._repr_methods)
             if pretty_repr is not None:
-                worker.log_event(channel, (unique_key, "display_expr", pretty_repr))
+                try:
+                    worker.log_event(channel, (unique_key, "display_expr", pretty_repr))
+                except Exception:
+                    exc_info = sys.exc_info()
+                    tb = traceback.format_exception(*exc_info)
+                    try:
+                        basic_repr = (repr(results.return_value), "__repr__", False)
+                        worker.log_event(channel, (unique_key, "display_expr", basic_repr))
+                    except Exception:
+                        exc_repr = (tb, pretty_repr[1], True)
+                        worker.log_event(channel, (unique_key, "display_expr", exc_repr))
                 send_finish = False
     finally:
         if capture_print and worker is not None and send_finish:

afar/_inspect.py

@@ -1,3 +1,4 @@
+"""Utilities to get the lines of the context body."""
 import dis
 from inspect import findsource
 

afar/_magic.py

@@ -1,3 +1,4 @@
+"""Define the IPython magic for using afar"""
 from textwrap import indent
 
 from dask.distributed import Client

afar/_printing.py

@@ -1,3 +1,4 @@
+"""Classes used to capture print statements within a Dask task."""
 import builtins
 import sys
 from io import StringIO
@@ -59,3 +60,6 @@ def __call__(self, *args, file=None, **kwargs):
                 pass
             else:
                 worker.log_event(self.channel, (self.key, stream_name, file.getvalue()))
+            # Print locally too
+            stream = sys.stdout if stream_name == "stdout" else sys.stderr
+            LocalPrint.printer(file.getvalue(), end="", file=stream)

afar/_reprs.py

@@ -1,3 +1,4 @@
+"""Utilities to calculate the (pretty) repr of objects remotely and display locally."""
 import sys
 import traceback