Unable to Run Parallel Taichi Simulations Using concurrent.futures

[akekic]

Hi Taichi team,

I am trying to run parallel simulations using the concurrent.futures library. However, I am encountering an issue related to pickling Taichi objects. Below is a minimal example that reproduces the error.

Directory Structure

multiprocessing_taichi
|- __init__.py
|- taichi_class.py
|- taichi_main.py

File Contents

__init__py

# This file is empty

taichi_class.py

import taichi as ti

ti.init(default_fp=ti.f32, random_seed=20, debug=True)
ti.set_logging_level(ti.WARN)

@ti.data_oriented
class MyTaichiClass:
    def __init__(self):
        super().__init__()

        self.x = ti.Vector.field(n=2, dtype=ti.f32)

    def run(self):
        # do something complicated
        # ...

        self.x = ti.Vector([1.0, 2.0])
        return self.x.to_numpy().mean()

taichi_main.py

from multiprocessing_taichi.taichi_class import MyTaichiClass

import concurrent.futures

def execute_run(taichi_instance):
    return taichi_instance.run()

if __name__ == "__main__":
    my_taichi_instance = MyTaichiClass()

    with concurrent.futures.ProcessPoolExecutor() as executor:
        n = 10
        futures = {
            executor.submit(
                execute_run,
                my_taichi_instance,
            ): idx
            for idx in range(n)
        }
        for future in concurrent.futures.as_completed(futures):
            try:
                future.result()
            except Exception as e:
                idx = futures[future]
                print(f"Failed to generate batch {idx}: {e}")

Logs

When I run taichi_main.py, I get the following logs/errors:

[Taichi] version 1.7.0, llvm 15.0.5, commit 2fd24490, osx, python 3.11.2
[Taichi] Starting on arch=x64
Failed to generate batch 1: cannot pickle 'taichi._lib.core.taichi_python.Expr' object
Failed to generate batch 0: cannot pickle 'taichi._lib.core.taichi_python.Expr' object
Failed to generate batch 2: cannot pickle 'taichi._lib.core.taichi_python.Expr' object
Failed to generate batch 3: cannot pickle 'taichi._lib.core.taichi_python.Expr' object
Failed to generate batch 4: cannot pickle 'taichi._lib.core.taichi_python.Expr' object
Failed to generate batch 5: cannot pickle 'taichi._lib.core.taichi_python.Expr' object
Failed to generate batch 6: cannot pickle 'taichi._lib.core.taichi_python.Expr' object
Failed to generate batch 7: cannot pickle 'taichi._lib.core.taichi_python.Expr' object
Failed to generate batch 8: cannot pickle 'taichi._lib.core.taichi_python.Expr' object
Failed to generate batch 9: cannot pickle 'taichi._lib.core.taichi_python.Expr' object
[Taichi] version 1.7.0, llvm 15.0.5, commit 2fd24490, osx, python 3.11.2
[Taichi] version 1.7.0, llvm 15.0.5, commit 2fd24490, osx, python 3.11.2
[Taichi] version 1.7.0, llvm 15.0.5, commit 2fd24490, osx, python 3.11.2
[Taichi] version 1.7.0, llvm 15.0.5, commit 2fd24490, osx, python 3.11.2
[Taichi] version 1.7.0, llvm 15.0.5, commit 2fd24490, osx, python 3.11.2
[Taichi] version 1.7.0, llvm 15.0.5, commit 2fd24490, osx, python 3.11.2
[Taichi] version 1.7.0, llvm 15.0.5, commit 2fd24490, osx, python 3.11.2
[Taichi] version 1.7.0, llvm 15.0.5, commit 2fd24490, osx, python 3.11.2
[Taichi] Starting on arch=x64
[Taichi] Starting on arch=x64
[Taichi] Starting on arch=x64
[Taichi] Starting on arch=x64
[Taichi] Starting on arch=x64
[Taichi] Starting on arch=x64
[Taichi] Starting on arch=x64
[Taichi] Starting on arch=x64

Environment

• Python Version: 3.11.2
• Taichi Version: 1.7.0

Issue and Question

It seems like I can't pickle a Taichi instance. Is there a way to execute separate taichi simulations in parallel? I want to generate many samples of the simulation under slightly different initial conditions. I wasn't able to find anything in the past issues or docs.

Any help or guidance would be greatly appreciated.

Thank you!

[BirchKwok]

I encountered a nearly identical issue when using concurrent.futures.ThreadPoolExecutor. Taichi cannot be used, and it throws the error [llvm_context.cpp:add_struct_module@674] Assertion failure: std::this_thread::get_id() == main_thread_id_. Does anyone know how to solve this?

[jarmitage]

#7980 related?

[BirchKwok]

#7980 related? 

Yes, is there a solution?