Add SLURM support

Project.toml

@@ -3,13 +3,15 @@ uuid = "9feedd95-f0e0-423f-a8dc-de0970eae6b3"
 version = "0.2.7"
 
 [deps]
+ClusterManagers = "34f1f09b-3a8b-5176-ab39-66d58a4d544e"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Format = "1fa38f19-a742-5d3f-a2b9-30dd87b9d5f8"
 Glob = "c27321d9-0574-5035-807b-f59d2c89b15c"
 IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MIMEs = "6c6e2e6c-3030-632d-7369-2d6c69616d65"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 Measurements = "eff96d63-e80a-5855-80a2-b1b0885c5ab7"
@@ -23,6 +25,7 @@ RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 StaticStrings = "4db0a0c5-418a-4e1d-8806-cb305fe13294"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
+ThreadPinning = "811555cd-349b-4f26-b7bc-1f208b848042"
 TypedTables = "9d95f2ec-7b3d-5a63-8d20-e2491e220bb9"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
@@ -36,13 +39,15 @@ LegendDataManagementSolidStateDetectorsExt = "SolidStateDetectors"
 LegendDataManagementRecipesBaseExt = "RecipesBase"
 
 [compat]
+ClusterManagers = "0.4.5"
 Dates = "1"
 Distributed = "1"
 Format = "1"
 Glob = "1"
 IntervalSets = "0.6, 0.7"
 JSON = "0.21, 1"
 LRUCache = "1.5"
+LinearAlgebra = "1"
 MIMEs = "0.1"
 Markdown = "1"
 Measurements = "2"
@@ -57,7 +62,8 @@ SolidStateDetectors = "0.8, 0.9"
 StaticStrings = "0.2"
 StructArrays = "0.4, 0.5, 0.6"
 Tables = "0.2, 1.0"
+ThreadPinning = "0.7"
 TypedTables = "1.4"
 UUIDs = "1"
 Unitful = "1"
-julia = "1.9"
+julia = "1.9"

src/LegendDataManagement.jl

@@ -8,6 +8,7 @@ using Dates
 using UUIDs
 
 import Distributed
+import LinearAlgebra
 import Pkg
 
 using Glob
@@ -32,6 +33,9 @@ using TypedTables
 import Markdown
 using MIMEs: mime_from_extension
 
+import ThreadPinning
+import ClusterManagers
+
 include("legend_report.jl")
 include("status_types.jl")
 include("atomic_fcreate.jl")
@@ -50,5 +54,6 @@ include("calibration_functions.jl")
 include("lprops.jl")
 include("utils/utils.jl")
 
+include("slurm.jl")
 
 end # module

src/slurm.jl

