perf(julia): precompile PySR's first fit in a bundled Julia package - #1396
Open
MilesCranmerBot wants to merge 1 commit into
Open
MilesCranmerBot wants to merge 1 commit into
MilesCranmerBot wants to merge 1 commit into
Conversation
Loading PythonCall invalidates SymbolicRegression.jl's precompiled code. PySRPrecompile loads both packages and precompiles PySR's calls, so that code survives. Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why
PySR always loads PythonCall. Loading PythonCall invalidates much of the native code that SymbolicRegression.jl caches from its precompile workload, so every PySR session recompiled the search on its first fit. In plain Julia, replaying SymbolicRegression.jl's precompiled workload takes 0.05 s. After
using PythonCall, the same replay takes 3.7 s.A package precompiled with PythonCall already loaded keeps its compiled code. This PR ships a small Julia package,
PySRPrecompile, inside PySR. It loads SymbolicRegression.jl and PythonCall, then runs a precompile workload shaped like PySR's own calls. SymbolicRegression.jl needs no changes. This replaces the approach in astroautomata/SymbolicRegression.jl#743, which added a PythonCall extension to SymbolicRegression.jl.Scope
pysr/julia/PySRPrecompile/: new Julia package.precompilestatements for the PythonCall conversions PySR uses to buildOptions,OperatorEnum,ExpressionSpecandExternalStop, and to copy NumPy arrays throughPyArray.@compile_workloadthat reruns SymbolicRegression.jl's genericdo_precompilation(Val(:compile)), then a search with PySR's default operators (+, -, /, *), the keyword arguments PySR passes toequation_search, progress output, and aSerializationround trip of the options and search state, as PySR's checkpoint does.pysr/juliapkg.json: addsPySRPrecompileby relativepathwith"dev": true. juliapkg resolves the path against the JSON file's directory.pysr/julia_import.py:using PySRPrecompilebefore the existingusing SymbolicRegression.Mainstill needsusing SymbolicRegression, because PySR evaluates user code such as custom operators and losses there.Tradeoffs
PySRPrecompilepinsPythonCall = "=0.9.36", because someprecompilestatements name PythonCall internals. Each PythonCall release that juliacall adopts needs a bump here. Without those statements, the first fit is about 1 s slower.PySRPrecompilesets no version bound on SymbolicRegression.jl, sopysr/juliapkg.jsonremains the only place that sets it, and custom SymbolicRegression.jl forks keep working. A fork that changes theequation_searchAPI would fail this package's precompile step.using PythonCallin PrecompileTools'@recompile_invalidationswas also tested. It changed first-fit time by less than 0.02 s and added 12 s to install, so it is not used.Blast Radius
Every PySR install precompiles one more Julia package, once per install or update. Runtime behavior does not change.
Verification
Measured on a 12-core Slurm node (Julia 1.13.1, Python 3.13), medians of three fresh processes, with the PySR changes from #1393 to #1395 applied:
Four alternating pairs of this PR against #743 put the README first fit 0.08 s faster (median paired difference) and the zero-iteration fit 0.02 s faster.
PySRPrecompilefromsite-packages, including with read-only package files. Three later imports took 4.76 to 4.86 s with no recompilation and byte-identical cache files.