guix-jupyter-kernel.scm 12.5 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
;;; Guix-kernel -- Guix kernel for Jupyter
;;; Copyright (C) 2018  Evgeny Panfilov <epanfilov@gmail.com>
;;; Copyright (C) 2018  Pierre-Antoine Rouby <pierre-antoine.rouby@inria.fr>
;;;
;;; 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 of the License, 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 this program.  If not, see <https://www.gnu.org/licenses/>.

18
19
(use-modules (json)
             (simple-zmq)
jerry40's avatar
jerry40 committed
20
	     (srfi srfi-1)
21
             (srfi srfi-11)
22
             (srfi srfi-13)
23
24
25
26
27
             (ice-9 vlist)
             (guix build utils)
             (guix-kernel tools)
             (guix-kernel hmac)
             (guix-kernel magic)
28
29
             (guix-kernel environ)
             (guix-kernel jupyter-client))
jerry40's avatar
jerry40 committed
30

31
32
33
34
;;
;; Notebook information.
;;

35
36
37
38
(define notebook-info
  (json->scm
   (open-input-file
    (car (last-pair (command-line))))))
jerry40's avatar
jerry40 committed
39

40
41
(define (get-notebook-info-atom name)
  (hash-ref notebook-info name))
jerry40's avatar
jerry40 committed
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
;; Parse json values.
(define notebook-info-control-port
  (get-notebook-info-atom "control_port"))

(define notebook-info-shell-port
  (get-notebook-info-atom "shell_port"))

(define notebook-info-transport
  (get-notebook-info-atom "transport"))

(define notebook-info-signature-scheme
  (get-notebook-info-atom "signature_scheme"))

(define notebook-info-stdin-port
  (get-notebook-info-atom "stdin_port"))

(define notebook-info-heartbeat-port
  (get-notebook-info-atom "hb_port"))

(define notebook-info-ip
  (get-notebook-info-atom "ip"))

(define notebook-info-iopub-port
  (get-notebook-info-atom "iopub_port"))

(define notebook-info-key
  (get-notebook-info-atom "key"))

71
72
73
74
(define (create-address port)
  (string-append notebook-info-transport
                 "://" notebook-info-ip ":"
                 (number->string port)))
jerry40's avatar
jerry40 committed
75

76
77
78
79
;;
;; ZeroMQ.
;;

80
;; ZeroMQ context.
jerry40's avatar
jerry40 committed
81
82
(define context (zmq-create-context))

83
;; Adresses.
jerry40's avatar
jerry40 committed
84
85
86
87
88
89
(define addr-heartbeat (create-address notebook-info-heartbeat-port))
(define addr-shell     (create-address notebook-info-shell-port))
(define addr-control   (create-address notebook-info-control-port))
(define addr-iopub     (create-address notebook-info-iopub-port))
(define addr-stdin     (create-address notebook-info-stdin-port))

90
;; Sockets.
91
92
93
94
95
(define socket-heartbeat (zmq-create-socket context ZMQ_REP))
(define socket-shell     (zmq-create-socket context ZMQ_ROUTER))
(define socket-control   (zmq-create-socket context ZMQ_ROUTER))
(define socket-iopub     (zmq-create-socket context ZMQ_PUB))
(define socket-stdin     (zmq-create-socket context ZMQ_ROUTER))
jerry40's avatar
jerry40 committed
96

97
;; Useful lists.
98
99
100
101
102
103
104
105
106
107
(define adresses (list addr-heartbeat
                       addr-shell
                       addr-control
                       addr-iopub
                       addr-stdin))
(define sockets  (list socket-heartbeat
                       socket-shell
                       socket-control
                       socket-iopub
                       socket-stdin))
jerry40's avatar
jerry40 committed
108

109
;; Bind sockets to addressess.
jerry40's avatar
jerry40 committed
110
111
(for-each zmq-bind-socket sockets adresses)

112
;;
113
;; Container tools.
114
115
;;

116
(define container-context (zmq-create-context))
117

