Skip to content

perf(julia): precompile PySR's first fit in a bundled Julia package - #1396

Open
MilesCranmerBot wants to merge 1 commit into
astroautomata:masterfrom
MilesCranmerBot:perf/pysr-precompile
Open

MilesCranmerBot wants to merge 1 commit into
astroautomata:masterfrom
MilesCranmerBot:perf/pysr-precompile

Conversation

@MilesCranmerBot

Copy link
Copy Markdown
Contributor

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.
    • precompile statements for the PythonCall conversions PySR uses to build Options, OperatorEnum, ExpressionSpec and ExternalStop, and to copy NumPy arrays through PyArray.
    • A @compile_workload that reruns SymbolicRegression.jl's generic do_precompilation(Val(:compile)), then a search with PySR's default operators (+, -, /, *), the keyword arguments PySR passes to equation_search, progress output, and a Serialization round trip of the options and search state, as PySR's checkpoint does.
  • pysr/juliapkg.json: adds PySRPrecompile by relative path with "dev": true. juliapkg resolves the path against the JSON file's directory.
  • pysr/julia_import.py: using PySRPrecompile before the existing using SymbolicRegression. Main still needs using SymbolicRegression, because PySR evaluates user code such as custom operators and losses there.

Tradeoffs

  • PySRPrecompile pins PythonCall = "=0.9.36", because some precompile statements name PythonCall internals. Each PythonCall release that juliacall adopts needs a bump here. Without those statements, the first fit is about 1 s slower.
  • PySRPrecompile sets no version bound on SymbolicRegression.jl, so pysr/juliapkg.json remains the only place that sets it, and custom SymbolicRegression.jl forks keep working. A fork that changes the equation_search API would fail this package's precompile step.
  • Wrapping using PythonCall in PrecompileTools' @recompile_invalidations was 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:

Release SymbolicRegression.jl#743 This PR
README example first fit 13.71 s 5.46 s 5.49 s
README example second fit 2.31 s 2.07 s 2.05 s
Zero-iteration first fit 10.12 s 2.15 s 2.13 s
Fresh install with PythonCall 117.9 s 132.1 s 136.1 s

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.

  • A wheel built from this branch and installed into a fresh virtual environment (not editable) resolved and precompiled PySRPrecompile from site-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.
  • The targeted PySR tests (constraints, weights, warm starts, checkpoints, pickling) pass: 10 of 10. Fixed-seed halls of fame for the README and default cases are byte-identical to The invalidation of pysr under some constraint conditions #743's.
  • A two-worker multiprocessing fit completes.
  • JuliaFormatter 2.4.0 (Blue style) formats the package.

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

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant