fix: issue with addprocs function undefined

MilesCranmer · Aug 13, 2024 · 09c2db5 · 09c2db5
1 parent 5e9b679
commit 09c2db5
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/pysr/julia_helpers.py b/pysr/julia_helpers.py
@@ -29,9 +29,11 @@ def _escape_filename(filename):
 
 def _load_cluster_manager(cluster_manager: str):
     if cluster_manager == "mpi":
+        jl.seval("using Distributed: addprocs")
         jl.seval("using MPIClusterManagers: MPIWorkerManager")
+
         return jl.seval(
-            "__pysr_mpi_addprocs(np; exeflags=``, kws...) = "
+            "(np; exeflags=``, kws...) -> "
             + "addprocs(MPIWorkerManager(np); exeflags=`$exeflags --project=$(Base.active_project())`, kws...)"
         )
     else:

diff --git a/pysr/test/test.py b/pysr/test/test.py
@@ -95,6 +95,7 @@ def _multiprocessing_turbo_custom_objective(self, cluster_manager):
             unary_operators=["sqrt"],
             procs=2,
             multithreading=False,
+            cluster_manager=cluster_manager,
             turbo=True,
             early_stop_condition="stop_if(loss, complexity) = loss < 1e-10 && complexity == 1",
             loss_function="""