guix-jupyter-kernel.scm 12.4 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
233
     ((magic-run? code)
      (let* ((list     (string-split magic #\ ))
             (env-name (list-ref list 2)))
234
235
        (values containers
                (proxy-exec-container containers env-name
236
                                      parts count))))
237
     (else
238
239
      (values containers
              (proxy-exec-container containers "default"
240
                                    parts count))))))
jerry40's avatar
jerry40 committed
241

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

250
251
252
253
;;
;; Dispatch route.
;;

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

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

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

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

(define (atexit)
296
297
298
299
300
301
302
303
304
305
306
307
  (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)

308
309
310
311
;;
;; Run.
;;

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

    (mkdir-p ipc-dir)

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

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

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

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

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

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