From cfd9d63b622aebd92bedc9c1f6319f1b9369feae Mon Sep 17 00:00:00 2001
From: Dilum Aluthge <dilum@aluthge.com>
Date: Sun, 9 Feb 2025 16:10:35 -0500
Subject: [PATCH] Factor the "construct `srun` command" functionality out into
 a separate `construct_srun_cmd()` function

---
 src/slurmmanager.jl | 47 ++++++++++++++++++++++++++++-----------------
 1 file changed, 29 insertions(+), 18 deletions(-)

diff --git a/src/slurmmanager.jl b/src/slurmmanager.jl
index ff6573e..a2c677c 100644
--- a/src/slurmmanager.jl
+++ b/src/slurmmanager.jl
@@ -171,31 +171,42 @@ function warn_if_unexpected_params(params::Dict)
     return nothing
 end
 
+@static if Base.VERSION >= v"1.6.0"
+  # Pass the key-value pairs from `params[:env]` to the `srun` command:
+  function construct_srun_cmd(; params::Dict)
+    exehome = params[:dir]
+    exename = params[:exename]
+    exeflags = params[:exeflags]
+    
+    _srun_cmd_without_env = `srun -D $exehome $exename $exeflags --worker`
+
+    env2 = _new_environment_additions(Dict{String,String}(params[:env]))
+    srun_cmd_with_env = addenv(_srun_cmd_without_env, env2)
+
+    return srun_cmd_with_env
+  end
+else
+  # See discussion above for why we don't support this functionality on Julia 1.5 and earlier.
+  function construct_srun_cmd(; params::Dict)
+    exehome = params[:dir]
+    exename = params[:exename]
+    exeflags = params[:exeflags]
+
+    _srun_cmd_without_env = `srun -D $exehome $exename $exeflags --worker`
+
+    return _srun_cmd_without_env
+  end
+end
+
 function Distributed.launch(manager::SlurmManager, params::Dict, instances_arr::Array, c::Condition)
     try
         warn_if_unexpected_params(params)
 
-        exehome = params[:dir]
-        exename = params[:exename]
-        exeflags = params[:exeflags]
-
-        _srun_cmd_without_env = `srun -D $exehome $exename $exeflags --worker`
-
-        @static if Base.VERSION >= v"1.6.0"
-          # Pass the key-value pairs from `params[:env]` to the `srun` command:
-          env2 = _new_environment_additions(Dict{String,String}(params[:env]))
-          srun_cmd_with_env = addenv(_srun_cmd_without_env, env2)
-        else
-          # See discussion above for why we don't support this functionality on Julia 1.5 and earlier.
-          if haskey(params, :env)
-            @warn "SlurmClusterManager.jl does not support params[:env] on Julia 1.5 and earlier" Base.VERSION
-          end
-          srun_cmd_with_env = _srun_cmd_without_env
-        end
+        srun_cmd = construct_srun_cmd(; params=params)
 
         # Pass cookie as stdin to srun; srun forwards stdin to process
         # This way the cookie won't be visible in ps, top, etc on the compute node
-        manager.srun_proc = open(srun_cmd_with_env, write=true, read=true)
+        manager.srun_proc = open(srun_cmd, write=true, read=true)
         write(manager.srun_proc, cluster_cookie())
         write(manager.srun_proc, "\n")