Using Dask throws ImportError when SageMath code is run in python

Question

This question is very similar to my earlier question and was prompted by one of the comments.

Recently, I have been trying to parallelize some code using Dask. The code involves computations in SageMath, but it seems that whenever I use Sage code in a function I am trying to parallelize it throws an ImportError even though Sage has been successfully loaded. I want to know why I am getting an ImportError even though Sage seems to have loaded successfully, and more importantly, how to fix it.

Here is a basic example of what I am running into. When I run this:

import time
from sage.all import *
from dask import delayed
from dask.distributed import Client


client = Client(n_workers=4)

#I can add Sage integers with no problem
#So Sage seems to be loaded
Integer(1)+Integer(1) 

def Hello():
    Integer(1)+Integer(1) #if I remove this line the code runs fine
    return 'Hello World'

z = delayed(Hello)()
z.compute()

I get this error

ImportError                               Traceback (most recent call last)
<timed eval> in <module>

~/.sage/local/lib/python3.9/site-packages/dask/base.py in compute(self, **kwargs)
    284         dask.base.compute
    285         """
--> 286         (result,) = compute(self, traverse=False, **kwargs)
    287         return result
    288 

~/.sage/local/lib/python3.9/site-packages/dask/base.py in compute(*args, **kwargs)
    566         postcomputes.append(x.__dask_postcompute__())
    567 
--> 568     results = schedule(dsk, keys, **kwargs)
    569     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    570 

~/.sage/local/lib/python3.9/site-packages/distributed/client.py in get(self, dsk, keys, workers, allow_other_workers, resources, sync, asynchronous, direct, retries, priority, fifo_timeout, actors, **kwargs)
   2669                     should_rejoin = False
   2670             try:
-> 2671                 results = self.gather(packed, asynchronous=asynchronous, direct=direct)
   2672             finally:
   2673                 for f in futures.values():

~/.sage/local/lib/python3.9/site-packages/distributed/client.py in gather(self, futures, errors, direct, asynchronous)
   1946             else:
   1947                 local_worker = None
-> 1948             return self.sync(
   1949                 self._gather,
   1950                 futures,

~/.sage/local/lib/python3.9/site-packages/distributed/client.py in sync(self, func, asynchronous, callback_timeout, *args, **kwargs)
    843             return future
    844         else:
--> 845             return sync(
    846                 self.loop, func, *args, callback_timeout=callback_timeout, **kwargs
    847             )

~/.sage/local/lib/python3.9/site-packages/distributed/utils.py in sync(loop, func, callback_timeout, *args, **kwargs)
    324     if error[0]:
    325         typ, exc, tb = error[0]
--> 326         raise exc.with_traceback(tb)
    327     else:
    328         return result[0]

~/.sage/local/lib/python3.9/site-packages/distributed/utils.py in f()
    307             if callback_timeout is not None:
    308                 future = asyncio.wait_for(future, callback_timeout)
--> 309             result[0] = yield future
    310         except Exception:
    311             error[0] = sys.exc_info()

/var/tmp/sage-9.4-current/local/lib/python3.9/site-packages/tornado/gen.py in run(self)
    733 
    734                     try:
--> 735                         value = future.result()
    736                     except Exception:
    737                         exc_info = sys.exc_info()

~/.sage/local/lib/python3.9/site-packages/distributed/client.py in _gather(self, futures, errors, direct, local_worker)
   1811                             exc = CancelledError(key)
   1812                         else:
-> 1813                             raise exception.with_traceback(traceback)
   1814                         raise exc
   1815                     if errors == "skip":

~/.sage/local/lib/python3.9/site-packages/distributed/protocol/pickle.py in loads()
     73             return pickle.loads(x, buffers=buffers)
     74         else:
---> 75             return pickle.loads(x)
     76     except Exception:
     77         logger.info("Failed to deserialize %s", x[:10000], exc_info=True)

/var/tmp/sage-9.4-current/local/lib/python3.9/site-packages/sage/rings/integer.pyx in init sage.rings.integer (build/cythonized/sage/rings/integer.c:54201)()
----> 1 r"""
      2 Elements of the ring `\ZZ` of integers
      3 
      4 Sage has highly optimized and extensive functionality for arithmetic with integers
      5 and the ring of integers.

