Skip to content

Commit 66c5403

Browse files
SylvainCorlaySylvainCorlay
committed
Nudge kernel on restart
1 parent 364f45a commit 66c5403

File tree

3 files changed

+232
-17
lines changed

3 files changed

+232
-17
lines changed

‎Untitled.ipynb

+185
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,185 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "code",
5+
"execution_count": 1,
6+
"metadata": {},
7+
"outputs": [
8+
{
9+
"name": "stdout",
10+
"output_type": "stream",
11+
"text": [
12+
"1\n"
13+
]
14+
}
15+
],
16+
"source": [
17+
"print(1)"
18+
]
19+
},
20+
{
21+
"cell_type": "code",
22+
"execution_count": 2,
23+
"metadata": {},
24+
"outputs": [
25+
{
26+
"name": "stdout",
27+
"output_type": "stream",
28+
"text": [
29+
"1\n"
30+
]
31+
}
32+
],
33+
"source": [
34+
"print(1)"
35+
]
36+
},
37+
{
38+
"cell_type": "code",
39+
"execution_count": 3,
40+
"metadata": {},
41+
"outputs": [
42+
{
43+
"name": "stdout",
44+
"output_type": "stream",
45+
"text": [
46+
"1\n"
47+
]
48+
}
49+
],
50+
"source": [
51+
"print(1)"
52+
]
53+
},
54+
{
55+
"cell_type": "code",
56+
"execution_count": 4,
57+
"metadata": {},
58+
"outputs": [
59+
{
60+
"name": "stdout",
61+
"output_type": "stream",
62+
"text": [
63+
"1\n"
64+
]
65+
}
66+
],
67+
"source": [
68+
"print(1)"
69+
]
70+
},
71+
{
72+
"cell_type": "code",
73+
"execution_count": 5,
74+
"metadata": {},
75+
"outputs": [
76+
{
77+
"name": "stdout",
78+
"output_type": "stream",
79+
"text": [
80+
"1\n"
81+
]
82+
}
83+
],
84+
"source": [
85+
"print(1)"
86+
]
87+
},
88+
{
89+
"cell_type": "code",
90+
"execution_count": 6,
91+
"metadata": {},
92+
"outputs": [
93+
{
94+
"name": "stdout",
95+
"output_type": "stream",
96+
"text": [
97+
"1\n"
98+
]
99+
}
100+
],
101+
"source": [
102+
"print(1)"
103+
]
104+
},
105+
{
106+
"cell_type": "code",
107+
"execution_count": 7,
108+
"metadata": {},
109+
"outputs": [
110+
{
111+
"name": "stdout",
112+
"output_type": "stream",
113+
"text": [
114+
"1\n"
115+
]
116+
}
117+
],
118+
"source": [
119+
"print(1)"
120+
]
121+
},
122+
{
123+
"cell_type": "code",
124+
"execution_count": 8,
125+
"metadata": {},
126+
"outputs": [
127+
{
128+
"name": "stdout",
129+
"output_type": "stream",
130+
"text": [
131+
"1\n"
132+
]
133+
}
134+
],
135+
"source": [
136+
"print(1)"
137+
]
138+
},
139+
{
140+
"cell_type": "code",
141+
"execution_count": 9,
142+
"metadata": {},
143+
"outputs": [
144+
{
145+
"name": "stdout",
146+
"output_type": "stream",
147+
"text": [
148+
"1\n"
149+
]
150+
}
151+
],
152+
"source": [
153+
"print(1)"
154+
]
155+
},
156+
{
157+
"cell_type": "code",
158+
"execution_count": null,
159+
"metadata": {},
160+
"outputs": [],
161+
"source": []
162+
}
163+
],
164+
"metadata": {
165+
"kernelspec": {
166+
"display_name": "Python 3",
167+
"language": "python",
168+
"name": "python3"
169+
},
170+
"language_info": {
171+
"codemirror_mode": {
172+
"name": "ipython",
173+
"version": 3
174+
},
175+
"file_extension": ".py",
176+
"mimetype": "text/x-python",
177+
"name": "python",
178+
"nbconvert_exporter": "python",
179+
"pygments_lexer": "ipython3",
180+
"version": "3.7.3"
181+
}
182+
},
183+
"nbformat": 4,
184+
"nbformat_minor": 4
185+
}

‎notebook/services/kernels/handlers.py

