Fail to apply the multiprocessing with gradient-related customized loss function

[HouJayson]

Hi,
Thank you for your great tool. I encountered a problem when I tried to use the multiprocessing calculation by setting the "multithreading=False", an error just occurred "KeyError: key DynamicExpressionsZygoteExt [98b68043-617e-5e87-ba04-89939eaf59ce] not found".

This issue arises possibly because I have customized the loss function (incorporating the gradient calculus), as I have tested a simple code provided in the manual (no customized loss function), and it works with the multiprocessing setup.

My devices info: Windows 10; Julia version: 1.9.3; pysr version is 0.16.3

My implementation:

from julia import MLJ
from julia import SymbolicRegression
from julia import Zygote
from julia import SymbolicUtils
model = PySRRegressor(
    procs=2,
    populations=4,
    multithreading=False,
    binary_operators=["+","*","pow"],
    full_objective=objective_all,
    constraints={"pow":(4,1)},
    nested_constraints={
      "pow": { "pow":0},			    			    
    }, 
    complexity_of_operators={"pow":2, "+":5,"*":5},		
    complexity_of_variables=4,      
    complexity_of_constants=1,
    maxsize=50,
    maxdepth=5,
    warmup_maxsize_by=0.2,
    turbo=True,
    enable_autodiff=True,
    niterations=2,
)

The customized loss function

function derivative_loss(tree, dataset::Dataset{T,L}, options, idx) where {T,L}
    
    # Column-major:
    X = idx === nothing ? dataset.X : view(dataset.X, :, idx)
    ∂y = idx === nothing ? dataset.y : view(dataset.y, idx)
    
    ŷ1, ∂ŷ1, completed = eval_grad_tree_array(tree, X, options; variable=true)
    
    ∂yx1_t = ∂ŷ1[1,1:41]
    ∂yx2_t = ∂ŷ1[1,42:82]
    ∂yx1_c = ∂ŷ1[1,83:123]
    ∂yx2_c = ∂ŷ1[1,124:164]
    ∂yx1_s = ∂ŷ1[1,165:205]
    ∂yx2_s = ∂ŷ1[1,206:246]
    
    ∂y_t   = ∂y[1:41]
    ∂y_c   = ∂y[83:123]
    ∂y_s   = ∂y[165:205]
    
    !completed && return L(Inf)
    diff_t = ∂yx1_t .- X[1,42:82] ./ X[1,1:41] .* ∂yx2_t	
    diff_c = ∂yx1_c .- X[1,124:164] ./ X[1,83:123] .* ∂yx2_c	
    diff_s = ((X[1,165:205]) .^ 2) ./ (1 .+ (X[1,165:205]) .^ 2) .* ∂yx1_s .- ((X[1,206:246]) .^ 2) ./ (1 .+ (X[1,206:246]) .^ 2) .* ∂yx2_s
    
    mse_t = sum(i -> (diff_t[i] - ∂y_t[i])^2, eachindex(∂y_t)) / length(∂y_t)
    mse_c = sum(i -> (diff_c[i] - ∂y_c[i])^2, eachindex(∂y_c)) / length(∂y_c)
    mse_s = sum(i -> (diff_s[i] - ∂y_s[i])^2, eachindex(∂y_s)) / length(∂y_s)
    
    return mse_t + mse_c + mse_s
end

[MilesCranmer]

I think the issue is that you are not starting Julia within the right environment for PySR, because you are calling it from julia import ... before the first PySR run happens. Thus, the worker cannot find the right Julia package extension. I think the worker is activating a different environment and loading a different version of SymbolicRegression.jl than the head worker.

Could you please try something like

from pysr.julia_helpers import init_julia
init_julia(julia_kwargs=dict(optimize=3, threads="auto"))

and only afterwards, run the from julia import ... stuff?

i.e., it would look like

from pysr.julia_helpers import init_julia
init_julia(julia_kwargs=dict(optimize=3, threads="auto"))

from julia import Zygote
...

Note you might need to do from julia import Pkg; Pkg.add("Zygote") after Julia has been loaded so that it installs Zygote in the correct environment.

[HouJayson]

Thank you so much for your prompt response! I followed your suggestion, but unfortunately, the error persists. Here's what I did based on your advice:

1. I made the recommended modification, but the issue remains.
   

2. I verified that these packages are installed in the same environment, as shown in the following screenshots:
   

3. I also attempted to reinstall Julia, rebuild the dependencies, and rerun the code with multithreading enabled, but the issue persisted.
   

4. I found that setting multi-threads for Julia with "JULIA_NUM_THREADS = 12" improved the prediction speed significantly, which greatly met my needs.

[MilesCranmer]

If multithreading works, you could just go with that (multithreading=True, with procs set to your number of cores)

I wonder if the issue has to do with the worker not having Zygote loaded? Perhaps try something like this:

objective_all = """
const zygote_loaded = Ref(false)
const zygote_lock = Threads.SpinLock()

function derivative_loss(tree, dataset::Dataset{T,L}, options, idx) where {T,L}
    zygote_loaded[] || lock(zygote_lock) do
        zygote_loaded[] || @eval using Zygote: Zygote
        zygote_loaded[] = true
    end

    # Column-major:
    X = idx === nothing ? dataset.X : view(dataset.X, :, idx)
    ∂y = idx === nothing ? dataset.y : view(dataset.y, idx)
    
    # [rest of loss function]
end
"""

This is a bit of a hack but it might work. Essentially what happens is that const zygote_loaded = Ref(false) will be created once on each process. Therefore, each process will be forced to run @eval using Zygote: Zygote exactly once, thus loading Zygote upon the first call.

However I'm not sure if this will work... A more robust solution would be to add another line here:
https://github.com/MilesCranmer/SymbolicRegression.jl/blob/141987a3587978039712f010172b142ad09d4757/src/Configure.jl#L217

that runs using Zygote after the activation if options.enable_autodiff is true and Zygote has already been loaded by the head worker.

[HouJayson]

It works by setting "multithreading=True", thank you so much!

[HouJayson]

Thank you so much, your suggestions are highly valued.

[MilesCranmer]

No problem. Could you mark the comment which solved your issue as the "Answer"? Then it will be helpful for other users scrolling these forums. Thanks! - Miles