@@ -0,0 +1,171 @@
+# ClusterManager for Slurm
+
+# This code is a modified version of ClusterManagers.SlurmManager, both
+# original code and modifications are licensed under the MIT License (MIT):
+# https://github.com/JuliaParallel/ClusterManagers.jl/blob/master/LICENSE.md
+
+# Modifications are to be upstreamed are planned to be upstreamed as far
+# as possible, once proved in the field.
+
+
+# ==================================================================
+using Base: ExponentialBackOff
+using Distributed: ClusterManager, WorkerConfig
+using ClusterManagers: worker_arg
+
+import Distributed: launch, manage
+# ==================================================================
+
+
+struct SlurmManager <: ClusterManager
+    np::Integer
+    retry_delays
+end
+
+struct SlurmException <: Exception
+    msg
+end
+
+function launch(manager::SlurmManager, params::Dict, instances_arr::Array,
+                c::Condition)
+
+    try
+        exehome = params[:dir]
+        exename = params[:exename]
+        exeflags = params[:exeflags]
+
+        stdkeys = keys(Distributed.default_addprocs_params())
+
+        p = filter(x->(!(x[1] in stdkeys) && x[1] != :job_file_loc), params)
+
+        srunargs = []
+        for k in keys(p)
+            if length(string(k)) == 1
+                push!(srunargs, "-$k")
+                val = p[k]
+                if length(val) > 0
+                    push!(srunargs, "$(p[k])")
+                end
+            else
+                k2 = replace(string(k), "_"=>"-")
+                val = p[k]
+                if length(val) > 0
+                    push!(srunargs, "--$(k2)=$(p[k])")
+                else
+                    push!(srunargs, "--$(k2)")
+                end
+            end
+        end
+
+        # Get job file location from parameter dictionary.
+        job_file_loc = joinpath(exehome, get(params, :job_file_loc, "."))
+
+        # Make directory if not already made.
+        if !isdir(job_file_loc)
+            mkdir(job_file_loc)
+        end
+
+    # Check for given output file name
+    jobname = "julia-$(getpid())"
+    has_output_name = ("-o" in srunargs) | ("--output" in srunargs)
+    if has_output_name
+        loc = findfirst(x-> x == "-o" || x == "--output", srunargs)
+        job_output_name = srunargs[loc+1]
+        job_output_template = joinpath(job_file_loc, job_output_name)
+        srunargs[loc+1] = job_output_template
+    else
+        job_output_name = "$(jobname)-$(trunc(Int, Base.time() * 10))"
+        make_job_output_path(task_num) = joinpath(job_file_loc, "$(job_output_name)-$(task_num).out")
+        job_output_template = make_job_output_path("%4t")
+        push!(srunargs, "-o", job_output_template)
+    end
+
+        np = manager.np
+        srun_cmd = `srun -J $jobname -n $np -D $exehome $(srunargs) $exename $exeflags $(worker_arg())`
+
+        global g_state = (;manager, params, instances_arr, c, srun_cmd, has_output_name, job_output_name, make_job_output_path, job_output_template)
+
+        @info "Starting SLURM job $jobname: $srun_cmd"
+        srun_proc = open(srun_cmd)
+
+        slurm_spec_regex = r"([\w]+):([\d]+)#(\d{1,3}.\d{1,3}.\d{1,3}.\d{1,3})"
+        could_not_connect_regex = r"could not connect"
+        exiting_regex = r"exiting."
+        retry_delays = manager.retry_delays
+
+        t_start = time()
+        t_waited = round(Int, time() - t_start)
+        for i = 0:np - 1
+            slurm_spec_match::Union{RegexMatch,Nothing} = nothing
+            worker_errors = String[]
+            if has_output_name
+                fn = job_output_template
+            else
+                fn = make_job_output_path(lpad(i, 4, "0"))
+            end
+            for retry_delay in push!(collect(retry_delays), 0)
+                t_waited = round(Int, time() - t_start)
+
+                # Wait for output log to be created and populated, then parse
+
+                if isfile(fn)
+                    if filesize(fn) > 0
+                        open(fn) do f
+                            # Due to error and warning messages, the specification
+                            # may not appear on the file's first line
+                            for line in eachline(f)
+                                re_match = match(slurm_spec_regex, line)
+                                if !isnothing(re_match)
+                                    slurm_spec_match = re_match
+                                end
+                                for expr in [could_not_connect_regex, exiting_regex]
+                                    if !isnothing(match(expr, line))
+                                        slurm_spec_match = nothing
+                                        push!(worker_errors, line)
+                                    end
+                                end
+                            end
+                        end
+                    end
+                    if !isempty(worker_errors) || !isnothing(slurm_spec_match)
+                        break   # break if error or specification found
+                    else
+                        @info "Worker $i (after $t_waited s): Output file found, but no connection details yet"
+                    end
+                else
+                    @info "Worker $i (after $t_waited s): No output file \"$fn\" yet"
+                end
+
+                # Sleep for some time to limit ressource usage while waiting for the job to start
+                sleep(retry_delay)
+            end
+
+            if !isempty(worker_errors)
+                throw(SlurmException("Worker $i failed after $t_waited s: $(join(worker_errors, " "))"))
+            elseif isnothing(slurm_spec_match)
+                throw(SlurmException("Timeout after $t_waited s while waiting for worker $i to get ready."))
+            end
+
+            config = WorkerConfig()
+            config.port = parse(Int, slurm_spec_match[2])
+            config.host = strip(slurm_spec_match[3])
+            @info "Worker $i ready after $t_waited s on host $(config.host), port $(config.port)"
+            # Keep a reference to the proc, so it's properly closed once
+            # the last worker exits.
+            config.userdata = srun_proc
+            push!(instances_arr, config)
+            notify(c)
+        end
+    catch e
+        println("Error launching Slurm job:")
+        rethrow(e)
+    end
+end
+
+function manage(manager::SlurmManager, id::Integer, config::WorkerConfig,
+                op::Symbol)
+    # This function needs to exist, but so far we don't do anything
+end
+
+SlurmManager(np::Integer) = SlurmManager(np, ExponentialBackOff(n=10, first_delay=1,
+                                                                max_delay=512, factor=2))