SSHManager: Add the custom_worker_flag option, but don't add it to the public API

Author: DilumAluthge

Diff for: src/slurmmanager.jl

@@ -44,22 +44,35 @@ end
 
 @static if Base.VERSION >= v"1.9.0"
   # In Julia 1.9 and later, the no-argument method `Distributed.default_addprocs_params()`
-  # includes :env, so we don't need to do anything.
+  # includes :env.
   # See also: https://github.com/JuliaLang/julia/blob/v1.9.0/stdlib/Distributed/src/cluster.jl#L526-L541
+  #
+  # So we don't need to add `:env`
+  # But we do still need to add `:custom_worker_flag`
 
-  Distributed.default_addprocs_params(::SlurmManager) = Distributed.default_addprocs_params()
+  function Distributed.default_addprocs_params(::SlurmManager)
+    our_stuff = Dict{Symbol,Any}(
+      custom_worker_flag = `--worker`,
+    )
+    upstreams_stuff = Distributed.default_addprocs_params()
+    total_stuff = merge(our_stuff, upstreams_stuff)
+    return total_stuff
+  end
 elseif v"1.6.0" <= Base.VERSION < v"1.9.0"
   # In Julia 1.6 through 1.8, the no-argument method `Distributed.default_addprocs_params()`
   # does not include :env. However, Julia does allow us to add a specialized method
   # `Distributed.default_addprocs_params(::SlurmManager)`, so we do so here.
   #
   # The ability to add the specialized `Distributed.default_addprocs_params(::SlurmManager)`
-  # method was added to Julia in https://github.com/JuliaLang/julia/pull/38570
+  # method was added to Julia in https://github.com/JuliaLang/julia/pull/38570,
+  # which corresponds to https://github.com/JuliaLang/julia/commit/687378ebfd1188ee7af302ec089cf3d31689647c,
+  # which is in Julia 1.6 and later.
   #
   # See also: https://github.com/JuliaLang/julia/blob/v1.8.0/stdlib/Distributed/src/cluster.jl#L526-L540
   function Distributed.default_addprocs_params(::SlurmManager)
     our_stuff = Dict{Symbol,Any}(
       :env => [],
+      custom_worker_flag = `--worker`,
     )
     upstreams_stuff = Distributed.default_addprocs_params()
     total_stuff = merge(our_stuff, upstreams_stuff)
@@ -69,6 +82,10 @@ elseif Base.VERSION < v"1.6.0"
   # In Julia 1.5 and earlier, Julia does not have the `addenv()` function.
   # I don't want to add a dependency on Compat.jl just for this one feature,
   # so we will just choose to not support `params[:env]` on Julia 1.5 and earlier.
+  #
+  # Also, in Julia 1.5 and earlier, Julia does not have the
+  # ability to add the specialized `Distributed.default_addprocs_params(::SlurmManager)`
+  # method.
 end
 
 function _new_environment_additions(params_env::Dict{String, String})
@@ -178,8 +195,9 @@ function Distributed.launch(manager::SlurmManager, params::Dict, instances_arr::
         exehome = params[:dir]
         exename = params[:exename]
         exeflags = params[:exeflags]
+        custom_worker_flag = params[:custom_worker_flag]
 
-        _srun_cmd_without_env = `srun -D $exehome $exename $exeflags --worker`
+        _srun_cmd_without_env = `srun -D $exehome $exename $exeflags $custom_worker_flag`
 
         @static if Base.VERSION >= v"1.6.0"
           # Pass the key-value pairs from `params[:env]` to the `srun` command: