Display pre-test time

src/ParallelTestRunner.jl

@@ -45,17 +45,20 @@ Malt.isrunning(wrkr::PTRWorker) = Malt.isrunning(wrkr.w)
 Malt.stop(wrkr::PTRWorker) = Malt.stop(wrkr.w)
 
 #Always set the max rss so that if tests add large global variables (which they do) we don't make the GC's life too hard
-if Sys.WORD_SIZE == 64
-    const JULIA_TEST_MAXRSS_MB = 3800
-else
-    # Assume that we only have 3.5GB available to a single process, and that a single
-    # test can take up to 2GB of RSS.  This means that we should instruct the test
-    # framework to restart any worker that comes into a test set with 1.5GB of RSS.
-    const JULIA_TEST_MAXRSS_MB = 1536
+function get_max_worker_rss()
+    mb = if haskey(ENV, "JULIA_TEST_MAXRSS_MB")
+        parse(Int, ENV["JULIA_TEST_MAXRSS_MB"])
+    elseif Sys.WORD_SIZE == 64
+        Sys.total_memory() > 8*2^30 ? 3800 : 3000
+    else
+        # Assume that we only have 3.5GB available to a single process, and that a single
+        # test can take up to 2GB of RSS.  This means that we should instruct the test
+        # framework to restart any worker that comes into a test set with 1.5GB of RSS.
+        1536
+    end
+    return mb * 2^20
 end
 
-const max_worker_rss = JULIA_TEST_MAXRSS_MB * 2^20
-
 function with_testset(f, testset)
     @static if VERSION >= v"1.13.0-DEV.1044"
         Test.@with_testset testset f()
@@ -84,6 +87,7 @@ struct TestRecord <: AbstractTestRecord
     bytes::UInt64
     gctime::Float64
     rss::UInt64
+    total_time::Float64
 end
 
 function memory_usage(rec::TestRecord)
@@ -103,6 +107,7 @@ struct TestIOContext
     stdout::IO
     stderr::IO
     color::Bool
+    debug_timing::Bool
     lock::ReentrantLock
     name_align::Int
     elapsed_align::Int
@@ -112,7 +117,7 @@ struct TestIOContext
     rss_align::Int
 end
 
-function test_IOContext(stdout::IO, stderr::IO, lock::ReentrantLock, name_align::Int)
+function test_IOContext(stdout::IO, stderr::IO, lock::ReentrantLock, name_align::Int, debug_timing::Bool)
     elapsed_align = textwidth("Time (s)")
     gc_align = textwidth("GC (s)")
     percent_align = textwidth("GC %")
