feat: add worker_timeout

Author: MilesCranmer

src/Configure.jl

@@ -368,6 +368,7 @@ function configure_workers(;
     procs::Union{Vector{Int},Nothing},
     numprocs::Int,
     addprocs_function::Function,
+    worker_timeout::Union{Float64,Nothing},
     options::AbstractOptions,
     @nospecialize(worker_imports::Union{Vector{Symbol},Nothing}),
     project_path,
@@ -378,7 +379,9 @@ function configure_workers(;
     runtests::Bool,
 )
     (procs, we_created_procs) = if procs === nothing
-        (addprocs_function(numprocs; lazy=false, exeflags), true)
+        withenv("JULIA_WORKER_TIMEOUT" => string(worker_timeout)) do
+            (addprocs_function(numprocs; lazy=false, exeflags), true)
+        end
     else
         (procs, false)
     end

src/MLJInterface.jl

@@ -77,6 +77,7 @@ function modelexpr(
             numprocs::Union{Int,Nothing} = nothing
             procs::Union{Vector{Int},Nothing} = nothing
             addprocs_function::Union{Function,Nothing} = nothing
+            worker_timeout::Union{Real,Nothing} = nothing
             heap_size_hint_in_bytes::Union{Integer,Nothing} = nothing
             worker_imports::Union{Vector{Symbol},Nothing} = nothing
             logger::Union{AbstractSRLogger,Nothing} = nothing
@@ -299,6 +300,7 @@ function _update(
         numprocs=m.numprocs,
         procs=m.procs,
         addprocs_function=m.addprocs_function,
+        worker_timeout=m.worker_timeout,
         heap_size_hint_in_bytes=m.heap_size_hint_in_bytes,
         worker_imports=m.worker_imports,
         runtests=m.runtests,

src/SearchUtils.jl

@@ -82,6 +82,7 @@ struct RuntimeOptions{PARALLELISM,DIM_OUT,RETURN_STATE,LOGGER} <: AbstractRuntim
     numprocs::Int64
     init_procs::Union{Vector{Int},Nothing}
     addprocs_function::Function
+    worker_timeout::Float64
     exeflags::Cmd
     worker_imports::Union{Vector{Symbol},Nothing}
     runtests::Bool
@@ -117,6 +118,7 @@ end
     numprocs::Union{Int,Nothing}=nothing,
     procs::Union{Vector{Int},Nothing}=nothing,
     addprocs_function::Union{Function,Nothing}=nothing,
+    worker_timeout::Union{Real,Nothing}=nothing,
     heap_size_hint_in_bytes::Union{Integer,Nothing}=nothing,
     worker_imports::Union{Vector{Symbol},Nothing}=nothing,
     runtests::Bool=true,
@@ -190,6 +192,13 @@ end
     _verbosity = something(verbosity, options_verbosity, 1)
     _progress = something(progress, options_progress, (_verbosity > 0) && nout == 1)
     _addprocs_function = something(addprocs_function, addprocs)
+    _worker_timeout = Float64(
+        @something(
+            worker_timeout,
+            tryparse(Float64, get(ENV, "JULIA_WORKER_TIMEOUT", "")),
+            min(60, _numprocs^2)
+        )
+    )
     _run_id = @something(run_id, generate_run_id())
 
     exeflags = if concurrency == :multiprocessing && isnothing(procs)
@@ -211,6 +220,7 @@ end
         _numprocs,
         procs,
         _addprocs_function,
+        _worker_timeout,
         exeflags,
         worker_imports,
         runtests,

src/SymbolicRegression.jl

@@ -406,6 +406,10 @@ which is useful for debugging and profiling.
     which is the number of processes to use, as well as the `lazy` keyword argument.
     For example, if set up on a slurm cluster, you could pass
     `addprocs_function = addprocs_slurm`, which will set up slurm processes.
+- `worker_timeout::Union{Real,Nothing}=nothing`: Timeout in seconds for worker processes
+    to establish connection with the master process. If `JULIA_WORKER_TIMEOUT` is already set,
+    that value is used. Otherwise defaults to `min(60, numprocs^2)`. When explicitly provided,
+    this temporarily overrides `JULIA_WORKER_TIMEOUT` only during worker creation.
 - `heap_size_hint_in_bytes::Union{Int,Nothing}=nothing`: On Julia 1.9+, you may set the `--heap-size-hint`
     flag on Julia processes, recommending garbage collection once a process
     is close to the recommended size. This is important for long-running distributed
@@ -469,6 +473,7 @@ function equation_search(
     numprocs::Union{Int,Nothing}=nothing,
     procs::Union{Vector{Int},Nothing}=nothing,
     addprocs_function::Union{Function,Nothing}=nothing,
+    worker_timeout::Union{Real,Nothing}=nothing,
     heap_size_hint_in_bytes::Union{Integer,Nothing}=nothing,
     worker_imports::Union{Vector{Symbol},Nothing}=nothing,
     runtests::Bool=true,
@@ -520,6 +525,7 @@ function equation_search(
         numprocs=numprocs,
         procs=procs,
         addprocs_function=addprocs_function,
+        worker_timeout=worker_timeout,
         heap_size_hint_in_bytes=heap_size_hint_in_bytes,
         worker_imports=worker_imports,
         runtests=runtests,
@@ -644,6 +650,7 @@ end
             procs=ropt.init_procs,
             ropt.numprocs,
             ropt.addprocs_function,
+            ropt.worker_timeout,
             options,
             worker_imports=ropt.worker_imports,
             project_path=splitdir(Pkg.project().path)[1],