Add init_worker_code for code that runs once per worker

docs/src/advanced.md

@@ -95,6 +95,38 @@ end # hide
 
 The `init_code` is evaluated in each test's sandbox module, so all definitions are available to your test files.
 
+## Worker Initialization
+
+For most situations, `init_code` described above should be used. However, if the common code takes so long to import that it makes a notable difference to run before every testset, you can use the `init_worker_code` keyword argument in [`runtests`](@ref) to have it run only once at worker initialization. However, you will also have to import the directly-used functionality in your testset module using `init_code` due to the way ParallelTestRunner.jl creates a temporary module for each testset.
+
+The example below is trivial and `init_worker_code` would not be necessary if this were used in a package, but it shows how it should be used. A real use-case of this is for tests using the GPUArrays.jl test suite; including it takes about 3s, so that 3s running before every testset can add a significant amount of runtime to the various GPU backend testsuites as opposed to running once when the runner is initally created.
+
+```@example mypackage
+using ParallelTestRunner
+using MyPackage
+
+const init_worker_code = quote
+    # Common code that's slow to import
+    function complex_common_test_helper(x)
+        return x * 2
+    end
+end
+
+const init_code = quote
+    # ParallelTestRunner creates a temporary module to run
+    #  each testset. `init_code` runs in this temporary module,
+    #  but code from `init_worker_code` that will be directly
+    #  called in a testset must be explicitly included in the
+    #  module namespace.
+    using Main: complex_common_test_helper
+end
+
+cd(test_dir) do # hide
+runtests(MyPackage, ARGS; init_worker_code, init_code)
+end # hide
+```
+The `init_worker_code` is evaluated once per worker, so all definitions can be imported for use by the test module.
+
 ## Custom Workers
 
 For tests that require specific environment variables or Julia flags, you can use the `test_worker` keyword argument to [`runtests`](@ref) to assign tests to custom workers:
@@ -200,7 +232,7 @@ function jltest {
 
 1. **Keep tests isolated**: Each test file runs in its own module, so avoid relying on global state between tests.
 
-1. **Use `init_code` for common setup**: Instead of duplicating setup code in each test file, use `init_code` to share common initialization.
+1. **Use `init_code` for common setup**: Instead of duplicating setup code in each test file, use `init_code` to share common initialization. For long-running initialization, consider using `init_worker_code` so that it is run only once per worker creation instead of before each test.
 
 1. **Filter tests appropriately**: Use [`filter_tests!`](@ref) to respect user-specified test filters while allowing additional programmatic filtering.
 

src/ParallelTestRunner.jl

@@ -432,33 +432,36 @@ function test_exe(color::Bool=false)
 end
 
 """
-    addworkers(; env=Vector{Pair{String, String}}(), exename=nothing, exeflags=nothing, color::Bool=false)
+    addworkers(; env=Vector{Pair{String, String}}(), init_worker_code = :(), exename=nothing, exeflags=nothing, color::Bool=false)
 
 Add `X` worker processes.
 To add a single worker, use [`addworker`](@ref).
 
 ## Arguments
 - `env`: Vector of environment variable pairs to set for the worker process.
+- `init_worker_code`: Code use to initialize each worker. This is run only once per worker instead of once per test.
 - `exename`: Custom executable to use for the worker process.
 - `exeflags`: Custom flags to pass to the worker process.
 - `color`: Boolean flag to decide whether to start `julia` with `--color=yes` (if `true`) or `--color=no` (if `false`).
 """
 addworkers(X; kwargs...) = [addworker(; kwargs...) for _ in 1:X]
 
 """
-    addworker(; env=Vector{Pair{String, String}}(), exename=nothing, exeflags=nothing; color::Bool=false)
+    addworker(; env=Vector{Pair{String, String}}(), init_worker_code = :(), exename=nothing, exeflags=nothing; color::Bool=false)
 
-Add a single worker process. 
+Add a single worker process.
 To add multiple workers, use [`addworkers`](@ref).
 
 ## Arguments
 - `env`: Vector of environment variable pairs to set for the worker process.
+- `init_worker_code`: Code use to initialize each worker. This is run only once per worker instead of once per test.
 - `exename`: Custom executable to use for the worker process.
 - `exeflags`: Custom flags to pass to the worker process.
 - `color`: Boolean flag to decide whether to start `julia` with `--color=yes` (if `true`) or `--color=no` (if `false`).
 """
 function addworker(;
         env = Vector{Pair{String, String}}(),
+        init_worker_code = :(),
         exename = nothing,
         exeflags = nothing,
         color::Bool = false,
@@ -476,7 +479,11 @@ function addworker(;
     push!(env, "JULIA_NUM_THREADS" => "1")
     # Malt already sets OPENBLAS_NUM_THREADS to 1
     push!(env, "OPENBLAS_NUM_THREADS" => "1")
-    return PTRWorker(; exename, exeflags, env)
+    wrkr =  PTRWorker(; exename, exeflags, env)
+    if init_worker_code != :()
+        Malt.remote_eval_wait(Main, wrkr.w, init_worker_code)
+    end
+    return wrkr
 end
 
 """
@@ -656,6 +663,7 @@ end
     runtests(mod::Module, args::Union{ParsedArgs,Array{String}};
              testsuite::Dict{String,Expr}=find_tests(pwd()),
              init_code = :(),
+             init_worker_code = :(),
              test_worker = Returns(nothing),
              stdout = Base.stdout,
              stderr = Base.stderr)
@@ -677,6 +685,7 @@ Several keyword arguments are also supported:
   By default, automatically discovers all `.jl` files in the test directory and its subdirectories.
 - `init_code`: Code use to initialize each test's sandbox module (e.g., import auxiliary
   packages, define constants, etc).
+- `init_worker_code`: Code use to initialize each worker. This is run only once per worker instead of once per test.
 - `test_worker`: Optional function that takes a test name 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`)
@@ -754,7 +763,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 = :(), test_worker = Returns(nothing),
+                  init_code = :(), init_worker_code = :(), test_worker = Returns(nothing),
                   stdout = Base.stdout, stderr = Base.stderr)
     #
     # set-up
@@ -993,7 +1002,7 @@ function runtests(mod::Module, args::ParsedArgs;
                     wrkr = p
                 end
                 if wrkr === nothing || !Malt.isrunning(wrkr)
-                    wrkr = p = addworker(; io_ctx.color)
+                    wrkr = p = addworker(; init_worker_code, io_ctx.color)
                 end
 
                 # run the test

test/runtests.jl

@@ -71,6 +71,34 @@ end
     @test contains(str, "SUCCESS")
 end
 
+@testset "init worker code" begin
+    init_worker_code = quote
+        should_be_defined() = true
+
+        macro should_also_be_defined()
+            return :(true)
+        end
+    end
+    init_code = quote
+        using Test
+        using Main: should_be_defined, @should_also_be_defined
+    end
+
+    testsuite = Dict(
+        "custom" => quote
+            @test should_be_defined()
+            @test @should_also_be_defined()
+        end
+    )
+
+    io = IOBuffer()
+    runtests(ParallelTestRunner, ["--verbose"]; init_code, init_worker_code, testsuite, stdout=io, stderr=io)
+
+    str = String(take!(io))
+    @test contains(str, r"custom .+ started at")
+    @test contains(str, "SUCCESS")
+end
+
 @testset "custom worker" begin
     function test_worker(name)
         if name == "needs env var"