Giant refactor to move all state into a Kernel struct #1145

JamesWrigley · 2025-02-05T21:57:16Z

Pretty much what the title says 😛 This is not ready yet but I'm opening it for transparency. There are two reasons for the refactor:

Easier testing, so that we can more confidently make changes like Make message receive and handling async #1140 and Jupyter messaging updates/correctness #1138.
Allow for creating a temporary kernel that can be used in a precompilation workload for Add precompile statements to reduce startup latency #1074.

Overview:

Almost all global state has been moved into the new Kernel struct, and all relevant internal methods now require a kernel be passed to them. That's the bulk of the diff. I've tried to minimize other changes as much as possible, so there should 🤞 be no change in behaviour from master.
There are certain fields like ans and n which are documented as globals and must stay global, so setproperty!(::Kernel,...) is implemented such that it will copy changes to those fields to their corresponding global variable. Similarly there are certain public methods like IJulia.history() that must still work without a kernel argument, so a new _default_kernel global is defined for normal usage and that's used by default for all public methods.

An edge-case are the dicts In and Out which are by default assigned to the corresponding fields of _default_kernel and will not contain changes from temporary test kernels. I don't think that's a particularly important discrepancy so I'm inclined to ignore it.
I implemented a basic test in waitloop.jl that spawns the IJulia waitloop and then makes a request to it using jupyter_client. This isn't working very well right now because blocking calls to the client causes hangs, I suspect there's a deadlock somewhere.

This will need a lot more work before it can be merged, but feel free to review the design already (cc @stevengj, @fredrikekre).

@MasonProtter, you should be able to do... something... with this for a precompilation workload. I would suggest creating a Kernel, connecting to some of the sockets with ZMQ.jl, and then sending some mock requests following the wire protocol: https://jupyter-client.readthedocs.io/en/latest/messaging.html#the-wire-protocol

codecov · 2025-02-07T00:48:38Z

Codecov Report

Attention: Patch coverage is 81.85185% with 49 lines in your changes missing coverage. Please review.

Project coverage is 59.61%. Comparing base (4390ef4) to head (20dc591).

Files with missing lines	Patch %	Lines
src/stdio.jl	68.75%	10 Missing ⚠️
src/handlers.jl	66.66%	8 Missing ⚠️
src/init.jl	82.92%	7 Missing ⚠️
src/IJulia.jl	92.68%	6 Missing ⚠️
src/comm_manager.jl	25.00%	6 Missing ⚠️
src/inline.jl	57.14%	6 Missing ⚠️
src/eventloop.jl	77.27%	5 Missing ⚠️
src/execute_request.jl	96.29%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #1145       +/-   ##
===========================================
+ Coverage   10.21%   59.61%   +49.40%     
===========================================
  Files          14       14               
  Lines         822      894       +72     
===========================================
+ Hits           84      533      +449     
+ Misses        738      361      -377

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

JamesWrigley · 2025-02-07T01:12:39Z

Figured out the deadlock, we can now start writing full tests with jupyter_client :)

This is the only change in 5.4, in previous versions it was sent on the shell socket.

This is what Jupyter expects.

The Jupyter messaging docs state that all messages on the shell (request) channel shall be ended with a `status: idle` message, and that: "This idle status message indicates that IOPub messages associated with a given request have all been received." Previously we did not attempt to ensure this, which meant that in some cases the `execute_result` message would be published after the `execute_reply` message had been sent. It's still not exactly guaranteed by the `yield()`, but it's something and it makes the `jupyter_kernel_test` tests pass.

This is required by the spec: https://jupyter-client.readthedocs.io/en/latest/messaging.html#introspection

stevengj · 2025-02-10T12:41:59Z

src/IJulia.jl

 using ZMQ, JSON, SoftGlobalScope
-import Base.invokelatest
+import Base: invokelatest, RefValue


Probably using rather than import since we only use these, not extend them? Also invokelatest is exported automatically these days.

(This code dates back to the days before using Foo: bar)

Agreed, but I think I'd prefer to do those sorts of changes in a later PR. This one is already quite large 😅

stevengj · 2025-02-10T12:43:59Z

src/IJulia.jl

+    setfield!(value, name, x)
+end
+
+function Base.close(kernel::Kernel)


should be called by a finalizer?

Unfortunately that wouldn't be easy because it wait's on the stdio tasks so it might yield. I think it's ok to rely on the do-constructor closing it in the tests.

stevengj · 2025-02-10T12:45:56Z

src/IJulia.jl

