initial commit

Author: kleinhenz

Diff for: LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2020 Joseph Kleinhenz
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

Diff for: Manifest.toml

@@ -0,0 +1,15 @@
+# This file is machine-generated - editing it directly is not advised
+
+[[Distributed]]
+deps = ["Random", "Serialization", "Sockets"]
+uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"
+
+[[Random]]
+deps = ["Serialization"]
+uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+
+[[Serialization]]
+uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
+
+[[Sockets]]
+uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

Diff for: Project.toml

@@ -0,0 +1,7 @@
+name = "SlurmClusterManager"
+uuid = "c82cd089-7bf7-41d7-976b-6b5d413cbe0a"
+authors = ["Joseph Kleinhenz <[email protected]>"]
+version = "0.1.0"
+
+[deps]
+Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"

Diff for: README.md

@@ -0,0 +1,34 @@
+# SlurmClusterManager.jl
+
+This package provides support for using julia within the Slurm cluster environment.
+The code is adapted from [ClusterManagers.jl](https://github.com/JuliaParallel/ClusterManagers.jl) with some modifications.
+
+## Usage
+
+This script uses all resources from a Slurm allocation as julia workers and prints the id and hostname on each one.
+
+```jl
+#!/usr/bin/env julia
+
+using Distributed, SlurmClusterManager
+addprocs(SlurmManager())
+@everywhere println("hello from $(myid()):$(gethostname())")
+```
+
+If the code is saved in `script.jl` it can be queued and executed on two nodes using 64 workers per node by running
+
+```
+sbatch -N 2 --ntasks-per-node=64 script.jl
+```
+
+## Differences from `ClusterManagers.jl`
+
+* Only supports Slurm (see this [issue](https://github.com/JuliaParallel/ClusterManagers.jl/issues/58) for some background).
+* Requires that `SlurmManager` be created inside a Slurm allocation created by sbatch/salloc.
+  Specifically `SLURM_JOBID` and `SLURM_NTASKS` must be defined in order to construct `SlurmManager`.
+  This matches typical HPC workflows where resources are requested using sbatch and then used by the application code.
+  In contrast `ClusterManagers.jl` will  *dynamically* request resources when run outside of an existing Slurm allocation.
+  I found that this was basically never what I wanted since this leaves the manager process running on a login node,
+  and makes the script wait until resources are granted which is better handled by the actual Slurm queueing system.
+* Does not take any Slurm arguments. All Slurm arguments are inherited from the external Slurm allocation created by sbatch/salloc.
+* Output from workers is redirected to the manager process instead of requiring a separate output file for every task.

Diff for: src/SlurmClusterManager.jl

@@ -0,0 +1,10 @@
+module SlurmClusterManager
+
+export SlurmManager, launch, manage
+
+using Distributed
+import Distributed: launch, manage
+
+include("slurmmanager.jl")
+
+end # module

Diff for: src/slurmmanager.jl

@@ -0,0 +1,71 @@
+"""
+ClusterManager for a Slurm allocation
+
+Represents the resources available within a slurm allocation created by salloc/sbatch.
+The environment variables `SLURM_JOBID` and `SLURM_NTASKS` must be defined to construct this object.
+"""
+struct SlurmManager <: ClusterManager 
+  jobid::Int
+  ntasks::Int
+  verbose::Bool
+  launch_timeout::Float64
+
+  function SlurmManager(;verbose=false, launch_timeout=60.0)
+    @assert "SLURM_JOBID" in keys(ENV)
+    @assert "SLURM_NTASKS" in keys(ENV)
+
+    jobid = parse(Int, ENV["SLURM_JOBID"])
+    ntasks = parse(Int, ENV["SLURM_NTASKS"])
+
+    new(jobid, ntasks, verbose, launch_timeout)
+  end
+end
+
+function launch(manager::SlurmManager, params::Dict, instances_arr::Array, c::Condition)
+    try
+        exehome = params[:dir]
+        exename = params[:exename]
+        exeflags = params[:exeflags]
+
+        srun_cmd = `srun -D $exehome $exename $exeflags --worker=$(cluster_cookie())`
+        srun_proc = open(srun_cmd)
+
+        t = @async for i in 1:manager.ntasks
+          manager.verbose && println("connecting to worker $i out of $(manager.ntasks)")
+
+          line = readline(srun_proc)
+          m = match(r".*:(\d*)#(.*)", line)
+          m === nothing && error("could not parse $line")
+
+          config = WorkerConfig()
+          config.port = parse(Int, m[1])
+          config.host = strip(m[2])
+
+          # 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
+
+        # workers must be launched before timeout otherwise interrupt
+        status = timedwait(() -> istaskdone(t), manager.launch_timeout)
+        if status !== :ok
+          @async Base.throwto(t, ErrorException("launch_timeout exceeded"))
+        end
+        wait(t)
+
+        # redirect output
+        @async while !eof(srun_proc)
+          line = readline(srun_proc)
+          println(line)
+        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