##// END OF EJS Templates
fix typo in write_connection_file
MinRK -
Show More
@@ -1,538 +1,538 b''
1 1 """Utilities for connecting to kernels
2 2
3 3 Authors:
4 4
5 5 * Min Ragan-Kelley
6 6
7 7 """
8 8
9 9 #-----------------------------------------------------------------------------
10 10 # Copyright (C) 2013 The IPython Development Team
11 11 #
12 12 # Distributed under the terms of the BSD License. The full license is in
13 13 # the file COPYING, distributed as part of this software.
14 14 #-----------------------------------------------------------------------------
15 15
16 16 #-----------------------------------------------------------------------------
17 17 # Imports
18 18 #-----------------------------------------------------------------------------
19 19
20 20 from __future__ import absolute_import
21 21
22 22 import glob
23 23 import json
24 24 import os
25 25 import socket
26 26 import sys
27 27 from getpass import getpass
28 28 from subprocess import Popen, PIPE
29 29 import tempfile
30 30
31 31 import zmq
32 32
33 33 # external imports
34 34 from IPython.external.ssh import tunnel
35 35
36 36 # IPython imports
37 37 # from IPython.config import Configurable
38 38 from IPython.core.profiledir import ProfileDir
39 39 from IPython.utils.localinterfaces import LOCALHOST
40 40 from IPython.utils.path import filefind, get_ipython_dir
41 41 from IPython.utils.py3compat import str_to_bytes, bytes_to_str
42 42 from IPython.utils.traitlets import (
43 43 Bool, Integer, Unicode, CaselessStrEnum,
44 44 HasTraits,
45 45 )
46 46
47 47
48 48 #-----------------------------------------------------------------------------
49 49 # Working with Connection Files
50 50 #-----------------------------------------------------------------------------
51 51
52 52 def write_connection_file(fname=None, shell_port=0, iopub_port=0, stdin_port=0, hb_port=0,
53 53 control_port=0, ip=LOCALHOST, key=b'', transport='tcp'):
54 54 """Generates a JSON config file, including the selection of random ports.
55 55
56 56 Parameters
57 57 ----------
58 58
59 59 fname : unicode
60 60 The path to the file to write
61 61
62 62 shell_port : int, optional
63 The port to use for ROUTER channel.
63 The port to use for ROUTER (shell) channel.
64 64
65 65 iopub_port : int, optional
66 66 The port to use for the SUB channel.
67 67
68 68 stdin_port : int, optional
69 69 The port to use for the ROUTER (raw input) channel.
70 70
71 71 control_port : int, optional
72 The port to use for the ROUTER (raw input) channel.
72 The port to use for the ROUTER (control) channel.
73 73
74 74 hb_port : int, optional
75 The port to use for the hearbeat REP channel.
75 The port to use for the heartbeat REP channel.
76 76
77 77 ip : str, optional
78 78 The ip address the kernel will bind to.
79 79
80 80 key : str, optional
81 81 The Session key used for HMAC authentication.
82 82
83 83 """
84 84 # default to temporary connector file
85 85 if not fname:
86 86 fname = tempfile.mktemp('.json')
87 87
88 88 # Find open ports as necessary.
89 89
90 90 ports = []
91 91 ports_needed = int(shell_port <= 0) + \
92 92 int(iopub_port <= 0) + \
93 93 int(stdin_port <= 0) + \
94 94 int(control_port <= 0) + \
95 95 int(hb_port <= 0)
96 96 if transport == 'tcp':
97 97 for i in range(ports_needed):
98 98 sock = socket.socket()
99 99 sock.bind(('', 0))
100 100 ports.append(sock)
101 101 for i, sock in enumerate(ports):
102 102 port = sock.getsockname()[1]
103 103 sock.close()
104 104 ports[i] = port
105 105 else:
106 106 N = 1
107 107 for i in range(ports_needed):
108 108 while os.path.exists("%s-%s" % (ip, str(N))):
109 109 N += 1
110 110 ports.append(N)
111 111 N += 1
112 112 if shell_port <= 0:
113 113 shell_port = ports.pop(0)
114 114 if iopub_port <= 0:
115 115 iopub_port = ports.pop(0)
116 116 if stdin_port <= 0:
117 117 stdin_port = ports.pop(0)
118 118 if control_port <= 0:
119 119 control_port = ports.pop(0)
120 120 if hb_port <= 0:
121 121 hb_port = ports.pop(0)
122 122
123 123 cfg = dict( shell_port=shell_port,
124 124 iopub_port=iopub_port,
125 125 stdin_port=stdin_port,
126 126 control_port=control_port,
127 127 hb_port=hb_port,
128 128 )
129 129 cfg['ip'] = ip
130 130 cfg['key'] = bytes_to_str(key)
131 131 cfg['transport'] = transport
132 132
133 133 with open(fname, 'w') as f:
134 134 f.write(json.dumps(cfg, indent=2))
135 135
136 136 return fname, cfg
137 137
138 138
139 139 def get_connection_file(app=None):
140 140 """Return the path to the connection file of an app
141 141
142 142 Parameters
143 143 ----------
144 144 app : IPKernelApp instance [optional]
145 145 If unspecified, the currently running app will be used
146 146 """
147 147 if app is None:
148 148 from IPython.kernel.zmq.kernelapp import IPKernelApp
149 149 if not IPKernelApp.initialized():
150 150 raise RuntimeError("app not specified, and not in a running Kernel")
151 151
152 152 app = IPKernelApp.instance()
153 153 return filefind(app.connection_file, ['.', app.profile_dir.security_dir])
154 154
155 155
156 156 def find_connection_file(filename, profile=None):
157 157 """find a connection file, and return its absolute path.
158 158
159 159 The current working directory and the profile's security
160 160 directory will be searched for the file if it is not given by
161 161 absolute path.
162 162
163 163 If profile is unspecified, then the current running application's
164 164 profile will be used, or 'default', if not run from IPython.
165 165
166 166 If the argument does not match an existing file, it will be interpreted as a
167 167 fileglob, and the matching file in the profile's security dir with
168 168 the latest access time will be used.
169 169
170 170 Parameters
171 171 ----------
172 172 filename : str
173 173 The connection file or fileglob to search for.
174 174 profile : str [optional]
175 175 The name of the profile to use when searching for the connection file,
176 176 if different from the current IPython session or 'default'.
177 177
178 178 Returns
179 179 -------
180 180 str : The absolute path of the connection file.
181 181 """
182 182 from IPython.core.application import BaseIPythonApplication as IPApp
183 183 try:
184 184 # quick check for absolute path, before going through logic
185 185 return filefind(filename)
186 186 except IOError:
187 187 pass
188 188
189 189 if profile is None:
190 190 # profile unspecified, check if running from an IPython app
191 191 if IPApp.initialized():
192 192 app = IPApp.instance()
193 193 profile_dir = app.profile_dir
194 194 else:
195 195 # not running in IPython, use default profile
196 196 profile_dir = ProfileDir.find_profile_dir_by_name(get_ipython_dir(), 'default')
197 197 else:
198 198 # find profiledir by profile name:
199 199 profile_dir = ProfileDir.find_profile_dir_by_name(get_ipython_dir(), profile)
200 200 security_dir = profile_dir.security_dir
201 201
202 202 try:
203 203 # first, try explicit name
204 204 return filefind(filename, ['.', security_dir])
205 205 except IOError:
206 206 pass
207 207
208 208 # not found by full name
209 209
210 210 if '*' in filename:
211 211 # given as a glob already
212 212 pat = filename
213 213 else:
214 214 # accept any substring match
215 215 pat = '*%s*' % filename
216 216 matches = glob.glob( os.path.join(security_dir, pat) )
217 217 if not matches:
218 218 raise IOError("Could not find %r in %r" % (filename, security_dir))
219 219 elif len(matches) == 1:
220 220 return matches[0]
221 221 else:
222 222 # get most recent match, by access time:
223 223 return sorted(matches, key=lambda f: os.stat(f).st_atime)[-1]
224 224
225 225
226 226 def get_connection_info(connection_file=None, unpack=False, profile=None):
227 227 """Return the connection information for the current Kernel.
228 228
229 229 Parameters
230 230 ----------
231 231 connection_file : str [optional]
232 232 The connection file to be used. Can be given by absolute path, or
233 233 IPython will search in the security directory of a given profile.
234 234 If run from IPython,
235 235
236 236 If unspecified, the connection file for the currently running
237 237 IPython Kernel will be used, which is only allowed from inside a kernel.
238 238 unpack : bool [default: False]
239 239 if True, return the unpacked dict, otherwise just the string contents
240 240 of the file.
241 241 profile : str [optional]
242 242 The name of the profile to use when searching for the connection file,
243 243 if different from the current IPython session or 'default'.
244 244
245 245
246 246 Returns
247 247 -------
248 248 The connection dictionary of the current kernel, as string or dict,
249 249 depending on `unpack`.
250 250 """
251 251 if connection_file is None:
252 252 # get connection file from current kernel
253 253 cf = get_connection_file()
254 254 else:
255 255 # connection file specified, allow shortnames:
256 256 cf = find_connection_file(connection_file, profile=profile)
257 257
258 258 with open(cf) as f:
259 259 info = f.read()
260 260
261 261 if unpack:
262 262 info = json.loads(info)
263 263 # ensure key is bytes:
264 264 info['key'] = str_to_bytes(info.get('key', ''))
265 265 return info
266 266
267 267
268 268 def connect_qtconsole(connection_file=None, argv=None, profile=None):
269 269 """Connect a qtconsole to the current kernel.
270 270
271 271 This is useful for connecting a second qtconsole to a kernel, or to a
272 272 local notebook.
273 273
274 274 Parameters
275 275 ----------
276 276 connection_file : str [optional]
277 277 The connection file to be used. Can be given by absolute path, or
278 278 IPython will search in the security directory of a given profile.
279 279 If run from IPython,
280 280
281 281 If unspecified, the connection file for the currently running
282 282 IPython Kernel will be used, which is only allowed from inside a kernel.
283 283 argv : list [optional]
284 284 Any extra args to be passed to the console.
285 285 profile : str [optional]
286 286 The name of the profile to use when searching for the connection file,
287 287 if different from the current IPython session or 'default'.
288 288
289 289
290 290 Returns
291 291 -------
292 292 subprocess.Popen instance running the qtconsole frontend
293 293 """
294 294 argv = [] if argv is None else argv
295 295
296 296 if connection_file is None:
297 297 # get connection file from current kernel
298 298 cf = get_connection_file()
299 299 else:
300 300 cf = find_connection_file(connection_file, profile=profile)
301 301
302 302 cmd = ';'.join([
303 303 "from IPython.frontend.qt.console import qtconsoleapp",
304 304 "qtconsoleapp.main()"
305 305 ])
306 306
307 307 return Popen([sys.executable, '-c', cmd, '--existing', cf] + argv, stdout=PIPE, stderr=PIPE)
308 308
309 309
310 310 def tunnel_to_kernel(connection_info, sshserver, sshkey=None):
311 311 """tunnel connections to a kernel via ssh
312 312
313 313 This will open four SSH tunnels from localhost on this machine to the
314 314 ports associated with the kernel. They can be either direct
315 315 localhost-localhost tunnels, or if an intermediate server is necessary,
316 316 the kernel must be listening on a public IP.
317 317
318 318 Parameters
319 319 ----------
320 320 connection_info : dict or str (path)
321 321 Either a connection dict, or the path to a JSON connection file
322 322 sshserver : str
323 323 The ssh sever to use to tunnel to the kernel. Can be a full
324 324 `user@server:port` string. ssh config aliases are respected.
325 325 sshkey : str [optional]
326 326 Path to file containing ssh key to use for authentication.
327 327 Only necessary if your ssh config does not already associate
328 328 a keyfile with the host.
329 329
330 330 Returns
331 331 -------
332 332
333 333 (shell, iopub, stdin, hb) : ints
334 334 The four ports on localhost that have been forwarded to the kernel.
335 335 """
336 336 if isinstance(connection_info, basestring):
337 337 # it's a path, unpack it
338 338 with open(connection_info) as f:
339 339 connection_info = json.loads(f.read())
340 340
341 341 cf = connection_info
342 342
343 343 lports = tunnel.select_random_ports(4)
344 344 rports = cf['shell_port'], cf['iopub_port'], cf['stdin_port'], cf['hb_port']
345 345
346 346 remote_ip = cf['ip']
347 347
348 348 if tunnel.try_passwordless_ssh(sshserver, sshkey):
349 349 password=False
350 350 else:
351 351 password = getpass("SSH Password for %s: "%sshserver)
352 352
353 353 for lp,rp in zip(lports, rports):
354 354 tunnel.ssh_tunnel(lp, rp, sshserver, remote_ip, sshkey, password)
355 355
356 356 return tuple(lports)
357 357
358 358
359 359 #-----------------------------------------------------------------------------
360 360 # Mixin for classes that work with connection files
361 361 #-----------------------------------------------------------------------------
362 362
363 363 channel_socket_types = {
364 364 'hb' : zmq.REQ,
365 365 'shell' : zmq.DEALER,
366 366 'iopub' : zmq.SUB,
367 367 'stdin' : zmq.DEALER,
368 368 'control': zmq.DEALER,
369 369 }
370 370
371 371 port_names = [ "%s_port" % channel for channel in ('shell', 'stdin', 'iopub', 'hb', 'control')]
372 372
373 373 class ConnectionFileMixin(HasTraits):
374 374 """Mixin for configurable classes that work with connection files"""
375 375
376 376 # The addresses for the communication channels
377 377 connection_file = Unicode('')
378 378 _connection_file_written = Bool(False)
379 379
380 380 transport = CaselessStrEnum(['tcp', 'ipc'], default_value='tcp', config=True)
381 381
382 382 ip = Unicode(LOCALHOST, config=True,
383 383 help="""Set the kernel\'s IP address [default localhost].
384 384 If the IP address is something other than localhost, then
385 385 Consoles on other machines will be able to connect
386 386 to the Kernel, so be careful!"""
387 387 )
388 388
389 389 def _ip_default(self):
390 390 if self.transport == 'ipc':
391 391 if self.connection_file:
392 392 return os.path.splitext(self.connection_file)[0] + '-ipc'
393 393 else:
394 394 return 'kernel-ipc'
395 395 else:
396 396 return LOCALHOST
397 397
398 398 def _ip_changed(self, name, old, new):
399 399 if new == '*':
400 400 self.ip = '0.0.0.0'
401 401
402 402 # protected traits
403 403
404 404 shell_port = Integer(0)
405 405 iopub_port = Integer(0)
406 406 stdin_port = Integer(0)
407 407 control_port = Integer(0)
408 408 hb_port = Integer(0)
409 409
410 410 @property
411 411 def ports(self):
412 412 return [ getattr(self, name) for name in port_names ]
413 413
414 414 #--------------------------------------------------------------------------
415 415 # Connection and ipc file management
416 416 #--------------------------------------------------------------------------
417 417
418 418 def get_connection_info(self):
419 419 """return the connection info as a dict"""
420 420 return dict(
421 421 transport=self.transport,
422 422 ip=self.ip,
423 423 shell_port=self.shell_port,
424 424 iopub_port=self.iopub_port,
425 425 stdin_port=self.stdin_port,
426 426 hb_port=self.hb_port,
427 427 control_port=self.control_port,
428 428 )
429 429
430 430 def cleanup_connection_file(self):
431 431 """Cleanup connection file *if we wrote it*
432 432
433 433 Will not raise if the connection file was already removed somehow.
434 434 """
435 435 if self._connection_file_written:
436 436 # cleanup connection files on full shutdown of kernel we started
437 437 self._connection_file_written = False
438 438 try:
439 439 os.remove(self.connection_file)
440 440 except (IOError, OSError, AttributeError):
441 441 pass
442 442
443 443 def cleanup_ipc_files(self):
444 444 """Cleanup ipc files if we wrote them."""
445 445 if self.transport != 'ipc':
446 446 return
447 447 for port in self.ports:
448 448 ipcfile = "%s-%i" % (self.ip, port)
449 449 try:
450 450 os.remove(ipcfile)
451 451 except (IOError, OSError):
452 452 pass
453 453
454 454 def write_connection_file(self):
455 455 """Write connection info to JSON dict in self.connection_file."""
456 456 if self._connection_file_written:
457 457 return
458 458
459 459 self.connection_file, cfg = write_connection_file(self.connection_file,
460 460 transport=self.transport, ip=self.ip, key=self.session.key,
461 461 stdin_port=self.stdin_port, iopub_port=self.iopub_port,
462 462 shell_port=self.shell_port, hb_port=self.hb_port,
463 463 control_port=self.control_port,
464 464 )
465 465 # write_connection_file also sets default ports:
466 466 for name in port_names:
467 467 setattr(self, name, cfg[name])
468 468
469 469 self._connection_file_written = True
470 470
471 471 def load_connection_file(self):
472 472 """Load connection info from JSON dict in self.connection_file."""
473 473 with open(self.connection_file) as f:
474 474 cfg = json.loads(f.read())
475 475
476 476 self.transport = cfg.get('transport', 'tcp')
477 477 self.ip = cfg['ip']
478 478 for name in port_names:
479 479 setattr(self, name, cfg[name])
480 480 self.session.key = str_to_bytes(cfg['key'])
481 481
482 482 #--------------------------------------------------------------------------
483 483 # Creating connected sockets
484 484 #--------------------------------------------------------------------------
485 485
486 486 def _make_url(self, channel):
487 487 """Make a ZeroMQ URL for a given channel."""
488 488 transport = self.transport
489 489 ip = self.ip
490 490 port = getattr(self, '%s_port' % channel)
491 491
492 492 if transport == 'tcp':
493 493 return "tcp://%s:%i" % (ip, port)
494 494 else:
495 495 return "%s://%s-%s" % (transport, ip, port)
496 496
497 497 def _create_connected_socket(self, channel, identity=None):
498 498 """Create a zmq Socket and connect it to the kernel."""
499 499 url = self._make_url(channel)
500 500 socket_type = channel_socket_types[channel]
501 501 self.log.info("Connecting to: %s" % url)
502 502 sock = self.context.socket(socket_type)
503 503 if identity:
504 504 sock.identity = identity
505 505 sock.connect(url)
506 506 return sock
507 507
508 508 def connect_iopub(self, identity=None):
509 509 """return zmq Socket connected to the IOPub channel"""
510 510 sock = self._create_connected_socket('iopub', identity=identity)
511 511 sock.setsockopt(zmq.SUBSCRIBE, b'')
512 512 return sock
513 513
514 514 def connect_shell(self, identity=None):
515 515 """return zmq Socket connected to the Shell channel"""
516 516 return self._create_connected_socket('shell', identity=identity)
517 517
518 518 def connect_stdin(self, identity=None):
519 519 """return zmq Socket connected to the StdIn channel"""
520 520 return self._create_connected_socket('stdin', identity=identity)
521 521
522 522 def connect_hb(self, identity=None):
523 523 """return zmq Socket connected to the Heartbeat channel"""
524 524 return self._create_connected_socket('hb', identity=identity)
525 525
526 526 def connect_control(self, identity=None):
527 527 """return zmq Socket connected to the Heartbeat channel"""
528 528 return self._create_connected_socket('control', identity=identity)
529 529
530 530
531 531 __all__ = [
532 532 'write_connection_file',
533 533 'get_connection_file',
534 534 'find_connection_file',
535 535 'get_connection_info',
536 536 'connect_qtconsole',
537 537 'tunnel_to_kernel',
538 538 ]
General Comments 0
You need to be logged in to leave comments. Login now