diff --git a/CHANGELOG.md b/CHANGELOG.md index 54a44e6f..f25fa1e9 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/Project.toml b/Project.toml index 1cd67639..9d825772 100644 --- a/Project.toml +++ b/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" diff --git a/src/Malt.jl b/src/Malt.jl index 03d966ef..b762cc67 100644 --- a/src/Malt.jl +++ b/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 diff --git a/src/QuartoNotebookRunner.jl b/src/QuartoNotebookRunner.jl index 2f7ac955..1ff8d260 100644 --- a/src/QuartoNotebookRunner.jl +++ b/src/QuartoNotebookRunner.jl @@ -32,6 +32,7 @@ import Preferences import ProgressLogging import REPL import Random +import Scratch import Sockets import SHA import TOML diff --git a/src/QuartoNotebookWorker/src/NotebookState.jl b/src/QuartoNotebookWorker/src/NotebookState.jl index 24cc8fa3..a0589275 100644 --- a/src/QuartoNotebookWorker/src/NotebookState.jl +++ b/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 diff --git a/src/QuartoNotebookWorker/src/PrecompileTools-post1.12/PrecompileTools.jl b/src/QuartoNotebookWorker/src/PrecompileTools-post1.12/PrecompileTools.jl new file mode 100644 index 00000000..6ea696bc --- /dev/null +++ b/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 diff --git a/src/QuartoNotebookWorker/src/PrecompileTools-post1.12/workloads.jl b/src/QuartoNotebookWorker/src/PrecompileTools-post1.12/workloads.jl new file mode 100644 index 00000000..047e3e0c --- /dev/null +++ b/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 diff --git a/src/QuartoNotebookWorker/src/PrecompileTools-pre1.12/PrecompileTools.jl b/src/QuartoNotebookWorker/src/PrecompileTools-pre1.12/PrecompileTools.jl new file mode 100644 index 00000000..ae901a16 --- /dev/null +++ b/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 diff --git a/src/QuartoNotebookWorker/src/PrecompileTools-pre1.12/workloads.jl b/src/QuartoNotebookWorker/src/PrecompileTools-pre1.12/workloads.jl new file mode 100644 index 00000000..19b304fa --- /dev/null +++ b/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 diff --git a/src/QuartoNotebookWorker/src/QuartoNotebookWorker.jl b/src/QuartoNotebookWorker/src/QuartoNotebookWorker.jl index 5868834b..f35594bf 100644 --- a/src/QuartoNotebookWorker/src/QuartoNotebookWorker.jl +++ b/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 diff --git a/src/QuartoNotebookWorker/src/precompile.jl b/src/QuartoNotebookWorker/src/precompile.jl new file mode 100644 index 00000000..d2f50847 --- /dev/null +++ b/src/QuartoNotebookWorker/src/precompile.jl @@ -0,0 +1,66 @@ +precompile(Tuple{typeof(QuartoNotebookWorker.Malt.main)}) +precompile( + Tuple{ + typeof(QuartoNotebookWorker.Malt._bson_deserialize), + QuartoNotebookWorker.Malt.Sockets.TCPSocket, + }, +) + +if VERSION >= v"1.9" + for type in [Int, Float64, String, Nothing, Missing] + precompile( + Tuple{ + typeof(Core.kwcall), + NamedTuple{(:inline,),Tuple{Bool}}, + typeof(QuartoNotebookWorker.render_mimetypes), + type, + Base.Dict{String,Any}, + }, + ) + end + precompile( + Tuple{ + typeof(Core.kwcall), + NamedTuple{ + (:file, :line, :cell_options), + Tuple{String,Int64,Base.Dict{String,Any}}, + }, + typeof(QuartoNotebookWorker.include_str), + Module, + String, + }, + ) +end + +precompile(Tuple{typeof(Base.Filesystem.mkpath),String}) +precompile(Tuple{typeof(QuartoNotebookWorker.refresh!),Base.Dict{String,Any}}) +precompile( + Tuple{ + typeof(QuartoNotebookWorker.refresh!), + Base.Dict{String,Any}, + Base.Dict{String,Any}, + }, +) + +module __PrecompilationModule end + +QuartoNotebookWorker.NotebookState.NotebookModuleForPrecompile[] = __PrecompilationModule + +PrecompileTools.@compile_workload begin + for code in ["1 + 1", "println(\"abc\")", "error()", "@info \"info text\" value=1"] + result = QuartoNotebookWorker.render( + code, + "some_file", + 1, + Dict{String,Any}("error" => "true"); + inline = false, + ) + io = IOBuffer() + bson = + QuartoNotebookWorker.Packages.BSON.bson(io, Dict{Symbol,Any}(:data => result)) + seekstart(io) + QuartoNotebookWorker.Packages.BSON.load(io)[:data] + end +end + +QuartoNotebookWorker.NotebookState.NotebookModuleForPrecompile[] = nothing diff --git a/src/server.jl b/src/server.jl index 3ed5c11f..b7e86060 100644 --- a/src/server.jl +++ b/src/server.jl @@ -28,11 +28,12 @@ mutable struct File timeout = _extract_timeout(merged_options) exe, _exeflags = _julia_exe(exeflags) + worker = cd( () -> Malt.Worker(; exe, exeflags = _exeflags, - env = vcat(env, quarto_env), + env = vcat(env, quarto_env, startup_env_variables()), ), dirname(path), ) @@ -66,6 +67,14 @@ mutable struct File end end +function startup_env_variables() + qnw_env_dir = Scratch.@get_scratch!("qnw-env-$(hash(Malt.worker_package))") + return [ + "QUARTONOTEBOOKWORKER_ENV_DIR=$qnw_env_dir", + "QUARTONOTEBOOKWORKER_PACKAGE_DIR=$(Malt.worker_package)", + ] +end + struct SourceRange file::Union{String,Nothing} lines::UnitRange{Int} @@ -237,7 +246,11 @@ function refresh!(file::File, options::Dict) Malt.stop(file.worker) exe, _exeflags = _julia_exe(exeflags) file.worker = cd( - () -> Malt.Worker(; exe, exeflags = _exeflags, env = vcat(env, quarto_env)), + () -> Malt.Worker(; + exe, + exeflags = _exeflags, + env = vcat(env, quarto_env, startup_env_variables()), + ), dirname(file.path), ) file.exe = exe @@ -1575,7 +1588,6 @@ function run!( # block until a decision is reached decision = take!(file.run_decision_channel) - # :forceclose might have been set from another task if decision === :forceclose close!(server, file.path) # this is in the same task, so reentrant lock allows access diff --git a/src/startup.jl b/src/startup.jl index 27b50f21..f00fbfd5 100644 --- a/src/startup.jl +++ b/src/startup.jl @@ -13,11 +13,15 @@ function capture(func) try return func() catch err + bt = catch_backtrace() errors_log_file = joinpath(ENV["MALT_WORKER_TEMP_DIR"], "errors.log") - open(errors_log_file, "w") do io + io = open(errors_log_file, "w") + try showerror(io, err) - Base.show_backtrace(io, catch_backtrace()) + Base.inferencebarrier(Base.show_backtrace)(io, bt) flush(io) + finally + close(io) end exit() end @@ -45,35 +49,38 @@ pushfirst!(LOAD_PATH, "@stdlib") # and if a user was to perform `Pkg` operations they may affect that # environment. Instead we provide a temporary sandbox environment that gets # discarded when the notebook process exits. + let temp = mktempdir() sandbox = joinpath(temp, "QuartoSandbox") - mkpath(sandbox) + mkdir(sandbox) # The empty project file is key to making this the active environment if # noting else is available if the rest of the `LOAD_PATH`. touch(joinpath(sandbox, "Project.toml")) push!(LOAD_PATH, sandbox) +end + +# Step 2b: +# +# We also need to ensure that the `QuartoNotebookWorker` package is +# available on the `LOAD_PATH`. - # Step 2b: - # - # We also need to ensure that the `QuartoNotebookWorker` package is - # available on the `LOAD_PATH`. This is done by creating another - # environment alongside the sandbox environment. We `Pkg.develop` the - # "local" `QuartoNotebookWorker` package into this environment. `Pkg` - # operations are logged to the `pkg.log` file that the server process can - # read to provide feedback to the user if needed. - # +qnw_env_dir = ENV["QUARTONOTEBOOKWORKER_ENV_DIR"] +qnw_package_dir = ENV["QUARTONOTEBOOKWORKER_PACKAGE_DIR"] + +env_path = joinpath(qnw_env_dir, string(VERSION)) +env_proj = joinpath(env_path, "Project.toml") +env_mani = joinpath(env_path, "Manifest.toml") + +if !(isfile(env_proj) && isfile(env_mani)) # `Pkg` is loaded outside of this closure otherwise the methods required do # not exist in a new enough world age to be callable. Pkg = Base.require(Base.PkgId(Base.UUID("44cfe95a-1eb2-52ea-b672-e2afdf69b78f"), "Pkg")) capture() do - worker = joinpath(temp, "QuartoNotebookWorker") - mkpath(worker) - push!(LOAD_PATH, worker) open(joinpath(ENV["MALT_WORKER_TEMP_DIR"], "pkg.log"), "w") do io ap = Base.active_project() try - Pkg.activate(worker; io) - Pkg.develop(; path = ENV["QUARTONOTEBOOKWORKER_PACKAGE"], io) + Pkg.activate(env_path; io) + Pkg.develop(; path = qnw_package_dir, io) finally # Ensure that we switch the active project back afterwards. Pkg.activate(ap; io) @@ -83,6 +90,8 @@ let temp = mktempdir() end end +push!(LOAD_PATH, env_path) + # Step 3: # # The parent process needs some additional metadata about this `julia` process to @@ -94,9 +103,10 @@ end # instead manually write the strings so that this can happen prior to any # stdlib loading, which could trigger errors that we would then want this # metadata to be able to properly inform the user about. -capture() do +function save_metadata() metadata_toml_file = joinpath(ENV["MALT_WORKER_TEMP_DIR"], "metadata.toml") - open(metadata_toml_file, "w") do io + io = open(metadata_toml_file, "w") + try project_toml_file = Base.active_project() if !isnothing(project_toml_file) && isfile(project_toml_file) println(io, "project = $(repr(project_toml_file))") @@ -107,8 +117,11 @@ capture() do end println(io, "julia_version = $(repr(string(VERSION)))") flush(io) + finally + close(io) end end +capture(save_metadata) # Step: 4 # @@ -119,18 +132,19 @@ end # should do a manual `import Revise` in their notebook if they need `Revise` # support. const QUARTO_ENABLE_REVISE = get(ENV, "QUARTO_ENABLE_REVISE", "false") == "true" -capture() do +function require_revise() if QUARTO_ENABLE_REVISE pkgid = Base.PkgId(Base.UUID("295af30f-e4ad-537b-8983-00126c2a3abe"), "Revise") Base.require(pkgid) end end +capture(require_revise) # Step 5: # # Now load in the worker package. This may trigger package precompilation on # first load, hence it is run under a `capture` should it fail to run. -const QuartoNotebookWorker = capture() do +function require_qnw() Base.require( Base.PkgId( Base.UUID("38328d9c-a911-4051-bc06-3f7f556ffeda"), @@ -138,6 +152,7 @@ const QuartoNotebookWorker = capture() do ), ) end +const QuartoNotebookWorker = capture(require_qnw) # Step 6: #