Add max_default_jobs to bound the default number of jobs

Project.toml

@@ -1,6 +1,6 @@
 name = "ParallelTestRunner"
 uuid = "d3525ed8-44d0-4b2c-a655-542cee43accc"
-version = "2.9.4"
+version = "2.10.0"
 authors = ["Valentin Churavy <v.churavy@gmail.com>"]
 
 [deps]

docs/src/advanced.md

@@ -238,7 +238,13 @@ runtests(MyPackage, ARGS; memory_per_worker = 2^30) # 1 GiB per worker
 Workers are recycled once they exceed the `JULIA_TEST_MAXRSS_MB` environment variable (or the `max_worker_rss` keyword).
 Lowering it keeps memory pressure in check, which can make it possible to run more workers than would otherwise be safe.
 
-The number of jobs can also be forced with `--jobs=N` or the `PTR_NUM_JOBS` environment variable, but this bypasses the memory check entirely.
+Resources that PTR does not know about, such as the memory of a GPU shared by all workers, can be accounted for by passing an upper bound on the default with `max_default_jobs`:
+
+```julia
+runtests(MyPackage, ARGS; max_default_jobs = gpu_free ÷ gpu_memory_per_worker)
+```
+
+The number of jobs can also be forced with `--jobs=N` or the `PTR_NUM_JOBS` environment variable, but this bypasses the memory check and `max_default_jobs` entirely.
 
 More jobs is not always faster: once the machine runs out of memory the whole suite slows down, sometimes enough to time out a CI job.
 This has been especially problematic on low-memory macOS machines (notably CI runners), where memory compression makes each garbage collection slower.

src/ParallelTestRunner.jl

@@ -56,7 +56,8 @@ include("execute.jl")
              stdout = Base.stdout,
              stderr = Base.stderr,
              max_worker_rss = get_max_worker_rss(),
-             memory_per_worker = 2 * 2^30)
+             memory_per_worker = 2 * 2^30,
+             max_default_jobs = typemax(Int),
              serial = String[],
              serial_position::Symbol = :before,
              recycle_on_failure::Bool = false,
@@ -111,6 +112,10 @@ Several keyword arguments are also supported:
   Packages whose tests use less memory can pass a smaller value to increase the default
   parallelism. Ignored when the number of jobs is set explicitly via `--jobs=N` or the
   `PTR_NUM_JOBS` environment variable.
+- `max_default_jobs`: Upper bound on the default number of jobs, for resources that PTR
+  does not account for itself, such as the memory of a GPU shared between workers
+  (default: no bound). Like `memory_per_worker`, it is ignored when the
+  number of jobs is set explicitly via `--jobs=N` or the `PTR_NUM_JOBS` environment variable.
 - `serial`: A vector of test names (keys of `testsuite`) that should be run one at a time
   instead of in parallel.
 - `serial_position`: When to run serial tests relative to the parallel batch.
