Kernel subshells (JEP91) implementation #1249

ianthomas23 · 2024-06-13T10:02:41Z

This is the implementation of the kernel subshells JEP (jupyter/enhancement-proposals#91). It follows the latest commit (1f1ad3d) with the addition of a %subshell magic command that is useful for debugging. To try this out I have a JupyterLab branch that talks to this branch and is most easily tried out using https://mybinder.org/v2/gh/ianthomas23/jupyterlab/jep91_demo?urlpath=lab; once the mybinder instance has started, open the subshell_demo_notebook.ipynb and follow the instructions therein.

The idea is that this is mergeable as it is now, it is backward compatible in that it does not break any existing use of ipykernel (subject to CI confirmation). There are some ramifications of the protocol additions (outlined below) that will need addressing eventually, but I consider these future work that can be in separate PRs.

Outline of changes

The parent subshell (i.e. the main shell) runs in the main thread.
Each new subshell runs in a separate thread.
There is a new thread that deals with all communication on the shell channel, previously this was performed in the main thread.
Communication between the shell channel thread and other threads is performed using ZMQ inproc pair sockets, which are essentially shared memory and avoid the use of thread synchronisation primitives.
Incoming shell messages are handled by the shell channel thread which extracts the subshell_id from the message and passes it on to the correct subshell.
Subshells are created and deleted via messages sent on the control channel. These are passed to the shell channel thread via inproc pair sockets so that the SubshellManager in the shell channel thread is responsible for subshell lifetimes.

Example scenario

Here is an example of the communication between threads when running a long task in the parent subshell (main thread) and whilst this is running a child subshell is created, used, and deleted.

sequenceDiagram
    participant client as Client
    participant control as Control thread
    participant shell as Shell channel thread
    participant main as Main thread

    client->>+shell: Execute request (main shell)
    shell->>-main: Execute request (inproc)
    activate main

    client->>+control: Create subshell request
    control->>-shell: Create subshell request (inproc)
    activate shell
    create participant subshell as Subshell thread
    shell-->>subshell: Create subshell thread

    shell->>control: Create subshell reply (inproc)
    deactivate shell
    activate control
    control->>-client: Create subshell reply

    client->>+shell: Execute request (subshell)
    shell->>-subshell: Execute request (inproc)
    activate subshell

    subshell->>shell: Execute reply (inproc)
    deactivate subshell
    activate shell
    shell->>-client: Execute reply (subshell)

    client->>+control: Delete subshell request
    control->>-shell: Delete subshell request (inproc)
    activate shell
    destroy subshell
    shell-->>subshell: Delete subshell thread

    shell->>control: Delete subshell reply (inproc)
    deactivate shell
    activate control
    control->>-client: Delete subshell reply

    main->>shell: Execute reply (inproc)
    deactivate main
    activate shell
    shell->>-client: Execute reply (main shell)

Future work

ipykernel

Shell channel thread deserialises ~~the whole~~ some of the message to get the subshell_id. Ideally it would only deserialise the header. May need changes in Jupyter Client.
Signalling a subshell to stop uses a threading.Event following the existing anyio implementation which requires an extra thread per Event. It would be nice if this could be changed so a subshell is a single thread not two.
Execution count. Should either be a separate count per subshell or a single count for a kernel. Needs a decision and changes in IPython as is currently not atomic.
History. Related to item 2 above.
input() on more than one subshell at the same time run but do not store correctly.
Debugger use needs investigating.
Busy/idle status needs investigating. Should there, as now, be separate status for each subshell, or the concept of kernel (i.e. any subshell) busy status? This issue is much wider than subshells as it includes status of the control channel, and how Jupyter Server should track status (Improve the busy/idle execution state tracking for kernels. jupyter-server/jupyter_server#1429).
Use of display hooks for e.g. Matplotlib. Should these be on the parent subshell, or child subshells too?

JupyterLab

The JupyterLab branch I am using to demo this isn't really intended to be merged. But if it was, it needs:

Check kernel_info to see if subshells are supported.
Delete subshell when close a subshell's ConsolePanel.
Report subshell IDs in tree view?
Display of subshell busy/idle status.

(Edited for clarity)

ianthomas23 · 2024-09-10T07:04:02Z

The kernel subshells JEP has been accepted and as this is the reference implementation for that I am asking for someone to review this.

davidbrochart · 2024-09-10T09:53:18Z

ipykernel/subshell_manager.py

+        while True:
+            for socket, _ in self._poller.poll(0):
+                msg = await socket.recv_multipart(copy=False)
+                self._shell_socket.send_multipart(msg)
+
+            # Yield to other tasks.
+            await sleep(0)


This loop leads to 100% CPU usage, we need to find another way.

Interesting. On my macOS dev machine this is fine, but I've confirmed that it is a problem on Linux.

I've changed the implementation in 97e3e91. I've removed the sleep(0) and poll(0) calls and now use an async zmq Poller in an anyio task. When a subshell is created or deleted this task is cancelled using an anyio.Event and rescheduled with the updated list of subshells. With this I no longer see 100% CPU load on Linux or macOS.

I've changed the implementation of this piece of code again. My use of await poller.poll() turned out to cause problems on python < 3.10 and on Windows. Now I am avoiding use of zmq.Poller and instead using a separate anyio task for each subshell to listen for reply messages and send them out to the client via the shell channel. This seems to be more robust on older python and Windows. I've also replaced the use of an anyio.Event with a memory object stream instead (essentially an anyio async queue).

for more information, see https://pre-commit.ci

ianthomas23 · 2024-09-17T16:47:12Z

I had missed a direct use of the shell_socket to send an abort_reply message when it should have been sent via the SubshellManager. With that fixed, all the local CI is now passing.

ipykernel/thread.py

Carreau

Thanks for your patience, I'm going to try reviewing this and pushing it forward.

I'm done a partial read – not in depth, and added a few comments – which are not mandatory request for changes but thoughts to myself, for when I'll come back to it.

I have to run some errands, so I'm going to post this as is for now but will come back to it later.

ipykernel/thread.py

Carreau · 2024-10-01T08:58:10Z

ipykernel/control.py

-
-    def set_task(self, task):
-        self._task = task
+        super().__init__(name=CONTROL_THREAD_NAME, **kwargs)

    def run(self):
        """Run the thread."""
        self.name = CONTROL_THREAD_NAME


I think this is unnecessary now Thread have setter/getter and super().__init__ should set the private name directly.

Maybe try to put an assert ==, and if test are passing, just remove the override of the run method ?

Removed the name duplication in b644f9e7 and 9109cc8 without causing any problems.

ipykernel/kernelbase.py

ipykernel/subshell.py

Carreau · 2024-10-01T09:16:15Z

ipykernel/subshell_manager.py

+
+        self._context: zmq.asyncio.Context = context
+        self._shell_socket = shell_socket
+        self._cache: dict[str, Subshell] = {}


Self curiosity:

check the difference between attribute typing in __init__ and at class level.

Is mypy smart enough, or does it say attribute may be unset if type set in __init__ ??

ipykernel/subshell_manager.py

Carreau · 2024-10-01T09:19:41Z

ipykernel/subshell_manager.py

+            elif type == "list":
+                reply["subshell_id"] = self.list_subshell()
+            else:
+                msg = f"Unrecognised message type {type}"


Suggested change

msg = f"Unrecognised message type {type}"

msg = f"Unrecognised message type {type!r}"

Fixed in d17c77f.

ipykernel/kernelbase.py

ipykernel/thread.py

ipykernel/zmqshell.py

tests/utils.py

Carreau · 2024-10-02T10:09:15Z

I'm +1 for this, a couple of stylistic notes – if you agree I'm happy to push some changes updates but don't want to do so without your approval as I know it can get unwieldy to have another maintainer push changes on your branch on big PRs.

I think we can merge this as "experimental" and then work on the related issues when using this features on various projects (IPython history etc...).

ianthomas23 · 2024-10-02T11:35:16Z

@Carreau Thanks for the review and for the offer of pushing changes. I'd rather make the changes myself, and I'll try to group them sensibly to avoid to much CI churn.

davidbrochart · 2024-10-02T13:26:48Z

It seems that you use zmq.PAIR sockets for inter-thread communication (sub-shell thread <---> main-shell thread, sub-shell thread <--> control thread), are those real sockets? If so, sockets are a scarce resource, and AnyIO has nice facilities for inter-thread communication, that could be used instead?
I see that you use AnyIO memory object streams in the sub-shell manager, could you explain your use of them, in particular if they are used for inter-thread communication?

ianthomas23 · 2024-10-02T15:29:50Z

It seems that you use zmq.PAIR sockets for inter-thread communication (sub-shell thread <---> main-shell thread, sub-shell thread <--> control thread), are those real sockets?

A ZMQ inproc pair is an area of memory shared between two threads and wrapped to look like ZMQ sockets.

I see that you use AnyIO memory object streams in the sub-shell manager, could you explain your use of them, in particular if they are used for inter-thread communication?

Here is where they are created in the source code which includes some explanation:

# anyio memory object stream for async queue-like communication between tasks.
# Used by _create_subshell to tell listen_from_subshells to spawn a new task.
self._send_stream, self._receive_stream = create_memory_object_stream[str]()

The communication is between two different async tasks running in the same thread.

Carreau

Some notes for myself and history, but all good on my side. Let's get that in and iterate if necessary.

Carreau · 2024-10-03T08:51:30Z

ipykernel/kernelbase.py

+        # socket=None is valid if kernel subshells are not supported.
+        try:
+            while True:
+                await self.process_shell_message(socket=socket)


Not for this PR:

I don't like that breaking out of this while loop requires an exception that is hard to find in the implementation of process_shell_message; we should likely have an explicit exit point.

Carreau · 2024-10-03T08:52:53Z

ipykernel/kernelbase.py

                if inspect.isawaitable(result):
                    await result


Not for this PR:

I think we should move to handler always returning an awaitable, or more likely handler always being a coroutine function.

Carreau · 2024-10-03T08:53:24Z

ipykernel/kernelbase.py

            "protocol_version": kernel_protocol_version,
            "implementation": self.implementation,
            "implementation_version": self.implementation_version,
            "language_info": self.language_info,
            "banner": self.banner,
            "help_links": self.help_links,
+            "supported_features": [],


Carreau · 2024-10-03T08:54:47Z

ipykernel/kernelbase.py

+            return
+        if not self._supports_kernel_subshells:
+            self.log.error("Subshells are not supported by this kernel")
+            return


And from the discussion bellow, maybe think in later maybe return something that indicate why the subshell was not created ?

Or maybe not as that would change the message spec.

I think I get that better now, that these handler return values likely go nowhere so why they do not raise exception. Let's maybe leave making this clearer to a subsequent refactor/maintenance.

Carreau · 2024-10-03T09:01:39Z

ipykernel/subshell_manager.py

+        except BaseException as err:
+            reply = {
+                "status": "error",
+                "evalue": str(err),


not for this PR:

Can this contain information that should not be send to the frontend ?

I don't believe so as the user is anyway the one that runs the kernel, but we had case where traceback would contain sensitive values. In other apps webapp this would be a likely issue (like revealing paths and filename of a server), but I don't think it's an issue here.

Carreau · 2024-10-03T09:06:06Z

I haven't merged something on this repo in quite some time, I think I would prefer a rebase-and-merge as you took time to craft the commits individually, but I think the policy here is squash-merged.

Any objections to squash-merge ?

ianthomas23 · 2024-10-03T09:34:46Z

Any objections to squash-merge ?

A squash-merge is fine by me.

davidbrochart

Thanks @ianthomas23 !

ianthomas23 added the enhancement label Jun 13, 2024

ianthomas23 mentioned this pull request Jun 13, 2024

Add JEP for sub-shells jupyter/enhancement-proposals#91

Merged

30 tasks

3coins mentioned this pull request Jun 20, 2024

Meeting Notes 2024 jupyter-server/team-compass#57

Open

ianthomas23 force-pushed the jep91 branch 2 times, most recently from 651299f to 9139669 Compare July 1, 2024 14:02

ianthomas23 force-pushed the jep91 branch from 9139669 to 3bfd93b Compare July 8, 2024 15:55

ianthomas23 force-pushed the jep91 branch from 3bfd93b to 938e844 Compare August 6, 2024 09:10

This was referenced Sep 10, 2024

Implement sub-shells #1062

Closed

Support subshells #955

Closed

davidbrochart requested changes Sep 10, 2024

View reviewed changes

jtpio mentioned this pull request Sep 11, 2024

Weekly Team Meetings: Jul–Dec 2024 jupyterlab/frontends-team-compass#251

Open

ianthomas23 and others added 18 commits September 12, 2024 14:02

Separate thread to handle incoming shell channel messages

4ea2d04

Add supported_features entry to kernel_info_reply

b801b79

Dummy subshell control message implementations

0ddc1b7

SubshellCache

17bd0f6

Subshell threads with correct lifetime but don't yet execute code

62e5d21

Extract subshell_id, and handle message in correct subshell

2f8273b

Add subshell test_execution_count

1cc8892

Use dataclass for Subshell, clean up correctly

f5b664d

Test can create subshell whilst main subshell is executing

c3fd89e

Add %subshell magic

7fc1d6e

Temporary mutex for execute_reply from subshells

b8b7af5

Process subshell control messages in SubshellCache

7dd3407

Send shell reply messages via shell channel thread

27ead02

Poll subshell inproc pair sockets in SubshellManager

1e79f4c

[pre-commit.ci] auto fixes from pre-commit.com hooks

365d231

for more information, see https://pre-commit.ci

Better error checking

90c3a20

Pin anyio <= 4.3.0

2fa7365

Remove anyio pin following fix in PR 1253

ad32ea7

ianthomas23 added 2 commits September 17, 2024 17:33

Add modules to docs API

e4fa75d

Increase timings of test_subshells::test_execution_counts

a708e0e

fleming79 mentioned this pull request Sep 29, 2024

Draft: Add bi-directional comms and run operations in tasks jtpio/ipylab#135

Closed

Carreau self-requested a review October 1, 2024 08:45

Carreau reviewed Oct 1, 2024

View reviewed changes

ipykernel/thread.py Show resolved Hide resolved

Carreau reviewed Oct 1, 2024

View reviewed changes

Carreau reviewed Oct 2, 2024

View reviewed changes

ianthomas23 added 8 commits October 2, 2024 13:05

Use _tasks_and_args in BaseThread

2a91bcd

Use super when calling base class constructors

4bcd59e

Improve thread classes

b644f9e

Always return supported_features in kernel_info

924015a

Make _supports_kernel_subshells a property

a1fb1de

Clarify socket=None in process_shell_message

be5c83a

Improved use of thread names

9109cc8

Small changes

d17c77f

ianthomas23 force-pushed the jep91 branch from a79f32a to b5b96a2 Compare October 2, 2024 13:25

More type annotations

5715dd0

ianthomas23 force-pushed the jep91 branch from b5b96a2 to 5715dd0 Compare October 2, 2024 14:06

Carreau approved these changes Oct 3, 2024

View reviewed changes

davidbrochart approved these changes Oct 3, 2024

View reviewed changes

davidbrochart merged commit 3089438 into ipython:main Oct 3, 2024
28 of 32 checks passed

ianthomas23 deleted the jep91 branch October 3, 2024 10:41

ianthomas23 mentioned this pull request Oct 21, 2024

Support kernel subshells (JEP91) jupyterlab/jupyterlab#16882

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Kernel subshells (JEP91) implementation #1249

Kernel subshells (JEP91) implementation #1249

ianthomas23 commented Jun 13, 2024 •

edited

Loading

ianthomas23 commented Sep 10, 2024

davidbrochart Sep 10, 2024

ianthomas23 Sep 12, 2024

ianthomas23 Sep 12, 2024 •

edited

Loading

ianthomas23 Sep 13, 2024

ianthomas23 commented Sep 17, 2024

Carreau left a comment

Carreau Oct 1, 2024 •

edited

Loading

ianthomas23 Oct 2, 2024

Carreau Oct 1, 2024

Carreau Oct 1, 2024

ianthomas23 Oct 2, 2024

Carreau commented Oct 2, 2024

ianthomas23 commented Oct 2, 2024

davidbrochart commented Oct 2, 2024

ianthomas23 commented Oct 2, 2024

Carreau left a comment

Carreau Oct 3, 2024

Carreau Oct 3, 2024

Carreau Oct 3, 2024

Carreau Oct 3, 2024

Carreau Oct 3, 2024

Carreau Oct 3, 2024

Carreau commented Oct 3, 2024

ianthomas23 commented Oct 3, 2024

davidbrochart left a comment

	msg = f"Unrecognised message type {type}"
	msg = f"Unrecognised message type {type!r}"

Kernel subshells (JEP91) implementation #1249

Kernel subshells (JEP91) implementation #1249

Conversation

ianthomas23 commented Jun 13, 2024 • edited Loading

Outline of changes

Example scenario

Future work

ipykernel

JupyterLab

ianthomas23 commented Sep 10, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ianthomas23 Sep 12, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ianthomas23 commented Sep 17, 2024

Carreau left a comment

Choose a reason for hiding this comment

Carreau Oct 1, 2024 • edited Loading

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

Carreau commented Oct 2, 2024

ianthomas23 commented Oct 2, 2024

davidbrochart commented Oct 2, 2024

ianthomas23 commented Oct 2, 2024

Carreau left a 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

Carreau commented Oct 3, 2024

ianthomas23 commented Oct 3, 2024

davidbrochart left a comment

Choose a reason for hiding this comment

ianthomas23 commented Jun 13, 2024 •

edited

Loading

ianthomas23 Sep 12, 2024 •

edited

Loading

Carreau Oct 1, 2024 •

edited

Loading