Add a precompilation workload

This makes use of the scoped context and in-process worker support to have both
the master and worker running in the same process for precompilation.

Also lowered the sleep time for the launch Condition notification to 0.1s. This
helps since the Condition is edge-triggered and launching an in-process worker
usually takes less than 1s.

Author: JamesWrigley

Project.toml

@@ -3,6 +3,7 @@ uuid = "fab6aee4-877b-4bac-a744-3eca44acbb6f"
 version = "1.2.0"
 
 [deps]
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
@@ -19,6 +20,7 @@ Aqua = "0.8"
 Distributed = "1"
 LibSSH = "0.7"
 LinearAlgebra = "1"
+PrecompileTools = "1"
 Random = "1"
 Revise = "3.7.0"
 ScopedValues = "1.6.0"

docs/src/_changelog.md

@@ -14,6 +14,7 @@ This documents notable changes in DistributedNext.jl. The format is based on
   single struct ([#61]). This should not be a user-visible change, but of course
   it's possible that some things slipped through the cracks so please open an
   issue if you encounter any bugs.
+- A precompilation workload was added to improve TTFX ([#62]).
 
 ## [v1.2.0] - 2026-03-21
 

ext/ReviseExt.jl

@@ -3,8 +3,11 @@ module ReviseExt
 import DistributedNext
 import DistributedNext: myid, workers, remotecall
 
-import Revise
+using PrecompileTools: @recompile_invalidations
 
+# Sadly Revise causes quite a few invalidations. TODO: make DistributedNext more
+# resistant to invalidations.
+@recompile_invalidations import Revise
 
 struct DistributedNextWorker <: Revise.AbstractWorker
     id::Int

src/DistributedNext.jl

@@ -146,7 +146,6 @@ include("macros.jl")      # @spawn and friends
 include("workerpool.jl")
 include("pmap.jl")
 include("managers.jl")    # LocalManager and SSHManager
-include("precompile.jl")
 
 # Bundles all mutable global state for a distributed cluster into a single
 # object. The active context is accessed via the `CTX` ScopedValue, allowing
@@ -212,6 +211,14 @@ include("precompile.jl")
     stdlib_watcher_timer::Union{Timer, Nothing} = nothing
 end
 
+function ClusterContext(f::Base.Callable; kwargs...)
+    ctx = ClusterContext(; kwargs...)
+    ret = @with CTX => ctx f()
+    close(ctx)
+
+    return ret
+end
+
 function Base.close(ctx::ClusterContext)
     ctx.shutting_down[] = true
     if !isnothing(ctx.gc_msgs_task)
@@ -231,6 +238,8 @@ end
 
 const CTX = ScopedValue(ClusterContext())
 
+include("precompile.jl")
+
 function __init__()
     init_parallel()
 

src/cluster.jl

@@ -539,7 +539,7 @@ function addprocs_locked(manager::ClusterManager, params)
             if isempty(launched)
                 istaskdone(t_launch) && break
                 @async begin
-                    sleep(1)
+                    sleep(0.1)
                     notify(launch_ntfy)
                 end
                 wait(launch_ntfy)

src/precompile.jl

@@ -1,40 +1,10 @@
-precompile(Tuple{typeof(DistributedNext.remotecall),Function,Int,Module,Vararg{Any, 100}})
-precompile(Tuple{typeof(DistributedNext.procs)})
-precompile(Tuple{typeof(DistributedNext.finalize_ref), DistributedNext.Future})
-precompile(Tuple{typeof(DistributedNext.setup_launched_worker), DistributedNext.LocalManager, DistributedNext.WorkerConfig, Vector{Int}})
-precompile(Tuple{typeof(DistributedNext.process_tcp_streams), Sockets.TCPSocket, Sockets.TCPSocket, Bool})
-precompile(Tuple{typeof(Serialization.serialize), DistributedNext.ClusterSerializer{Sockets.TCPSocket}, Array{Any, 1}})
-precompile(Tuple{typeof(DistributedNext.parse_connection_info), String})
-precompile(Tuple{typeof(DistributedNext._run_callbacks_concurrently), String, Base.Dict{Any, Union{Function, Type}}, Int64, Array{Tuple{DistributedNext.LocalManager, Base.Dict{Symbol, Any}}, 1}})
-precompile(Tuple{typeof(DistributedNext._run_callbacks_concurrently), String, Base.Dict{Any, Union{Function, Type}}, Int64, Array{Int64, 1}})
-precompile(Tuple{typeof(DistributedNext.read_worker_host_port), Base.PipeEndpoint})
-precompile(Tuple{typeof(Base.put!), Base.Channel{Any}, Int64})
-precompile(Tuple{typeof(Base.put!), Base.Channel{Any}, WeakRef})
+using PrecompileTools: @compile_workload
 
-precompile(Tuple{typeof(DistributedNext.handle_msg), DistributedNext.IdentifySocketMsg, DistributedNext.MsgHeader, Sockets.TCPSocket, Sockets.TCPSocket, VersionNumber})
-precompile(Tuple{typeof(DistributedNext.handle_msg), DistributedNext.CallWaitMsg, DistributedNext.MsgHeader, Sockets.TCPSocket, Sockets.TCPSocket, VersionNumber})
-precompile(Tuple{typeof(DistributedNext.handle_msg), DistributedNext.CallMsg{:call}, DistributedNext.MsgHeader, Sockets.TCPSocket, Sockets.TCPSocket, VersionNumber})
-precompile(Tuple{typeof(DistributedNext.handle_msg), DistributedNext.CallMsg{:call_fetch}, DistributedNext.MsgHeader, Sockets.TCPSocket, Sockets.TCPSocket, VersionNumber})
-precompile(Tuple{typeof(DistributedNext.handle_msg), DistributedNext.ResultMsg, DistributedNext.MsgHeader, Sockets.TCPSocket, Sockets.TCPSocket, VersionNumber})
-precompile(Tuple{typeof(DistributedNext.handle_msg), DistributedNext.IdentifySocketAckMsg, DistributedNext.MsgHeader, Sockets.TCPSocket, Sockets.TCPSocket, VersionNumber})
-precompile(Tuple{typeof(DistributedNext.handle_msg), DistributedNext.RemoteDoMsg, DistributedNext.MsgHeader, Sockets.TCPSocket, Sockets.TCPSocket, VersionNumber})
-precompile(Tuple{typeof(DistributedNext.handle_msg), DistributedNext.JoinPGRPMsg, DistributedNext.MsgHeader, Sockets.TCPSocket, Sockets.TCPSocket, VersionNumber})
-precompile(Tuple{typeof(DistributedNext.handle_msg), DistributedNext.JoinCompleteMsg, DistributedNext.MsgHeader, Sockets.TCPSocket, Sockets.TCPSocket, VersionNumber})
-
-# These lines were obtained from `addprocs(1; exeflags=`--trace-compile=compilation.txt --trace-compile-timing`)`
-precompile(Tuple{typeof(DistributedNext.start_worker)})
-precompile(Tuple{typeof(DistributedNext.start_worker), Base.PipeEndpoint, String})
-precompile(Tuple{typeof(DistributedNext.set_valid_processes), Array{Int64, 1}})
-precompile(Tuple{typeof(DistributedNext.terminate_all_workers)})
-
-# This is disabled because it doesn't give much benefit
-# and the code in Distributed is poorly typed causing many invalidations
-# TODO: Maybe reenable now that Distributed is not in sysimage.
-#=
-    precompile_script *= """
-    using Distributed
-    addprocs(2)
-    pmap(x->iseven(x) ? 1 : 0, 1:4)
-    @distributed (+) for i = 1:100 Int(rand(Bool)) end
-    """
-=#
+@compile_workload begin
+    # Run the workload in a separate ClusterContext so the default one stays clean
+    ClusterContext() do
+        # Use an in-process worker to avoid spawning a real process during precompilation
+        pid = only(addprocs(LocalManager(1, true, true)))
+        rmprocs(pid)
+    end
+end