Skip to content

Add wrapper magic: %remove_kernel #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
OutlyingWest wants to merge 1 commit into from
Closed

Add wrapper magic: %remove_kernel #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,18 @@ Output:
Successfully wrapped kernel: python3
```

#### `%remove_kernel <kernel_name>`
Remove a wrapper-created kernel spec. If the kernel is currently wrapped, it will be shut down first.

```python
%remove_kernel my_wrapped_kernel
```

Output:
```
Removed kernel spec: my_wrapped_kernel
```

### Jumper Extension Commands

All jumper-extension magic commands are available and executed locally:
Expand Down
184 changes: 152 additions & 32 deletions jumper_wrapper_kernel/kernel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
while keeping jumper-extension magic commands local.
"""

import os
import sys
from ipykernel.ipkernel import IPythonKernel
from jupyter_client import KernelManager
Expand Down Expand Up @@ -53,6 +54,15 @@ def wrap_kernel(self, line):
"""
self._kernel._wrap_kernel(line.strip())

@line_magic
def remove_kernel(self, line):
"""Remove a Jupyter kernel spec created by the wrapper.

Usage:
%remove_kernel <kernel_name>
"""
self._kernel._remove_kernel(line.strip())


class JumperWrapperKernel(IPythonKernel):
"""A Jupyter kernel that wraps other kernels."""
Expand Down Expand Up @@ -85,6 +95,7 @@ def __init__(self, **kwargs):
self._kernel_manager = None
self._kernel_client = None
self._kernel_spec_manager = KernelSpecManager()
self._wrapper_kernel_spec_name = os.environ.get("JUMPER_WRAPPER_KERNEL_SPEC", "")

# Set of magic commands registered (populated after loading extensions)
self._jumper_magic_commands = set()
Expand Down Expand Up @@ -319,6 +330,9 @@ def _save_wrapped_kernel_spec(self, wrapped_kernel_name, new_kernel_name):
],
'display_name': new_kernel_name,
'language': wrapped_spec.get('language', 'python'),
'env': {
'JUMPER_WRAPPER_KERNEL_SPEC': new_kernel_name,
},
'metadata': {
'debugger': False,
'jumper_wrapper': {
Expand All @@ -334,27 +348,133 @@ def _save_wrapped_kernel_spec(self, wrapped_kernel_name, new_kernel_name):
return True
except Exception as e:
return False


def _remove_kernel(self, args):
"""Remove a kernel spec created by the wrapper."""
kernel_name = args.strip()

if not kernel_name:
error_msg = "Usage: %remove_kernel <kernel_name>\nUse %list_kernels to see available kernels."
self.send_response(self.iopub_socket, 'stream', {'name': 'stderr', 'text': error_msg})
return {
'status': 'error',
'ename': 'ValueError',
'evalue': 'No kernel name specified',
'traceback': [error_msg],
'execution_count': self.execution_count,
}

# Shutdown wrapped kernel if it matches the kernel name
if self._wrapped_kernel_name == kernel_name:
self._shutdown_wrapped_kernel()

available_kernels = self._get_available_kernels()
if kernel_name not in available_kernels:
error_msg = f"Kernel '{kernel_name}' not found. Use %list_kernels to see available kernels."
self.send_response(self.iopub_socket, 'stream', {'name': 'stderr', 'text': error_msg})
return {
'status': 'error',
'ename': 'ValueError',
'evalue': f'Kernel not found: {kernel_name}',
'traceback': [error_msg],
'execution_count': self.execution_count,
}

spec = available_kernels.get(kernel_name, {}).get('spec', {})
metadata = spec.get('metadata', {})
wrapper_metadata = metadata.get('jumper_wrapper', {})
if not wrapper_metadata:
error_msg = (
f"Kernel '{kernel_name}' was not created by the wrapper; "
"refusing to remove. Use %list_kernels to inspect available kernels."
)
self.send_response(self.iopub_socket, 'stream', {'name': 'stderr', 'text': error_msg})
return {
'status': 'error',
'ename': 'ValueError',
'evalue': 'Not a wrapper-created kernel',
'traceback': [error_msg],
'execution_count': self.execution_count,
}

# Case 1: We were started from a named wrapper kernel spec.
# Shut down this wrapper kernel only if the removed kernel name
# matches our own wrapper kernel spec name.
if self._wrapper_kernel_spec_name:
should_shutdown_self = (self._wrapper_kernel_spec_name == kernel_name)
else:
# Case 2: We were started via auto-wrap without a wrapper spec name.
# Shut down this wrapper kernel if the removed kernel is the same
# as our auto-wrapped target and there is no other wrapper spec
# for that target (i.e., this wrapper is the only one).
wrapped_kernel = (wrapper_metadata or {}).get("wrapped_kernel")
should_shutdown_self = bool(
wrapped_kernel
and self.auto_wrap_kernel == wrapped_kernel
and self._is_only_wrapper_spec(available_kernels, wrapped_kernel)
)

return self._try_remove(kernel_name, should_shutdown_self)


def _is_only_wrapper_spec(self, available_kernels, wrapped_kernel) -> bool:
"""Return True if there is exactly one wrapper spec for wrapped_kernel."""
return sum(
1
for spec_info in available_kernels.values()
if spec_info.get("spec", {})
.get("metadata", {})
.get("jumper_wrapper", {})
.get("wrapped_kernel") == wrapped_kernel
) == 1

def _try_remove(self,kernel_name: str, should_shutdown_self: bool):
try:
self._kernel_spec_manager.remove_kernel_spec(kernel_name)
success_msg = f"Removed kernel spec: {kernel_name}\n"
if should_shutdown_self:
success_msg += "Current kernel will shut down.\n"
self.send_response(self.iopub_socket, 'stream', {'name': 'stdout', 'text': success_msg})
if should_shutdown_self:
self._shutdown_wrapped_kernel()
self.shell.exit_now = True
return {
'status': 'ok',
'execution_count': self.execution_count,
'payload': [],
'user_expressions': {},
}
except Exception as e:
error_msg = f"Failed to remove kernel '{kernel_name}': {str(e)}"
self.send_response(self.iopub_socket, 'stream', {'name': 'stderr', 'text': error_msg})
return {
'status': 'error',
'ename': type(e).__name__,
'evalue': str(e),
'traceback': [error_msg],
'execution_count': self.execution_count,
}

def _update_language_info_from_wrapped_kernel(self):
"""Get language info from the wrapped kernel and update our language_info."""
if self._kernel_client is None:
return

try:
# Request kernel_info from wrapped kernel
msg_id = self._kernel_client.kernel_info()
reply = self._kernel_client.get_shell_msg(timeout=10)

if reply['content'].get('status') == 'ok':
wrapped_language_info = reply['content'].get('language_info', {})

# Update our language info to match the wrapped kernel
if wrapped_language_info:
self.language = wrapped_language_info.get('name', 'python')
self.language_info = wrapped_language_info.copy()
except Exception:
pass

def _notify_frontend_language_change(self):
"""Send JavaScript to frontend to trigger kernel info refresh for syntax highlighting."""
js_code = """
Expand Down Expand Up @@ -382,19 +502,19 @@ def _notify_frontend_language_change(self):
'metadata': {},
}
)

def _shutdown_wrapped_kernel(self):
"""Shutdown the currently wrapped kernel."""
if self._kernel_client is not None:
self._kernel_client.stop_channels()
self._kernel_client = None

if self._kernel_manager is not None:
self._kernel_manager.shutdown_kernel(now=True)
self._kernel_manager = None

self._wrapped_kernel_name = None

def _trigger_pre_run_cell(self, code, silent, store_history):
"""Trigger pre_run_cell event for jumper-extension hooks."""
info = ExecutionInfo(
Expand All @@ -406,7 +526,7 @@ def _trigger_pre_run_cell(self, code, silent, store_history):
)
self.shell.events.trigger('pre_run_cell', info)
return info

def _trigger_post_run_cell(self, info, success=True, result_value=None, error=None):
"""Trigger post_run_cell event for jumper-extension hooks."""
exec_result = ExecutionResult(info)
Expand All @@ -415,7 +535,7 @@ def _trigger_post_run_cell(self, info, success=True, result_value=None, error=No
exec_result.error_in_exec = error
self.shell.events.trigger('post_run_cell', exec_result)
return exec_result

def _forward_to_wrapped_kernel(self, code, silent, store_history, user_expressions, allow_stdin):
"""Forward code execution to the wrapped kernel."""
if self._kernel_client is None:
Expand All @@ -428,10 +548,10 @@ def _forward_to_wrapped_kernel(self, code, silent, store_history, user_expressio
'traceback': [error_msg],
'execution_count': self.execution_count,
}

# Trigger pre_run_cell event for jumper-extension
exec_info = self._trigger_pre_run_cell(code, silent, store_history)

# Execute on wrapped kernel
msg_id = self._kernel_client.execute(
code,
Expand All @@ -440,25 +560,25 @@ def _forward_to_wrapped_kernel(self, code, silent, store_history, user_expressio
user_expressions=user_expressions,
allow_stdin=allow_stdin,
)

# Process messages from the wrapped kernel until we get the shell reply
# We need to process iopub messages while waiting for the reply
execution_error = None
got_idle = False
result = None

# Process iopub messages until we see idle status
while not got_idle:
try:
msg = self._kernel_client.get_iopub_msg(timeout=0.1)
msg_type = msg['header']['msg_type']
content = msg['content']

# Only process messages for our execution request
parent_msg_id = msg.get('parent_header', {}).get('msg_id')
if parent_msg_id != msg_id:
continue

if msg_type == 'status':
if content.get('execution_state') == 'idle':
got_idle = True
Expand All @@ -467,15 +587,15 @@ def _forward_to_wrapped_kernel(self, code, silent, store_history, user_expressio
self.send_response(self.iopub_socket, msg_type, content)
elif msg_type in ('stream', 'display_data', 'execute_result'):
self.send_response(self.iopub_socket, msg_type, content)

except Exception:
# Timeout - check if shell reply is available
try:
reply = self._kernel_client.get_shell_msg(timeout=0.1)
result = reply['content']
except Exception:
pass

# If we didn't get the shell reply yet, get it now
if result is None:
try:
Expand All @@ -490,24 +610,24 @@ def _forward_to_wrapped_kernel(self, code, silent, store_history, user_expressio
'traceback': [str(e)],
'execution_count': self.execution_count,
}

# Trigger post_run_cell event for jumper-extension
self._trigger_post_run_cell(
exec_info,
success=(result.get('status') == 'ok'),
error=execution_error
)

return result

def _execute_local_magic(self, code):
"""Execute a magic command locally using IPython."""
try:
result = self.shell.run_cell(code)

if result.success:
if result.result is not None:
self.send_response(self.iopub_socket, 'stream',
self.send_response(self.iopub_socket, 'stream',
{'name': 'stdout', 'text': str(result.result) + '\n'})
return {
'status': 'ok',
Expand Down Expand Up @@ -540,11 +660,11 @@ def _execute_local_magic(self, code):
'traceback': [str(e)],
'execution_count': self.execution_count,
}

def _is_local_magic(self, code):
"""Check if code is a magic command that should be executed locally."""
return is_local_magic_cell(code, self._get_local_magics())

def do_execute(self, code, silent, store_history=True, user_expressions=None, allow_stdin=False):
"""Execute code - either locally or forwarded to wrapped kernel."""
user_expressions = user_expressions or {}
Expand All @@ -560,12 +680,12 @@ def do_execute(self, code, silent, store_history=True, user_expressions=None, al

# Forward everything else to the wrapped kernel
return self._forward_to_wrapped_kernel(code, silent, store_history, user_expressions, allow_stdin)

def do_shutdown(self, restart):
"""Shutdown the kernel."""
self._shutdown_wrapped_kernel()
return {'status': 'ok', 'restart': restart}

def do_complete(self, code, cursor_pos):
"""Handle code completion - forward to wrapped kernel if available."""
# Perform deferred auto-wrap if pending
Expand All @@ -580,15 +700,15 @@ def do_complete(self, code, cursor_pos):
return reply['content']
except Exception:
pass

return {
'status': 'ok',
'matches': [],
'cursor_start': cursor_pos,
'cursor_end': cursor_pos,
'metadata': {},
}

def do_inspect(self, code, cursor_pos, detail_level=0):
"""Handle object inspection - forward to wrapped kernel if available."""
# Perform deferred auto-wrap if pending
Expand All @@ -603,14 +723,14 @@ def do_inspect(self, code, cursor_pos, detail_level=0):
return reply['content']
except Exception:
pass

return {
'status': 'ok',
'found': False,
'data': {},
'metadata': {},
}

@property
def kernel_info(self):
"""Return kernel info with current language_info (may be from wrapped kernel)."""
Expand Down