Merge pull request #488 from MilesCranmer/timeout_parameter

feat: add `worker_timeout` to allow workers more time to connect

Author: MilesCranmer

src/Configure.jl

@@ -368,6 +368,7 @@ function configure_workers(;
     procs::Union{Vector{Int},Nothing},
     numprocs::Int,
     addprocs_function::Function,
+    worker_timeout::Float64,
     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

@@ -78,6 +78,7 @@ function modelexpr(
             procs::Union{Vector{Int},Nothing} = nothing
             addprocs_function::Union{Function,Nothing} = nothing
             heap_size_hint_in_bytes::Union{Integer,Nothing} = nothing
+            worker_timeout::Union{Real,Nothing} = nothing
             worker_imports::Union{Vector{Symbol},Nothing} = nothing
             logger::Union{AbstractSRLogger,Nothing} = nothing
             runtests::Bool = true
@@ -300,6 +301,7 @@ function _update(
         procs=m.procs,
         addprocs_function=m.addprocs_function,
         heap_size_hint_in_bytes=m.heap_size_hint_in_bytes,
+        worker_timeout=m.worker_timeout,
         worker_imports=m.worker_imports,
         runtests=m.runtests,
         saved_state=(old_fitresult === nothing ? nothing : old_fitresult.state),

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
@@ -118,6 +119,7 @@ end
     procs::Union{Vector{Int},Nothing}=nothing,
     addprocs_function::Union{Function,Nothing}=nothing,
     heap_size_hint_in_bytes::Union{Integer,Nothing}=nothing,
+    worker_timeout::Union{Real,Nothing}=nothing,
     worker_imports::Union{Vector{Symbol},Nothing}=nothing,
     runtests::Bool=true,
     return_state::VRS=nothing,
@@ -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", "")),
+            max(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

@@ -411,6 +411,9 @@ which is useful for debugging and profiling.
     is close to the recommended size. This is important for long-running distributed
     jobs where each process has an independent memory, and can help avoid
     out-of-memory errors. By default, this is set to `Sys.free_memory() / numprocs`.
+- `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 `max(60, numprocs^2)`.
 - `worker_imports::Union{Vector{Symbol},Nothing}=nothing`: If you want to import
     additional modules on each worker, pass them here as a vector of symbols.
     By default some of the extensions will automatically be loaded when needed.
@@ -470,6 +473,7 @@ function equation_search(
     procs::Union{Vector{Int},Nothing}=nothing,
     addprocs_function::Union{Function,Nothing}=nothing,
     heap_size_hint_in_bytes::Union{Integer,Nothing}=nothing,
+    worker_timeout::Union{Real,Nothing}=nothing,
     worker_imports::Union{Vector{Symbol},Nothing}=nothing,
     runtests::Bool=true,
     saved_state=nothing,
@@ -521,6 +525,7 @@ function equation_search(
         procs=procs,
         addprocs_function=addprocs_function,
         heap_size_hint_in_bytes=heap_size_hint_in_bytes,
+        worker_timeout=worker_timeout,
         worker_imports=worker_imports,
         runtests=runtests,
         saved_state=saved_state,
@@ -644,6 +649,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],