@@ -122,19 +127,26 @@ function test_IOContext(stdout::IO, stderr::IO, lock::ReentrantLock, name_align:
     color = get(stdout, :color, false)
 
     return TestIOContext(
-        stdout, stderr, color, lock, name_align, elapsed_align, gc_align, percent_align,
+        stdout, stderr, color, debug_timing, lock, name_align, elapsed_align, gc_align, percent_align,
         alloc_align, rss_align
     )
 end
 
 function print_header(ctx::TestIOContext, testgroupheader, workerheader)
     lock(ctx.lock)
     try
+        # header top
         printstyled(ctx.stdout, " "^(ctx.name_align + textwidth(testgroupheader) - 3), " │ ")
-        printstyled(ctx.stdout, "         │ ──────────────── CPU ──────────────── │\n", color = :white)
+        printstyled(ctx.stdout, "         |", color = :white)
+        ctx.debug_timing && printstyled(ctx.stdout, "   Init   │", color = :white)
+        printstyled(ctx.stdout, " ──────────────── CPU ──────────────── │\n", color = :white)
+
+        # header bottom
         printstyled(ctx.stdout, testgroupheader, color = :white)
         printstyled(ctx.stdout, lpad(workerheader, ctx.name_align - textwidth(testgroupheader) + 1), " │ ", color = :white)
-        printstyled(ctx.stdout, "Time (s) │ GC (s) │ GC % │ Alloc (MB) │ RSS (MB) │\n", color = :white)
+        printstyled(ctx.stdout, "Time (s) │", color = :white)
+        ctx.debug_timing && printstyled(ctx.stdout, " time (s) │", color = :white)
+        printstyled(ctx.stdout, " GC (s) │ GC % │ Alloc (MB) │ RSS (MB) │\n", color = :white)
         flush(ctx.stdout)
     finally
         unlock(ctx.lock)
@@ -160,9 +172,15 @@ function print_test_finished(record::TestRecord, wrkr, test, ctx::TestIOContext)
     try
         printstyled(ctx.stdout, test, color = :white)
         printstyled(ctx.stdout, lpad("($wrkr)", ctx.name_align - textwidth(test) + 1, " "), " │ ", color = :white)
-        time_str = @sprintf("%7.2f", record.time)
+        time = record.time
+        time_str = @sprintf("%7.2f", time)
         printstyled(ctx.stdout, lpad(time_str, ctx.elapsed_align, " "), " │ ", color = :white)
 
+        if ctx.debug_timing
+            init_time_str = @sprintf("%7.2f", record.total_time - time)
+            printstyled(ctx.stdout, lpad(init_time_str, ctx.elapsed_align, " "), " │ ", color = :white)
+        end
+
         gc_str = @sprintf("%5.2f", record.gctime)
         printstyled(ctx.stdout, lpad(gc_str, ctx.gc_align, " "), " │ ", color = :white)
         percent_str = @sprintf("%4.1f", 100 * record.gctime / record.time)
@@ -188,8 +206,11 @@ function print_test_failed(record::TestRecord, wrkr, test, ctx::TestIOContext)
             lpad("($wrkr)", ctx.name_align - textwidth(test) + 1, " "), " |"
             , color = :red
         )
-        time_str = @sprintf("%7.2f", record.time)
+        time = record.time
+        time_str = @sprintf("%7.2f", time)
         printstyled(ctx.stderr, lpad(time_str, ctx.elapsed_align + 1, " "), " │", color = :red)
+        init_time_str = @sprintf("%7.2f", record.total_time - time)
+        printstyled(ctx.stdout, lpad(init_time_str, ctx.elapsed_align, " "), " │ ", color = :red)
 
         failed_str = "failed at $(now())\n"
         # 11 -> 3 from " | " 3x and 2 for each " " on either side
@@ -275,7 +296,7 @@ function Test.finish(ts::WorkerTestSet)
     return ts.wrapped_ts
 end
 
-function runtest(f, name, init_code)
+function runtest(f, name, init_code, start_time)
     function inner()
         # generate a temporary module to execute the tests in
         mod = @eval(Main, module $(gensym(name)) end)
@@ -304,7 +325,7 @@ function runtest(f, name, init_code)
 
         # process results
         rss = Sys.maxrss()
-        record = TestRecord(data..., rss)
+        record = TestRecord(data..., rss, time() - start_time)
 
         GC.gc(true)
         return record
@@ -539,6 +560,7 @@ Fields are
 
 * `jobs::Union{Some{Int}, Nothing}`: the number of jobs
 * `verbose::Union{Some{Nothing}, Nothing}`: whether verbose printing was enabled
+* `debug_timing::Union{Some{Nothing}, Nothing}`: whether debug timing printing was enabled
 * `quickfail::Union{Some{Nothing}, Nothing}`: whether quick fail was enabled
 * `list::Union{Some{Nothing}, Nothing}`: whether tests should be listed
 * `custom::Dict{String,Any}`: a dictionary of custom arguments
@@ -547,6 +569,7 @@ Fields are
 struct ParsedArgs
     jobs::Union{Some{Int}, Nothing}
     verbose::Union{Some{Nothing}, Nothing}
+    debug_timing::Union{Some{Nothing}, Nothing}
     quickfail::Union{Some{Nothing}, Nothing}
     list::Union{Some{Nothing}, Nothing}
 
@@ -604,7 +627,8 @@ function parse_args(args; custom::Array{String} = String[])
                --list             List all available tests.
                --verbose          Print more information during testing.
                --quickfail        Fail the entire run as soon as a single test errored.
