[BUG] AutoLeanServer is not robust to crashes induced by stack overflows

[jonathan-laurent]

Describe the bug

When the server crashes due to a stack-overflow (which can easily happen, for example when using norm_num on expressions containing large numbers), AutoLeanServer fails to properly restart it, probably due to corrupted session caches. (Not using session caches and manually rerunning initialisation code solves the problem).

To Reproduce
Your environment:

• OS: Linux, macOS,
• Python version: 3.12
• LeanInteract version: 0.9.2
  Then the steps to reproduce the behavior.

Running this program:

from lean_interact import TempRequireProject, LeanREPLConfig, Command, AutoLeanServer

project = TempRequireProject(lean_version="v4.23.0", require="mathlib")
config = LeanREPLConfig(project=project, memory_hard_limit_mb=8000)

command = """
theorem mathd_numbertheory_328 : (5^999999) % 7 = 6 := by norm_num
"""

server = AutoLeanServer(config)
env = server.run(Command(cmd="import Mathlib"), add_to_session_cache=True).env

try:
    # Larger timeout values sometimes fail to replicate the bug, especially on MacOS.
    server.run(Command(cmd=command, env=env), timeout=1e-2)
except TimeoutError:
    print("Timeout!")
    
print("Server not alive, so will be restarted.")
server.run(Command(cmd=command, env=env))

leads to the following output:

Timeout!
Server not alive, so will be restarted.

and then exception:

ConnectionAbortedError: The Lean server closed unexpectedly.
--------------------------------------------------
stdout:

--------------------------------------------------
stderr:
libc++abi: terminating due to uncaught exception of type lean::exception: cannot evaluate `[init]` declaration 'Mathlib.Meta.NormNum.normNumExt' in the same module

--------------------------------------------------If stdout and stderr are empty or obscure, here is a list of possible reasons (not exhaustive):
- Not enough memory and/or compute available
- The cached Lean REPL is corrupted. In this case, clear the cache using the `clear-lean-cache` command.
- An uncaught exception in the Lean REPL

[augustepoiroux]

Hello!
Thanks for reporting this issue with details.
I did some testing, and, in short, the issue is coming from the pickling feature (with LeanInteract >= 0.9.0). More details below. Since the pickling feature is also introducing other kinds of bugs I am personally not using session caches anymore. While this is something I would like to fix, this is not my priority at the moment.
Instead, I recommend using incremental elaboration with AutoLeanServer and no session cache. Incremental elaboration makes interaction with Lean slightly different from the classical REPL but takes care of optimizing shared computations automatically. Instead of running commands one by one and recording environment states for imports for example, one simply runs whole file content at once. Incremental elaboration takes care of caching computation and imports and will automatically reuse the best environment state across the whole interaction history (in constant time thanks to the use of a trie, in practice negligible overhead).
On restarts, the current pickling feature was anyway replaying constants and reloading the imports, leading to no gain in computation time in comparison to rerunning the commands.

More details about the issue

We can get the same exception by replacing

try:
    # Larger timeout values sometimes fail to replicate the bug, especially on MacOS.
    server.run(Command(cmd=command, env=env), timeout=1e-2)
except TimeoutError:
    print("Timeout!")

with

server.kill()

Replacing this with the underlying pickling logic and AutoLeanServer with LeanServer, the following code raises the same exception:

from lean_interact import TempRequireProject, LeanREPLConfig, Command, LeanServer, PickleEnvironment, UnpickleEnvironment

project = TempRequireProject(lean_version="v4.23.0", require="mathlib")
config = LeanREPLConfig(project=project)

command = """
theorem mathd_numbertheory_328 : (5^999999) % 7 = 6 := by norm_num
"""

server = LeanServer(config)
env = server.run(Command(cmd="import Mathlib")).env
server.run(PickleEnvironment(env=env, pickle_to="mathlib.olean"))
server.restart()
env2 = server.run(UnpickleEnvironment(unpickle_env_from="mathlib.olean")).env
print(server.run(Command(cmd=command, env=env2)))

Commenting the server.restart() line makes the exception disappear. Some data from the REPL state which cannot be recovered after a restart is therefore probably missing in the pickle file.

[augustepoiroux]

I fixed AutoLeanServer by implementing a more reliable session cache, ReplaySessionCache . It reruns commands put in the session cache instead of using the pickling feature. It is also thread-safe, meaning that you can share one session cache instance across multiple AutoLeanServer instances (in a multi-threaded context, not a multi-processing context).
Your snippet should now be running fine in v0.10.0.
I also added more explanations about the incremental elaboration feature in the documentation.

[jonathan-laurent]

This incremental elaboration feature seems like a big usability boost. Thanks for your work on this package!