guix-jupyter-kernel.scm 14.7 KB
Newer Older
1
;;; Guix-kernel -- Guix kernel for Jupyter
2
3
4
;;; Copyright (C) 2018 Evgeny Panfilov <epanfilov@gmail.com>
;;; Copyright (C) 2018 Pierre-Antoine Rouby <pierre-antoine.rouby@inria.fr>
;;; Copyright (C) 2018 Inria
5
6
7
8
9
10
11
12
13
14
15
16
17
18
;;;
;;; 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/>.

19
20
(use-modules (json)
             (simple-zmq)
jerry40's avatar
jerry40 committed
21
	     (srfi srfi-1)
22
             (srfi srfi-11)
23
             (srfi srfi-13)
24
             (rnrs bytevectors)
25
             (ice-9 vlist)
26
27
             (ice-9 futures)
             (ice-9 match)
28
             (sxml simple)
29
30
31
             (gnu build linux-container)
             (gnu system file-systems)
             (guix build syscalls)
32
33
34
35
             (guix build utils)
             (guix-kernel tools)
             (guix-kernel hmac)
             (guix-kernel magic)
36
37
             (guix-kernel environ)
             (guix-kernel jupyter-client))
jerry40's avatar
jerry40 committed
38

39
40
41
(define session-id (random (* 255 255) (seed->random-state
                                        (current-time))))

jerry40's avatar
jerry40 committed
42

43
;;
44
;; Container tools.
45
46
;;

47
(define container-context (zmq-create-context))
48

