Closed MilesCranmer closed 1 year ago
Where are you expecting to compile?
Have you tried https://github.com/timholy/SnoopCompile.jl/tree/master/SnoopPrecompile ?
Also see this: https://sciml.ai/news/2022/09/21/compile_time/
The current version of pysr.create_sysimage
here runs a single model.fit()
with default settings, and saves a sysimage of the Julia runtime. I'm thinking this could be done automatically as part of pysr.install()
- it should help startup time quite a bit and be kind of like any C-backed Python library.
I played around with SnoopPrecompile but by itself it didn't seem to boost startup time that much. But I'm assuming that is just for precompilation, rather than building a whole sysimage - the latter of which should help startup time more, right? I was a bit confused where the boundary is between precompilation and sysimages. If I improve precompilation, do I get a better sysimage, than if I just run a full model pass and save the runtime? @ChrisRackauckas? (Great blogpost btw)
PySR might also be a bit different than ODE solvers here in that I don't need to re-compile if the input data/settings change (unless its a new operator or loss function), like I would have to with a new ODE. i.e., since the user provides data in a standard array format, I think a simple sysimage might be good enough...?
A system image refers to native code caching. However, we still need to put work in to figure out which code to compile. Note that we really need to address the compiled_modules=False
issue.
A system image refers to native code caching.
Wait, why can't I just do this, and cache the native code for a standard run?
Note that we really need to address the
compiled_modules=False
issue.
Do you mean make PyJulia do this automatically like PySR does (easy-ish), or actually fix the underlying pre-compilation issue (seems impossibly difficult)?
Think about when native code is generated in Julia. It is not generated from the function definitions alone. Rather it is generated when you call the function.
In order to generate the native code, you need to either run an example workflow or generate precompile statements. SnoopPrecompile helps with both aspects.
One thing that SnoopPrecompile does is figure I out when we are caching type inference or native code to disk. When it detects that, then it executes thr demo workflow and/or precompile statements.
Thanks for trying to explain this. I don't completely follow the underlying reasons why it would improve over simply caching a typical PySR run, but it sounds like SnoopCompile will help with startup time regardless. So I will try setting it up with a @precompile_all_calls
in SymbolicRegression.jl. Is there anything I can do to improve the PyCall.jl startup too (specifically for the SymbolicRegression.jl interface?)?
Okay it's actually working now! You can try it out with: python -c 'import pysr; pysr.install()'; python -c 'import pysr; pysr.compile()'
(need to be separate sessions for some reasons). There are some tricky things to work out still (the custom DEPOT test is currently failing), but it's cool that it's working. I get about a 33% improved speed overall in running example.py
.
By default it tries to create a pysr.so
and put it in the pysr-0.11.12
shared project directory. That way it can simply check for the existence of that sysimage at startup, and load it if it's there.
Interesting. It seems like most of the startup time is now due to getting the Julia version and project path, with, e.g.,
cmds = [
"julia",
"-e using Pkg; print(Pkg.project().path)",
]
julia_project_dir_str = subprocess.run(
cmds,
capture_output=True,
env=os.environ,
).stdout.decode()
is there a way to get this information much faster? This uses JULIA_PROJECT
to set the project, although sometimes it will be, e.g., @pysr-0.11.12
, so Python by itself can't figure out the path.
I think we might be able to extract this from pyjulia earlier than we boot up PyCall.jl or even without PyCall.jl.
I got it a bit faster, now it's not so bad:
time julia --startup-file=no --compile=min -O0 -g0 --project=@pysr-0.11.12 -e 'import Pkg: project; print(project().path)'
# 0.15s user 0.21s system 287% cpu 0.128 total
I think it might have just been the startup file to be honest.
It's now down to 13 seconds to run example.py
, whereas before it was 32 seconds! Not bad at all. It's still a bit faster to run .fit()
the second time though; I wonder what else is missing.
cc @Quarlos
It's still a bit faster to run
.fit()
the second time though
This suggests there are more precompile statements to generate or write.
One quick way would be to enable trace compile. SnoopCompile.jl also works.
julia --trace-compile=stderr
Good idea.
One interesting thing is that the startup time in pure-Julia is pretty much the same as the second run. But for some reason the Python version experiences extra startup penalty. I wonder how much is due to PyCall.jl compiling interfaces?
For future reference, since I didn't know how to set --trace-compile
, you can do it with the low-level interface:
# within `init_julia`, hence `julia_kwargs` is a dict of CLI options:
from julia.api import LibJulia
api = LibJulia.load(julia="julia")
api.sysimage = sysimage_name
cmds = [
f"--{k}={v}" if v is not None else f"--{k}" for k, v in julia_kwargs.items()
]
cmds += ["--trace-compile=stderr"]
print("Running with", cmds)
api.init_julia(cmds)
I also had to modify example.py
with verbosity=0, progress=False
, so that the progress bar is muted (otherwise it prints to stderr
).
Looks like most of the calls are PyCall.jl - I guess it needs better precompilation. Some of the LoopVectorization.jl and DynamicExpressions.jl are minor but I should be able to hit them with a modified precompilation in SymbolicRegression.jl
We're currently running with compiled-modules=no
. That's a problem. We are not caching code that we could be caching.
That's only for conda installs, though. The paste I sent above is with compiled_modules=True
.
When building the system image, you could also exercise PyCall
Do you mean call PySRRegressor().fit()
on some example data, before saving the sysimage? Maybe I'm doing something wrong but it didn't seem to help much, for whatever reason. I wonder if much of PyCall.jl is unable to be cached, since it can't tell if the Python function calling it is the same or not?
Part of the initial latency might simply be resolving pointers to symbols within libpython. That unfortunately cannot be cached.
Thanks for the note, that makes sense to me. I guess that few seconds isn't a major deal anyways.
Btw, I made a PR to PyCall.jl here: https://github.com/JuliaPy/PyCall.jl/pull/1015 which addresses some of the missed precompiles. I think it would be better to use SnoopCompile.jl, but it seems to not play well with Python. So I think manual precompile
is good enough for now.
I think this would be good to try merging under the heading that this is very experimental feature, and is only available opt-in. I'd personally be interested in using it, and I probably know of some others. One other idea would also be to hash the Manifest.toml
file in the env, and throw an error loading the sys image if the hash does not match.
We probably should see if we can fix the issue in PyCall... this is going to be a real problem when Julia 1.9 lands.
Do you mean the static libpython issue?
I mean the compiled_modules=False
thing
What is the main thread for that again? I see there is https://github.com/JuliaPy/PyCall.jl/pull/945 and https://github.com/JuliaPy/pyjulia/issues/496
Also, do you mean it is a problem for speed? Or would it e.g., cause crashes on Julia 1.9?
I'm going to close this for now. It just doesn't get much of a performance improvement, because the precompilation is already good enough, especially on the latest Julia versions. The increased dev complexity wouldn't justify a ~3% improvement in startup time.
This is an early WIP to integrate PackageCompiler.jl into PySR. No more startup time!
cc @mkitti