forked from jupyter/jupyter_client
-
Notifications
You must be signed in to change notification settings - Fork 0
/
threaded.py
350 lines (286 loc) · 10.9 KB
/
threaded.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
""" Defines a KernelClient that provides thread-safe sockets with async callbacks on message
replies.
"""
import asyncio
import atexit
import time
from concurrent.futures import Future
from functools import partial
from threading import Thread
from typing import Any, Dict, List, Optional
import zmq
from tornado.ioloop import IOLoop
from traitlets import Instance, Type
from traitlets.log import get_logger
from zmq.eventloop import zmqstream
from .channels import HBChannel
from .client import KernelClient
from .session import Session
# Local imports
# import ZMQError in top-level namespace, to avoid ugly attribute-error messages
# during garbage collection of threads at exit
class ThreadedZMQSocketChannel:
"""A ZMQ socket invoking a callback in the ioloop"""
session = None
socket = None
ioloop = None
stream = None
_inspect = None
def __init__(
self,
socket: Optional[zmq.Socket],
session: Optional[Session],
loop: Optional[IOLoop],
) -> None:
"""Create a channel.
Parameters
----------
socket : :class:`zmq.Socket`
The ZMQ socket to use.
session : :class:`session.Session`
The session to use.
loop
A tornado ioloop to connect the socket to using a ZMQStream
"""
super().__init__()
self.socket = socket
self.session = session
self.ioloop = loop
f: Future = Future()
def setup_stream():
try:
assert self.socket is not None
self.stream = zmqstream.ZMQStream(self.socket, self.ioloop)
self.stream.on_recv(self._handle_recv)
except Exception as e:
f.set_exception(e)
else:
f.set_result(None)
assert self.ioloop is not None
self.ioloop.add_callback(setup_stream)
# don't wait forever, raise any errors
f.result(timeout=10)
_is_alive = False
def is_alive(self) -> bool:
"""Whether the channel is alive."""
return self._is_alive
def start(self) -> None:
"""Start the channel."""
self._is_alive = True
def stop(self) -> None:
"""Stop the channel."""
self._is_alive = False
def close(self) -> None:
"""Close the channel."""
if self.stream is not None and self.ioloop is not None:
# c.f.Future for threadsafe results
f: Future = Future()
def close_stream():
try:
if self.stream is not None:
self.stream.close(linger=0)
self.stream = None
except Exception as e:
f.set_exception(e)
else:
f.set_result(None)
self.ioloop.add_callback(close_stream)
# wait for result
try:
f.result(timeout=5)
except Exception as e:
log = get_logger()
msg = f"Error closing stream {self.stream}: {e}"
log.warning(msg, RuntimeWarning, stacklevel=2)
if self.socket is not None:
try:
self.socket.close(linger=0)
except Exception:
pass
self.socket = None
def send(self, msg: Dict[str, Any]) -> None:
"""Queue a message to be sent from the IOLoop's thread.
Parameters
----------
msg : message to send
This is threadsafe, as it uses IOLoop.add_callback to give the loop's
thread control of the action.
"""
def thread_send():
assert self.session is not None
self.session.send(self.stream, msg)
assert self.ioloop is not None
self.ioloop.add_callback(thread_send)
def _handle_recv(self, msg_list: List) -> None:
"""Callback for stream.on_recv.
Unpacks message, and calls handlers with it.
"""
assert self.ioloop is not None
assert self.session is not None
ident, smsg = self.session.feed_identities(msg_list)
msg = self.session.deserialize(smsg)
# let client inspect messages
if self._inspect:
self._inspect(msg) # type:ignore[unreachable]
self.call_handlers(msg)
def call_handlers(self, msg: Dict[str, Any]) -> None:
"""This method is called in the ioloop thread when a message arrives.
Subclasses should override this method to handle incoming messages.
It is important to remember that this method is called in the thread
so that some logic must be done to ensure that the application level
handlers are called in the application thread.
"""
pass
def process_events(self) -> None:
"""Subclasses should override this with a method
processing any pending GUI events.
"""
pass
def flush(self, timeout: float = 1.0) -> None:
"""Immediately processes all pending messages on this channel.
This is only used for the IOPub channel.
Callers should use this method to ensure that :meth:`call_handlers`
has been called for all messages that have been received on the
0MQ SUB socket of this channel.
This method is thread safe.
Parameters
----------
timeout : float, optional
The maximum amount of time to spend flushing, in seconds. The
default is one second.
"""
# We do the IOLoop callback process twice to ensure that the IOLoop
# gets to perform at least one full poll.
stop_time = time.monotonic() + timeout
assert self.ioloop is not None
if self.stream is None or self.stream.closed():
# don't bother scheduling flush on a thread if we're closed
_msg = "Attempt to flush closed stream"
raise OSError(_msg)
def flush(f):
try:
self._flush()
except Exception as e:
f.set_exception(e)
else:
f.set_result(None)
for _ in range(2):
f: Future = Future()
self.ioloop.add_callback(partial(flush, f))
# wait for async flush, re-raise any errors
timeout = max(stop_time - time.monotonic(), 0)
try:
f.result(max(stop_time - time.monotonic(), 0))
except TimeoutError:
# flush with a timeout means stop waiting, not raise
return
def _flush(self) -> None:
"""Callback for :method:`self.flush`."""
assert self.stream is not None
self.stream.flush()
self._flushed = True
class IOLoopThread(Thread):
"""Run a pyzmq ioloop in a thread to send and receive messages"""
_exiting = False
ioloop = None
def __init__(self):
"""Initialize an io loop thread."""
super().__init__()
self.daemon = True
@staticmethod
@atexit.register
def _notice_exit() -> None:
# Class definitions can be torn down during interpreter shutdown.
# We only need to set _exiting flag if this hasn't happened.
if IOLoopThread is not None:
IOLoopThread._exiting = True
def start(self) -> None:
"""Start the IOLoop thread
Don't return until self.ioloop is defined,
which is created in the thread
"""
self._start_future: Future = Future()
Thread.start(self)
# wait for start, re-raise any errors
self._start_future.result(timeout=10)
def run(self) -> None:
"""Run my loop, ignoring EINTR events in the poller"""
try:
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
async def assign_ioloop():
self.ioloop = IOLoop.current()
loop.run_until_complete(assign_ioloop())
except Exception as e:
self._start_future.set_exception(e)
else:
self._start_future.set_result(None)
loop.run_until_complete(self._async_run())
async def _async_run(self):
"""Run forever (until self._exiting is set)"""
while not self._exiting:
await asyncio.sleep(1)
def stop(self) -> None:
"""Stop the channel's event loop and join its thread.
This calls :meth:`~threading.Thread.join` and returns when the thread
terminates. :class:`RuntimeError` will be raised if
:meth:`~threading.Thread.start` is called again.
"""
self._exiting = True
self.join()
self.close()
self.ioloop = None
def __del__(self):
self.close()
def close(self) -> None:
"""Close the io loop thread."""
if self.ioloop is not None:
try:
self.ioloop.close(all_fds=True)
except Exception:
pass
class ThreadedKernelClient(KernelClient):
"""A KernelClient that provides thread-safe sockets with async callbacks on message replies."""
@property
def ioloop(self):
if self.ioloop_thread:
return self.ioloop_thread.ioloop
ioloop_thread = Instance(IOLoopThread, allow_none=True)
def start_channels(
self,
shell: bool = True,
iopub: bool = True,
stdin: bool = True,
hb: bool = True,
control: bool = True,
) -> None:
"""Start the channels on the client."""
self.ioloop_thread = IOLoopThread()
self.ioloop_thread.start()
if shell:
self.shell_channel._inspect = self._check_kernel_info_reply
super().start_channels(shell, iopub, stdin, hb, control)
def _check_kernel_info_reply(self, msg: Dict[str, Any]) -> None:
"""This is run in the ioloop thread when the kernel info reply is received"""
if msg["msg_type"] == "kernel_info_reply":
self._handle_kernel_info_reply(msg)
self.shell_channel._inspect = None
def stop_channels(self) -> None:
"""Stop the channels on the client."""
super().stop_channels()
if self.ioloop_thread and self.ioloop_thread.is_alive():
self.ioloop_thread.stop()
iopub_channel_class = Type(ThreadedZMQSocketChannel) # type:ignore[arg-type]
shell_channel_class = Type(ThreadedZMQSocketChannel) # type:ignore[arg-type]
stdin_channel_class = Type(ThreadedZMQSocketChannel) # type:ignore[arg-type]
hb_channel_class = Type(HBChannel) # type:ignore[arg-type]
control_channel_class = Type(ThreadedZMQSocketChannel) # type:ignore[arg-type]
def is_alive(self) -> bool:
"""Is the kernel process still running?"""
if self._hb_channel is not None:
# We don't have access to the KernelManager,
# so we use the heartbeat.
return self._hb_channel.is_beating()
# no heartbeat and not local, we can't tell if it's running,
# so naively return True
return True