PrecompileTools doesn't run __init__() so some functionality may not work during package compilation?

[NHDaly]

If I have a package's workfile file that runs functionality in the package, from what I understand @compile_workload doesn't run the package's __init__() function, so the package's module will still be uninitialized during the workload?

For example:

module TestPackage

using PrecompileTools

const PORT = Ref(0)

function __init__()
    PORT[] = rand(2222:8888)
end

function test()
    # start up a server and run the tests...
    @assert PORT[] != 0
    # ...
end

@compile_workload begin
    test()
end

end # module TestPackage

Then:

julia> using TestPackage
[ Info: Precompiling TestPackage [13cf7a4f-c364-47a8-9398-f77e91fc5657]
ERROR: LoadError: AssertionError: PORT[] != 0
Stacktrace:
 [1] test
   @ ~/tmp/TestPackage/src/TestPackage.jl:13 [inlined]
 [2] macro expansion
   @ ~/tmp/TestPackage/src/TestPackage.jl:18 [inlined]
 [3] top-level scope
   @ ~/.julia/packages/PrecompileTools/kmH5L/src/workloads.jl:78
 [4] include
   @ ./Base.jl:457 [inlined]
 [5] include_package_for_output(pkg::Base.PkgId, input::String, depot_path::Vector{String}, dl_load_path::Vector{String}, load_path::Vector{String}, concrete_deps::Vector{Pair{Base.PkgId, UInt128}}, source::Nothing)
   @ Base ./loading.jl:2049
 [6] top-level scope
   @ stdin:3
in expression starting at /Users/nathandaly/tmp/TestPackage/src/TestPackage.jl:1

Does that make sense?
What do other packages usually do here? Are we meant to call __init__() manually? I didn't see anything about this in the docs.
Thanks :)

[NHDaly]

And if there are nested modules in the package, how do we ensure that we can find an init all of those? Is the recommendation to do all of this in like a wrapper package? Like move the actual package to a subpackage and then have the "main package" be nothing other than a wrapper around the subpackage + the @compile_workload statement?

EDIT: I don't think this would work either

[NHDaly]

😁 this is really dumb, but this seems to at least unblock me for now lol

recursively_init_modules(m::Module) = _recursively_init_modules(m, m)
function _recursively_init_modules(root::Module, m::Module)
    # iterate all recursively reachable modules from m and call __init__() on them:
    for name in names(m, all=true)
        # @show name
        if name isa Symbol && name != :Base && name != :Core && name != nameof(m)
            m2 = try Core.eval(m, name) catch ; continue ; end
            # @show m2
            if m2 isa Module && fullname(m2)[1] == nameof(root)
                _recursively_init_modules(root, m2)
            end
        end
    end
    @info m
    if isdefined(m, :__init__)
        @info "init $m"
        @eval m __init__()
    end
end
@setup_workload begin
    recursively_init_modules(TestPackage)
    @info PORT[]
    @compile_workload begin
        test()
    end
end

EDIT: This is the slightly more robust terrible really dumb thing that I'm trying now

[vchuravy]

Precompilation doesn't run __init__ of the current module at all. This is not specific to PrecompileTools.
You need to do the setup twice, once at top level and then in init for resetting it

[NHDaly]

@vchuravy and I talked about this on slack, and it seems like a nice change to make to julia would be to move the compile_workload step to after the Module is closed, so that we can run after the module is initialized, and we don't need to worry about modifying state in the module.

Can we do something where packages can register a compilation callback / function that runs after we've finalized the module for serialization? Something like this?:

module MyPackage
...
end

Base.precompilation(MyPackage) do
    @setup_workload begin
        # ...
        @compile_workload begin
            # ...
        end
    end
end

and then Base would deepcopy the MyPackage module, run the callback, and then only keep the new compilations but serialize the original copy of the module? It's a bit convoluted, but it would ensure that precompilation snoop scripts don't mutate the state of the module, and that they can run after the init() is called, so the expected state is present.

(Or maybe the PrecompileTools macros would expand to setting this callback or something.)

Thoughts?