118
119
(define container-path
  (string-replace (current-filename) "guix-jupyter-container.scm"
120
121
                  (+ (string-rindex (current-filename) #\/) 1)))

122
(define (register-container vhash name values)
123
  (vhash-cons name values vhash))
124

125
(define (container-by-name vhash name)
126
  (vhash-assoc name vhash))
127

128
(define (container-pid-by-name vhash name)
129
  (assoc-ref (vhash-assoc name vhash) "pid"))
130

131
(define (container-socket-by-name vhash name)
132
  (assoc-ref (vhash-assoc name vhash) "socket"))
133

134
135
136
137
;;
;; Send.
;;

138
(define (send-to-container containers socket name parts)
139
  (zmq-send-msg-parts socket (append (list name) parts)))
140

141
142
143
144
(define (proxy-exec-container containers name parts count)
  (let ((socket  (container-socket-by-name containers name)))
    ;; Send Jupyter command tu default container.
    (send-to-container containers socket name parts)
145
146
147
148
149
150
151
152
153
154
155
156
157

    ;; Get results.
    (let ((execute-input   (zmq-get-msg-parts socket))
          (error-or-result (zmq-get-msg-parts socket))
          (execute-reply   (zmq-get-msg-parts socket)))

      ;; Resend results to Jupyter.
      (zmq-send-msg-parts socket-iopub execute-input)
      (zmq-send-msg-parts socket-shell execute-reply)
      (zmq-send-msg-parts socket-iopub error-or-result)

      (display ".")
      (+ count 1))))
158

159
;;
160
;; Handlers.
161
162
;;

jerry40's avatar
jerry40 committed
163
(define (heartbeat-handler)
164
165
166
  (zmq-message-send socket-heartbeat
                    (zmq-message-receive socket-heartbeat
                                         (zmq-msg-init)))
jerry40's avatar
jerry40 committed
167
168
  (heartbeat-handler))

169
(define (general-handler socket containers count)
jerry40's avatar
jerry40 committed
170
  (let* ((parts (zmq-get-msg-parts socket))
171
172
173
174
175
176
         (wire-uuid          (car parts))
         (wire-delimiter     (cadr parts))
         (wire-signature     (caddr parts))
         (wire-header        (json-string->scm(list-ref parts 3)))
         (wire-parent-header (list-ref parts 4))
         (wire-metadata      (json-string->scm(list-ref parts 5)))
177
178
179
180
181
182
         (wire-content       (json-string->scm(list-ref parts 6)))
         (msg-type      (hash-ref wire-header "msg_type"))
         (msg-username  (hash-ref wire-header "username"))
         (msg-session   (hash-ref wire-header "session"))
         (msg-version   (hash-ref wire-header "version"))
         (header- (make-header msg-username msg-session msg-version)))
183
184
    (pub-busy socket-iopub header- (scm->json-string wire-header)
              notebook-info-key)
185
    (let-values (((containers count) ((dispatch msg-type) socket wire-uuid
186
187
188
                                      header- wire-parent-header
                                      (scm->json-string wire-metadata)
                                      wire-content parts
189
                                      containers count)))
190
191
      (pub-idle socket-iopub header- (scm->json-string wire-header)
                notebook-info-key)
192
      (general-handler socket containers count))))
jerry40's avatar
jerry40 committed
193

194
;; Unknown request type, ignore it.
195
(define (ignore-request socket uuid header- parent-header
196
                        metadata content parts containers
197
                        count)
198
  (values containers count))
jerry40's avatar
jerry40 committed
199

200
;; Send kernel-info.
201
(define (reply-kernel-info-request socket uuid header- parent-header
202
                                   metadata content parts containers
203
                                   count)
204
205
206
  (send-to-jupyter socket uuid (header- "kernel_info_reply")
                   parent-header metadata (scm->json-string KERNEL-INFO)
                   notebook-info-key)
207
  (values containers
208
          count))
209
210

(define (reply-execute-request socket uuid header- parent-header
211
                               metadata content parts containers
212
                               count)
213
214
  (let* ((code  (hash-ref content "code"))
         (magic (get-magic-line code)))
215
216
    (cond
     ((magic-env? code)                ;Guix environment magic command
217
218
219
      (let* ((list     (string-split magic #\ ))
             (env-name (list-ref list 2))
             (env      (list-cdr-ref list 3)))
220
        ;; Create new environment if not exist
221
222
223
        (if (container-by-name containers env-name)
            (values containers
                    (proxy-exec-container containers env-name
224
                                          parts count))
225
226
227
228
229
            (let* ((pid       (run-new-container env-name env))
                   (container (new-container-connect env-name pid))
                   (new       (register-container containers env-name
                                                  container)))
              (values new (proxy-exec-container new env-name
230
                                                parts count))))))
231
232
     ((or (magic-run? code)
          (magic-kernel? code))
233
234
      (let* ((list     (string-split magic #\ ))
             (env-name (list-ref list 2)))
235
236
        (values containers
                (proxy-exec-container containers env-name
237
                                      parts count))))
238
     (else
239
240
      (values containers
              (proxy-exec-container containers "default"
241
                                    parts count))))))
jerry40's avatar
jerry40 committed
242

243
(define (shutdown socket uuid header- parent-header
244
                  metadata content parts containers
245
                  count)
246
  (display "Shutdown kernel !\n")
247
  (kill-containers containers)
248
249
  (atexit)
  (exit #t))
jerry40's avatar
jerry40 committed
250

251
252
253
254
;;
;; Dispatch route.
;;

jerry40's avatar
jerry40 committed
255
256
257
258
(define dispatch-route
  `(("kernel_info_request" . ,reply-kernel-info-request)
    ("execute_request"     . ,reply-execute-request)
    ("shutdown_request"    . ,shutdown)
259
    ("comm_info_request"   . ,ignore-request)))
jerry40's avatar
jerry40 committed
260
261
262
263

(define (dispatch msg-type)
  (let ((res (assoc-ref dispatch-route msg-type)))
    (unless res
jerry40's avatar
jerry40 committed
264
      (display
265
266
       (string-append "\n(WW) unknown message type: "
                      msg-type "\n\n")))
jerry40's avatar
jerry40 committed
267
    (if res res ignore-request)))
jerry40's avatar
jerry40 committed
268

269
270
271
272
;;
;; Process exit and signals handler.
;;

273
274
275
276
(define (kill-containers containers)
  (display "Kill container... ")
  (let ((list (vlist->list containers)))
    (define (loop container)
277
      (cond
278
       ((null? container)
279
        (begin
280
          (zmq-destroy-context container-context)
281
282
283
          (display "Done.\n")
          #t))
       (else
284
285
286
        (let* ((container (car container))
               (name      (car container))
               (pid       (assoc-ref container "pid")))
287
288
          (kill pid SIGTERM)
          (waitpid pid)
289
          (zmq-close-socket (assoc-ref (assoc-ref container "sockets")
290
                                       "shell"))
291
          (zmq-close-socket (assoc-ref (assoc-ref container "sockets")
292
                                       "shell"))
293
          (loop (cdr container))))))
294
295
296
    (loop list)))

(define (atexit)
297
298
299
300
301
302
303
304
305
306
307
308
  (display "Close connections... ")
  (for-each zmq-close-socket sockets)
  (zmq-destroy-context context)
  (display "Done.\n"))

(define (sig-exit-handler signum)
  (atexit)
  (quit))

(sigaction SIGTERM sig-exit-handler)
(sigaction SIGINT  sig-exit-handler)

309
310
311
312
;;
;; Run.
;;

313
(define (new-container-connect name pid)
314
  (let* ((ipc-dir  (string-append "/tmp/guix-kernel/ipc/" name "-"
315
                                  (number->string pid) "/"))
316
317
         (container-addr (string-append "ipc://" ipc-dir))
         (container-addr-sock (string-append container-addr "shell")))
318
319
320

    (mkdir-p ipc-dir)

321
322
    (let ((container-socket (zmq-create-socket container-context ZMQ_DEALER)))
      (zmq-set-socket-option container-socket ZMQ_IDENTITY "guix-kernel")
323

324
      (zmq-bind-socket container-socket container-addr-sock)
325
      (display "+")
326
      
327
328
      (let ((new-container `(("pid"    . ,pid)
                             ("socket" . ,container-socket)
329
                             ("uptime" . ,(current-time)))))
330
        new-container))))
331

332
333
(define (start-container-kernel name env)
  (display (string-append "container: pid: " (number->string (getpid))
334
                          "(" (number->string (getpgrp)) ")\n"))
335
  (let ((guile-version (if (null? env)
336
                           (guile-current-version->path)
337
338
339
340
341
342
343
                           (if (not (package-in-list->path env "guile"))
                               (guile-current-version->path)
                               (guile->bin-path
                                (package-in-list->path env "guile")))))
        (new-env (if (null? env)
                     (environ)
                     (make-new-environment (environ) name env))))
344
    (apply execle guile-version
345
           new-env
346
347
348
349
           (append  (list "guile")
                    (guile-current-load-path->args-list)
                    (guile-current-load-compiled-path->args-list)
                    (list "--no-auto-compile" "-s"
350
                          container-path
351
                          name notebook-info-key)))))
352
353
354
355

(define (start-kernel pid)
  (display (string-append "kernel: pid: " (number->string (getpid))
                          "(" (number->string (getpgrp)) ")\n"))
356
  (let ((default-container (alist->vhash
357
                            `(("default" .
358
359
                               ,(new-container-connect "default" pid))))))
    (general-handler socket-shell default-container 0)))
360

361
(define (run-new-container name env)
362
363
  (let ((pid (primitive-fork)))
    (if (zero? pid)
364
        (start-container-kernel name env)
365
366
        pid)))

367
(start-kernel (run-new-container "default" '()))