Aesara is not threadsafe

[ferrine]

Aesara is not threadsafe

Please provide a minimal, self-contained, and reproducible example.

import threading
import aesara
import numpy as np
import aesara.tensor as at
import aesara.compile.compiledir
aesara.compile.compiledir.cleanup()
np.random.seed(32)
def run_aesara(i, out):
    s = np.random.randint(10, 50)
    a = at.random.normal(size=(s, s))
    fn = aesara.function([],a)
    for _ in range(10):
        fn()
    out[i] = True
    
T = 32
out = [None] * T

threads = [threading.Thread(target=run_aesara, args=(i, out)) for i in range(T)]
[t.start() for t in threads]
[t.join() for t in threads]
assert all(out)

Please provide the full traceback of any errors.

ERROR (aesara.graph.opt): Optimization failure due to: constant_folding
ERROR (aesara.graph.opt): node: Subtensor{int64}(TensorConstant{(2,) of 36}, ScalarConstant{1})
ERROR (aesara.graph.opt): TRACEBACK:
ERROR (aesara.graph.opt): Traceback (most recent call last):
  File "/home/ferres/.miniconda3/envs/pymc/lib/python3.9/site-packages/aesara/graph/opt.py", line 1850, in process_node
    replacements = lopt.transform(fgraph, node)
  File "/home/ferres/.miniconda3/envs/pymc/lib/python3.9/site-packages/aesara/graph/opt.py", line 1055, in transform
    return self.fn(fgraph, node)
  File "/home/ferres/.miniconda3/envs/pymc/lib/python3.9/site-packages/aesara/tensor/basic_opt.py", line 2943, in constant_folding
    thunk = node.op.make_thunk(node, storage_map, compute_map, no_recycling=[])
  File "/home/ferres/.miniconda3/envs/pymc/lib/python3.9/site-packages/aesara/link/c/op.py", line 131, in make_thunk
    return self.make_c_thunk(node, storage_map, compute_map, no_recycling)
  File "/home/ferres/.miniconda3/envs/pymc/lib/python3.9/site-packages/aesara/link/c/op.py", line 96, in make_c_thunk
    outputs = cl.make_thunk(
  File "/home/ferres/.miniconda3/envs/pymc/lib/python3.9/site-packages/aesara/link/c/basic.py", line 1192, in make_thunk
    cthunk, module, in_storage, out_storage, error_storage = self.__compile__(
  File "/home/ferres/.miniconda3/envs/pymc/lib/python3.9/site-packages/aesara/link/c/basic.py", line 1127, in __compile__
    thunk, module = self.cthunk_factory(
  File "/home/ferres/.miniconda3/envs/pymc/lib/python3.9/site-packages/aesara/link/c/basic.py", line 1623, in cthunk_factory
    module = get_module_cache().module_from_key(key=key, lnk=self)
  File "/home/ferres/.miniconda3/envs/pymc/lib/python3.9/site-packages/aesara/link/c/cmodule.py", line 1183, in module_from_key
    module = self._get_from_hash(module_hash, key)
  File "/home/ferres/.miniconda3/envs/pymc/lib/python3.9/site-packages/aesara/link/c/cmodule.py", line 1086, in _get_from_hash
    key_data.add_key(key, save_pkl=bool(key[0]))
  File "/home/ferres/.miniconda3/envs/pymc/lib/python3.9/site-packages/aesara/link/c/cmodule.py", line 526, in add_key
    assert key not in self.keys
AssertionError

Please provide any additional information below.

Versions and main components

• Aesara version: '2.6.6'

[fonnesbeck]

Can confirm this on Python 3.10 on Debian

[ferrine]

I think the major pain here, as I dive deeper in the code, compile lock is called recursively. And compilation really relies on the fact there is a single process.

[ricardoV94]

Theano always claimed to not be thread safe. You can however copy a function before threading, which also avoids costly recompilation

https://theano-pymc.readthedocs.io/en/latest/tutorial/examples.html#copying-functions

[ferrine]

I got the idea of thread local + filelock.
https://github.com/aesara-devs/aesara/blob/main/aesara/compile/compilelock.py#L71
Filelock is used to be the actual lock, and threadlocal allows the current thread to skip the lock. A concurrent process would wait the filelock to be released. The intention was indeed to be threadsafe, but the implementation is not robust. I'll look into that again soon

[ricardoV94]

Sounds like a bug then

[ferrine]

should we convert this to a bug report then?

[ferrine]

This can be a long standing bug, but with proper attention, we can solve this. I already found couple of not thread safe places, but I think the list is longer

[brandonwillard]

Yes, if there are specific thread-safety changes to be made, then issues can be opened for those.