[BUG]: TemplateExpressionSpec and loss_function_expression not compatible with multiprocessing

[ibengtsson]

What happened?

Combining TemplateExpressionSpec and a custom loss function (loss_function_expression) doesn't work when running in the multiprocessing setting. It works fine when using multithreading, however.

Version

1.5.8

Operating System

macOS

Package Manager

pip

Interface

Script (i.e., python my_script.py)

Relevant log output

❯ prp minimal_example.py
Compiling Julia backend...
[ Info: Automatically setting `--heap-size-hint=139M` on each Julia process. You can configure this with the `heap_size_hint_in_bytes` parameter.
[ Info: Importing SymbolicRegression on workers.
[ Info: Finished!
[ Info: Copying definition of __sr_template_11363233297447453560 to workers...
[ Info: Finished!
[ Info: Copying definition of eval_loss to workers...
[ Info: Finished!
Traceback (most recent call last):
  File "/Users/isakbe/Dev/modelling/il-sr/il_sr/scripts/minimal_example.py", line 52, in <module>
  File "/Users/isakbe/Dev/modelling/il-sr/il_sr/scripts/minimal_example.py", line 43, in main
    model.fit(x, y)
  File "/Users/isakbe/Library/Caches/pypoetry/virtualenvs/il-sr-9TFUWRsR-py3.11/lib/python3.11/site-packages/pysr/sr.py", line 2318, in fit
    self._run(X, y, runtime_params, weights=weights, seed=seed, category=category)
  File "/Users/isakbe/Library/Caches/pypoetry/virtualenvs/il-sr-9TFUWRsR-py3.11/lib/python3.11/site-packages/pysr/sr.py", line 2109, in _run
    out = SymbolicRegression.equation_search(
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/isakbe/.julia/packages/PythonCall/WMWY0/src/JlWrap/any.jl", line 262, in __call__
    return self._jl_callmethod($(pyjl_methodnum(pyjlany_call)), args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
juliacall.JuliaError: On worker 2:
MethodError: no method matching eval_loss(::TemplateExpression{Float32, TemplateStructure{(:f, :g), (), typeof(__sr_template_11363233297447453560), @NamedTuple{f::Int64, g::Int64}, @NamedTuple{}}, Node{Float32}, ComposableExpression{Float32, Node{Float32}, @NamedTuple{operators::OperatorEnum{Tuple{typeof(+), typeof(-), typeof(*), typeof(/), typeof(safe_pow)}, Tuple{}}, variable_names::Nothing, eval_options::EvalOptions{false, false, true, Nothing}}}, @NamedTuple{f::ComposableExpression{Float32, Node{Float32}, @NamedTuple{operators::OperatorEnum{Tuple{typeof(+), typeof(-), typeof(*), typeof(/), typeof(safe_pow)}, Tuple{}}, variable_names::Nothing, eval_options::EvalOptions{false, false, true, Nothing}}}, g::ComposableExpression{Float32, Node{Float32}, @NamedTuple{operators::OperatorEnum{Tuple{typeof(+), typeof(-), typeof(*), typeof(/), typeof(safe_pow)}, Tuple{}}, variable_names::Nothing, eval_options::EvalOptions{false, false, true, Nothing}}}}, @NamedTuple{structure::TemplateStructure{(:f, :g), (), typeof(__sr_template_11363233297447453560), @NamedTuple{f::Int64, g::Int64}, @NamedTuple{}}, operators::OperatorEnum{Tuple{typeof(+), typeof(-), typeof(*), typeof(/), typeof(safe_pow)}, Tuple{}}, variable_names::Nothing, parameters::@NamedTuple{}}}, ::SymbolicRegression.CoreModule.DatasetModule.BasicDataset{Float32, Float32, Matrix{Float32}, Vector{Float32}, Nothing, @NamedTuple{}, Nothing, Nothing, Nothing, Nothing}, ::Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, OperatorEnum{Tuple{typeof(+), typeof(-), typeof(*), typeof(/), typeof(safe_pow)}, Tuple{}}, Node, TemplateExpression, @NamedTuple{structure::TemplateStructure{(:f, :g), (), typeof(__sr_template_11363233297447453560), @NamedTuple{f::Int64, g::Int64}, @NamedTuple{}}}, MutationWeights, false, false, nothing, Nothing, 5})

Closest candidates are:
  eval_loss(::Any, ::Dataset{T, L}, ::Any, !Matched::Any) where {T, L} (method too new to be called from this world context.)
   @ Main none:2

Stacktrace:
 [1] #12
   @ /opt/homebrew/Cellar/julia/1.10.4/share/julia/stdlib/v1.10/Distributed/src/macros.jl:83
 [2] #invokelatest#2
   @ ./essentials.jl:892 [inlined]
 [3] invokelatest
   @ ./essentials.jl:889
 [4] #107
   @ /opt/homebrew/Cellar/julia/1.10.4/share/julia/stdlib/v1.10/Distributed/src/process_messages.jl:283
 [5] run_work_thunk
   @ /opt/homebrew/Cellar/julia/1.10.4/share/julia/stdlib/v1.10/Distributed/src/process_messages.jl:70
 [6] run_work_thunk
   @ /opt/homebrew/Cellar/julia/1.10.4/share/julia/stdlib/v1.10/Distributed/src/process_messages.jl:79
 [7] #100
   @ /opt/homebrew/Cellar/julia/1.10.4/share/julia/stdlib/v1.10/Distributed/src/process_messages.jl:88
Stacktrace:
  [1] remotecall_fetch(f::Function, w::Distributed.Worker, args::Distributed.RRID; kwargs::@Kwargs{})
    @ Distributed /opt/homebrew/Cellar/julia/1.10.4/share/julia/stdlib/v1.10/Distributed/src/remotecall.jl:465
  [2] remotecall_fetch(f::Function, w::Distributed.Worker, args::Distributed.RRID)
    @ Distributed /opt/homebrew/Cellar/julia/1.10.4/share/julia/stdlib/v1.10/Distributed/src/remotecall.jl:454
  [3] remotecall_fetch
    @ /opt/homebrew/Cellar/julia/1.10.4/share/julia/stdlib/v1.10/Distributed/src/remotecall.jl:492 [inlined]
  [4] call_on_owner
    @ /opt/homebrew/Cellar/julia/1.10.4/share/julia/stdlib/v1.10/Distributed/src/remotecall.jl:565 [inlined]
  [5] fetch(r::Distributed.Future)
    @ Distributed /opt/homebrew/Cellar/julia/1.10.4/share/julia/stdlib/v1.10/Distributed/src/remotecall.jl:619
  [6] test_function_on_workers(example_inputs::Tuple{TemplateExpression{Float32, TemplateStructure{(:f, :g), (), typeof(__sr_template_11363233297447453560), @NamedTuple{f::Int64, g::Int64}, @NamedTuple{}}, Node{Float32}, ComposableExpression{Float32, Node{Float32}, @NamedTuple{operators::OperatorEnum{Tuple{typeof(+), typeof(-), typeof(*), typeof(/), typeof(safe_pow)}, Tuple{}}, variable_names::Nothing, eval_options::EvalOptions{false, false, true, Nothing}}}, @NamedTuple{f::ComposableExpression{Float32, Node{Float32}, @NamedTuple{operators::OperatorEnum{Tuple{typeof(+), typeof(-), typeof(*), typeof(/), typeof(safe_pow)}, Tuple{}}, variable_names::Nothing, eval_options::EvalOptions{false, false, true, Nothing}}}, g::ComposableExpression{Float32, Node{Float32}, @NamedTuple{operators::OperatorEnum{Tuple{typeof(+), typeof(-), typeof(*), typeof(/), typeof(safe_pow)}, Tuple{}}, variable_names::Nothing, eval_options::EvalOptions{false, false, true, Nothing}}}}, @NamedTuple{structure::TemplateStructure{(:f, :g), (), typeof(__sr_template_11363233297447453560), @NamedTuple{f::Int64, g::Int64}, @NamedTuple{}}, operators::OperatorEnum{Tuple{typeof(+), typeof(-), typeof(*), typeof(/), typeof(safe_pow)}, Tuple{}}, variable_names::Nothing, parameters::@NamedTuple{}}}, SymbolicRegression.CoreModule.DatasetModule.BasicDataset{Float32, Float32, Matrix{Float32}, Vector{Float32}, Nothing, @NamedTuple{}, Nothing, Nothing, Nothing, Nothing}, Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, OperatorEnum{Tuple{typeof(+), typeof(-), typeof(*), typeof(/), typeof(safe_pow)}, Tuple{}}, Node, TemplateExpression, @NamedTuple{structure::TemplateStructure{(:f, :g), (), typeof(__sr_template_11363233297447453560), @NamedTuple{f::Int64, g::Int64}, @NamedTuple{}}}, MutationWeights, false, false, nothing, Nothing, 5}}, op::Function, procs::Vector{Int64})
    @ SymbolicRegression ~/.julia/packages/SymbolicRegression/MdISO/src/Configure.jl:237
  [7] move_functions_to_workers(procs::Vector{Int64}, options::Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, OperatorEnum{Tuple{typeof(+), typeof(-), typeof(*), typeof(/), typeof(safe_pow)}, Tuple{}}, Node, TemplateExpression, @NamedTuple{structure::TemplateStructure{(:f, :g), (), typeof(__sr_template_11363233297447453560), @NamedTuple{f::Int64, g::Int64}, @NamedTuple{}}}, MutationWeights, false, false, nothing, Nothing, 5}, dataset::SymbolicRegression.CoreModule.DatasetModule.BasicDataset{Float32, Float32, Matrix{Float32}, Vector{Float32}, Nothing, @NamedTuple{}, Nothing, Nothing, Nothing, Nothing}, verbosity::Int64)
    @ SymbolicRegression ~/.julia/packages/SymbolicRegression/MdISO/src/Configure.jl:211
  [8] configure_workers(; procs::Nothing, numprocs::Int64, addprocs_function::typeof(Distributed.addprocs), options::Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, OperatorEnum{Tuple{typeof(+), typeof(-), typeof(*), typeof(/), typeof(safe_pow)}, Tuple{}}, Node, TemplateExpression, @NamedTuple{structure::TemplateStructure{(:f, :g), (), typeof(__sr_template_11363233297447453560), @NamedTuple{f::Int64, g::Int64}, @NamedTuple{}}}, MutationWeights, false, false, nothing, Nothing, 5}, worker_imports::Union{Nothing, Vector{Symbol}}, project_path::String, file::String, exeflags::Cmd, verbosity::Int64, example_dataset::SymbolicRegression.CoreModule.DatasetModule.BasicDataset{Float32, Float32, Matrix{Float32}, Vector{Float32}, Nothing, @NamedTuple{}, Nothing, Nothing, Nothing, Nothing}, runtests::Bool)
    @ SymbolicRegression ~/.julia/packages/SymbolicRegression/MdISO/src/Configure.jl:389
  [9] _create_workers(datasets::Vector{SymbolicRegression.CoreModule.DatasetModule.BasicDataset{Float32, Float32, Matrix{Float32}, Vector{Float32}, Nothing, @NamedTuple{}, Nothing, Nothing, Nothing, Nothing}}, ropt::SymbolicRegression.SearchUtilsModule.RuntimeOptions{:multiprocessing, 1, true, Nothing}, options::Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, OperatorEnum{Tuple{typeof(+), typeof(-), typeof(*), typeof(/), typeof(safe_pow)}, Tuple{}}, Node, TemplateExpression, @NamedTuple{structure::TemplateStructure{(:f, :g), (), typeof(__sr_template_11363233297447453560), @NamedTuple{f::Int64, g::Int64}, @NamedTuple{}}}, MutationWeights, false, false, nothing, Nothing, 5})
    @ SymbolicRegression ~/.julia/packages/SymbolicRegression/MdISO/src/SymbolicRegression.jl:633
 [10] _equation_search(datasets::Vector{SymbolicRegression.CoreModule.DatasetModule.BasicDataset{Float32, Float32, Matrix{Float32}, Vector{Float32}, Nothing, @NamedTuple{}, Nothing, Nothing, Nothing, Nothing}}, ropt::SymbolicRegression.SearchUtilsModule.RuntimeOptions{:multiprocessing, 1, true, Nothing}, options::Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, OperatorEnum{Tuple{typeof(+), typeof(-), typeof(*), typeof(/), typeof(safe_pow)}, Tuple{}}, Node, TemplateExpression, @NamedTuple{structure::TemplateStructure{(:f, :g), (), typeof(__sr_template_11363233297447453560), @NamedTuple{f::Int64, g::Int64}, @NamedTuple{}}}, MutationWeights, false, false, nothing, Nothing, 5}, saved_state::Nothing)
    @ SymbolicRegression ~/.julia/packages/SymbolicRegression/MdISO/src/SymbolicRegression.jl:568
 [11] equation_search(datasets::Vector{SymbolicRegression.CoreModule.DatasetModule.BasicDataset{Float32, Float32, Matrix{Float32}, Vector{Float32}, Nothing, @NamedTuple{}, Nothing, Nothing, Nothing, Nothing}}; options::Options{SymbolicRegression.CoreModule.OptionsStructModule.ComplexityMapping{Int64, Int64}, OperatorEnum{Tuple{typeof(+), typeof(-), typeof(*), typeof(/), typeof(safe_pow)}, Tuple{}}, Node, TemplateExpression, @NamedTuple{structure::TemplateStructure{(:f, :g), (), typeof(__sr_template_11363233297447453560), @NamedTuple{f::Int64, g::Int64}, @NamedTuple{}}}, MutationWeights, false, false, nothing, Nothing, 5}, saved_state::Nothing, runtime_options::Nothing, runtime_options_kws::@Kwargs{niterations::Int64, parallelism::String, numprocs::Int64, procs::Nothing, addprocs_function::Nothing, heap_size_hint_in_bytes::Nothing, worker_imports::Nothing, runtests::Bool, return_state::Bool, run_id::String, verbosity::Int64, logger::Nothing, progress::Bool, v_dim_out::Val{1}})
    @ SymbolicRegression ~/.julia/packages/SymbolicRegression/MdISO/src/SymbolicRegression.jl:561
 [12] equation_search
    @ ~/.julia/packages/SymbolicRegression/MdISO/src/SymbolicRegression.jl:542 [inlined]
 [13] #equation_search#23
    @ ~/.julia/packages/SymbolicRegression/MdISO/src/SymbolicRegression.jl:511 [inlined]
 [14] equation_search
    @ ~/.julia/packages/SymbolicRegression/MdISO/src/SymbolicRegression.jl:456 [inlined]
 [15] #equation_search#24
    @ ~/.julia/packages/SymbolicRegression/MdISO/src/SymbolicRegression.jl:535 [inlined]
 [16] pyjlany_call(self::typeof(equation_search), args_::Py, kwargs_::Py)
    @ PythonCall.JlWrap ~/.julia/packages/PythonCall/WMWY0/src/JlWrap/any.jl:44
 [17] _pyjl_callmethod(f::Any, self_::Ptr{PythonCall.C.PyObject}, args_::Ptr{PythonCall.C.PyObject}, nargs::Int64)
    @ PythonCall.JlWrap ~/.julia/packages/PythonCall/WMWY0/src/JlWrap/base.jl:73
 [18] _pyjl_callmethod(o::Ptr{PythonCall.C.PyObject}, args::Ptr{PythonCall.C.PyObject})
    @ PythonCall.JlWrap.Cjl ~/.julia/packages/PythonCall/WMWY0/src/JlWrap/C.jl:63

Extra Info

A minimal example:

from pysr import PySRRegressor, TemplateExpressionSpec
import numpy as np
from sklearn.metrics import r2_score

def main():

    x = np.random.uniform(low=[-1, 3, 10], high=[6, 10, 100], size=(100, 3))
    y = np.sin(x[:, 0] - x[:, 2] ** 3) * np.exp(x[:, 1] / 5 + x[:, 0]) 
    
    template = TemplateExpressionSpec(
        expressions=["f", "g"],
        combine="sin(f(x1, x3)) * exp(g(x1, x2))",
        variable_names=["x1", "x2", "x3"]
    )

    loss = """ 
        using Statistics
        function eval_loss(tree, dataset::Dataset{T,L}, options, idx)::L where {T,L}

            x = idx === nothing ? dataset.X : view(dataset.X, :, idx)
            y = idx === nothing ? dataset.y : view(dataset.y, idx)
            
            prediction, flag = eval_tree_array(tree, x, options)

            if !flag
                return L(Inf)
            end
            
            errors = (prediction .- y) .^2
            return mean(errors)
        end
    """

    model = PySRRegressor(
        binary_operators=["+", "-", "*", "/", "^"],
        constraints={"^": (1, 9)},
        expression_spec=template,
        parallelism="multiprocessing",
        loss_function_expression=loss,
    )

    
    model.fit(x, y)
    
    preds = model.predict(x)
    score = r2_score(y, preds)
    print(score)

if __name__ == "__main__":
    main()

[MilesCranmer]

Ah. This error is not due to multiprocessing. It's actually due to a couple unrelated issues:

1. You need to chose a new function name. eval_loss is already used as the internal name, and this appears to be overwriting it somehow (I think). And it seems like it is not overwriting it on the worker process, which leads to a conflict.
2. It needs to have a 3-arg version available. You should declare it with idx=nothing rather than just idx.

You can fix it with the following patch:

- function eval_loss(tree, dataset::Dataset{T,L}, options, idx)::L where {T,L}
+ function my_custom_loss(tree, dataset::Dataset{T,L}, options, idx=nothing)::L where {T,L}

Also, in the latest PySR/SymbolicRegression.jl, you don't need the idx argument at all, it will do this all automatically for you. So you could actually use the following patch instead of the above one:

- function eval_loss(tree, dataset::Dataset{T,L}, options, idx)::L where {T,L}
+ function my_custom_loss(tree, dataset::Dataset{T,L}, options)::L where {T,L}

-     x = idx === nothing ? dataset.X : view(dataset.X, :, idx)
-     y = idx === nothing ? dataset.y : view(dataset.y, idx)
+     x = dataset.X
+     y = dataset.y
      
      prediction, flag = eval_tree_array(tree, x, options)

      if !flag
          return L(Inf)
      end
      
      errors = (prediction .- y) .^2
      return mean(errors)
  end

[ibengtsson]

Ah, should've guessed that the error was on my end... I really should spend some more time in Julia :) Thanks for the swift reply and suggestions, I really appreciate it.

[MilesCranmer]

No problem! Glad it works