@@ -233,6 +238,7 @@ function runtests(mod::Module, args::ParsedArgs;
                   stderr = Base.stderr,
                   max_worker_rss = get_max_worker_rss(),
                   memory_per_worker = DEFAULT_MEMORY_PER_WORKER,
+                  max_default_jobs = DEFAULT_MAX_DEFAULT_JOBS,
                   recycle_on_failure::Bool = false,
                   retries::Integer = 0,
                   history_flush_every::Integer = 20,
@@ -301,6 +307,7 @@ function runtests(mod::Module, args::ParsedArgs;
         stderr,
         max_worker_rss,
         memory_per_worker,
+        max_default_jobs,
         recycle_on_failure,
         retries,
         history_flush_every,
@@ -334,6 +341,7 @@ function _runtests(mod::Module, args::ParsedArgs;
                    stderr = Base.stderr,
                    max_worker_rss = get_max_worker_rss(),
                    memory_per_worker = DEFAULT_MEMORY_PER_WORKER,
+                   max_default_jobs = DEFAULT_MAX_DEFAULT_JOBS,
                    recycle_on_failure::Bool = false,
                    retries::Integer = 0,
                    history_flush_every::Integer = 20,
@@ -344,7 +352,8 @@ function _runtests(mod::Module, args::ParsedArgs;
 
     # determine parallelism
     env_jobs = tryparse(Int, get(ENV, "PTR_NUM_JOBS", ""))
-    _jobs = @something args.jobs env_jobs default_njobs(; memory_per_worker)
+
+    _jobs = @something args.jobs env_jobs default_njobs(; memory_per_worker, max_default_jobs)
     jobs = clamp(_jobs, 1, max(1, length(parallel_tests)))
     worker_pool = Channel{Union{Nothing, PTRWorker}}(jobs)
     for _ in 1:jobs

src/resources.jl

@@ -35,20 +35,31 @@ end
 # memory). Packages whose tests are heavier can pass a larger
 # `memory_per_worker` to `runtests`.
 const DEFAULT_MEMORY_PER_WORKER = 2 * Int64(2)^30
+const DEFAULT_MAX_DEFAULT_JOBS = typemax(Int)
 
 """
-    default_njobs(; memory_per_worker = 2 * 2^30)
+    default_njobs(; memory_per_worker = 2 * 2^30, [max_default_jobs])
 
 A function used to determine the default number of parallel jobs. Calculated as the
 number of CPU threads, clamped such that each worker can be assumed to use `memory_per_worker`
-bytes of the available system memory.
+bytes of the available system memory, bounded by `max_default_jobs` if specified.
+
+Always returns at least `1`.
 """
 function default_njobs(;
-        memory_per_worker = DEFAULT_MEMORY_PER_WORKER,
+        memory_per_worker::Real = DEFAULT_MEMORY_PER_WORKER,
+        max_default_jobs::Integer = DEFAULT_MAX_DEFAULT_JOBS,
         # Just use Sys.EFFECTIVE_CPU_THREADS when min VERSION >= v"1.13"
         _cpu_threads = (@static isdefined(Sys, :EFFECTIVE_CPU_THREADS) ? Sys.EFFECTIVE_CPU_THREADS : Sys.CPU_THREADS),
         _free_memory = available_memory(),
     )
+    # Accept max_default_jobs = 0 since a common user pattern
+    #  is to limit the number of jobs with a flooring division
+    max_default_jobs >= 0 ||
+        throw(ArgumentError("`max_default_jobs` cannot be negative, got $(repr(max_default_jobs))"))
+    memory_per_worker >= 1 ||
+        throw(ArgumentError("`memory_per_worker` must be >= 1, got $(repr(memory_per_worker))"))
+
     memory_jobs = Int64(_free_memory) ÷ memory_per_worker
-    return max(1, min(_cpu_threads, memory_jobs))
+    return Int(max(1, min(_cpu_threads, memory_jobs, max_default_jobs)))
 end

test/core.jl

@@ -47,6 +47,29 @@ end
     # lighter per-worker memory estimate increases the default
     @test ParallelTestRunner.default_njobs(; _cpu_threads=4, _free_memory=UInt64(2) ^ 32,
                                              memory_per_worker=1 * Int64(2) ^ 30) == 4
+
+    # max_default_jobs respected even when other parameters would select more
+    @test ParallelTestRunner.default_njobs(; _cpu_threads=4, _free_memory=UInt64(2) ^ 32,
+                                             memory_per_worker=1 * Int64(2) ^ 30,
+                                             max_default_jobs=1) == 1
+    @test ParallelTestRunner.default_njobs(; _cpu_threads=4, _free_memory=UInt64(2) ^ 32,
+                                             memory_per_worker=1 * Int64(2) ^ 30,
+                                             max_default_jobs=2) == 2
+
+    # max_default_jobs = 0 returns 1
+    @test ParallelTestRunner.default_njobs(; max_default_jobs=0) == 1
+
+    # always returns an Int, regardless of the argument types
+    @test ParallelTestRunner.default_njobs(; memory_per_worker=1.5 * 2 ^ 30) isa Int
+
+    # invalid argument values
+    @test_throws ArgumentError ParallelTestRunner.default_njobs(; memory_per_worker=0)
+    @test_throws ArgumentError ParallelTestRunner.default_njobs(; memory_per_worker=-1)
+    @test_throws TypeError ParallelTestRunner.default_njobs(; memory_per_worker="2")
+
+    @test_throws ArgumentError ParallelTestRunner.default_njobs(; max_default_jobs=-1)
+    @test_throws TypeError ParallelTestRunner.default_njobs(; max_default_jobs=1.5)
+    @test_throws TypeError ParallelTestRunner.default_njobs(; max_default_jobs="2")
 end
 
 @testset "number of jobs" begin
@@ -69,6 +92,43 @@ end
         runtests(ParallelTestRunner, ["--jobs=1"]; testsuite, stdout=io, stderr=io)
     end
     @test contains(String(take!(io)), "using 1 parallel jobs")
+
+    # max_default_jobs bounds the default, but not explicit settings
+    if ParallelTestRunner.default_njobs() >= 2
+        io = IOBuffer()
+        withenv("PTR_NUM_JOBS" => nothing) do
+            runtests(ParallelTestRunner, String[]; testsuite, stdout=io, stderr=io)
+        end
+        @test !contains(String(take!(io)), "using 1 parallel jobs")
+    end
+    io = IOBuffer()
+    withenv("PTR_NUM_JOBS" => nothing) do
+        runtests(ParallelTestRunner, String[]; testsuite, stdout=io, stderr=io,
+                 max_default_jobs=1)
+    end
+    @test contains(String(take!(io)), "using 1 parallel jobs")
+
+    io = IOBuffer()
+    withenv("PTR_NUM_JOBS" => nothing) do
+        runtests(ParallelTestRunner, ["--jobs=2"]; testsuite, stdout=io, stderr=io,
+                 max_default_jobs=1)
+    end
+    @test contains(String(take!(io)), "using 2 parallel jobs")
+
+    io = IOBuffer()
+    withenv("PTR_NUM_JOBS" => "2") do
+        runtests(ParallelTestRunner, String[]; testsuite, stdout=io, stderr=io,
+                 max_default_jobs=1)
+    end
+    @test contains(String(take!(io)), "using 2 parallel jobs")
+
+    # Run with 1 job despite a system not having enough resources for 1 job
+    io = IOBuffer()
+    withenv("PTR_NUM_JOBS" => nothing) do
+        runtests(ParallelTestRunner, String[]; testsuite, stdout=io, stderr=io,
+                 max_default_jobs=2÷4)
+    end
+    @test contains(String(take!(io)), "using 1 parallel jobs")
 end
 
 @testset "subdir use" begin