Merge pull request #1210 from JuliaLang/precompile

Improve precompilation

Author: JamesWrigley

docs/src/_changelog.md

@@ -28,6 +28,8 @@ Changelog](https://keepachangelog.com).
   needing to rebuild IJulia to update the kernel after every patch release of
   Julia, but it does mean that IJulia will only create kernels for each Julia
   minor release instead of each patch release.
+- Extended the precompilation workload to cover more calls, and added a
+  workaround to minimize TTFX when Revise is used ([#1210]).
 
 ## [v1.31.1] - 2025-10-20
 

src/IJulia.jl

@@ -140,6 +140,14 @@ REPL.REPLDisplay(repl::MiniREPL) = repl.display
     stop_event::Base.Event = Base.Event()
     waitloop_task::RefValue{Task} = Ref{Task}()
 
+    # This is a bit strange, but IJulia invalidates a lot of Revise because of
+    # its definition of Base.lock(::IJuliaStdio). This contributes significantly
+    # to TTFX because Revise.revise() is automatically set as a preexecute hook,
+    # so if run_kernel() enables the preexecute hook it also spawns a task to
+    # precompile Revise.revise() asynchronously. We store the task here to
+    # prevent any lingering tasks upon shutdown that could cause hangs.
+    revise_precompile_task::RefValue{Task} = Ref{Task}()
+
     requests_task::RefValue{Task} = Ref{Task}()
     watch_stdout_task::RefValue{Task} = Ref{Task}()
     watch_stderr_task::RefValue{Task} = Ref{Task}()
@@ -199,6 +207,10 @@ function start_shutdown(kernel::Kernel)
 end
 
 function Base.close(kernel::Kernel)
+    if isassigned(kernel.revise_precompile_task)
+        wait(kernel.revise_precompile_task[])
+    end
+
     # Reset the IO streams first so that any later errors get printed
     if kernel.capture_stdout
         redirect_stdout(orig_stdout[])

src/comm_manager.jl

@@ -29,11 +29,11 @@ function Comm(target,
     return comm
 end
 
-comm_target(comm :: Comm{target}) where {target} = target
+comm_target(comm :: Comm{target}) where {target} = target::Symbol
 
 function comm_info_request(sock, kernel, msg)
     reply = if haskey(msg.content, "target_name")
-        t = Symbol(msg.content["target_name"])
+        t = Symbol(msg.content["target_name"]::String)
         filter(kv -> comm_target(kv.second) == t, kernel.comms)
     else
         # reply with all comms.
@@ -87,9 +87,9 @@ end
 
 function comm_open(sock, kernel, msg)
     if haskey(msg.content, "comm_id")
-        comm_id = msg.content["comm_id"]
+        comm_id = msg.content["comm_id"]::String
         if haskey(msg.content, "target_name")
-            target = msg.content["target_name"]
+            target = msg.content["target_name"]::String
             if !haskey(msg.content, "data")
                 msg.content["data"] = Dict()
             end
@@ -104,11 +104,13 @@ function comm_open(sock, kernel, msg)
                                   msg, "comm_close"))
         end
     end
+
+    nothing
 end
 
 function comm_msg(sock, kernel, msg)
     if haskey(msg.content, "comm_id")
-        comm_id = msg.content["comm_id"]
+        comm_id = msg.content["comm_id"]::String
         if haskey(kernel.comms, comm_id)
             comm = kernel.comms[comm_id]
         else
@@ -121,11 +123,13 @@ function comm_msg(sock, kernel, msg)
         end
         comm.on_msg(msg)
     end
+
+    nothing
 end
 
 function comm_close(sock, kernel, msg)
     if haskey(msg.content, "comm_id")
-        comm_id = msg.content["comm_id"]
+        comm_id = msg.content["comm_id"]::String
         comm = kernel.comms[comm_id]
 
         if !haskey(msg.content, "data")
@@ -135,6 +139,8 @@ function comm_close(sock, kernel, msg)
 
         delete!(kernel.comms, comm.id)
     end
+
+    nothing
 end
 
 end # module

src/execute_request.jl

@@ -75,6 +75,9 @@ function execute_request(socket, kernel, msg)
     hcode = replace(code, r"^\s*\?" => "")
 
     try
+        if isassigned(kernel.revise_precompile_task)
+            wait(kernel.revise_precompile_task[])
+        end
         foreach(invokelatest, IJulia._preexecute_hooks)
 
         kernel.ans = result = if hcode != code # help request

src/kernel.jl

@@ -8,7 +8,8 @@ function run_kernel()
     ENV["COLUMNS"] = get(ENV, "COLUMNS", "80")
 
     println(Core.stdout, "Starting kernel event loops.")
-    IJulia.init(ARGS, IJulia.Kernel())
+    kernel = Kernel()
+    IJulia.init(ARGS, kernel)
 
     let startupfile = !isempty(DEPOT_PATH) ? abspath(DEPOT_PATH[1], "config", "startup_ijulia.jl") : ""
         isfile(startupfile) && Base.JLOptions().startupfile != 2 && Base.include(Main, startupfile)
@@ -19,11 +20,12 @@ function run_kernel()
 
     # check whether Revise is running and as needed configure it to run before every prompt
     if isdefined(Main, :Revise)
-        let mode = get(ENV, "JULIA_REVISE", "auto")
-            mode == "auto" && IJulia.push_preexecute_hook(Main.Revise.revise)
+        if get(ENV, "JULIA_REVISE", "auto") == "auto"
+            IJulia.push_preexecute_hook(Main.Revise.revise)
+            kernel.revise_precompile_task[] = Threads.@spawn precompile(Main.Revise.revise, ())
         end
     end
 
-    wait(IJulia._default_kernel::Kernel)
-    close(IJulia._default_kernel::Kernel)
+    wait(kernel)
+    close(kernel)
 end

src/precompile.jl

@@ -6,22 +6,13 @@ const _TEST_KEY = "a0436f6c-1916-498b-8eb9-e81ab9368e84"
 
 # How to update the precompilation workload:
 # 1. Uncomment the `@show` expressions in `recv_ipython()` in msg.jl.
-# 2. Copy this workload into tests/kernel.jl and update as desired:
-#
-#    Kernel(profile; capture_stdout=false, capture_stderr=false) do kernel
-#        jupyter_client(profile) do client
-#            kernel_info(client)
-#            execute(client, "42")
-#            execute(client, "?import")
-#            execute(client, """error("foo")""")
-#        end
-#    end
-#
-# 3. When the above runs it will print out the contents of the received messages
-#    as strings. You can copy these verbatim into the precompilation workload
-#    below. Note that if you modify any step of the workload you will need to
-#    update *all* the messages to ensure they have the right parent
-#    headers/signatures.
+# 2. Uncomment the call to `run_precompile()` in tests/kernel.jl (and comment
+#    out the rest of the tests if you like).
+# 3. When the `run_precompile()` runs it will print out the contents of the
+#    received messages as strings. You can copy these verbatim into the
+#    precompilation workload below. Note that if you modify any step of the
+#    workload you will need to update *all* the messages to ensure they have the
+#    right parent headers/signatures.
 @compile_workload begin
     local profile = create_profile(45_000; key=_TEST_KEY)
 
@@ -38,21 +29,32 @@ const _TEST_KEY = "a0436f6c-1916-498b-8eb9-e81ab9368e84"
             port = profile["shell_port"]
             ZMQ.connect(requests_socket, "tcp://$(ip):$(port)")
 
-            # kernel_info
-            idents = ["a1cd3a77-85e3881309b0cc3701e51156"]
-            signature = "af03588308cec89e76d0568134c0eaf24e9fe869aac729d5b9aed6baac6d369b"
-            header = "{\"msg_id\": \"a1cd3a77-85e3881309b0cc3701e51156_1694526_0\", \"msg_type\": \"kernel_info_request\", \"username\": \"james\", \"session\": \"a1cd3a77-85e3881309b0cc3701e51156\", \"date\": \"2025-08-29T09:54:47.389494Z\", \"version\": \"5.4\"}"
+            # Kernel info
+            idents = ["626c4427-479d61edd6b98ccca470f2d6"]
+            signature = "306b616a72292e9a736fe42b3c7d6fd51e10653ea2c5bc8f33810a06d33df8b5"
+            header = "{\"msg_id\": \"626c4427-479d61edd6b98ccca470f2d6_3346283_0\", \"msg_type\": \"kernel_info_request\", \"username\": \"james\", \"session\": \"626c4427-479d61edd6b98ccca470f2d6\", \"date\": \"2025-11-02T18:59:07.097698Z\", \"version\": \"5.4\"}"
             parent_header = "{}"
             metadata = "{}"
             content = "{}"
 
             ZMQ.send_multipart(requests_socket, [only(idents), "<IDS|MSG>", signature, header, parent_header, metadata, content])
             ZMQ.recv_multipart(requests_socket, String)
 
+            # Completion request
+            idents = ["626c4427-479d61edd6b98ccca470f2d6"]
+            signature = "721d3988d167417f9f0f03a823d3870470a5d924184a9a07dff422d8993eefa8"
+            header = "{\"msg_id\": \"626c4427-479d61edd6b98ccca470f2d6_3346283_1\", \"msg_type\": \"complete_request\", \"username\": \"james\", \"session\": \"626c4427-479d61edd6b98ccca470f2d6\", \"date\": \"2025-11-02T18:59:07.587641Z\", \"version\": \"5.4\"}"
+            parent_header = "{}"
+            metadata = "{}"
+            content = "{\"code\": \"mk\", \"cursor_pos\": 2}"
+
+            ZMQ.send_multipart(requests_socket, [only(idents), "<IDS|MSG>", signature, header, parent_header, metadata, content])
+            ZMQ.recv_multipart(requests_socket, String)
+
             # Execute `42`
-            idents = ["a1cd3a77-85e3881309b0cc3701e51156"]
-            signature = "23df4f581ab69b5b249caced71fc0a77bcb8f5c1f4eeb88d44ca49456db16e0d"
-            header = "{\"msg_id\": \"a1cd3a77-85e3881309b0cc3701e51156_1694526_1\", \"msg_type\": \"execute_request\", \"username\": \"james\", \"session\": \"a1cd3a77-85e3881309b0cc3701e51156\", \"date\": \"2025-08-29T09:54:49.546467Z\", \"version\": \"5.4\"}"
+            idents = ["626c4427-479d61edd6b98ccca470f2d6"]
+            signature = "01e135a939f00a11708ca906c6a16ce1226ad07669b7dad49c12c4b792ef6242"
+            header = "{\"msg_id\": \"626c4427-479d61edd6b98ccca470f2d6_3346283_2\", \"msg_type\": \"execute_request\", \"username\": \"james\", \"session\": \"626c4427-479d61edd6b98ccca470f2d6\", \"date\": \"2025-11-02T18:59:14.947227Z\", \"version\": \"5.4\"}"
             parent_header = "{}"
             metadata = "{}"
             content = "{\"code\": \"42\", \"silent\": false, \"store_history\": true, \"user_expressions\": {}, \"allow_stdin\": true, \"stop_on_error\": true}"
@@ -61,9 +63,9 @@ const _TEST_KEY = "a0436f6c-1916-498b-8eb9-e81ab9368e84"
             ZMQ.recv_multipart(requests_socket, String)
 
             # Execute `?import`
-            idents = ["a1cd3a77-85e3881309b0cc3701e51156"]
-            signature = "8f2b31cc4751d17bbd4d1216180e129b0759fdc9f05e0de1e91ca03db85fc5e1"
-            header = "{\"msg_id\": \"a1cd3a77-85e3881309b0cc3701e51156_1694526_2\", \"msg_type\": \"execute_request\", \"username\": \"james\", \"session\": \"a1cd3a77-85e3881309b0cc3701e51156\", \"date\": \"2025-08-29T09:54:49.951328Z\", \"version\": \"5.4\"}"
+            idents = ["626c4427-479d61edd6b98ccca470f2d6"]
+            signature = "a0fc54c9ed7250e8a151a766c7293843a1b8f91b1d766579a4761c18f34862df"
+            header = "{\"msg_id\": \"626c4427-479d61edd6b98ccca470f2d6_3346283_3\", \"msg_type\": \"execute_request\", \"username\": \"james\", \"session\": \"626c4427-479d61edd6b98ccca470f2d6\", \"date\": \"2025-11-02T18:59:15.621963Z\", \"version\": \"5.4\"}"
             parent_header = "{}"
             metadata = "{}"
             content = "{\"code\": \"?import\", \"silent\": false, \"store_history\": true, \"user_expressions\": {}, \"allow_stdin\": true, \"stop_on_error\": true}"
@@ -72,25 +74,52 @@ const _TEST_KEY = "a0436f6c-1916-498b-8eb9-e81ab9368e84"
             ZMQ.recv_multipart(requests_socket, String)
 
             # Execute `error("foo")`
-            idents = ["a1cd3a77-85e3881309b0cc3701e51156"]
-            signature = "c8415a60d32d231b582128a8f85ecce0996ee76734ab6aecf93af850d7e19e4a"
-            header = "{\"msg_id\": \"a1cd3a77-85e3881309b0cc3701e51156_1694526_3\", \"msg_type\": \"execute_request\", \"username\": \"james\", \"session\": \"a1cd3a77-85e3881309b0cc3701e51156\", \"date\": \"2025-08-29T09:54:50.755285Z\", \"version\": \"5.4\"}"
+            idents = ["626c4427-479d61edd6b98ccca470f2d6"]
+            signature = "7c7810bce6ff448c3ad8b0c928822d03ad8baf838e417ea5c29ec620a2cc36a7"
+            header = "{\"msg_id\": \"626c4427-479d61edd6b98ccca470f2d6_3346283_4\", \"msg_type\": \"execute_request\", \"username\": \"james\", \"session\": \"626c4427-479d61edd6b98ccca470f2d6\", \"date\": \"2025-11-02T18:59:18.648766Z\", \"version\": \"5.4\"}"
             parent_header = "{}"
             metadata = "{}"
             content = "{\"code\": \"error(\\\"foo\\\")\", \"silent\": false, \"store_history\": true, \"user_expressions\": {}, \"allow_stdin\": true, \"stop_on_error\": true}"
 
             ZMQ.send_multipart(requests_socket, [only(idents), "<IDS|MSG>", signature, header, parent_header, metadata, content])
             ZMQ.recv_multipart(requests_socket, String)
 
+            # Get history
+            idents = ["626c4427-479d61edd6b98ccca470f2d6"]
+            signature = "1d4bc84a2efb28efa0e8efe8eba32bc6b84decfd55d95e765b062a1adae2ae62"
+            header = "{\"msg_id\": \"626c4427-479d61edd6b98ccca470f2d6_3346283_5\", \"msg_type\": \"history_request\", \"username\": \"james\", \"session\": \"626c4427-479d61edd6b98ccca470f2d6\", \"date\": \"2025-11-02T18:59:18.813677Z\", \"version\": \"5.4\"}"
+            parent_header = "{}"
+            metadata = "{}"
+            content = "{\"raw\": true, \"output\": false, \"hist_access_type\": \"range\", \"session\": 0, \"start\": 0}"
+
+            ZMQ.send_multipart(requests_socket, [only(idents), "<IDS|MSG>", signature, header, parent_header, metadata, content])
+            ZMQ.recv_multipart(requests_socket, String)
+
+            # Get comm info
+            idents = ["626c4427-479d61edd6b98ccca470f2d6"]
+            signature = "ff8e9a57a81f5cbdebd12a34a09aad953d63e24e0c866fc3af936a4ef764a181"
+            header = "{\"msg_id\": \"626c4427-479d61edd6b98ccca470f2d6_3346283_6\", \"msg_type\": \"comm_info_request\", \"username\": \"james\", \"session\": \"626c4427-479d61edd6b98ccca470f2d6\", \"date\": \"2025-11-02T18:59:18.953835Z\", \"version\": \"5.4\"}"
+            parent_header = "{}"
+            metadata = "{}"
+            content = "{}"
+
+            ZMQ.send_multipart(requests_socket, [only(idents), "<IDS|MSG>", signature, header, parent_header, metadata, content])
+            ZMQ.recv_multipart(requests_socket, String)
+
             close(requests_socket)
         end
     end
+
+    # Clear global variables
+    empty!(_preexecute_hooks)
+    empty!(_postexecute_hooks)
+    empty!(_posterror_hooks)
 end
 
 # This function is executed by Jupyter so make sure that it's precompiled
 precompile(run_kernel, ())
 
 # Precompile all the handlers
-for f in handlers
+for f in values(handlers)
     precompile(f, (ZMQ.Socket, Kernel, Msg))
 end

test/kernel.jl

@@ -91,6 +91,7 @@ history(client) = make_request(client.history, client.get_shell_msg)
 shutdown(client; wait=true) = make_request(client.shutdown, client.get_control_msg; wait)
 execute(client, code) = make_request(client.execute, client.get_shell_msg; code)
 inspect(client, code) = make_request(client.inspect, client.get_shell_msg; code)
+complete(client, code) = make_request(client.complete, client.get_shell_msg; code)
 get_stdin_msg(client) = make_request(Returns(nothing), client.get_stdin_msg)
 get_iopub_msg(client) = make_request(Returns(nothing), client.get_iopub_msg)
 
@@ -130,6 +131,51 @@ function jupyter_client(f, profile)
     end
 end
 
+function run_precompile()
+    profile = IJulia.create_profile(; key=IJulia._TEST_KEY)
+
+    zmq_recv = """
+
+    ZMQ.send_multipart(requests_socket, [only(idents), "<IDS|MSG>", signature, header, parent_header, metadata, content])
+    ZMQ.recv_multipart(requests_socket, String)
+    """
+
+    Kernel(profile; capture_stdout=false, capture_stderr=false) do kernel
+        jupyter_client(profile) do client
+            println("# Kernel info")
+            kernel_info(client)
+            println(zmq_recv)
+
+            println("# Completion request")
+            complete(client, "mk")
+            println(zmq_recv)
+
+            println("# Execute `42`")
+            execute(client, "42")
+            println(zmq_recv)
+
+            println("# Execute `?import`")
+            execute(client, "?import")
+            println(zmq_recv)
+
+            println("""# Execute `error("foo")`""")
+            execute(client, """error("foo")""")
+            println(zmq_recv)
+
+            println("# Get history")
+            history(client)
+            println(zmq_recv)
+
+            println("# Get comm info")
+            comm_info(client)
+            println(zmq_recv)
+        end
+    end
+end
+
+# Uncomment this to run the precompilation workload
+# run_precompile()
+
 @testset "Kernel" begin
     profile = IJulia.create_profile(; key=IJulia._TEST_KEY)
     profile_kwargs = Dict([Symbol(key) => value for (key, value) in profile])