Dealing with small recursion limits

[hoodmane]

I am trying to use Jedi in Pyodide (CPython compiled to wasm) and in chrome the wasm stack seems to be very small. In particular, the recursion_limit in chrome must be set to around 150. This is small enough that Jedi frequently causes stack overflows. The variable jedi.inference.recursion.recursion_limit seems to be relevant, but I see the same maximum stack depth no matter what value I set this to. What am I doing wrong?

Here is the script I am using to test. The result I see is that no matter what value I set for recursion_limit the script always prints "112". Of course 112 is less than 150 and this example causes no problem in my use case but setting a low value for recursion_limit has no effect either in examples that exceed my stack budget of 150. What am I supposed to do to communicate this limit to Jedi?

from contextlib import contextmanager
import sys
import jedi

max_depth = 0
def f(frame, event, args):
    if event == "return":
        return
    frame.f_trace_lines = False
    cur_frame = frame
    depth = 0
    while cur_frame is not None:
        cur_frame = cur_frame.f_back
        depth += 1
    global max_depth
    max_depth = max(depth, max_depth)
    return f

@contextmanager
def count_depth():
    max_depth = 0
    sys.settrace(f)
    try:
        yield
    finally:
        sys.settrace(None)

def test_recursion_level(n):
    jedi.inference.recursion.recursion_limit = n
    with count_depth():
        jedi.Interpreter("import pathlib; pathlib.Path().", [globals()]).complete(1,31)
    return max_depth

if __name__ == "__main__":
    print(test_recursion_level(int(sys.argv[1])))

[davidhalter]

Yeah, I'm sorry for that, but I guess you are out of luck. Jedi uses an insane amount of recursion internally for everything. It's not possible to change that. Type inference is implemented with recursion.

Of course it would be nicer if that was not the case, but it's actually not that easy to change it. It was very straight-forward to implement a lot of internals with recursion - and I also don't think it's fundamentally wrong.

jedi.inference.recursion.recursion_limit won't really help with what you are trying. It's just a setting that lets you limit recursion a bit for type inference, if type inference calls the same function again and again.

What are you trying to achieve with WASM?

[hoodmane]

Thanks for your fast reply. Jedi is a fantastic package.

it's actually not that easy to change it. It was very straight-forward to implement a lot of internals with recursion - and I also don't think it's fundamentally wrong.

I agree: it's intrinsic to the problem domain. Presumably JEDI is recursively descending the AST of the code. You could only avoid recursion if you made you moved the data off the call stack onto your own stack which you could update manually. This would be error prone, annoying, and create a huge mess in your code. For what benefit?

If you ever did get around to rewriting Jedi in Rust, that would fix my problem: a single Python call frame takes up a huge amount of space on the native stack, heuristically equivalent to around 50 normal sized native function calls. This is really the chrome people's fault for making their WASM stack so tiny, the stack in firefox is three times larger and everything works fine.

I'm using this to assist with discoverability in a repl, so it'd be nice if it works in more complicated situations, but my main goal is to cover things like some_class.some_method().some_field, and only for a limited number of simple packages. For things like this, Jedi is hovering right over / under hitting a stack overflow. It seems to often stack overflow on the first query but then the second time it sees the same query it has cached enough stuff that the inference works. So I was thinking if I could prepopulate the cache with stuff about the particular packages I'm interested in, maybe that'd help it work better... Or if I could give Jedi specific advice to shortcut certain situations... All of these options are probably more effort than they are worth.

I'm making a Python repl for use by math researchers to create charts of certain math computations. Trying to do this entirely in the browser is mainly for the sake of personal interest, but I expect to have a small group of dedicated users. I think it's a very cool idea to "do data science entirely in your browser" but the browser platform is not quite mature enough yet.

[davidhalter]

If you ever did get around to rewriting Jedi in Rust, that would fix my problem: a single Python call frame takes up a huge amount of space on the native stack, heuristically equivalent to around 50 normal sized native function calls. This is really the chrome people's fault for making their WASM stack so tiny, the stack in firefox is three times larger and everything works fine.

It's probably also a problem of running Python that way. They could probably also avoid quite a space on the stack if they really optimized it. But I know how hard that can be. The rust solution is unfortunately quite far away and I'm not even 100% sure it will work on WASM (even though that would be awesome).

I'm making a Python repl for use by math researchers to create charts of certain math computations.

Sounds awesome :) Is it related to Jupyter notebooks somehow? Does it use similar libraries? They are using Jedi in a Python process on the server, which is probably easier to get working than having it in the browser.

[hoodmane]

It's probably also a problem of running Python that way. They could probably also avoid quite a space on the stack if they really optimized it.

