Skip to content

Commit 68c4b7e

Browse files
committed
wip
1 parent 1a12302 commit 68c4b7e

File tree

3 files changed

+31
-15
lines changed

3 files changed

+31
-15
lines changed

jupyter_server/pytest_plugin.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -472,11 +472,11 @@ def inner(nbpath):
472472

473473

474474
@pytest.fixture(autouse=True)
475-
def jp_server_cleanup():
475+
def jp_server_cleanup(io_loop):
476476
yield
477477
app: ServerApp = ServerApp.instance()
478-
loop = asyncio.get_event_loop_policy().get_event_loop()
479-
loop.run_until_complete(app._stop())
478+
loop = io_loop.asyncio_loop
479+
loop.run_until_complete(app._cleanup())
480480
ServerApp.clear_instance()
481481

482482

jupyter_server/services/kernels/kernelmanager.py

Lines changed: 26 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,7 @@ def _default_kernel_buffers(self):
149149

150150
def __init__(self, **kwargs):
151151
self.pinned_superclass = MultiKernelManager
152+
self._pending_kernel_tasks = {}
152153
self.pinned_superclass.__init__(self, **kwargs)
153154
self.last_kernel_activity = utcnow()
154155

@@ -216,9 +217,11 @@ async def start_kernel(self, kernel_id=None, path=None, **kwargs):
216217
kwargs["kernel_id"] = kernel_id
217218
kernel_id = await ensure_async(self.pinned_superclass.start_kernel(self, **kwargs))
218219
self._kernel_connections[kernel_id] = 0
219-
fut = asyncio.ensure_future(self._finish_kernel_start(kernel_id))
220+
task = asyncio.create_task(self._finish_kernel_start(kernel_id))
220221
if not getattr(self, "use_pending_kernels", None):
221-
await fut
222+
await task
223+
else:
224+
self._pending_kernel_tasks[kernel_id] = task
222225
# add busy/activity markers:
223226
kernel = self.get_kernel(kernel_id)
224227
kernel.execution_state = "starting"
@@ -372,7 +375,7 @@ def stop_buffering(self, kernel_id):
372375
buffer_info = self._kernel_buffers.pop(kernel_id)
373376
# close buffering streams
374377
for stream in buffer_info["channels"].values():
375-
if not stream.closed():
378+
if not stream.socket.closed:
376379
stream.on_recv(None)
377380
stream.close()
378381

@@ -387,13 +390,19 @@ def stop_buffering(self, kernel_id):
387390
def shutdown_kernel(self, kernel_id, now=False, restart=False):
388391
"""Shutdown a kernel by kernel_id"""
389392
self._check_kernel_id(kernel_id)
390-
self.stop_watching_activity(kernel_id)
391-
self.stop_buffering(kernel_id)
392393

393394
# Decrease the metric of number of kernels
394395
# running for the relevant kernel type by 1
395396
KERNEL_CURRENTLY_RUNNING_TOTAL.labels(type=self._kernels[kernel_id].kernel_name).dec()
396397

398+
if kernel_id in self._pending_kernel_tasks:
399+
task = self._pending_kernel_tasks.pop(kernel_id)
400+
task.cancel()
401+
return
402+
403+
self.stop_watching_activity(kernel_id)
404+
self.stop_buffering(kernel_id)
405+
397406
self.pinned_superclass.shutdown_kernel(self, kernel_id, now=now, restart=restart)
398407

399408
async def restart_kernel(self, kernel_id, now=False):
@@ -533,7 +542,8 @@ def stop_watching_activity(self, kernel_id):
533542
"""Stop watching IOPub messages on a kernel for activity."""
534543
kernel = self._kernels[kernel_id]
535544
if getattr(kernel, "_activity_stream", None):
536-
kernel._activity_stream.close()
545+
if not kernel._activity_stream.socket.closed:
546+
kernel._activity_stream.close()
537547
kernel._activity_stream = None
538548

539549
def initialize_culler(self):
@@ -638,19 +648,25 @@ def __init__(self, **kwargs):
638648
self.pinned_superclass = AsyncMultiKernelManager
639649
self.pinned_superclass.__init__(self, **kwargs)
640650
self.last_kernel_activity = utcnow()
651+
self._pending_kernel_tasks = {}
641652

642653
async def shutdown_kernel(self, kernel_id, now=False, restart=False):
643654
"""Shutdown a kernel by kernel_id"""
644655
self._check_kernel_id(kernel_id)
645-
self.stop_watching_activity(kernel_id)
646-
self.stop_buffering(kernel_id)
647656

648657
# Decrease the metric of number of kernels
649658
# running for the relevant kernel type by 1
650659
KERNEL_CURRENTLY_RUNNING_TOTAL.labels(type=self._kernels[kernel_id].kernel_name).dec()
651660

661+
if kernel_id in self._pending_kernel_tasks:
662+
task = self._pending_kernel_tasks.pop(kernel_id)
663+
task.cancel()
664+
return
665+
666+
self.stop_watching_activity(kernel_id)
667+
self.stop_buffering(kernel_id)
668+
652669
# Finish shutting down the kernel before clearing state to avoid a race condition.
653-
ret = await self.pinned_superclass.shutdown_kernel(
670+
return await self.pinned_superclass.shutdown_kernel(
654671
self, kernel_id, now=now, restart=restart
655672
)
656-
return ret

pyproject.toml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,9 @@ addopts = "-raXs --durations 10 --color=yes --doctest-modules"
1313
testpaths = [
1414
"tests/"
1515
]
16-
timeout = 300
16+
timeout = 30
1717
# Restore this setting to debug failures
18-
# timeout_method = "thread"
18+
timeout_method = "thread"
1919
filterwarnings = [
2020
"error",
2121
"ignore:There is no current event loop:DeprecationWarning",

0 commit comments

Comments
 (0)