WIP JET integration idea

[Drvi]

Very WIP and ugly, but consider the following tests with a very contrived bug in the foo function:

# file: jet_tests.jl
@testsetup module Constants
    const DEF = 1
    export DEF
end

@testitem "Test with JET on" jet=:typo setup=[Constants] begin
    using .Threads: @spawn

    @test DEF+1 == 2


    function foo()
        r = Ref(2)
        @spawn ($(r)[] += UNDEF)
        sleep(0.5)
        return r[]
    end
    @test foo() == 2
end

@testitem "Test with JET off" setup=[Constants] begin
    using .Threads: @spawn

    @test DEF+1 == 2


    function foo()
        r = Ref(2)
        @spawn ($(r)[] += UNDEF)
        sleep(0.5)
        return r[]
    end
    @test foo() == 2
end

Without JET integration, the test item would happily report 2 passes. With JET integration we do get a failed test about the typo in foo:

12:57:54 | START (1/2) test item "Test with JET on" at jet_tests.jl:6
12:57:54 | DONE  (1/2) test item "Test with JET on" 0.6 secs (3.5% compile), 1.12 M allocs (65.187 MB)

Error in testset "JET :typo mode" on worker 12811:
...

12:57:55 | START (2/2) test item "Test with JET off" at jet_tests.jl:21
12:57:55 | DONE  (2/2) test item "Test with JET off" 0.5 secs (1.6% compile), 16.59 K allocs (1.197 MB)
Test Summary:           | Pass  Fail  Total  Time
ReTestItems             |    4     1      5  1.6s
  .                     |    4     1      5      
    jet_tests.jl        |    4     1      5      
      Test with JET on  |    2     1      3  0.6s
        JET :typo mode  |          1      1  0.1s
      Test with JET off |    2            2  0.5s

When I tried using JET directly, the analysis failed to discover the issue and evaluated the code... I'm not sure how to invoke JET to avoid that issue.

JET.report_file("jet_tests.jl", context=ReTestItems, concretization_patterns = [:(x_)], mode=:typo) # evaluates the code and fails to find issues

TODO: tests
TODO: Should we automagically load JET on the workers iff it is available in the test environment? Should there be a toggle for that? EDIT: Yes 🪄

[nickrobinson251]

integration with JET could be really nice. We'll need to work out how the dependency is going to work...

i wonder if we can/should use package extensions for this? I've not used them yet myself, so not sure. https://pkgdocs.julialang.org/v1/creating-packages/#Conditional-loading-of-code-in-packages-(Extensions)

I'd be very hesitant to depend on JET directly/unconditionally: ReTestItems.jl having dependencies means we require users' code be compatible with those dependencies, specifically the version of those dependencies which ReTestItems.jl is using, and this compatibility issue is amplified by how unstable the dependency is... and JET is very unstable (we've seen that its often best to pin to a specific JET patch version, and the JET version is very tied to the Julia version). I think the VS Code extension vendors their dependencies to avoid this issue (and it's also why we use our own "Worker" code rather than depend on ConcurrentUtilities.jl)

Currently we only have 1 (non-stdlib) dependency, which is the very stable LoggingExtras.jl v1 and honestly i think i'd rather drop even that (what do we use it for? is it just @debugv? if so, it'd be trivial for us to roll our own version of that).

[Drvi]

@nickrobinson251 I agree it would be better if we didn't have to depend on JET unconditionally. Maybe this will finally force me to try out writing package extensions (but I worry this might not be the best use case for them -- IIUC, extensions shine when you simply add new methods for dispatch on public API, but adding a new keyword to a macro feels slightly different... will have to try it out I guess)

[nickrobinson251]

IIUC, extensions shine when you simply add new methods for dispatch on public API, but adding a new keyword to a macro feels slightly different... will have to try it out I guess

Perhaps we can have the macro always accept the jet=... keyword, and always call the jet_test function, but then by default we just define a jet_test(::Any) function e.g.

function jet_test(mode::Any)
    if !isnothing(mode)
        @testset "JET $(repr(mode)) mode" begin
            throw("JET tests are not supported unless the JET package is loaded. Got jet=$repr(mode)")
        end
    end
end 

then have the package extension define a more specific method that works actually runs tests if JET is loaded, something like

function jet_test(mode::Symbol)
    @testset "JET $(repr(mode)) mode" begin
            JET.test(...; mode)
    end
end

[Drvi]

This is to illustrate that for package extensions to kick in, the weak dependency must be loaded. Here is how the two invocations look:

The first one doesn't run with JET, because on the workers, JET is not loaded.

[Drvi]

@nickrobinson251 What do you think about recording a broken test instead of an error? We'll hit this method whenever JET is not loaded in the session evaluating the test items, so e.g. if the user wants a faster run locally without JET analysis, they'd get just a bunch of broken tests instead of noisy errors.

[nickrobinson251]

if running via runtests, that'll activate the test environment which will then presumably load JET (if that test suite has Jet tests), right? How would you run a test suite that has JET as a test dependency without causing JET to get loaded?

i wonder if we add a runtests keyword skip_jet::Bool=false, which when true causes all JET tests to be skipped (i.e. any testitem with jet=... will record a skipped test, regardless of whether or not JET is loaded)

[Drvi]

Activating the environment doesn't load its packages, one needs to import JET or using JET on the worker before we try to use the extension, so even if the suite has JET tests, if those are not eval'd first, we won't have the extension available for all test items. This is why I wonder -- do we want the user to manually preload JET via init_worker_expr or do we want to preload the package for them?

The runtest-level keyword is a good idea 👍 and I think that apart skip_jet we might also want default_jet::Symbol=:none (with corresponding ENV variables), so that the user doesn't have to edit a bunch of test items to run JET on all of them.

[nickrobinson251]

hmmm, interesting. what if, when a testitem sets the jet keyword, the runtestitem code adds using JET to the test code (just like how it automatically adds using Test and using $PackageBeingTested)? Would that work.

we might also want default_jet::Symbol=:none (with corresponding ENV variables),

Agreed. I'd be inclined to just name this jet::Symbol... maybe even we allow jet=:skip as the way to skip JET tests?

[Drvi]

what if, when a testitem sets the jet keyword, the runtestitem code adds using JET to the test code (just like how it automatically adds using Test and using $PackageBeingTested)? Would that work.

Precisely! Ok, thanks, I'll look into it 👍

jet::Symbol... maybe even we allow jet=:skip

💯

[Drvi]

Another cute design problem: JET failures alone shouldn't trigger retries (since there is no chance they'll recover). I'm experimenting with adding an "is_retriable" field on the TestItemResult, but the API for / our usage of Test.AbstractTestSet is quite stateful and exception-centric, making it a bit challenging.