Disable PythonCall GC during equation search

pysr/julia_helpers.py

@@ -1,6 +1,16 @@
 """Functions for initializing the Julia environment and installing deps."""
+import os
 import warnings
 
+# Required to avoid segfaults (https://juliapy.github.io/PythonCall.jl/dev/faq/)
+if os.environ["PYTHON_JULIACALL_HANDLE_SIGNALS"] not in {"yes", ""}:
+    warnings.warn(
+        "PYTHON_JULIACALL_HANDLE_SIGNALS environment variable is set to something other than 'yes' or ''. "
+        + "You will experience segfaults if running with multithreading."
+    )
+
+os.environ["PYTHON_JULIACALL_HANDLE_SIGNALS"] = "yes"
+
 import juliacall
 import juliapkg
 
@@ -8,6 +18,9 @@ jl = juliacall.newmodule("PySR")
 
 from juliacall import convert as jl_convert
 
+jl.seval("using PythonCall: PythonCall")
+PythonCall = jl.PythonCall
+
 juliainfo = None
 julia_initialized = False
 julia_kwargs_at_initialization = None

pysr/sr.py

@@ -32,7 +32,13 @@ from .export_numpy import sympy2numpy
 from .export_sympy import assert_valid_sympy_symbol, create_sympy_symbols, pysr2sympy
 from .export_torch import sympy2torch
 from .feature_selection import run_feature_selection
-from .julia_helpers import _escape_filename, _load_cluster_manager, jl, jl_convert
+from .julia_helpers import (
+    PythonCall,
+    _escape_filename,
+    _load_cluster_manager,
+    jl,
+    jl_convert,
+)
 from .utils import (
     _csv_filename_to_pkl_filename,
     _preprocess_julia_floats,
@@ -1716,6 +1722,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
             jl.Vector, self.display_feature_names_in_.tolist()
         )
 
+        PythonCall.GC.disable()
         # Call to Julia backend.
         # See https://github.com/MilesCranmer/SymbolicRegression.jl/blob/master/src/SymbolicRegression.jl
         self.raw_julia_state_ = SymbolicRegression.equation_search(
@@ -1738,6 +1745,7 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator):
             progress=progress and self.verbosity > 0 and len(y.shape) == 1,
             verbosity=int(self.verbosity),
         )
+        PythonCall.GC.enable()
 
         # Set attributes
         self.equations_ = self.get_hof()