forked from emacs-jupyter/jupyter
-
Notifications
You must be signed in to change notification settings - Fork 0
/
jupyter-server-kernel.el
375 lines (325 loc) · 15.4 KB
/
jupyter-server-kernel.el
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
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
;;; jupyter-server-kernel.el --- Working with kernels behind a Jupyter server -*- lexical-binding: t -*-
;; Copyright (C) 2020-2024 Nathaniel Nicandro
;; Author: Nathaniel Nicandro <[email protected]>
;; Created: 23 Apr 2020
;; This program is free software; you can redistribute it and/or
;; modify it under the terms of the GNU General Public License as
;; published by the Free Software Foundation; either version 3, or (at
;; your option) any later version.
;; This program is distributed in the hope that it will be useful, but
;; WITHOUT ANY WARRANTY; without even the implied warranty of
;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
;; General Public License for more details.
;; You should have received a copy of the GNU General Public License
;; along with GNU Emacs; see the file COPYING. If not, write to the
;; Free Software Foundation, Inc., 59 Temple Place - Suite 330,
;; Boston, MA 02111-1307, USA.
;;; Commentary:
;; Holds the definitions of `jupyter-server', what communicates to the
;; Jupyter server using the REST API, and `jupyter-kernel-server' a
;; representation of a kernel on a server.
;;; Code:
(require 'jupyter-kernel)
(require 'jupyter-rest-api)
(require 'jupyter-monads)
(require 'websocket)
(declare-function jupyter-encode-raw-message "jupyter-messages")
(declare-function jupyter-tramp-server-from-file-name "jupyter-tramp")
(declare-function jupyter-tramp-file-name-p "jupyter-tramp")
(declare-function jupyter-server-kernel-id-from-name "jupyter-server")
(defgroup jupyter-server-kernel nil
"Kernel behind a Jupyter server"
:group 'jupyter)
;;; `jupyter-server'
(defvar-local jupyter-current-server nil
"The `jupyter-server' associated with the current buffer.
Used in, e.g. a `jupyter-server-kernel-list-mode' buffer.")
(put 'jupyter-current-server 'permanent-local t)
(defvar jupyter--servers nil)
;; TODO: We should really rename `jupyter-server' to something like
;; `jupyter-server-client' since it isn't a representation of a server, but a
;; communication channel with one.
(defclass jupyter-server (jupyter-rest-client eieio-instance-tracker)
((tracking-symbol :initform 'jupyter--servers)
(kernelspecs
:type json-plist
:initform nil
:documentation "Kernelspecs for the kernels available behind
this gateway. Access them through `jupyter-kernelspecs'.")))
(cl-defmethod make-instance ((_class (subclass jupyter-server)) &rest slots)
(cl-assert (plist-get slots :url))
(or (cl-loop
with url = (plist-get slots :url)
for server in jupyter--servers
if (equal url (oref server url)) return server)
(cl-call-next-method)))
(defun jupyter-servers ()
"Return a list of all `jupyter-server's."
(jupyter-gc-servers)
jupyter--servers)
(defun jupyter-gc-servers ()
"Delete `jupyter-server' instances that are no longer accessible."
(dolist (server jupyter--servers)
(unless (jupyter-api-server-exists-p server)
(jupyter-api-delete-cookies (oref server url))
(delete-instance server))))
(cl-defmethod jupyter-api-request :around ((server jupyter-server) _method &rest _plist)
(condition-case nil
(cl-call-next-method)
(jupyter-api-unauthenticated
(if (memq jupyter-api-authentication-method '(ask token password))
(oset server auth jupyter-api-authentication-method)
(error "Unauthenticated request, can't attempt re-authentication \
with default `jupyter-api-authentication-method'"))
(prog1 (cl-call-next-method)
(jupyter-reauthenticate-websockets server)))))
(cl-defmethod jupyter-kernelspecs ((client jupyter-rest-client) &optional _refresh)
(or (jupyter-api-get-kernelspec client)
(error "Can't retrieve kernelspecs from server @ %s"
(oref client url))))
(cl-defmethod jupyter-kernelspecs ((server jupyter-server) &optional refresh)
"Return the kernelspecs on SERVER.
By default the available kernelspecs are cached. To force an
update of the cached kernelspecs, give a non-nil value to
REFRESH."
(when (or refresh (null (oref server kernelspecs)))
(let ((specs (cl-call-next-method)))
(plist-put specs :kernelspecs
(cl-loop
for (_ spec) on (plist-get specs :kernelspecs) by #'cddr
for name = (plist-get spec :name)
collect (make-jupyter-kernelspec
:name name
:plist (plist-get spec :spec))))
(oset server kernelspecs specs)))
(plist-get (oref server kernelspecs) :kernelspecs))
(cl-defmethod jupyter-kernelspecs :extra "server" ((host string) &optional refresh)
(if (jupyter-tramp-file-name-p host)
(jupyter-kernelspecs (jupyter-tramp-server-from-file-name host) refresh)
(cl-call-next-method)))
(cl-defmethod jupyter-server-has-kernelspec-p ((server jupyter-server) name)
"Return non-nil if SERVER can launch kernels with kernelspec NAME."
(jupyter-guess-kernelspec name (jupyter-kernelspecs server)))
;;; Kernel definition
(cl-defstruct (jupyter-server-kernel
(:include jupyter-kernel))
(server jupyter-current-server
:read-only t
:documentation "The kernel server.")
;; TODO: Make this read only by only allowing creating
;; representations of kernels that have already been launched and
;; have a connection to the kernel.
(id nil
:type (or null string)
:documentation "The kernel ID."))
(cl-defmethod jupyter-alive-p ((kernel jupyter-server-kernel))
(pcase-let (((cl-struct jupyter-server-kernel server id) kernel))
(and id server
;; TODO: Cache this call
(condition-case err
(jupyter-api-get-kernel server id)
(file-error nil) ; Non-existent server
(jupyter-api-http-error
(unless (= (nth 1 err) 404) ; Not Found
(signal (car err) (cdr err)))))
(cl-call-next-method))))
(defun jupyter-server-kernel (&rest args)
"Return a `jupyter-server-kernel' initialized with ARGS."
(apply #'make-jupyter-server-kernel args))
(cl-defmethod jupyter-kernel :extra "server" (&rest args)
"Return a representation of a kernel on a Jupyter server.
If ARGS has a :server key, return a `jupyter-server-kernel'
initialized using ARGS. If ARGS also has a :spec key, whose
value is the name of a kernelspec, the returned kernel's spec
slot will be the corresponding `jupyter-kernelspec'.
Call the next method if ARGS does not contain :server."
(let ((server (plist-get args :server)))
(if (not server) (cl-call-next-method)
(cl-assert (object-of-class-p server 'jupyter-server))
(let ((spec (plist-get args :spec)))
(when (stringp spec)
(plist-put args :spec
;; TODO: (jupyter-server-kernelspec server "python3")
;; which returns an I/O action and then arrange
;; for that action to be bound by mlet* and set
;; as the spec value. Or better yet, have
;; `jupyter-kernel' return a delayed kernel with
;; the server connection already open and
;; kernelspecs already retrieved.
(or (jupyter-guess-kernelspec
spec (jupyter-kernelspecs server))
;; TODO: Return the error to the I/O context.
(error "No kernelspec matching %s @ %s" spec
(oref server url))))))
(apply #'jupyter-server-kernel args))))
;;; Websocket IO
(defvar jupyter--reauth-subscribers (make-hash-table :weakness 'key :test 'eq))
(defun jupyter-reauthenticate-websockets (server)
"Re-authenticate WebSocket connections of SERVER."
(when-let* ((pub (gethash server jupyter--reauth-subscribers)))
(jupyter-run-with-io pub
(jupyter-publish 'reauthenticate))))
(cl-defmethod jupyter-websocket-io ((kernel jupyter-server-kernel))
"Return a list representing an IO connection to KERNEL.
The list is composed of two elements (IO-PUB ACTION-SUB), IO-PUB
is a publisher used to send/receive messages to/from KERNEL and
ACTION-SUB is a subscriber of kernel actions to perform on
KERNEL.
To send a message to KERNEL, publish a list of the form
(list \='send CHANNEL MSG-TYPE CONTENT MSG-ID)
to IO-PUB, e.g.
(jupyter-run-with-io IO-PUB
(jupyter-publish (list \='send CHANNEL MSG-TYPE CONTENT MSG-ID)))
To receive messages from KERNEL, subscribe to IO-PUB e.g.
(jupyter-run-with-io IO-PUB
(jupyter-subscribe
(jupyter-subscriber
(lambda (msg)
...))))
The value \='interrupt or \='shutdown can be published to ACTION-SUB
to interrupt or shutdown KERNEL. The value (list \='action FN)
where FN is a single argument function can also be published, in
this case FN will be evaluated on KERNEL."
(jupyter-launch kernel)
(pcase-let* (((cl-struct jupyter-server-kernel server id) kernel))
(letrec ((status-pub (jupyter-publisher))
(reauth-pub (or (gethash server jupyter--reauth-subscribers)
(setf (gethash server jupyter--reauth-subscribers)
(jupyter-publisher))))
(shutdown nil)
(kernel-io
(jupyter-publisher
(lambda (event)
(pcase event
(`(message . ,rest) (jupyter-content rest))
(`(send ,channel ,msg-type ,content ,msg-id)
(when shutdown
(error "Attempting to send message to shutdown kernel"))
(let ((send
(lambda ()
(websocket-send-text
ws (let* ((cd (websocket-client-data ws))
(session (plist-get cd :session)))
(jupyter-encode-raw-message session msg-type
:channel channel
:msg-id msg-id
:content content))))))
(condition-case nil
(funcall send)
(websocket-closed
(setq ws (funcall make-websocket))
(funcall send)))))
('start
(when shutdown
(error "Can't start I/O connection to shutdown kernel"))
(unless (websocket-openp ws)
(setq ws (funcall make-websocket))))
('stop (websocket-close ws))))))
(ws-failed-to-open t)
(make-websocket
(lambda ()
(jupyter-api-kernel-websocket
server id
:custom-header-alist (jupyter-api-auth-headers server)
:on-open
(lambda (_ws)
(setq ws-failed-to-open nil))
:on-close
(lambda (_ws)
(if ws-failed-to-open
;; TODO: Retry?
(error "Kernel connection could not be established")
(setq ws-failed-to-open t)))
;; TODO: on-error publishes to status-pub
:on-message
(lambda (_ws frame)
(pcase (websocket-frame-opcode frame)
((or 'text 'binary)
(let ((msg (jupyter-read-plist-from-string
(websocket-frame-payload frame))))
(jupyter-run-with-io kernel-io
(jupyter-publish (cons 'message msg)))))
(_
(jupyter-run-with-io status-pub
(jupyter-publish
(list 'error (websocket-frame-opcode frame))))))))))
(ws (prog1 (funcall make-websocket)
(jupyter-run-with-io reauth-pub
(jupyter-subscribe
(jupyter-subscriber
(lambda (_reauth)
(if shutdown (jupyter-unsubscribe)
(jupyter-run-with-io kernel-io
(jupyter-do
(jupyter-publish 'stop)
(jupyter-publish 'start)))))))))))
(list kernel-io
(jupyter-subscriber
(lambda (action)
(pcase action
('interrupt
(jupyter-interrupt kernel))
('shutdown
(jupyter-shutdown kernel)
(setq shutdown t)
(when (websocket-openp ws)
(websocket-close ws)))
('restart
(jupyter-restart kernel))
(`(action ,fn)
(funcall fn kernel)))))))))
(cl-defmethod jupyter-io ((kernel jupyter-server-kernel))
(jupyter-websocket-io kernel))
;;; Kernel management
;; The KERNEL argument is optional here so that `jupyter-launch'
;; does not require more than one argument just to handle this case.
(cl-defmethod jupyter-launch ((server jupyter-server) &optional kernel)
(cl-check-type kernel string)
(let* ((spec (jupyter-guess-kernelspec
kernel (jupyter-kernelspecs server)))
(plist (jupyter-api-start-kernel
server (jupyter-kernelspec-name spec))))
(jupyter-kernel :server server :id (plist-get plist :id) :spec spec)))
;; FIXME: Don't allow creating kernels without them being launched.
(cl-defmethod jupyter-launch ((kernel jupyter-server-kernel))
"Launch KERNEL based on its kernelspec.
When KERNEL does not have an ID yet, launch KERNEL on SERVER
using its SPEC."
(pcase-let (((cl-struct jupyter-server-kernel server id spec session) kernel))
(unless session
(and id (setq id (or (jupyter-server-kernel-id-from-name server id) id)))
(if id
;; When KERNEL already has an ID before it has a session,
;; assume we are connecting to an already launched kernel. In
;; this case, make sure the KERNEL's SPEC is the same as the
;; one being connected to.
;;
;; Note, this also has the side effect of raising an error
;; when the ID does not match one on the server.
(unless spec
(let ((model (jupyter-api-get-kernel server id)))
(setf (jupyter-kernel-spec kernel)
(jupyter-guess-kernelspec
(plist-get model :name)
(jupyter-kernelspecs server)))))
(let ((plist (jupyter-api-start-kernel
server (jupyter-kernelspec-name spec))))
(setf (jupyter-server-kernel-id kernel) (plist-get plist :id))
(sit-for 1)))
;; TODO: Replace with the real session object
(setf (jupyter-kernel-session kernel) (jupyter-session))))
(cl-call-next-method))
(cl-defmethod jupyter-shutdown ((kernel jupyter-server-kernel))
(pcase-let (((cl-struct jupyter-server-kernel server id session) kernel))
(cl-call-next-method)
(when session
(jupyter-api-shutdown-kernel server id))))
(cl-defmethod jupyter-restart ((kernel jupyter-server-kernel))
(pcase-let (((cl-struct jupyter-server-kernel server id session) kernel))
(when session
(jupyter-api-restart-kernel server id))))
(cl-defmethod jupyter-interrupt ((kernel jupyter-server-kernel))
(pcase-let (((cl-struct jupyter-server-kernel server id) kernel))
(jupyter-api-interrupt-kernel server id)))
(provide 'jupyter-server-kernel)
;;; jupyter-server-kernel.el ends here