ipengineapp.py
398 lines
| 13.1 KiB
| text/x-python
|
PythonLexer
|
r3604 | #!/usr/bin/env python | ||
# encoding: utf-8 | ||||
""" | ||||
The IPython engine application | ||||
|
r4018 | |||
Authors: | ||||
* Brian Granger | ||||
* MinRK | ||||
|
r3604 | """ | ||
#----------------------------------------------------------------------------- | ||||
|
r4018 | # Copyright (C) 2008-2011 The IPython Development Team | ||
|
r3604 | # | ||
# Distributed under the terms of the BSD License. The full license is in | ||||
# the file COPYING, distributed as part of this software. | ||||
#----------------------------------------------------------------------------- | ||||
#----------------------------------------------------------------------------- | ||||
# Imports | ||||
#----------------------------------------------------------------------------- | ||||
|
r3631 | import json | ||
|
r3604 | import os | ||
import sys | ||||
|
r4120 | import time | ||
|
r3604 | |||
import zmq | ||||
from zmq.eventloop import ioloop | ||||
|
r4024 | from IPython.core.profiledir import ProfileDir | ||
|
r4115 | from IPython.parallel.apps.baseapp import ( | ||
BaseParallelApplication, | ||||
base_aliases, | ||||
base_flags, | ||||
|
r5214 | catch_config_error, | ||
|
r4115 | ) | ||
|
r3604 | from IPython.zmq.log import EnginePUBHandler | ||
|
r6887 | from IPython.zmq.ipkernel import Kernel, IPKernelApp | ||
|
r4962 | from IPython.zmq.session import ( | ||
Session, session_aliases, session_flags | ||||
) | ||||
|
r3604 | |||
|
r3985 | from IPython.config.configurable import Configurable | ||
|
r4962 | |||
|
r3673 | from IPython.parallel.engine.engine import EngineFactory | ||
|
r7890 | from IPython.parallel.util import disambiguate_ip_address | ||
|
r3673 | |||
|
r3604 | from IPython.utils.importstring import import_item | ||
|
r6813 | from IPython.utils.py3compat import cast_bytes | ||
|
r6887 | from IPython.utils.traitlets import Bool, Unicode, Dict, List, Float, Instance | ||
|
r3604 | |||
|
r3614 | |||
|
r3604 | #----------------------------------------------------------------------------- | ||
# Module level variables | ||||
#----------------------------------------------------------------------------- | ||||
#: The default config file name for this application | ||||
|
r3672 | default_config_file_name = u'ipengine_config.py' | ||
|
r3604 | |||
|
r3990 | _description = """Start an IPython engine for parallel computing. | ||
|
r3985 | |||
IPython engines run in parallel and perform computations on behalf of a client | ||||
and controller. A controller needs to be started before the engines. The | ||||
engine can be configured using command line options or using a cluster | ||||
directory. Cluster directories contain config, log and security files and are | ||||
|
r4024 | usually located in your ipython directory and named as "profile_name". | ||
|
r4218 | See the `profile` and `profile-dir` options for details. | ||
|
r3985 | """ | ||
|
r4216 | _examples = """ | ||
ipengine --ip=192.168.0.1 --port=1000 # connect to hub at ip and port | ||||
|
r4218 | ipengine --log-to-file --log-level=DEBUG # log to a file with DEBUG verbosity | ||
|
r4216 | """ | ||
|
r3985 | |||
#----------------------------------------------------------------------------- | ||||
# MPI configuration | ||||
#----------------------------------------------------------------------------- | ||||
|
r3604 | |||
mpi4py_init = """from mpi4py import MPI as mpi | ||||
mpi.size = mpi.COMM_WORLD.Get_size() | ||||
mpi.rank = mpi.COMM_WORLD.Get_rank() | ||||
""" | ||||
pytrilinos_init = """from PyTrilinos import Epetra | ||||
class SimpleStruct: | ||||
pass | ||||
mpi = SimpleStruct() | ||||
mpi.rank = 0 | ||||
mpi.size = 0 | ||||
""" | ||||
|
r3985 | class MPI(Configurable): | ||
"""Configurable for MPI initialization""" | ||||
|
r3988 | use = Unicode('', config=True, | ||
|
r3985 | help='How to enable MPI (mpi4py, pytrilinos, or empty string to disable).' | ||
) | ||||
|
r3604 | |||
|
r4850 | def _use_changed(self, name, old, new): | ||
|
r3985 | # load default init script if it's not set | ||
if not self.init_script: | ||||
self.init_script = self.default_inits.get(new, '') | ||||
|
r3988 | init_script = Unicode('', config=True, | ||
|
r3985 | help="Initialization code for MPI") | ||
default_inits = Dict({'mpi4py' : mpi4py_init, 'pytrilinos':pytrilinos_init}, | ||||
config=True) | ||||
|
r3604 | |||
#----------------------------------------------------------------------------- | ||||
|
r3985 | # Main application | ||
|
r3604 | #----------------------------------------------------------------------------- | ||
|
r4115 | aliases = dict( | ||
file = 'IPEngineApp.url_file', | ||||
c = 'IPEngineApp.startup_command', | ||||
s = 'IPEngineApp.startup_script', | ||||
|
r3604 | |||
|
r4115 | url = 'EngineFactory.url', | ||
|
r4585 | ssh = 'EngineFactory.sshserver', | ||
sshkey = 'EngineFactory.sshkey', | ||||
|
r4115 | ip = 'EngineFactory.ip', | ||
transport = 'EngineFactory.transport', | ||||
port = 'EngineFactory.regport', | ||||
location = 'EngineFactory.location', | ||||
timeout = 'EngineFactory.timeout', | ||||
mpi = 'MPI.use', | ||||
) | ||||
aliases.update(base_aliases) | ||||
|
r4962 | aliases.update(session_aliases) | ||
flags = {} | ||||
flags.update(base_flags) | ||||
flags.update(session_flags) | ||||
|
r4216 | |||
|
r3992 | class IPEngineApp(BaseParallelApplication): | ||
|
r3985 | |||
|
r4847 | name = 'ipengine' | ||
description = _description | ||||
|
r4216 | examples = _examples | ||
|
r3991 | config_file_name = Unicode(default_config_file_name) | ||
|
r4006 | classes = List([ProfileDir, Session, EngineFactory, Kernel, MPI]) | ||
|
r3985 | |||
startup_script = Unicode(u'', config=True, | ||||
help='specify a script to be run at startup') | ||||
|
r3988 | startup_command = Unicode('', config=True, | ||
|
r3604 | help='specify a command to be run at startup') | ||
|
r3985 | url_file = Unicode(u'', config=True, | ||
help="""The full location of the file containing the connection information for | ||||
the controller. If this is not given, the file must be in the | ||||
security directory of the cluster directory. This location is | ||||
|
r3992 | resolved using the `profile` or `profile_dir` options.""", | ||
|
r3985 | ) | ||
|
r4120 | wait_for_url_file = Float(5, config=True, | ||
help="""The maximum number of seconds to wait for url_file to exist. | ||||
This is useful for batch-systems and shared-filesystems where the | ||||
controller and engine are started at the same time and it | ||||
may take a moment for the controller to write the connector files.""") | ||||
|
r3604 | |||
|
r4847 | url_file_name = Unicode(u'ipcontroller-engine.json', config=True) | ||
def _cluster_id_changed(self, name, old, new): | ||||
if new: | ||||
|
r4850 | base = 'ipcontroller-%s' % new | ||
|
r4847 | else: | ||
base = 'ipcontroller' | ||||
|
r4850 | self.url_file_name = "%s-engine.json" % base | ||
|
r4847 | |||
|
r3989 | log_url = Unicode('', config=True, | ||
help="""The URL for the iploggerapp instance, for forwarding | ||||
logging to a central location.""") | ||||
|
r6887 | |||
# an IPKernelApp instance, used to setup listening for shell frontends | ||||
kernel_app = Instance(IPKernelApp) | ||||
|
r3604 | |||
|
r4115 | aliases = Dict(aliases) | ||
|
r4962 | flags = Dict(flags) | ||
|
r6879 | |||
@property | ||||
def kernel(self): | ||||
"""allow access to the Kernel object, so I look like IPKernelApp""" | ||||
return self.engine.kernel | ||||
|
r3604 | |||
|
r3614 | def find_url_file(self): | ||
|
r4120 | """Set the url file. | ||
|
r3604 | |||
Here we don't try to actually see if it exists for is valid as that | ||||
is hadled by the connection logic. | ||||
""" | ||||
|
r3985 | config = self.config | ||
|
r3604 | # Find the actual controller key file | ||
|
r3985 | if not self.url_file: | ||
self.url_file = os.path.join( | ||||
|
r3992 | self.profile_dir.security_dir, | ||
|
r3985 | self.url_file_name | ||
|
r3604 | ) | ||
|
r4585 | |||
def load_connector_file(self): | ||||
"""load config from a JSON connector file, | ||||
at a *lower* priority than command-line/config files. | ||||
""" | ||||
|
r5483 | self.log.info("Loading url_file %r", self.url_file) | ||
|
r4585 | config = self.config | ||
with open(self.url_file) as f: | ||||
d = json.loads(f.read()) | ||||
|
r7889 | # allow hand-override of location for disambiguation | ||
# and ssh-server | ||||
|
r4585 | try: | ||
config.EngineFactory.location | ||||
except AttributeError: | ||||
config.EngineFactory.location = d['location'] | ||||
try: | ||||
config.EngineFactory.sshserver | ||||
except AttributeError: | ||||
|
r7889 | config.EngineFactory.sshserver = d.get('ssh') | ||
location = config.EngineFactory.location | ||||
|
r7890 | proto, ip = d['interface'].split('://') | ||
|
r8156 | ip = disambiguate_ip_address(ip, location) | ||
|
r7890 | d['interface'] = '%s://%s' % (proto, ip) | ||
|
r7889 | |||
# DO NOT allow override of basic URLs, serialization, or exec_key | ||||
# JSON file takes top priority there | ||||
|
r7890 | config.Session.key = cast_bytes(d['exec_key']) | ||
|
r7889 | |||
|
r7890 | config.EngineFactory.url = d['interface'] + ':%i' % d['registration'] | ||
|
r7889 | |||
config.Session.packer = d['pack'] | ||||
config.Session.unpacker = d['unpack'] | ||||
self.log.debug("Config changed:") | ||||
self.log.debug("%r", config) | ||||
self.connection_info = d | ||||
|
r6887 | |||
|
r6893 | def bind_kernel(self, **kwargs): | ||
"""Promote engine to listening kernel, accessible to frontends.""" | ||||
|
r6887 | if self.kernel_app is not None: | ||
return | ||||
|
r6893 | self.log.info("Opening ports for direct connections as an IPython kernel") | ||
|
r6887 | |||
kernel = self.kernel | ||||
|
r6893 | kwargs.setdefault('config', self.config) | ||
kwargs.setdefault('log', self.log) | ||||
kwargs.setdefault('profile_dir', self.profile_dir) | ||||
kwargs.setdefault('session', self.engine.session) | ||||
app = self.kernel_app = IPKernelApp(**kwargs) | ||||
# allow IPKernelApp.instance(): | ||||
IPKernelApp._instance = app | ||||
|
r6887 | app.init_connection_file() | ||
# relevant contents of init_sockets: | ||||
app.shell_port = app._bind_socket(kernel.shell_streams[0], app.shell_port) | ||||
app.log.debug("shell ROUTER Channel on port: %i", app.shell_port) | ||||
|
r4585 | |||
|
r6887 | app.iopub_port = app._bind_socket(kernel.iopub_socket, app.iopub_port) | ||
app.log.debug("iopub PUB Channel on port: %i", app.iopub_port) | ||||
kernel.stdin_socket = self.engine.context.socket(zmq.ROUTER) | ||||
app.stdin_port = app._bind_socket(kernel.stdin_socket, app.stdin_port) | ||||
app.log.debug("stdin ROUTER Channel on port: %i", app.stdin_port) | ||||
# start the heartbeat, and log connection info: | ||||
app.init_heartbeat() | ||||
app.log_connection_info() | ||||
app.write_connection_file() | ||||
|
r3985 | def init_engine(self): | ||
|
r3604 | # This is the working dir by now. | ||
sys.path.insert(0, '') | ||||
|
r3985 | config = self.config | ||
# print config | ||||
self.find_url_file() | ||||
|
r4120 | |||
# was the url manually specified? | ||||
keys = set(self.config.EngineFactory.keys()) | ||||
keys = keys.union(set(self.config.RegistrationFactory.keys())) | ||||
if keys.intersection(set(['ip', 'url', 'port'])): | ||||
# Connection info was specified, don't wait for the file | ||||
url_specified = True | ||||
self.wait_for_url_file = 0 | ||||
else: | ||||
url_specified = False | ||||
if self.wait_for_url_file and not os.path.exists(self.url_file): | ||||
|
r5483 | self.log.warn("url_file %r not found", self.url_file) | ||
self.log.warn("Waiting up to %.1f seconds for it to arrive.", self.wait_for_url_file) | ||||
|
r4120 | tic = time.time() | ||
while not os.path.exists(self.url_file) and (time.time()-tic < self.wait_for_url_file): | ||||
|
r5483 | # wait for url_file to exist, or until time limit | ||
|
r4120 | time.sleep(0.1) | ||
|
r3985 | if os.path.exists(self.url_file): | ||
|
r4585 | self.load_connector_file() | ||
|
r4120 | elif not url_specified: | ||
|
r5483 | self.log.fatal("Fatal: url file never arrived: %s", self.url_file) | ||
|
r4120 | self.exit(1) | ||
|
r3614 | |||
|
r3985 | try: | ||
|
r8062 | exec_lines = config.IPKernelApp.exec_lines | ||
|
r3985 | except AttributeError: | ||
|
r8062 | try: | ||
exec_lines = config.InteractiveShellApp.exec_lines | ||||
except AttributeError: | ||||
exec_lines = config.IPKernelApp.exec_lines = [] | ||||
try: | ||||
exec_files = config.IPKernelApp.exec_files | ||||
except AttributeError: | ||||
try: | ||||
exec_files = config.InteractiveShellApp.exec_files | ||||
except AttributeError: | ||||
exec_files = config.IPKernelApp.exec_files = [] | ||||
|
r3614 | |||
|
r3985 | if self.startup_script: | ||
|
r8062 | exec_files.append(self.startup_script) | ||
|
r3985 | if self.startup_command: | ||
exec_lines.append(self.startup_command) | ||||
|
r3604 | |||
|
r3985 | # Create the underlying shell class and Engine | ||
|
r3604 | # shell_class = import_item(self.master_config.Global.shell_class) | ||
|
r3985 | # print self.config | ||
|
r3604 | try: | ||
|
r7889 | self.engine = EngineFactory(config=config, log=self.log, | ||
connection_info=self.connection_info, | ||||
) | ||||
|
r3604 | except: | ||
self.log.error("Couldn't start the Engine", exc_info=True) | ||||
self.exit(1) | ||||
|
r4585 | |||
|
r3989 | def forward_logging(self): | ||
if self.log_url: | ||||
|
r5483 | self.log.info("Forwarding logging to %s", self.log_url) | ||
|
r3989 | context = self.engine.context | ||
lsock = context.socket(zmq.PUB) | ||||
lsock.connect(self.log_url) | ||||
handler = EnginePUBHandler(self.engine, lsock) | ||||
handler.setLevel(self.log_level) | ||||
self.log.addHandler(handler) | ||||
|
r4585 | |||
|
r3985 | def init_mpi(self): | ||
|
r3604 | global mpi | ||
|
r3985 | self.mpi = MPI(config=self.config) | ||
mpi_import_statement = self.mpi.init_script | ||||
if mpi_import_statement: | ||||
|
r3604 | try: | ||
self.log.info("Initializing MPI:") | ||||
self.log.info(mpi_import_statement) | ||||
exec mpi_import_statement in globals() | ||||
except: | ||||
mpi = None | ||||
else: | ||||
mpi = None | ||||
|
r5214 | @catch_config_error | ||
|
r3986 | def initialize(self, argv=None): | ||
super(IPEngineApp, self).initialize(argv) | ||||
self.init_mpi() | ||||
self.init_engine() | ||||
|
r3989 | self.forward_logging() | ||
|
r3986 | |||
|
r3985 | def start(self): | ||
|
r3604 | self.engine.start() | ||
try: | ||||
self.engine.loop.start() | ||||
except KeyboardInterrupt: | ||||
self.log.critical("Engine Interrupted, shutting down...\n") | ||||
def launch_new_instance(): | ||||
|
r3985 | """Create and run the IPython engine""" | ||
|
r3999 | app = IPEngineApp.instance() | ||
|
r3986 | app.initialize() | ||
|
r3604 | app.start() | ||
if __name__ == '__main__': | ||||
launch_new_instance() | ||||