+1-2
Original file line numberDiff line numberDiff line change
@@ -78,9 +78,8 @@ def post(self, kernel_id, action):
7878
yield maybe_future(km.interrupt_kernel(kernel_id))
7979
self.set_status(204)
8080
if action == 'restart':
81-
8281
try:
83-
yield maybe_future(km.restart_kernel(kernel_id))
82+
yield maybe_future(km.restart_kernel(kernel_id, channels))
8483
except Exception as e:
8584
self.log.error("Exception restarting kernel", exc_info=True)
8685
self.set_status(500)

‎notebook/services/kernels/kernelmanager.py

+46-15
Original file line numberDiff line numberDiff line change
@@ -304,33 +304,54 @@ def shutdown_kernel(self, kernel_id, now=False, restart=False):
304304

305305
return self.pinned_superclass.shutdown_kernel(self, kernel_id, now=now, restart=restart)
306306

307-
async def restart_kernel(self, kernel_id, now=False):
307+
async def restart_kernel(self, kernel_id, channels, now=False):
308308
"""Restart a kernel by kernel_id"""
309309
self._check_kernel_id(kernel_id)
310310
await maybe_future(self.pinned_superclass.restart_kernel(self, kernel_id, now=now))
311311
kernel = self.get_kernel(kernel_id)
312312
# return a Future that will resolve when the kernel has successfully restarted
313-
channel = kernel.connect_shell()
313+
shell_channel = self.channels['shell']
314+
iopub_channel = self.channels['iopub']
315+
314316
future = Future()
317+
info_future = Future()
318+
iopub_future = Future()
315319

316320
def finish():
317-
"""Common cleanup when restart finishes/fails for any reason."""
318-
if not channel.closed():
319-
channel.close()
321+
"""Common cleanup"""
320322
loop.remove_timeout(timeout)
323+
loop.remove_timeout(nudge_handle)
324+
iopub_channel.stop_on_recv()
325+
shell_channel.stop_on_recv()
321326
kernel.remove_restart_callback(on_restart_failed, 'dead')
322327

323-
def on_reply(msg):
324-
self.log.debug("Kernel info reply received: %s", kernel_id)
325-
finish()
326-
if not future.done():
327-
future.set_result(msg)
328+
def on_shell_reply(msg):
329+
if not info_future.done():
330+
self.log.debug("Nudge: shell info reply received: %s", self.kernel_id)
331+
shell_channel.stop_on_recv()
332+
self.log.debug("Nudge: resolving shell future")
333+
info_future.set_result(msg)
334+
if iopub_future.done():
335+
finish()
336+
self.log.debug("Nudge: resolving main future in shell handler")
337+
future.set_result(info_future.result())
338+
339+
def on_iopub(msg):
340+
if not iopub_future.done():
341+
self.log.debug("Nudge: first IOPub received: %s", self.kernel_id)
342+
iopub_channel.stop_on_recv()
343+
self.log.debug("Nudge: resolving iopub future")
344+
iopub_future.set_result(None)
345+
if info_future.done():
346+
finish()
347+
self.log.debug("Nudge: resolving main future in iopub handler")
348+
future.set_result(info_future.result())
328349

329350
def on_timeout():
330-
self.log.warning("Timeout waiting for kernel_info_reply: %s", kernel_id)
351+
self.log.warning("Nudge: Timeout waiting for kernel_info_reply: %s", self.kernel_id)
331352
finish()
332353
if not future.done():
333-
future.set_exception(TimeoutError("Timeout waiting for restart"))
354+
future.set_exception(TimeoutError("Timeout waiting for nudge"))
334355

335356
def on_restart_failed():
336357
self.log.warning("Restarting kernel failed: %s", kernel_id)
@@ -339,10 +360,20 @@ def on_restart_failed():
339360
future.set_exception(RuntimeError("Restart failed"))
340361

341362
kernel.add_restart_callback(on_restart_failed, 'dead')
342-
kernel.session.send(channel, "kernel_info_request")
343-
channel.on_recv(on_reply)
363+
364+
iopub_channel.on_recv(on_iopub)
365+
shell_channel.on_recv(on_shell_reply)
344366
loop = IOLoop.current()
345-
timeout = loop.add_timeout(loop.time() + self.kernel_info_timeout, on_timeout)
367+
368+
# Nudge the kernel with kernel info requests until we get an IOPub message
369+
def nudge():
370+
self.log.debug("Nudge")
371+
if not future.done():
372+
self.log.debug("nudging")
373+
self.session.send(shell_channel, "kernel_info_request")
374+
nudge_handle = loop.call_later(0.5, nudge)
375+
nudge_handle = loop.call_later(0, nudge)
376+
346377
return future
347378

348379
def notify_connect(self, kernel_id):

0 commit comments

Comments
 (0)