-               --jobs=N           Launch `N` processes to perform tests."""
+               --jobs=N           Launch `N` processes to perform tests.
+               --debug-timing     Print testset initialization timings."""
 
         if !isempty(custom)
             usage *= "\n\nCustom arguments:"
@@ -619,6 +643,7 @@ function parse_args(args; custom::Array{String} = String[])
 
     jobs = extract_flag!(args, "--jobs"; typ = Int)
     verbose = extract_flag!(args, "--verbose")
+    debug_timing = extract_flag!(args, "--debug-timing")
     quickfail = extract_flag!(args, "--quickfail")
     list = extract_flag!(args, "--list")
 
@@ -633,7 +658,7 @@ function parse_args(args; custom::Array{String} = String[])
         error("Unknown test options `$(join(optlike_args, " "))` (try `--help` for usage instructions)")
     end
 
-    return ParsedArgs(jobs, verbose, quickfail, list, custom_args, args)
+    return ParsedArgs(jobs, verbose, debug_timing, quickfail, list, custom_args, args)
 end
 
 """
@@ -666,7 +691,8 @@ end
              init_worker_code = :(),
              test_worker = Returns(nothing),
              stdout = Base.stdout,
-             stderr = Base.stderr)
+             stderr = Base.stderr,
+             max_worker_rss = get_max_worker_rss())
     runtests(mod::Module, ARGS; ...)
 
 Run Julia tests in parallel across multiple worker processes.
@@ -689,6 +715,7 @@ Several keyword arguments are also supported:
 - `test_worker`: Optional function that takes a test name and `init_worker_code` if `init_worker_code` is defined and returns a specific worker.
   When returning `nothing`, the test will be assigned to any available default worker.
 - `stdout` and `stderr`: I/O streams to write to (default: `Base.stdout` and `Base.stderr`)
+- `max_worker_rss`: RSS threshold where a worker will be restarted once it is reached.
 
 ## Command Line Options
 
@@ -764,7 +791,7 @@ issues during long test runs. The memory limit is set based on system architectu
 function runtests(mod::Module, args::ParsedArgs;
                   testsuite::Dict{String,Expr} = find_tests(pwd()),
                   init_code = :(), init_worker_code = :(), test_worker = Returns(nothing),
-                  stdout = Base.stdout, stderr = Base.stderr)
+                  stdout = Base.stdout, stderr = Base.stderr, max_worker_rss = get_max_worker_rss())
     #
     # set-up
     #
@@ -831,7 +858,7 @@ function runtests(mod::Module, args::ParsedArgs;
         stderr.lock = print_lock
     end
 
-    io_ctx = test_IOContext(stdout, stderr, print_lock, name_align)
+    io_ctx = test_IOContext(stdout, stderr, print_lock, name_align, !isnothing(args.debug_timing))
     print_header(io_ctx, testgroupheader, workerheader)
 
     status_lines_visible = Ref(0)
@@ -1015,7 +1042,7 @@ function runtests(mod::Module, args::ParsedArgs;
                 result = try
                     Malt.remote_eval_wait(Main, wrkr.w, :(import ParallelTestRunner))
                     Malt.remote_call_fetch(invokelatest, wrkr.w, runtest,
-                                           testsuite[test], test, init_code)
+                                           testsuite[test], test, init_code, test_t0)
                 catch ex
                     if isa(ex, InterruptException)
                         # the worker got interrupted, signal other tasks to stop

test/runtests.jl

@@ -24,6 +24,17 @@ cd(@__DIR__)
     @test isfile(ParallelTestRunner.get_history_file(ParallelTestRunner))
 end
 
+@testset "debug timing" begin
+    io = IOBuffer()
+    io_color = IOContext(io, :color => true)
+    runtests(ParallelTestRunner, ["--debug-timing"]; stdout=io_color, stderr=io_color)
+    str = String(take!(io))
+
+    @test contains(str, "Init") # debug timing
+    @test contains(str, "time (s)") # debug timing
+    @test contains(str, "Time (s)") # normal timing
+end
+
 @testset "subdir use" begin
     d = @__DIR__
     testsuite = find_tests(d)