Closed MilesCranmer closed 1 year ago
mkitti
commented
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/
MilesCranmer
commented
1 year ago 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...?
mkitti
commented
1 year ago 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.
MilesCranmer
commented
1 year ago 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=Falseissue.
Do you mean make PyJulia do this automatically like PySR does (easy-ish), or actually fix the underlying pre-compilation issue (seems impossibly difficult)?
mkitti
commented
1 year ago 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.
MilesCranmer
commented
1 year ago 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?)?
MilesCranmer
commented
1 year ago 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.
MilesCranmer
commented
1 year ago 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.
mkitti
commented
1 year ago I think we might be able to extract this from pyjulia earlier than we boot up PyCall.jl or even without PyCall.jl.
MilesCranmer
commented
1 year ago 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 totalI think it might have just been the startup file to be honest.
MilesCranmer
commented
1 year ago 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.
MilesCranmer
commented
1 year ago cc @Quarlos
mkitti
commented
1 year ago 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=stderrGood 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?
MilesCranmer
commented
1 year ago 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).
MilesCranmer
commented
1 year ago 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
mkitti
commented
1 year ago We're currently running with compiled-modules=no. That's a problem. We are not caching code that we could be caching.
MilesCranmer
commented
1 year ago That's only for conda installs, though. The paste I sent above is with compiled_modules=True.
mkitti
commented
1 year ago When building the system image, you could also exercise PyCall
MilesCranmer
commented
1 year ago 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?
mkitti
commented
1 year ago Part of the initial latency might simply be resolving pointers to symbols within libpython. That unfortunately cannot be cached.
MilesCranmer
commented
1 year ago 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.
MilesCranmer
commented
1 year ago 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.
mkitti
commented
1 year ago 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.
MilesCranmer
commented
1 year ago Do you mean the static libpython issue?
mkitti
commented
1 year ago I mean the compiled_modules=False thing
MilesCranmer
commented
1 year ago 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
MilesCranmer
commented
1 year ago Also, do you mean it is a problem for speed? Or would it e.g., cause crashes on Julia 1.9?
MilesCranmer
commented
1 year ago 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