49
(define container-path
50
51
  (string-replace (car (command-line)) "guix-jupyter-container.scm"
                  (+ (string-rindex (car (command-line)) #\/) 1)))
52

53
(define (register-container vhash name values)
54
  (vhash-cons name values vhash))
55

56
(define (container-by-name vhash name)
57
  (vhash-assoc name vhash))
58

59
(define (container-socket-by-name vhash name)
60
  (assoc-ref (vhash-assoc name vhash) "socket"))
61

62
63
64
65
;;
;; Send.
;;

66
(define (send-to-container containers socket name parts)
67
68
  (zmq-send-msg-parts-bytevector socket
                                 (cons (string->utf8 name) parts)))
69

70
71
(define* (proxy-exec-container containers parts count
                              #:key name notebook)
72
73
74
75
76
77
78
  (define (stream-loop socket streams)
    (cond
     ((null? streams) #t)
     (else
      (zmq-send-msg-parts socket
                          (car streams))
      (stream-loop socket (cdr streams)))))
79

80
81
82
83
  (define (get-messages socket
                        execute-input error-or-result
                        execute-reply streams)
    (let* ((new-parts (zmq-get-msg-parts socket))
84

85
86
           (header    (json-string->scm (list-ref new-parts 3)))
           (content   (json-string->scm (list-ref new-parts 6)))
87

88
           (msg-type  (hash-ref header "msg_type")))
89
      (display ".")
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
      (cond
       ((or (equal? msg-type "execute_result")
            (equal? msg-type "error"))
        (get-messages socket execute-input
                      new-parts execute-reply streams))
       ((equal? msg-type "execute_input")
        (get-messages socket new-parts
                      error-or-result execute-reply streams))
       ((equal? msg-type "execute_reply")
        (get-messages socket execute-input
                      error-or-result new-parts streams))
       ((equal? msg-type "stream")
        (get-messages socket execute-input
                      error-or-result execute-reply
                      (append (list new-parts)
                              streams)))
       (else
        (values execute-input error-or-result execute-reply streams)))))

  (let* ((container  (container-socket-by-name containers name)))
    (send-to-container containers container name parts)

    (let-values (((execute-input error-or-result execute-reply streams)
                  (get-messages container '() '() '() '())))
114
      (zmq-send-msg-parts (notebook-iopub notebook) execute-input)
115
116

      (unless (null? error-or-result)
117
        (zmq-send-msg-parts (notebook-iopub notebook) error-or-result))
118
      (unless (null? streams)
119
        (stream-loop (notebook-iopub notebook) streams))
120
      (unless (null? execute-reply)
121
        (zmq-send-msg-parts (notebook-shell notebook) execute-reply))
122

123
      (+ count 1))))
124

125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
;;
;; Html.
;;

(define (error->shtml error)
  (match error
    (('guix-kernel msg errno)
     `(h3 (@ (style "color: red;"))
          ,(string-append "Error: " msg)))
    (_
     `(h3 (@ (style "color: red;"))
          "Error !"))))

(define (error->html error)
  (call-with-output-string
    (lambda (port)
      (sxml->xml (error->shtml error) port))))

143
;;
144
;; Handlers.
145
146
;;

147
148
(define (heartbeat-handler notebook)
  (let ((message (zmq-message-receive-bytevector (notebook-heartbeat notebook)
149
                                                 (zmq-msg-init))))
150
151
    (zmq-message-send-bytevector (notebook-heartbeat notebook) message))
  (heartbeat-handler notebook))
jerry40's avatar
jerry40 committed
152

153

154
(define (general-handler notebook containers count)
155
156
157
158
  ;; FIXME Use 'zmq_poll'.
  (define (waiting-data socket next)
    (catch 'zmq-error
      (λ _
159
        (zmq-get-msg-parts-bytevector socket))
160
161
162
163
164
165
      (λ stuff
        (cond
         ((equal? EAGAIN (cadr stuff))
          (waiting-data next socket))
         (else #f)))))

166
167
  (let ((parts (waiting-data (notebook-shell notebook)
                             (notebook-control notebook))))
168
    (match parts
169
170
171
172
173
174
      ((wire-uuid wire-delimiter wire-signature
                  (= utf8->string wire-header)
                  wire-parent-header
                  (= utf8->string wire-metadata)
                  (= utf8->string wire-content))
       (let* ((header        (json-string->scm wire-header))
175
176
177
178
              (msg-type      (hash-ref header "msg_type"))
              (msg-username  (hash-ref header "username"))
              (msg-session   (hash-ref header "session"))
              (msg-version   (hash-ref header "version"))
179
              (header- (make-header msg-username msg-session msg-version)))
180
181
182
         (pub-busy (notebook-iopub notebook) header- wire-header
                   (notebook-key notebook))
         (let-values (((containers count) ((dispatch msg-type) notebook wire-uuid
183
184
                                           header- wire-header
                                           wire-metadata
185
186
                                           wire-content parts
                                           containers count)))
187
188
189
           (pub-idle (notebook-iopub notebook) header- wire-header
                     (notebook-key notebook))
           (general-handler notebook containers count))))
190
      (_
191
       (general-handler notebook containers count)))))
jerry40's avatar
jerry40 committed
192

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

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

209
(define (reply-execute-request notebook uuid header- parent-header
210
                               metadata content parts containers
211
                               count)
212
  (let* ((code  (hash-ref (json-string->scm content) "code"))
213
         (magic (get-magic-line code)))
214
215
216
217
218
219
220
221
222
    (catch 'guix-kernel
      (λ ()
        (cond
         ((magic-env? code)
          (let* ((list     (string-split magic #\ ))
                 (env-name (list-ref list 2))
                 (env      (list-cdr-ref list 4)))
            (if (container-by-name containers env-name)
                (values containers
223
224
225
226
227
                        (proxy-exec-container containers
                                              parts count
                                              #:name env-name
                                              #:notebook notebook))
                (let* ((id        (start-container notebook env-name env))
228
229
230
231
                       (container (new-container-connect env-name id))
                       (new       (register-container containers env-name
                                                      container)))
                  (values new
232
233
234
                          (proxy-exec-container new parts count
                                                #:name env-name
                                                #:notebook notebook))))))
235
236
237
238
         ((or (magic-run? code)
              (magic-kernel? code))
          (let* ((list     (string-split magic #\ ))
                 (env-name (list-ref list 2)))
239
            (values containers
240
241
242
                    (proxy-exec-container containers parts count
                                          #:name env-name
                                          #:notebook notebook))))
243
244
         ((magic-html? code)
          (values containers
245
246
                  (reply-html (notebook-shell notebook)
                              (notebook-iopub notebook)
247
248
249
                              uuid header- parent-header
                              metadata content
                              (delete-magic code)
250
                              (notebook-key notebook) count)))
251
252
         (else
          (values containers
253
254
255
                  (proxy-exec-container containers parts count
                                        #:name "default"
                                        #:notebook notebook)))))
256
      (λ error
257
        (values containers
258
259
                (reply-html (notebook-shell notebook)
                            (notebook-iopub notebook)
260
261
262
                            uuid header- parent-header
                            metadata content
                            (error->html error)
263
                            (notebook-key notebook) count))))))
jerry40's avatar
jerry40 committed
264

265
(define (shutdown notebook uuid header- parent-header
266
                  metadata content parts containers
267
                  count)
268
  (kill-containers containers parts)
269
  (exit #t))
jerry40's avatar
jerry40 committed
270

271
272
273
274
;;
;; Dispatch route.
;;

jerry40's avatar
jerry40 committed
275
276
277
278
(define dispatch-route
  `(("kernel_info_request" . ,reply-kernel-info-request)
    ("execute_request"     . ,reply-execute-request)
    ("shutdown_request"    . ,shutdown)
279
    ("comm_info_request"   . ,ignore-request)))
jerry40's avatar
jerry40 committed
280
281
282
283

(define (dispatch msg-type)
  (let ((res (assoc-ref dispatch-route msg-type)))
    (unless res
jerry40's avatar
jerry40 committed
284
      (display
285
286
       (string-append "\n(WW) unknown message type: "
                      msg-type "\n\n")))
jerry40's avatar
jerry40 committed
287
    (if res res ignore-request)))
jerry40's avatar
jerry40 committed
288

289
290
291
292
;;
;; Process exit and signals handler.
;;

293
294
(define (kill-containers containers parts)
  (display "\nKill container ")
295
  (let ((list (vlist->list containers)))
296
    (define (loop list)
297
      (cond
298
       ((null? list) #t)
299
       (else
300
        (let* ((container (car list))
301
               (name      (car container))
302
303
304
305
306
307
308
309
310
               (pid       (assoc-ref container "pid"))
               (sock      (assoc-ref container "socket")))
          (send-to-container containers sock name parts)
          (zmq-close-socket sock)
          (display ".")
          (loop (cdr list))))))

    (loop list)
    (display " Done.\n")))
311

312

313
314
315
316
;;
;; Run.
;;

317
(define (new-container-connect name pid)
318
  (let* ((ipc-dir  (string-append "/tmp/guix-kernel/ipc/" name "-"
319
                                  (number->string session-id) "/"))
320
321
         (container-addr (string-append "ipc://" ipc-dir))
         (container-addr-sock (string-append container-addr "shell")))
322
323
324

    (mkdir-p ipc-dir)

325
326
    (let ((container-socket (zmq-create-socket container-context ZMQ_DEALER)))
      (zmq-set-socket-option container-socket ZMQ_IDENTITY "guix-kernel")
327

328
      (zmq-bind-socket container-socket container-addr-sock)
329
      (display "+")
330

331
332
      (let ((new-container `(("pid"    . ,pid)
                             ("socket" . ,container-socket)
333
                             ("uptime" . ,(current-time)))))
334
        new-container))))
335

336
337
338
339
340
341
342
343
(define %network-file-systems
  ;; The list of <file-system> objects corresponding to bind-mounts required
  ;; for networking.
  (filter-map (lambda (mapping)
                (let ((fs (file-system-mapping->bind-mount mapping)))
                  (and (file-exists? (file-system-device fs))
                       fs)))
              %network-file-mappings))
344

345
(define (start-container notebook name env)
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
  (let* ((guile-version (if (null? env)
                            (guile-current-version->path)
                            (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)))
         (exec (lambda ()
                 (set-network-interface-up "lo") ;up lo interface

                 (apply execle guile-version
                        new-env
                        (append  (list "guile")
                                 (guile-current-load-path->args-list)
                                 (guile-current-load-compiled-path->args-list)
                                 (list "--no-auto-compile" "-s"
                                       container-path
                                       name (number->string session-id)
366
                                       (notebook-key notebook))))))
367
368
         (root (string-append "/tmp/guix-kernel/container/"
                              name "-" (number->string session-id)))
369
370
371
372
373
374
375
376
377
         (fs   (cons* %immutable-store
                      (file-system
                        (device "/tmp/guix-kernel")
                        (mount-point "/tmp/guix-kernel")
                        (type "none")
                        (check? #f)
                        (flags '(bind-mount)))
                      (append %container-file-systems
                              %network-file-systems))))
378
379
380
381
    (mkdir-p root)
    (run-container root fs
                   %namespaces 1
                   exec)))
382

383
(define (start-kernel pid notebook)
384
  (let ((default-container (alist->vhash
385
                            `(("default" .
386
                               ,(new-container-connect "default" pid))))))
387
    (general-handler notebook default-container 0)))
388

389

390
391
392
393
394
(define (exit-handler notebook)
  (lambda _
    (close-notebook notebook)
    (exit 1)))

395
396
397
398
399
400
401
;; Start!
(let ((notebook (call-with-input-file (car (last-pair (command-line)))
                  json->notebook)))
  (sigaction SIGTERM (exit-handler notebook))
  (sigaction SIGINT  (exit-handler notebook))
  (start-kernel (start-container notebook "default" '())
                notebook))