ipengineapp.py
276 lines
| 8.9 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 | ||||
import zmq | ||||
from zmq.eventloop import ioloop | ||||
|
r4024 | from IPython.core.profiledir import ProfileDir | ||
|
r3993 | from IPython.parallel.apps.baseapp import BaseParallelApplication | ||
|
r3604 | from IPython.zmq.log import EnginePUBHandler | ||
|
r3985 | from IPython.config.configurable import Configurable | ||
|
r4006 | from IPython.zmq.session import Session | ||
|
r3673 | from IPython.parallel.engine.engine import EngineFactory | ||
from IPython.parallel.engine.streamkernel import Kernel | ||||
|
r3666 | from IPython.parallel.util import disambiguate_url | ||
|
r3673 | |||
|
r3604 | from IPython.utils.importstring import import_item | ||
|
r3989 | from IPython.utils.traitlets import Bool, Unicode, Dict, List | ||
|
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". | ||
|
r3992 | See the `profile` and `profile_dir` options for details. | ||
|
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 | |||
|
r3985 | def _on_use_changed(self, old, new): | ||
# 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 | #----------------------------------------------------------------------------- | ||
|
r3992 | class IPEngineApp(BaseParallelApplication): | ||
|
r3985 | |||
|
r4025 | name = Unicode(u'ipengine') | ||
|
r3985 | description = Unicode(_description) | ||
|
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 | ) | ||
|
r3604 | |||
|
r3985 | url_file_name = Unicode(u'ipcontroller-engine.json') | ||
|
r3989 | log_url = Unicode('', config=True, | ||
help="""The URL for the iploggerapp instance, for forwarding | ||||
logging to a central location.""") | ||||
|
r3604 | |||
|
r3985 | aliases = Dict(dict( | ||
file = 'IPEngineApp.url_file', | ||||
c = 'IPEngineApp.startup_command', | ||||
s = 'IPEngineApp.startup_script', | ||||
|
r3604 | |||
|
r4006 | ident = 'Session.session', | ||
user = 'Session.username', | ||||
exec_key = 'Session.keyfile', | ||||
|
r3604 | |||
|
r3985 | url = 'EngineFactory.url', | ||
ip = 'EngineFactory.ip', | ||||
transport = 'EngineFactory.transport', | ||||
port = 'EngineFactory.regport', | ||||
location = 'EngineFactory.location', | ||||
timeout = 'EngineFactory.timeout', | ||||
|
r3992 | profile = "IPEngineApp.profile", | ||
profile_dir = 'ProfileDir.location', | ||||
|
r3985 | |||
mpi = 'MPI.use', | ||||
log_level = 'IPEngineApp.log_level', | ||||
|
r3989 | log_url = 'IPEngineApp.log_url' | ||
|
r3985 | )) | ||
|
r3604 | |||
|
r3614 | # def find_key_file(self): | ||
# """Set the key file. | ||||
# | ||||
# Here we don't try to actually see if it exists for is valid as that | ||||
# is hadled by the connection logic. | ||||
# """ | ||||
# config = self.master_config | ||||
# # Find the actual controller key file | ||||
# if not config.Global.key_file: | ||||
# try_this = os.path.join( | ||||
|
r3992 | # config.Global.profile_dir, | ||
|
r3614 | # config.Global.security_dir, | ||
# config.Global.key_file_name | ||||
# ) | ||||
# config.Global.key_file = try_this | ||||
def find_url_file(self): | ||||
|
r3604 | """Set the key file. | ||
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 | ) | ||
|
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() | ||||
|
r3614 | # if os.path.exists(config.Global.key_file) and config.Global.secure: | ||
# config.SessionFactory.exec_key = config.Global.key_file | ||||
|
r3985 | if os.path.exists(self.url_file): | ||
with open(self.url_file) as f: | ||||
|
r3614 | d = json.loads(f.read()) | ||
for k,v in d.iteritems(): | ||||
if isinstance(v, unicode): | ||||
d[k] = v.encode() | ||||
if d['exec_key']: | ||||
|
r4006 | config.Session.key = d['exec_key'] | ||
|
r3614 | d['url'] = disambiguate_url(d['url'], d['location']) | ||
|
r3985 | config.EngineFactory.url = d['url'] | ||
|
r3614 | config.EngineFactory.location = d['location'] | ||
|
r3985 | try: | ||
exec_lines = config.Kernel.exec_lines | ||||
except AttributeError: | ||||
config.Kernel.exec_lines = [] | ||||
exec_lines = config.Kernel.exec_lines | ||||
|
r3614 | |||
|
r3985 | if self.startup_script: | ||
enc = sys.getfilesystemencoding() or 'utf8' | ||||
cmd="execfile(%r)"%self.startup_script.encode(enc) | ||||
exec_lines.append(cmd) | ||||
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: | ||
|
r3985 | self.engine = EngineFactory(config=config, log=self.log) | ||
|
r3604 | except: | ||
self.log.error("Couldn't start the Engine", exc_info=True) | ||||
self.exit(1) | ||||
|
r3989 | def forward_logging(self): | ||
if self.log_url: | ||||
self.log.info("Forwarding logging to %s"%self.log_url) | ||||
context = self.engine.context | ||||
lsock = context.socket(zmq.PUB) | ||||
lsock.connect(self.log_url) | ||||
self.log.removeHandler(self._log_handler) | ||||
handler = EnginePUBHandler(self.engine, lsock) | ||||
handler.setLevel(self.log_level) | ||||
self.log.addHandler(handler) | ||||
self._log_handler = handler | ||||
|
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 | ||||
|
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() | ||||