It's hard for me to see that there would be any better way to do it: They want full low level compatibility with CPython so that they can run Python packages written in C. It's difficult to do this because the native platform has an operating system with preemptive multitasking whereas the browser only offers single threaded webworkers. Emscripten is a C to wasm compiler combined with shims for enough of the standard OS calls to make existing C code work miraculously decently. They compile CPython with Emscripten and make as few changes as possible to make it work. If they start optimizing stack usage, they would no longer be using CPython and they would have a much harder time maintaining compatibility with the CPython ecosystem.

The rust solution is unfortunately quite far away and I'm not even 100% sure it will work on WASM (even though that would be awesome).

Yeah I have Rust code with PyO3 bindings that I would be interested to run on Pyodide, but I don't really know how to get started with that. I think this is an area where it's less an issue of what is technically possible than an issue of what people have energy to implement.

I'm curious how you are going to set up the Rust implementation. Do you have a prototype of a small piece of it? Are you planning to use PyO3? Or does that not expose enough of the Python internals? The other key design detail is how to work with the ownership memory model. I don't have a particularly good understanding of what the Jedi data type is like, but it seems like the whole thing is a huge jumble of references, which is hard to encode in Rust. If you have a lot of data types that are immutable except for a cache field that is only updated once, you can use OnceCell for those instead of RwLock. Perhaps Shredder's garbage collected pointers would be appropriate? Or maybe have one main slotmap that owns everything and store ids into the slotmap for the graph?

Is it related to Jupyter notebooks somehow? Does it use similar libraries?

It's related in behavior to Jupyter notebooks, but I haven't looked into Jupyter's source code nor into its UI decisions. The repl is based on monaco, it connects to Pyodide running on a webworker to evaluate python code and do Jedi completions. The chart is drawn with a WebGL2 backend I wrote in Rust compiled with wasm-bindgen. I use a service worker to emulate a webserver and manage communication between the chart and the repl.

They are using Jedi in a Python process on the server, which is probably easier to get working than having it in the browser.

Communicating with a native Python process via a websocket and using the browser for the front end is the sane approach, but I thought it would be really cool to make a "client-only" design that can run without installation. It works surprisingly well, this stack issue is the only major downside (it's also much slower than native but that could be an acceptable trade for convenience).

[davidhalter]

I'm curious how you are going to set up the Rust implementation. Do you have a prototype of a small piece of it? Are you planning to use PyO3? Or does that not expose enough of the Python internals? The other key design detail is how to work with the ownership memory model. I don't have a particularly good understanding of what the Jedi data type is like, but it seems like the whole thing is a huge jumble of references, which is hard to encode in Rust. If you have a lot of data types that are immutable except for a cache field that is only updated once, you can use OnceCell for those instead of RwLock. Perhaps Shredder's garbage collected pointers would be appropriate? Or maybe have one main slotmap that owns everything and store ids into the slotmap for the graph?

It's pretty much not clear how that is going to work. I'm not worrying right now. I'm pretty sure I will create an internal API in the beginning (Jedi also has that: inference folder vs api folder). At some point I will then start exposing an API that users can reasonably use. A Python wrapper will then probably need to extend that API somehow. However, I'm not even sure if there is ever going to be a Jedi written in Rust. It might just be a different project that has no connection to Jedi. It might not even have a Python wrapper, because most People don't really care about that, they just want to use it in their editors, where a language server implementation would be good enough.

So for now it's all very unclear, I'm writing the parser at the moment, which is quite a lot of work in itself (and fast, I would like it to be up to 100 times faster than the Python version). It's also when I realized, that it's not as easy as just exposing the Rust API in Parso, because Rust does not have GC and reference counting everywhere feels pretty annoying.

but I thought it would be really cool to make a "client-only" design that can run without installation. It works surprisingly well, this stack issue is the only major downside

I agree, it would definitely be cool! I'm surprised you got it running :)

[hoodmane]

Just as an update, I opened an issue here asking for the chrome wasm stack to be enlarged in order to resolve this.

[leopsidom]

Not sure if it's a reasonable ask. @davidhalter do you know if it's possible to convert the recursion implementation to iterative approach and how hard it might be ? Just so it can still run on a platform even if it has a limited stack size.

[davidhalter]

@leopsidom

It is unfortunately impossible to just use an iterative approach. The recursion is scattered all over the code base. So it's probably easier to just work on the existing constants in recursion.py and have at least less recursion. But no recursion is probably impossible.

[leopsidom]

@davidhalter Thanks for the response. I've played with these settings:

jedi.inference.recursion.recursion_limit = 15
jedi.inference.recursion.total_function_execution_limit = 50
jedi.inference.recursion.per_function_execution_limit = 3
jedi.inference.recursion.per_function_recursion_limit = 2

But it seems like it still uses up the stack space sometimes. I'll try to play around with them a bit more. Anyway thanks for the pointer.

[hoodmane]

As an update, since we merged https://github.com/pyodide/pyodide/pull/2019, we can now support a recursion limit of 1000+ on Pyodide main branch, so Jedi is fully functional in Pyodide now.

[leopsidom]

@hoodmane this is amazing. Can't wait to test this out in the next version of pyodide.