/var/tmp/sage-9.4-current/local/lib/python3.9/site-packages/sage/rings/rational.pyx in init sage.rings.rational (build/cythonized/sage/rings/rational.cpp:40442)()
     98 
     99 
--> 100 import sage.rings.real_mpfr
    101 import sage.rings.real_double
    102 from libc.stdint cimport uint64_t

/var/tmp/sage-9.4-current/local/lib/python3.9/site-packages/sage/rings/real_mpfr.pyx in init sage.rings.real_mpfr (build/cythonized/sage/rings/real_mpfr.c:46795)()
----> 1 r"""
      2 Arbitrary Precision Real Numbers
      3 
      4 AUTHORS:
      5 

/var/tmp/sage-9.4-current/local/lib/python3.9/site-packages/sage/libs/mpmath/utils.pyx in init sage.libs.mpmath.utils (build/cythonized/sage/libs/mpmath/utils.c:9062)()
----> 1 """
      2 Utilities for Sage-mpmath interaction
      3 
      4 Also patches some mpmath functions for speed
      5 """

/var/tmp/sage-9.4-current/local/lib/python3.9/site-packages/sage/rings/complex_mpfr.pyx in init sage.rings.complex_mpfr (build/cythonized/sage/rings/complex_mpfr.c:34594)()
----> 1 """
      2 Arbitrary Precision Floating Point Complex Numbers
      3 
      4 AUTHORS:
      5 

/var/tmp/sage-9.4-current/local/lib/python3.9/site-packages/sage/rings/complex_double.pyx in init sage.rings.complex_double (build/cythonized/sage/rings/complex_double.c:25284)()
     96 from cypari2.convert cimport new_gen_from_double, new_t_COMPLEX_from_double
     97 
---> 98 from . import complex_mpfr
     99 
    100 from .complex_mpfr import ComplexField

ImportError: cannot import name complex_mpfr

Perhaps this has something to do with Dask not importing Sage when it goes to parallelize things

Answer 1

Unfortunately, you might be out of luck here (somewhat). It looks like sage is not developed with threaded execution driven by another language in mind - their root level modules modify key elements of the python environment and really try to take control of low-level functionality by default. For example, sage.__init__ modifies the way that both inspect and sqllite work (gross!)

The specific issue you're running into is that importing sage invokes the signal module, which cannot be run from a thread other than the main one. The issue isn't in sage operations, but simply the import statement:

In [8]: def hello_sage():
   ...:     from sage.all import Integer
   ...:     return 'Hello World'
   ...:

In [9]: futures = client.submit(hello_sage)

In [10]: distributed.worker - WARNING - Compute Failed
Function:  hello_sage
args:      ()
kwargs:    {}
Exception: ValueError('signal only works in main thread of the main interpreter')

Unfortunately, this is fairly incompatible with dask, which runs all delayed jobs within threads. It's not that dask can't import modules locally to a remote function (it definitely can), it's that those functions can't use signal to control execution.

Because of the way sage is written, as far as multithreading goes I think your only choice is to go with the parallelization options their developers have provided. That said, you can trick sage into thinking it's in a world of its own by having threads start their own subprocesses:

In [1]: import dask.distributed as dd

In [2]: from subprocess import Popen, PIPE

In [3]: def invoke_sage_cli():
   ...:     cmd = ["sage", "-c", "print(factor(35))"]
   ...:     p = Popen(cmd, stdout=PIPE, stderr=PIPE, text=True)
   ...:     o, e = p.communicate()
   ...:
   ...:     if e:
   ...:         raise SystemError(e)
   ...:
   ...:     return o
   ...:

In [4]: client = dd.Client(n_workers=4)

In [5]: future = client.submit(invoke_sage_cli)

In [6]: print(future.result())
5 * 7

This is a pretty hacky way of getting around this issue, and I think it's unlikely to offer any performance benefits over the native sage parallelization options as long as you're working on a single machine. If you're using dask to scale up a Kubernetes cluster or work with nodes on an HPC or something, then you could definitely use this route to schedule distributed jobs and then have sage manage multithreading within each node.