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 84.19118% with 43 lines in your changes missing coverage. Please review.

Project coverage is 61.14%. Comparing base (4390ef4) to head (d491800).

Files with missing lines	Patch %	Lines
src/handlers.jl	66.66%	8 Missing ⚠️
src/init.jl	82.92%	7 Missing ⚠️
src/IJulia.jl	92.59%	6 Missing ⚠️
src/comm_manager.jl	25.00%	6 Missing ⚠️
src/inline.jl	57.14%	6 Missing ⚠️
src/stdio.jl	82.35%	6 Missing ⚠️
src/eventloop.jl	86.36%	3 Missing ⚠️
src/execute_request.jl	96.29%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #1145       +/-   ##
===========================================
+ Coverage   10.21%   61.14%   +50.92%     
===========================================
  Files          14       14               
  Lines         822      893       +71     
===========================================
+ Hits           84      546      +462     
+ Misses        738      347      -391

☔ 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 :)

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.

Fixed the tests for ans/In/Out being propagated to kernel.current_module in 2fd607f, also moved verbose fully into Kernel.

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.~~

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

JamesWrigley · 2025-02-13T23:19:32Z

I believe this is ready to be reviewed now. One thorny issue is segfaults on windows that appear to be coming from ZMQ, using @threadcall for the heartbeat seemed to fix it: d491800
But I'm not convinced I understand the problem fully, will try to debug it a bit more before merging.

For reviewers, I've tried to keep the commits atomic so I would recommend reviewing each commit one-by-one.

JamesWrigley self-assigned this Feb 5, 2025

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

stevengj reviewed Feb 10, 2025

View reviewed changes

JamesWrigley force-pushed the kernel-struct branch from 1dd655d to 212ff17 Compare February 13, 2025 22:01

JamesWrigley added 7 commits February 13, 2025 23:09

Giant refactor to move all state into a Kernel struct

f7e22bd

Handle ZMQ.StateError's in the eventloop

ab94ecc

Send shutdown_reply on the control socket for protocol 5.4

0f0fa25

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

4343060

This is what Jupyter expects.

Send a metadata key in inspect_reply messages

576d2a3

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

Make history() print each entry on a new line

3c393ba

JamesWrigley force-pushed the kernel-struct branch from 212ff17 to 3c393ba Compare February 13, 2025 22:09

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

d491800

JamesWrigley marked this pull request as ready for review February 13, 2025 23:16

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

JamesWrigley Feb 12, 2025

stevengj Feb 10, 2025

JamesWrigley Feb 10, 2025 •

edited

Loading

JamesWrigley commented Feb 10, 2025 •

edited

Loading

JamesWrigley commented Feb 13, 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

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 • edited Loading

JamesWrigley commented Feb 13, 2025

codecov bot commented Feb 7, 2025 •

edited

Loading

JamesWrigley Feb 10, 2025 •

edited

Loading

JamesWrigley commented Feb 10, 2025 •

edited

Loading