Improvements to worker startup latency

CHANGELOG.md

@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+- Improved precompilation of QuartoNotebookWorker and avoided Pkg startup costs by using a scratch space to cache QNW project files [#341].
+
 ## [v0.17.4] - 2025-09-26
 
 ### Added
@@ -473,3 +475,4 @@ caching is enabled. Delete this folder to clear the cache. [#259]
 [#306]: https://github.com/PumasAI/QuartoNotebookRunner.jl/issues/306
 [#335]: https://github.com/PumasAI/QuartoNotebookRunner.jl/issues/335
 [#339]: https://github.com/PumasAI/QuartoNotebookRunner.jl/issues/339
+[#341]: https://github.com/PumasAI/QuartoNotebookRunner.jl/issues/341

Project.toml

@@ -22,6 +22,7 @@ REPL = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RelocatableFolders = "05181044-ff0b-4ac5-8273-598c1e38db00"
 SHA = "ea8e919c-243c-51af-8825-aaa63cd721ce"
+Scratch = "6c6a2e73-6563-6170-7368-637461726353"
 Sockets = "6462fe0b-24de-5631-8697-dd941f90decc"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6"
@@ -45,6 +46,7 @@ REPL = "1.6"
 Random = "1.6"
 RelocatableFolders = "1"
 SHA = "0.7, 1.6"
+Scratch = "1"
 Sockets = "1.6"
 TOML = "1"
 YAML = "0.4"

src/Malt.jl

@@ -388,10 +388,7 @@ const empty_file = RelocatableFolders.@path joinpath(@__DIR__, "empty.jl")
 const worker_package = RelocatableFolders.@path joinpath(@__DIR__, "QuartoNotebookWorker")
 
 function _get_worker_cmd(; exe, env, exeflags, file = String(startup_file))
-    defaults = Dict(
-        "OPENBLAS_NUM_THREADS" => "1",
-        "QUARTONOTEBOOKWORKER_PACKAGE" => String(worker_package),
-    )
+    defaults = Dict("OPENBLAS_NUM_THREADS" => "1")
     env = vcat(Base.byteenv(defaults), Base.byteenv(env))
     return addenv(`$exe --startup-file=no $exeflags $file`, env)
 end

src/QuartoNotebookRunner.jl

@@ -32,6 +32,7 @@ import Preferences
 import ProgressLogging
 import REPL
 import Random
+import Scratch
 import Sockets
 import SHA
 import TOML

src/QuartoNotebookWorker/src/NotebookState.jl

@@ -47,7 +47,11 @@ function define_notebook_module!(root = Main)
     return mod
 end
 
+const NotebookModuleForPrecompile = Base.RefValue{Union{Nothing,Module}}(nothing)
+
 # `getfield` ends up throwing a segfault here, `getproperty` works fine though.
-notebook_module() = Base.getproperty(Main, :Notebook)::Module
+notebook_module() =
+    NotebookModuleForPrecompile[] === nothing ? Base.getproperty(Main, :Notebook)::Module :
+    NotebookModuleForPrecompile[]
 
 end

src/QuartoNotebookWorker/src/PrecompileTools-post1.12/PrecompileTools.jl

@@ -0,0 +1,15 @@
+module PrecompileTools
+
+export @setup_workload, @compile_workload, @recompile_invalidations
+
+const verbose = Ref(false)    # if true, prints all the precompiles
+
+function precompile_mi(mi::Core.MethodInstance)
+    precompile(mi)
+    verbose[] && println(mi)
+    return
+end
+
+include("workloads.jl")
+
+end

src/QuartoNotebookWorker/src/PrecompileTools-post1.12/workloads.jl

@@ -0,0 +1,132 @@
+const newly_inferred = Core.CodeInstance[]   # only used to support verbose[]
+
+function workload_enabled(mod::Module)
+    # try
+    #     if load_preference(@__MODULE__, "precompile_workloads", true)
+    #         return load_preference(mod, "precompile_workload", true)
+    #     else
+    #         return false
+    #     end
+    # catch
+    #     true
+    # end
+    true # not using Preferences
+end
+
+@noinline is_generating_output() = ccall(:jl_generating_output, Cint, ()) == 1
+
+macro latestworld_if_toplevel()
+    Expr(Symbol("latestworld-if-toplevel"))
+end
+
+function tag_newly_inferred_enable()
+    ccall(:jl_tag_newly_inferred_enable, Cvoid, ())
+    if !Base.generating_output()   # for verbose[]
+        ccall(:jl_set_newly_inferred, Cvoid, (Any,), newly_inferred)
+    end
+end
+function tag_newly_inferred_disable()
+    ccall(:jl_tag_newly_inferred_disable, Cvoid, ())
+    if !Base.generating_output()   # for verbose[]
+        ccall(:jl_set_newly_inferred, Cvoid, (Any,), nothing)
+    end
+    if verbose[]
+        for ci in newly_inferred
+            println(ci.def)
+        end
+    end
+    return nothing
+end
+
+"""
+    @compile_workload f(args...)
+
+`precompile` (and save in the `compile_workload` file) any method-calls that occur inside the expression. All calls (direct or indirect) inside a
+`@compile_workload` block will be cached.
+
+`@compile_workload` has three key features:
+
+1. code inside runs only when the package is being precompiled (i.e., a `*.ji`
+   precompile `compile_workload` file is being written)
+2. the interpreter is disabled, ensuring your calls will be compiled
+3. both direct and indirect callees will be precompiled, even for methods defined in other packages
+   and even for runtime-dispatched callees (requires Julia 1.8 and above).
+
+!!! note
+    For comprehensive precompilation, ensure the first usage of a given method/argument-type combination
+    occurs inside `@compile_workload`.
+
+    In detail: runtime-dispatched callees are captured only when type-inference is executed, and they
+    are inferred only on first usage. Inferrable calls that trace back to a method defined in your package,
+    and their *inferrable* callees, will be precompiled regardless of "ownership" of the callees
+    (Julia 1.8 and higher).
+
+    Consequently, this recommendation matters only for:
+
+        - direct calls to methods defined in Base or other packages OR
+        - indirect runtime-dispatched calls to such methods.
+"""
+macro compile_workload(ex::Expr)
+    local iscompiling = :(
+        $PrecompileTools.is_generating_output() &&
+        $PrecompileTools.workload_enabled(@__MODULE__)
+    )
+    ex = quote
+        begin
+            $PrecompileTools.@latestworld_if_toplevel  # block inference from proceeding beyond this point (xref https://github.com/JuliaLang/julia/issues/57957)
+            $(esc(ex))
+        end
+    end
+    ex = quote
+        $PrecompileTools.tag_newly_inferred_enable()
+        try
+            $ex
+        finally
+            $PrecompileTools.tag_newly_inferred_disable()
+        end
+    end
+    return quote
+        if $iscompiling || $PrecompileTools.verbose[]
+            $ex
+        end
+    end
+end
+
+"""
+    @setup_workload begin
+        vars = ...
+        ⋮
+    end
+
+Run the code block only during package precompilation. `@setup_workload` is often used in combination
+with [`@compile_workload`](@ref), for example:
+
+    @setup_workload begin
+        vars = ...
+        @compile_workload begin
+            y = f(vars...)
+            g(y)
+            ⋮
+        end
+    end
+
+`@setup_workload` does not force compilation (though it may happen anyway) nor intentionally capture
+runtime dispatches (though they will be precompiled anyway if the runtime-callee is for a method belonging
+to your package).
+"""
+macro setup_workload(ex::Expr)
+    local iscompiling = :((
+        ccall(:jl_generating_output, Cint, ()) == 1 &&
+        $PrecompileTools.workload_enabled(@__MODULE__)
+    ))
+    # Ideally we'd like a `let` around this to prevent namespace pollution, but that seem to
+    # trigger inference & codegen in undesirable ways (see #16).
+    return quote
+        if $iscompiling || $PrecompileTools.verbose[]
+            let
+                $PrecompileTools.@latestworld_if_toplevel  # block inference from proceeding beyond this point (xref https://github.com/JuliaLang/julia/issues/57957)
+                $(esc(ex))
+            end
+        end
+    end
+end

src/QuartoNotebookWorker/src/PrecompileTools-pre1.12/PrecompileTools.jl

@@ -0,0 +1,19 @@
+module PrecompileTools
+
+export @setup_workload, @compile_workload, @recompile_invalidations
+
+const verbose = Ref(false)    # if true, prints all the precompiles
+const have_inference_tracking = isdefined(Core.Compiler, :__set_measure_typeinf)
+const have_force_compile =
+    isdefined(Base, :Experimental) &&
+    isdefined(Base.Experimental, Symbol("#@force_compile"))
+
+function precompile_mi(mi)
+    precompile(mi.specTypes) # TODO: Julia should allow one to pass `mi` directly (would handle `invoke` properly)
+    verbose[] && println(mi)
+    return
+end
+
+include("workloads.jl")
+
+end

src/QuartoNotebookWorker/src/PrecompileTools-pre1.12/workloads.jl

@@ -0,0 +1,151 @@
+
+function workload_enabled(mod::Module)
+    # try
+    #     if load_preference(@__MODULE__, "precompile_workloads", true)
+    #         return load_preference(mod, "precompile_workload", true)
+    #     else
+    #         return false
+    #     end
+    # catch
+    #     true
+    # end
+    return true
+end
+
+"""
+    check_edges(node)
+
+Recursively ensure that all callees of `node` are precompiled. This is (rarely) necessary
+because sometimes there is no backedge from callee to caller (xref https://github.com/JuliaLang/julia/issues/49617),
+and `staticdata.c` relies on the backedge to trace back to a MethodInstance that is tagged `mi.precompiled`.
+"""
+function check_edges(node)
+    parentmi = node.mi_info.mi
+    for child in node.children
+        childmi = child.mi_info.mi
+        if !(isdefined(childmi, :backedges) && parentmi ∈ childmi.backedges)
+            precompile_mi(childmi)
+        end
+        check_edges(child)
+    end
+end
+
+function precompile_roots(roots)
+    @assert have_inference_tracking
+    for child in roots
+        precompile_mi(child.mi_info.mi)
+        check_edges(child)
+    end
+end
+
+"""
+    @compile_workload f(args...)
+
+`precompile` (and save in the compile_workload file) any method-calls that occur inside the expression. All calls (direct or indirect) inside a
+`@compile_workload` block will be cached.
+
+`@compile_workload` has three key features:
+
+1. code inside runs only when the package is being precompiled (i.e., a `*.ji`
+   precompile compile_workload file is being written)
+2. the interpreter is disabled, ensuring your calls will be compiled
+3. both direct and indirect callees will be precompiled, even for methods defined in other packages
+   and even for runtime-dispatched callees (requires Julia 1.8 and above).
+
+!!! note
+    For comprehensive precompilation, ensure the first usage of a given method/argument-type combination
+    occurs inside `@compile_workload`.
+
+    In detail: runtime-dispatched callees are captured only when type-inference is executed, and they
+    are inferred only on first usage. Inferrable calls that trace back to a method defined in your package,
+    and their *inferrable* callees, will be precompiled regardless of "ownership" of the callees
+    (Julia 1.8 and higher).
+
+    Consequently, this recommendation matters only for:
+
+        - direct calls to methods defined in Base or other packages OR
+        - indirect runtime-dispatched calls to such methods.
+"""
+macro compile_workload(ex::Expr)
+    local iscompiling = if Base.VERSION < v"1.6"
+        :(ccall(:jl_generating_output, Cint, ()) == 1)
+    else
+        :((
+            ccall(:jl_generating_output, Cint, ()) == 1 &&
+            $PrecompileTools.workload_enabled(@__MODULE__)
+        ))
+    end
+    if have_force_compile
+        ex = quote
+            begin
+                Base.Experimental.@force_compile
+                $(esc(ex))
+            end
+        end
+    else
+        # Use the hack on earlier Julia versions that blocks the interpreter
+        ex = quote
+            while false
+            end
+            $(esc(ex))
+        end
+    end
+    if have_inference_tracking
+        ex = quote
+            Core.Compiler.Timings.reset_timings()
+            Core.Compiler.__set_measure_typeinf(true)
+            try
+                $ex
+            finally
+                Core.Compiler.__set_measure_typeinf(false)
+                Core.Compiler.Timings.close_current_timer()
+            end
+            $PrecompileTools.precompile_roots(Core.Compiler.Timings._timings[1].children)
+        end
+    end
+    return quote
+        if $iscompiling || $PrecompileTools.verbose[]
+            $ex
+        end
+    end
+end
+
+"""
+    @setup_workload begin
+        vars = ...
+        ⋮
+    end
+
+Run the code block only during package precompilation. `@setup_workload` is often used in combination
+with [`@compile_workload`](@ref), for example:
+
+    @setup_workload begin
+        vars = ...
+        @compile_workload begin
+            y = f(vars...)
+            g(y)
+            ⋮
+        end
+    end
+
+`@setup_workload` does not force compilation (though it may happen anyway) nor intentionally capture
+runtime dispatches (though they will be precompiled anyway if the runtime-callee is for a method belonging
+to your package).
+"""
+macro setup_workload(ex::Expr)
+    local iscompiling = if Base.VERSION < v"1.6"
+        :(ccall(:jl_generating_output, Cint, ()) == 1)
+    else
+        :((
+            ccall(:jl_generating_output, Cint, ()) == 1 &&
+            $PrecompileTools.workload_enabled(@__MODULE__)
+        ))
+    end
+    # Ideally we'd like a `let` around this to prevent namespace pollution, but that seem to
+    # trigger inference & codegen in undesirable ways (see #16).
+    return quote
+        if $iscompiling || $PrecompileTools.verbose[]
+            $(esc(ex))
+        end
+    end
+end

src/QuartoNotebookWorker/src/QuartoNotebookWorker.jl

@@ -101,4 +101,11 @@ include("notebook_metadata.jl")
 include("manifest_validation.jl")
 include("python.jl")
 
+if VERSION >= v"1.12-rc1"
+    include("PrecompileTools-post1.12/PrecompileTools.jl")
+else
+    include("PrecompileTools-pre1.12/PrecompileTools.jl")
+end
+include("precompile.jl")
+
 end