Attempt to make PySR process a daemon

Author: MilesCranmer

gui/processing.py

@@ -4,7 +4,6 @@
 import time
 from pathlib import Path
 
-import numpy as np
 import pandas as pd
 from data import generate_data, read_csv
 
@@ -17,6 +16,37 @@
 )
 
 
+def pysr_fit(queue: mp.Queue, out_queue: mp.Queue):
+    import pysr
+
+    while True:
+        # Get the arguments from the queue, if available
+        args = queue.get()
+        if args is None:
+            break
+        X = args["X"]
+        y = args["y"]
+        kwargs = args["kwargs"]
+        model = pysr.PySRRegressor(
+            progress=False,
+            timeout_in_seconds=1000,
+            **kwargs,
+        )
+        model.fit(X, y)
+        out_queue.put(None)
+
+
+class PySRProcess:
+    def __init__(self):
+        self.queue = mp.Queue()
+        self.out_queue = mp.Queue()
+        self.process = mp.Process(target=pysr_fit, args=(self.queue, self.out_queue))
+        self.process.start()
+
+
+PERSISTENT_WRITER = None
+
+
 def processing(
     file_input,
     force_run,
@@ -41,6 +71,11 @@ def processing(
     batch_size,
 ):
     """Load data, then spawn a process to run the greet function."""
+    global PERSISTENT_WRITER
+    if PERSISTENT_WRITER is None:
+        print("Starting PySR process")
+        PERSISTENT_WRITER = PySRProcess()
+
     if file_input is not None:
         try:
             X, y = read_csv(file_input, force_run)
@@ -53,31 +88,41 @@ def processing(
         base = Path(tmpdirname)
         equation_file = base / "hall_of_fame.csv"
         equation_file_bkup = base / "hall_of_fame.csv.bkup"
-        process = mp.Process(
-            target=pysr_fit,
-            kwargs=dict(
+        # Check if queue is empty, if not, kill the process
+        # and start a new one
+        if not PERSISTENT_WRITER.queue.empty():
+            print("Restarting PySR process")
+            if PERSISTENT_WRITER.process.is_alive():
+                PERSISTENT_WRITER.process.terminate()
+                PERSISTENT_WRITER.process.join()
+
+            PERSISTENT_WRITER = PySRProcess()
+        # Write these to queue instead:
+        PERSISTENT_WRITER.queue.put(
+            dict(
                 X=X,
                 y=y,
-                niterations=niterations,
-                maxsize=maxsize,
-                binary_operators=binary_operators,
-                unary_operators=unary_operators,
-                equation_file=equation_file,
-                parsimony=parsimony,
-                populations=populations,
-                population_size=population_size,
-                ncycles_per_iteration=ncycles_per_iteration,
-                elementwise_loss=elementwise_loss,
-                adaptive_parsimony_scaling=adaptive_parsimony_scaling,
-                optimizer_algorithm=optimizer_algorithm,
-                optimizer_iterations=optimizer_iterations,
-                batching=batching,
-                batch_size=batch_size,
-            ),
+                kwargs=dict(
+                    niterations=niterations,
+                    maxsize=maxsize,
+                    binary_operators=binary_operators,
+                    unary_operators=unary_operators,
+                    equation_file=equation_file,
+                    parsimony=parsimony,
+                    populations=populations,
+                    population_size=population_size,
+                    ncycles_per_iteration=ncycles_per_iteration,
+                    elementwise_loss=elementwise_loss,
+                    adaptive_parsimony_scaling=adaptive_parsimony_scaling,
+                    optimizer_algorithm=optimizer_algorithm,
+                    optimizer_iterations=optimizer_iterations,
+                    batching=batching,
+                    batch_size=batch_size,
+                ),
+            )
         )
-        process.start()
         last_yield_time = None
-        while process.is_alive():
+        while PERSISTENT_WRITER.out_queue.empty():
             if equation_file_bkup.exists():
                 try:
                     # First, copy the file to a the copy file
@@ -109,21 +154,3 @@ def processing(
                     last_yield_time = time.time()
                 except pd.errors.EmptyDataError:
                     pass
-
-        process.join()
-
-
-def pysr_fit(
-    *,
-    X,
-    y,
-    **pysr_kwargs,
-):
-    import pysr
-
-    model = pysr.PySRRegressor(
-        progress=False,
-        timeout_in_seconds=1000,
-        **pysr_kwargs,
-    )
-    model.fit(X, y)