+    # global counterparts.
+    if name ∈ (:ans, :n, :In, :Out, :inited, :verbose)
+        setproperty!(IJulia, name, x)
+    end


aren't ans, In, Out global in Main (not IJulia)?

Maybe this kind of global should be copied to/from Main before/after each execute request?

whereas n and verbose at least can probably be moved in to the Kernel?

You're quite right 🤔 I need to figure out a good way to test that, for now I've added a @test_broken test in d9acb6e.

whereas n and verbose at least can probably be moved in to the Kernel?

n is documented to be module-global (https://julialang.github.io/IJulia.jl/dev/library/public/#IJulia.n) so I don't think we can move that fully into Kernel, but verbose should be possible. I think only @vprintln/@verror_show will need tweaking to look at _default_kernel.verbose instead of the global.

stevengj · 2025-02-10T12:51:45Z

src/stdio.jl

@@ -274,8 +260,8 @@ function oslibuv_flush()
 end

 import Base.flush
-function flush(io::IJuliaStdio)
+function flush(io::IJuliaStdio, kernel)


This is a user-visible function, e.g. if the user calls flush(stdout).

It would be better to store a reference to the kernel in the IJuliaStdio object.

Good idea, fixed in d9acb6e and added a test for it.

JamesWrigley · 2025-02-10T21:24:34Z

Progress update:

Moved the waitloop task into Kernel to make testing easier, and added a wait(::Kernel) method to wait for it to finish.
Added/expanded tests for autocompletion, MIME display, and clear_output() using jupyter_kernel_test. I think that covers everything we can use jupyter_kernel_test for.
Moved the calls to watch_stdio() and pushdisplay() into init(), now close(::Kernel) will ensure they're cleaned up properly.
Added a bunch of tests for some of the public API: clear_history(), load(), the hook system, and history().
Tweaked history() to print all the entries on new lines in 20dc591. From the docstring, I think that's what it should be doing?

To-do:

Address remaining review comments.
Add a test for kernel.jl to ensure that the kernel works outside of the test suite.
Figure out why the load() test is failing on Windows.
Go through the coverage reports to find more low-hanging fruit to test.

JamesWrigley self-assigned this Feb 5, 2025

JamesWrigley force-pushed the kernel-struct branch 3 times, most recently from 4aa3633 to 64ec133 Compare February 10, 2025 00:01

JamesWrigley added 6 commits February 10, 2025 01:10

Giant refactor to move all state into a Kernel struct

fee074a

Handle ZMQ.StateError's in the eventloop

d42e0d9

Send shutdown_reply on the control socket for protocol 5.4

92af7ca

This is the only change in 5.4, in previous versions it was sent on the shell socket.

Convert the package version to a string before sending it

3c41114

This is what Jupyter expects.

Send a metadata key in inspect_reply messages

b19c5aa

This is required by the spec: https://jupyter-client.readthedocs.io/en/latest/messaging.html#introspection

JamesWrigley force-pushed the kernel-struct branch from 64ec133 to b19c5aa Compare February 10, 2025 00:11

stevengj reviewed Feb 10, 2025

View reviewed changes

JamesWrigley added 2 commits February 10, 2025 22:05

fixup! Giant refactor to move all state into a Kernel struct

d9acb6e

Make history() print each entry on a new line

20dc591

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Giant refactor to move all state into a Kernel struct #1145

Giant refactor to move all state into a Kernel struct #1145

JamesWrigley commented Feb 5, 2025

codecov bot commented Feb 7, 2025 •

edited

Loading

JamesWrigley commented Feb 7, 2025

stevengj Feb 10, 2025

JamesWrigley Feb 10, 2025

stevengj Feb 10, 2025

JamesWrigley Feb 10, 2025

stevengj Feb 10, 2025

stevengj Feb 10, 2025

JamesWrigley Feb 10, 2025

stevengj Feb 10, 2025

JamesWrigley Feb 10, 2025 •

edited

Loading

JamesWrigley commented Feb 10, 2025

Giant refactor to move all state into a Kernel struct #1145

Are you sure you want to change the base?

Giant refactor to move all state into a Kernel struct #1145

Conversation

JamesWrigley commented Feb 5, 2025

codecov bot commented Feb 7, 2025 • edited Loading

Codecov Report

JamesWrigley commented Feb 7, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

JamesWrigley Feb 10, 2025 • edited Loading

Choose a reason for hiding this comment

JamesWrigley commented Feb 10, 2025

codecov bot commented Feb 7, 2025 •

edited

Loading

JamesWrigley Feb 10, 2025 •

edited

Loading