Update slurmmanager.jl

Author: DilumAluthge

src/slurmmanager.jl

@@ -12,34 +12,40 @@ mutable struct SlurmManager <: ClusterManager
   srun_proc
 
   function SlurmManager(; launch_timeout=60.0, srun_post_exit_sleep=0.01)
+    ntasks_int = get_slurm_ntasks_int()
+    jobid_int = get_slurm_jobid_int()
 
-    jobid =
-    if "SLURM_JOB_ID" in keys(ENV)
-        ENV["SLURM_JOB_ID"]
-    elseif "SLURM_JOBID" in keys(ENV)
-        ENV["SLURM_JOBID"]
-    else
-        error("""
-              SlurmManager must be constructed inside a slurm allocation environemnt.
-              SLURM_JOB_ID or SLURM_JOBID must be defined.
-              """)
-    end
-
-    ntasks =
-    if "SLURM_NTASKS" in keys(ENV)
-      ENV["SLURM_NTASKS"]
-    else
-      error("""
-            SlurmManager must be constructed inside a slurm environment with a specified number of tasks.
-            SLURM_NTASKS must be defined.
-            """)
-    end
+    new(jobid_int, ntasks_int, launch_timeout, srun_post_exit_sleep, nothing)
+  end
+end
 
-    jobid = parse(Int, jobid)
-    ntasks = parse(Int, ntasks)
+function get_slurm_ntasks_int()
+  if "SLURM_JOB_ID" in keys(ENV)
+    ntasks_str = ENV["SLURM_JOB_ID"]
+  elseif "SLURM_JOBID" in keys(ENV)
+    ntasks_str = ENV["SLURM_JOBID"]
+  else
+    msg = "SlurmManager must be constructed inside a Slurm allocation environment." *
+          "SLURM_JOB_ID or SLURM_JOBID must be defined."
+    error(msg)
+  end
+  ntasks_int = parse(Int, ntasks_str)::Int
+  return ntasks_int
+end
 
-    new(jobid, ntasks, launch_timeout, srun_post_exit_sleep, nothing)
+function get_slurm_jobid_int()
+  if "SLURM_JOB_ID" in keys(ENV)
+      jobid_str = ENV["SLURM_JOB_ID"]
+  elseif "SLURM_JOBID" in keys(ENV)
+      jobid_str = ENV["SLURM_JOBID"]
+  else
+      msg = "SlurmManager must be constructed inside a Slurm allocation environment." *
+            "SLURM_JOB_ID or SLURM_JOBID must be defined."
+      error(msg)
   end
+
+  jobid_int = parse(Int, jobid_str)::Int
+  return jobid_int
 end
 
 @static if Base.VERSION >= v"1.9.0"