Merge pull request #174 from CliMA/ne/addprocs

Add default function to add workers

Author: nefrathenrici

docs/src/api.md

@@ -11,6 +11,7 @@ ClimaCalibrate.postprocess_g_ensemble
 
 ## Worker Interface
 ```@docs
+ClimaCalibrate.add_workers
 ClimaCalibrate.WorkerBackend
 ClimaCalibrate.SlurmManager
 ClimaCalibrate.PBSManager

src/backends.jl

@@ -34,6 +34,8 @@ Used for CliMA's private GPU server.
 """
 struct ClimaGPUBackend <: SlurmBackend end
 
+struct GCPBackend <: SlurmBackend end
+
 """
     DerechoBackend
 
@@ -61,6 +63,8 @@ function get_backend()
         (r"^clima.gps.caltech.edu$", ClimaGPUBackend),
         (r"^login[1-4].cm.cluster$", CaltechHPCBackend),
         (r"^hpc-(\d\d)-(\d\d).cm.cluster$", CaltechHPCBackend),
+        (r"^hpc\d+-slurm-login-\d+$", GCPBackend),
+        (r"^hpc\d+-a\d+nodeset-\d+$", GCPBackend),
         (r"derecho([1-8])$", DerechoBackend),
         (r"deg(\d\d\d\d)$", DerechoBackend), # This should be more specific
     ]
@@ -434,3 +438,9 @@ function model_run(
     end
     return job_id
 end
+
+backend_worker_kwargs(::Type{DerechoBackend}) = (; q = "main", A = "UCIT0011")
+
+backend_worker_kwargs(::Type{GCPBackend}) = (; partition = "a3")
+
+backend_worker_kwargs(::Type{<:AbstractBackend}) = (;)

src/workers.jl

@@ -1,7 +1,8 @@
 using Distributed
 using Logging
 
-export SlurmManager, PBSManager, default_worker_pool, set_worker_loggers
+export add_workers,
+    SlurmManager, PBSManager, default_worker_pool, set_worker_loggers
 
 # Set the time limit for the Julia worker to be contacted by the main process, default = "60.0s"
 # https://docs.julialang.org/en/v1/manual/environment-variables/#JULIA_WORKER_TIMEOUT
@@ -444,3 +445,103 @@ function set_worker_loggers(workers = workers())
         end
     end
 end
+
+
+function is_pbs_available()
+    return all([
+        !isnothing(Sys.which("qstat")),
+        !isnothing(Sys.which("pbsnodes")),
+        !isnothing(Sys.which("qsub")),
+    ])
+end
+
+
+function is_slurm_available()
+    return all([
+        !isnothing(Sys.which("sinfo")),
+        !isnothing(Sys.which("srun")),
+        !isnothing(Sys.which("sbatch")),
+    ])
+end
+
+function is_cluster_environment()
+    return is_pbs_available() || is_slurm_available()
+end
+
+const DEFAULT_WALLTIME = 60
+
+default_cpu_kwargs(::SlurmManager) = (;
+    cpus_per_task = 1,
+    time = format_slurm_time(DEFAULT_WALLTIME),
+    backend_worker_kwargs(get_backend())...,
+)
+default_cpu_kwargs(::PBSManager) = (;
+    l_select = "ncpus=1",
+    l_walltime = format_pbs_time(DEFAULT_WALLTIME),
+    backend_worker_kwargs(get_backend())...,
+)
+
+default_gpu_kwargs(::SlurmManager) = (;
+    gpus_per_task = 1,
+    cpus_per_task = 4,
+    time = format_slurm_time(DEFAULT_WALLTIME),
+    backend_worker_kwargs(get_backend())...,
+)
+default_gpu_kwargs(::PBSManager) = (;
+    l_select = "ngpus=1:ncpus=4",
+    l_walltime = format_pbs_time(DEFAULT_WALLTIME),
+    backend_worker_kwargs(get_backend())...,
+)
+
+function get_manager(cluster = :auto, nworkers = 1)
+    if cluster == :slurm || (cluster == :auto && is_slurm_available())
+        SlurmManager(nworkers)
+    elseif cluster == :pbs || (cluster == :auto && is_pbs_available())
+        PBSManager(nworkers)
+    else
+        error(
+            "Unknown cluster type: $cluster. Valid options are :auto, :pbs, :slurm, or :local",
+        )
+    end
+end
+
+"""
+    add_workers(
+        nworkers;
+        device = :gpu,
+        cluster = :auto,
+        kwargs...
+    )
+
+Add `nworkers` worker processes to the current Julia session, automatically detecting and configuring for the available computing environment.
+
+# Arguments
+- `nworkers::Int`: The number of worker processes to add.
+- `device::Symbol = :gpu`: The target compute device type, either `:gpu` (1 GPU, 4 CPU cores) or `:cpu` (1 CPU core).
+- `cluster::Symbol = :auto`: The cluster management system to use. Options:
+  * `:auto`: Auto-detect available cluster environment (SLURM, PBS, or local)
+  * `:slurm`: Force use of SLURM scheduler
+  * `:pbs`: Force use of PBS scheduler
+  * `:local`: Force use of local processing (standard `addprocs`)
+- `kwargs`: Other kwargs can be passed directly through to `addprocs`.
+"""
+function add_workers(nworkers::Int; device = :gpu, cluster = :auto, kwargs...)
+    if cluster == :local || (cluster == :auto && !is_cluster_environment())
+        # Use standard addprocs for local computation
+        @info "Using local processing mode, adding $nworkers worker$(nworkers == 1 ? "" : "s")"
+        return addprocs(nworkers)
+    else
+        # Select the manager based on environment or explicit selection
+        manager = get_manager(cluster, nworkers)
+        @info "Using $(nameof(typeof(manager))) to add $nworkers workers"
+
+        default_kwargs =
+            device == :gpu ? default_gpu_kwargs(manager) :
+            default_cpu_kwargs(manager)
+
+        # Merge the default kwargs with the user-provided kwargs, user kwargs take precedence
+        merged_kwargs = merge(default_kwargs, kwargs)
+
+        return addprocs(manager; merged_kwargs...)
+    end
+end

test/hpc_backend.jl

@@ -17,7 +17,7 @@ if backend == DerechoBackend
 end
 
 original_model_interface = model_interface
-interruption_model_interface, io = mktemp()
+interruption_model_interface, io = mktemp(@__DIR__)
 model_interface_str = """
 import ClimaCalibrate
 ClimaCalibrate.forward_model(iter, member) = member == 1 && exit()

test/pbs_manager_unit_tests.jl

@@ -1,13 +1,8 @@
 using Test, ClimaCalibrate, Distributed, Logging
 
 @testset "PBSManager Unit Tests" begin
-    p = addprocs(
-        PBSManager(1),
-        q = "main",
-        A = "UCIT0011",
-        l_select = "ngpus=1",
-        l_walltime = "00:05:00",
-    )
+    @test ClimaCalibrate.get_manager() == PBSManager(1)
+    p = add_workers(1; l_walltime = "00:05:00")
     @test nprocs() == length(p) + 1
     @test workers() == p
     @test remotecall_fetch(myid, 2) == 2

test/slurm_manager_unit_tests.jl

@@ -1,8 +1,9 @@
 using Test, ClimaCalibrate, Distributed, Logging
 
 @testset "SlurmManager Unit Tests" begin
-    out_file = "my_slurm_job.out"
-    p = addprocs(SlurmManager(1); o = out_file)
+    @test ClimaCalibrate.get_manager() == SlurmManager(1)
+    out_file = tempname()
+    p = add_workers(1; device = :cpu, o = out_file)
     @test nprocs() == 2
     @test workers() == p
     @test fetch(@spawnat :any myid()) == p[1]
@@ -16,7 +17,6 @@ using Test, ClimaCalibrate, Distributed, Logging
     @test workers() == [1]
     # Check output file creation
     @test isfile(out_file)
-    rm(out_file)
 
     # Test incorrect generic arguments
     @test_throws TaskFailedException p = addprocs(SlurmManager(1), time = "w")