Add init_worker_code for code that runs once per worker (#85)

Author: christiangnrd

Project.toml

@@ -1,7 +1,7 @@
 name = "ParallelTestRunner"
 uuid = "d3525ed8-44d0-4b2c-a655-542cee43accc"
 authors = ["Valentin Churavy <v.churavy@gmail.com>"]
-version = "2.1.0"
+version = "2.2.0"
 
 [deps]
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"

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.
+    import ..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:
@@ -134,6 +166,11 @@ The `test_worker` function receives the test name and should return either:
 - A worker object (from [`addworker`](@ref)) for tests that need special configuration
 - `nothing` to use the default worker pool
 
+!!! note
+    If your test suite uses both a `test_worker` function and `init_worker_code` as described in a prior section,
+    `test_worker` must also take in `init_worker_code` as a second argument. You are responsible for passing it to
+    [`addworker`](@ref) if your `init_code` depends on any `init_worker_code` definitions.
+
 ## Custom Arguments
 
 If your package needs to accept its own command-line arguments in addition to `ParallelTestRunner`'s options, use [`parse_args`](@ref) with custom flags:
@@ -200,7 +237,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,7 +685,8 @@ 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).
-- `test_worker`: Optional function that takes a test name and returns a specific worker.
+- `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 `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`)
 
@@ -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
@@ -987,13 +996,18 @@ function runtests(mod::Module, args::ParsedArgs;
                     test, test_t0
                 end
 
-                # if a worker failed, spawn a new one
-                wrkr = test_worker(test)
-                if wrkr ===  nothing
+                # pass in init_worker_code to custom worker function if defined
+                wrkr = if init_worker_code == :()
+                    test_worker(test)
+                else
+                    test_worker(test, init_worker_code)
+                end
+                if wrkr === nothing
                     wrkr = p
                 end
+                # if a worker failed, spawn a new one
                 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
+        import ..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"
@@ -106,6 +134,52 @@ end
     @test contains(str, "SUCCESS")
 end
 
+@testset "custom worker with `init_worker_code`" begin
+    init_worker_code = quote
+        should_be_defined() = true
+    end
+    init_code = quote
+        using Test
+        import ..should_be_defined
+    end
+    function test_worker(name, init_worker_code)
+        if name == "needs env var"
+            return addworker(env = ["SPECIAL_ENV_VAR" => "42"]; init_worker_code)
+        elseif name == "threads/2"
+            return addworker(exeflags = ["--threads=2"]; init_worker_code)
+        end
+        return nothing
+    end
+    testsuite = Dict(
+        "needs env var" => quote
+            @test ENV["SPECIAL_ENV_VAR"] == "42"
+            @test should_be_defined()
+        end,
+        "doesn't need env var" => quote
+            @test !haskey(ENV, "SPECIAL_ENV_VAR")
+            @test should_be_defined()
+        end,
+        "threads/1" => quote
+            @test Base.Threads.nthreads() == 1
+            @test should_be_defined()
+        end,
+        "threads/2" => quote
+            @test Base.Threads.nthreads() == 2
+            @test should_be_defined()
+        end
+    )
+
+    io = IOBuffer()
+    runtests(ParallelTestRunner, ["--verbose"]; test_worker, init_code, init_worker_code, testsuite, stdout=io, stderr=io)
+
+    str = String(take!(io))
+    @test contains(str, r"needs env var .+ started at")
+    @test contains(str, r"doesn't need env var .+ started at")
+    @test contains(str, r"threads/1 .+ started at")
+    @test contains(str, r"threads/2 .+ started at")
+    @test contains(str, "SUCCESS")
+end
+
 @testset "failing test" begin
     testsuite = Dict(
         "failing test" => quote