Switched to multi-threading, added export

[alexunderch]

Dear @lanctot,

With python3.14 approaching, there is some room to upgrade with AZ. Python 3.14 changed the default start method on POSIX from fork to forkserver which is not the same on Linux, for example. This would require a lot of monkey patching or there is a risk of the race condition with the import of spawn.py. That's why I decided to propose to switch to multi-threading instead with Python 3.13+.

Current implementation has N actors and M evaluators. Each actor/evaluator is a separate spawned process with its own JAX runtime. This means:

• N+M copies of compiled XLA programs (one per process)
• N+M copies of GPU memory allocations (JAX pre-allocates ~90% of GPU memory by default)

The learner saves to disk, then broadcasts the path string. Every actor then can probably retrigger a lot of computations: recompile, for example. Also, it cause disk I/O + deserialisation bottlenecks. To alleviate the issue, I did a single-node adjustments (because we have an ed implementation after all). There is single learner thread that takes all the computation and two queues of actors and evaluators, with shared model state to avoid I/O. and instantiate the model once (see README).

Metric	Before (processes)	After (threads)
JAX compilation caches	N + M	1
GPU memory overhead	~90% × (N+M)	~90% × 1
Weight sync latency	Disk I/O + deserialization	Atomic reference swap
Actor startup time	Model init + JIT compile per process	Instant (thread start)
Python 3.14 compatibility	Fragile (forkserver/spawn conflicts)	Native

What do you think about this?

Additionally, I modified the export_model part, and it should allow to export the model to onyx/jit for web or on-device demoes. At least it now has some purpose.

Looking forward to reading your thoughts.