class EngineConnector(HasTraits): """A simple object for accessing the various zmq connections of an object. Attributes are: id (int): engine ID uuid (str): uuid (unused?) queue (str): identity of queue's XREQ socket registration (str): identity of registration XREQ socket heartbeat (str): identity of heartbeat XREQ socket """ id=Integer(0) queue=CBytes() control=CBytes() registration=CBytes() heartbeat=CBytes() pending=Set()
class ZMQDataPublisher(Configurable): topic = topic = CBytes(b'datapub') session = Instance(Session) pub_socket = Instance('zmq.Socket') parent_header = Dict({}) def set_parent(self, parent): """Set the parent for outbound messages.""" self.parent_header = extract_header(parent) def publish_data(self, data): """publish a data_message on the IOPub channel Parameters ---------- data : dict The data to be published. Think of it as a namespace. """ session = self.session buffers = serialize_object(data, buffer_threshold=session.buffer_threshold, item_threshold=session.item_threshold, ) content = json_clean(dict(keys=data.keys())) session.send(self.pub_socket, 'data_message', content=content, parent=self.parent_header, buffers=buffers, ident=self.topic, )
class ZMQDisplayPublisher(DisplayPublisher): """A display publisher that publishes data using a ZeroMQ PUB socket.""" session = Instance(Session) pub_socket = Instance(SocketABC) parent_header = Dict({}) topic = CBytes(b'displaypub') def set_parent(self, parent): """Set the parent for outbound messages.""" self.parent_header = extract_header(parent) def _flush_streams(self): """flush IO Streams prior to display""" sys.stdout.flush() sys.stderr.flush() def publish(self, source, data, metadata=None): self._flush_streams() if metadata is None: metadata = {} self._validate_data(source, data, metadata) content = {} content['source'] = source content['data'] = encode_images(data) content['metadata'] = metadata self.session.send( self.pub_socket, u'display_data', json_clean(content), parent=self.parent_header, ident=self.topic, ) def clear_output(self, stdout=True, stderr=True, other=True): content = dict(stdout=stdout, stderr=stderr, other=other) if stdout: print('\r', file=sys.stdout, end='') if stderr: print('\r', file=sys.stderr, end='') self._flush_streams() self.session.send( self.pub_socket, u'clear_output', content, parent=self.parent_header, ident=self.topic, )
class ZMQDisplayPublisher(DisplayPublisher): """A display publisher that publishes data using a ZeroMQ PUB socket.""" session = Instance(Session, allow_none=True) pub_socket = Instance(SocketABC, allow_none=True) parent_header = Dict({}) topic = CBytes(b'display_data') def set_parent(self, parent): """Set the parent for outbound messages.""" self.parent_header = extract_header(parent) def _flush_streams(self): """flush IO Streams prior to display""" sys.stdout.flush() sys.stderr.flush() def publish(self, data, metadata=None, source=None): self._flush_streams() if metadata is None: metadata = {} self._validate_data(data, metadata) content = {} content['data'] = encode_images(data) content['metadata'] = metadata self.session.send( self.pub_socket, u'display_data', json_clean(content), parent=self.parent_header, ident=self.topic, ) def clear_output(self, wait=False): content = dict(wait=wait) self._flush_streams() self.session.send( self.pub_socket, u'clear_output', content, parent=self.parent_header, ident=self.topic, )
class Session(Configurable): """Object for handling serialization and sending of messages. The Session object handles building messages and sending them with ZMQ sockets or ZMQStream objects. Objects can communicate with each other over the network via Session objects, and only need to work with the dict-based IPython message spec. The Session will handle serialization/deserialization, security, and metadata. Sessions support configurable serialiization via packer/unpacker traits, and signing with HMAC digests via the key/keyfile traits. Parameters ---------- debug : bool whether to trigger extra debugging statements packer/unpacker : str : 'json', 'pickle' or import_string importstrings for methods to serialize message parts. If just 'json' or 'pickle', predefined JSON and pickle packers will be used. Otherwise, the entire importstring must be used. The functions must accept at least valid JSON input, and output *bytes*. For example, to use msgpack: packer = 'msgpack.packb', unpacker='msgpack.unpackb' pack/unpack : callables You can also set the pack/unpack callables for serialization directly. session : bytes the ID of this Session object. The default is to generate a new UUID. username : unicode username added to message headers. The default is to ask the OS. key : bytes The key used to initialize an HMAC signature. If unset, messages will not be signed or checked. keyfile : filepath The file containing a key. If this is set, `key` will be initialized to the contents of the file. """ debug = Bool(False, config=True, help="""Debug output in the Session""") packer = DottedObjectName( 'json', config=True, help="""The name of the packer for serializing messages. Should be one of 'json', 'pickle', or an import name for a custom callable serializer.""") def _packer_changed(self, name, old, new): if new.lower() == 'json': self.pack = json_packer self.unpack = json_unpacker elif new.lower() == 'pickle': self.pack = pickle_packer self.unpack = pickle_unpacker else: self.pack = import_item(str(new)) unpacker = DottedObjectName( 'json', config=True, help="""The name of the unpacker for unserializing messages. Only used with custom functions for `packer`.""") def _unpacker_changed(self, name, old, new): if new.lower() == 'json': self.pack = json_packer self.unpack = json_unpacker elif new.lower() == 'pickle': self.pack = pickle_packer self.unpack = pickle_unpacker else: self.unpack = import_item(str(new)) session = CBytes(b'', config=True, help="""The UUID identifying this session.""") def _session_default(self): return bytes(uuid.uuid4()) username = Unicode( os.environ.get('USER', 'username'), config=True, help="""Username for the Session. Default is your system username.""") # message signature related traits: key = CBytes(b'', config=True, help="""execution key, for extra authentication.""") def _key_changed(self, name, old, new): if new: self.auth = hmac.HMAC(new) else: self.auth = None auth = Instance(hmac.HMAC) digest_history = Set() keyfile = Unicode('', config=True, help="""path to file containing execution key.""") def _keyfile_changed(self, name, old, new): with open(new, 'rb') as f: self.key = f.read().strip() pack = Any(default_packer) # the actual packer function def _pack_changed(self, name, old, new): if not callable(new): raise TypeError("packer must be callable, not %s" % type(new)) unpack = Any(default_unpacker) # the actual packer function def _unpack_changed(self, name, old, new): # unpacker is not checked - it is assumed to be if not callable(new): raise TypeError("unpacker must be callable, not %s" % type(new)) def __init__(self, **kwargs): """create a Session object Parameters ---------- debug : bool whether to trigger extra debugging statements packer/unpacker : str : 'json', 'pickle' or import_string importstrings for methods to serialize message parts. If just 'json' or 'pickle', predefined JSON and pickle packers will be used. Otherwise, the entire importstring must be used. The functions must accept at least valid JSON input, and output *bytes*. For example, to use msgpack: packer = 'msgpack.packb', unpacker='msgpack.unpackb' pack/unpack : callables You can also set the pack/unpack callables for serialization directly. session : bytes the ID of this Session object. The default is to generate a new UUID. username : unicode username added to message headers. The default is to ask the OS. key : bytes The key used to initialize an HMAC signature. If unset, messages will not be signed or checked. keyfile : filepath The file containing a key. If this is set, `key` will be initialized to the contents of the file. """ super(Session, self).__init__(**kwargs) self._check_packers() self.none = self.pack({}) @property def msg_id(self): """always return new uuid""" return str(uuid.uuid4()) def _check_packers(self): """check packers for binary data and datetime support.""" pack = self.pack unpack = self.unpack # check simple serialization msg = dict(a=[1, 'hi']) try: packed = pack(msg) except Exception: raise ValueError("packer could not serialize a simple message") # ensure packed message is bytes if not isinstance(packed, bytes): raise ValueError("message packed to %r, but bytes are required" % type(packed)) # check that unpack is pack's inverse try: unpacked = unpack(packed) except Exception: raise ValueError("unpacker could not handle the packer's output") # check datetime support msg = dict(t=datetime.now()) try: unpacked = unpack(pack(msg)) except Exception: self.pack = lambda o: pack(squash_dates(o)) self.unpack = lambda s: extract_dates(unpack(s)) def msg_header(self, msg_type): return msg_header(self.msg_id, msg_type, self.username, self.session) def msg(self, msg_type, content=None, parent=None, subheader=None): """Return the nested message dict. This format is different from what is sent over the wire. The self.serialize method converts this nested message dict to the wire format, which uses a message list. """ msg = {} msg['header'] = self.msg_header(msg_type) msg['msg_id'] = msg['header']['msg_id'] msg['parent_header'] = {} if parent is None else extract_header(parent) msg['msg_type'] = msg_type msg['content'] = {} if content is None else content sub = {} if subheader is None else subheader msg['header'].update(sub) return msg def sign(self, msg_list): """Sign a message with HMAC digest. If no auth, return b''. Parameters ---------- msg_list : list The [p_header,p_parent,p_content] part of the message list. """ if self.auth is None: return b'' h = self.auth.copy() for m in msg_list: h.update(m) return h.hexdigest() def serialize(self, msg, ident=None): """Serialize the message components to bytes. Parameters ---------- msg : dict or Message The nexted message dict as returned by the self.msg method. Returns ------- msg_list : list The list of bytes objects to be sent with the format: [ident1,ident2,...,DELIM,HMAC,p_header,p_parent,p_content, buffer1,buffer2,...]. In this list, the p_* entities are the packed or serialized versions, so if JSON is used, these are uft8 encoded JSON strings. """ content = msg.get('content', {}) if content is None: content = self.none elif isinstance(content, dict): content = self.pack(content) elif isinstance(content, bytes): # content is already packed, as in a relayed message pass elif isinstance(content, unicode): # should be bytes, but JSON often spits out unicode content = content.encode('utf8') else: raise TypeError("Content incorrect type: %s" % type(content)) real_message = [ self.pack(msg['header']), self.pack(msg['parent_header']), content ] to_send = [] if isinstance(ident, list): # accept list of idents to_send.extend(ident) elif ident is not None: to_send.append(ident) to_send.append(DELIM) signature = self.sign(real_message) to_send.append(signature) to_send.extend(real_message) return to_send def send(self, stream, msg_or_type, content=None, parent=None, ident=None, buffers=None, subheader=None, track=False): """Build and send a message via stream or socket. The message format used by this function internally is as follows: [ident1,ident2,...,DELIM,HMAC,p_header,p_parent,p_content, buffer1,buffer2,...] The self.serialize method converts the nested message dict into this format. Parameters ---------- stream : zmq.Socket or ZMQStream the socket-like object used to send the data msg_or_type : str or Message/dict Normally, msg_or_type will be a msg_type unless a message is being sent more than once. content : dict or None the content of the message (ignored if msg_or_type is a message) parent : Message or dict or None the parent or parent header describing the parent of this message ident : bytes or list of bytes the zmq.IDENTITY routing path subheader : dict or None extra header keys for this message's header buffers : list or None the already-serialized buffers to be appended to the message track : bool whether to track. Only for use with Sockets, because ZMQStream objects cannot track messages. Returns ------- msg : message dict the constructed message (msg,tracker) : (message dict, MessageTracker) if track=True, then a 2-tuple will be returned, the first element being the constructed message, and the second being the MessageTracker """ if not isinstance(stream, (zmq.Socket, ZMQStream)): raise TypeError("stream must be Socket or ZMQStream, not %r" % type(stream)) elif track and isinstance(stream, ZMQStream): raise TypeError("ZMQStream cannot track messages") if isinstance(msg_or_type, (Message, dict)): # we got a Message, not a msg_type # don't build a new Message msg = msg_or_type else: msg = self.msg(msg_or_type, content, parent, subheader) buffers = [] if buffers is None else buffers to_send = self.serialize(msg, ident) flag = 0 if buffers: flag = zmq.SNDMORE _track = False else: _track = track if track: tracker = stream.send_multipart(to_send, flag, copy=False, track=_track) else: tracker = stream.send_multipart(to_send, flag, copy=False) for b in buffers[:-1]: stream.send(b, flag, copy=False) if buffers: if track: tracker = stream.send(buffers[-1], copy=False, track=track) else: tracker = stream.send(buffers[-1], copy=False) # omsg = Message(msg) if self.debug: pprint.pprint(msg) pprint.pprint(to_send) pprint.pprint(buffers) msg['tracker'] = tracker return msg def send_raw(self, stream, msg_list, flags=0, copy=True, ident=None): """Send a raw message via ident path. This method is used to send a already serialized message. Parameters ---------- stream : ZMQStream or Socket The ZMQ stream or socket to use for sending the message. msg_list : list The serialized list of messages to send. This only includes the [p_header,p_parent,p_content,buffer1,buffer2,...] portion of the message. ident : ident or list A single ident or a list of idents to use in sending. """ to_send = [] if isinstance(ident, bytes): ident = [ident] if ident is not None: to_send.extend(ident) to_send.append(DELIM) to_send.append(self.sign(msg_list)) to_send.extend(msg_list) stream.send_multipart(msg_list, flags, copy=copy) def recv(self, socket, mode=zmq.NOBLOCK, content=True, copy=True): """Receive and unpack a message. Parameters ---------- socket : ZMQStream or Socket The socket or stream to use in receiving. Returns ------- [idents], msg [idents] is a list of idents and msg is a nested message dict of same format as self.msg returns. """ if isinstance(socket, ZMQStream): socket = socket.socket try: msg_list = socket.recv_multipart(mode) except zmq.ZMQError as e: if e.errno == zmq.EAGAIN: # We can convert EAGAIN to None as we know in this case # recv_multipart won't return None. return None, None else: raise # split multipart message into identity list and message dict # invalid large messages can cause very expensive string comparisons idents, msg_list = self.feed_identities(msg_list, copy) try: return idents, self.unpack_message(msg_list, content=content, copy=copy) except Exception as e: print(idents, msg_list) # TODO: handle it raise e def feed_identities(self, msg_list, copy=True): """Split the identities from the rest of the message. Feed until DELIM is reached, then return the prefix as idents and remainder as msg_list. This is easily broken by setting an IDENT to DELIM, but that would be silly. Parameters ---------- msg_list : a list of Message or bytes objects The message to be split. copy : bool flag determining whether the arguments are bytes or Messages Returns ------- (idents,msg_list) : two lists idents will always be a list of bytes - the indentity prefix msg_list will be a list of bytes or Messages, unchanged from input msg_list should be unpackable via self.unpack_message at this point. """ if copy: idx = msg_list.index(DELIM) return msg_list[:idx], msg_list[idx + 1:] else: failed = True for idx, m in enumerate(msg_list): if m.bytes == DELIM: failed = False break if failed: raise ValueError("DELIM not in msg_list") idents, msg_list = msg_list[:idx], msg_list[idx + 1:] return [m.bytes for m in idents], msg_list def unpack_message(self, msg_list, content=True, copy=True): """Return a message object from the format sent by self.send. Parameters: ----------- content : bool (True) whether to unpack the content dict (True), or leave it serialized (False) copy : bool (True) whether to return the bytes (True), or the non-copying Message object in each place (False) """ minlen = 4 message = {} if not copy: for i in range(minlen): msg_list[i] = msg_list[i].bytes if self.auth is not None: signature = msg_list[0] if signature in self.digest_history: raise ValueError("Duplicate Signature: %r" % signature) self.digest_history.add(signature) check = self.sign(msg_list[1:4]) if not signature == check: raise ValueError("Invalid Signature: %r" % signature) if not len(msg_list) >= minlen: raise TypeError( "malformed message, must have at least %i elements" % minlen) message['header'] = self.unpack(msg_list[1]) message['msg_type'] = message['header']['msg_type'] message['parent_header'] = self.unpack(msg_list[2]) if content: message['content'] = self.unpack(msg_list[3]) else: message['content'] = msg_list[3] message['buffers'] = msg_list[4:] return message
class Session(Configurable): """Object for handling serialization and sending of messages. The Session object handles building messages and sending them with ZMQ sockets or ZMQStream objects. Objects can communicate with each other over the network via Session objects, and only need to work with the dict-based IPython message spec. The Session will handle serialization/deserialization, security, and metadata. Sessions support configurable serialization via packer/unpacker traits, and signing with HMAC digests via the key/keyfile traits. Parameters ---------- debug : bool whether to trigger extra debugging statements packer/unpacker : str : 'json', 'pickle' or import_string importstrings for methods to serialize message parts. If just 'json' or 'pickle', predefined JSON and pickle packers will be used. Otherwise, the entire importstring must be used. The functions must accept at least valid JSON input, and output *bytes*. For example, to use msgpack: packer = 'msgpack.packb', unpacker='msgpack.unpackb' pack/unpack : callables You can also set the pack/unpack callables for serialization directly. session : bytes the ID of this Session object. The default is to generate a new UUID. username : unicode username added to message headers. The default is to ask the OS. key : bytes The key used to initialize an HMAC signature. If unset, messages will not be signed or checked. keyfile : filepath The file containing a key. If this is set, `key` will be initialized to the contents of the file. """ debug = Bool(False, config=True, help="""Debug output in the Session""") packer = DottedObjectName( 'json', config=True, help="""The name of the packer for serializing messages. Should be one of 'json', 'pickle', or an import name for a custom callable serializer.""") def _packer_changed(self, name, old, new): if new.lower() == 'json': self.pack = json_packer self.unpack = json_unpacker self.unpacker = new elif new.lower() == 'pickle': self.pack = pickle_packer self.unpack = pickle_unpacker self.unpacker = new else: self.pack = import_item(str(new)) unpacker = DottedObjectName( 'json', config=True, help="""The name of the unpacker for unserializing messages. Only used with custom functions for `packer`.""") def _unpacker_changed(self, name, old, new): if new.lower() == 'json': self.pack = json_packer self.unpack = json_unpacker self.packer = new elif new.lower() == 'pickle': self.pack = pickle_packer self.unpack = pickle_unpacker self.packer = new else: self.unpack = import_item(str(new)) session = CUnicode(u'', config=True, help="""The UUID identifying this session.""") def _session_default(self): u = unicode_type(uuid.uuid4()) self.bsession = u.encode('ascii') return u def _session_changed(self, name, old, new): self.bsession = self.session.encode('ascii') # bsession is the session as bytes bsession = CBytes(b'') username = Unicode( str_to_unicode(os.environ.get('USER', 'username')), help="""Username for the Session. Default is your system username.""", config=True) metadata = Dict( {}, config=True, help= """Metadata dictionary, which serves as the default top-level metadata dict for each message.""" ) # if 0, no adapting to do. adapt_version = Integer(0) # message signature related traits: key = CBytes(b'', config=True, help="""execution key, for extra authentication.""") def _key_changed(self): self._new_auth() signature_scheme = Unicode( 'hmac-sha256', config=True, help="""The digest scheme used to construct the message signatures. Must have the form 'hmac-HASH'.""") def _signature_scheme_changed(self, name, old, new): if not new.startswith('hmac-'): raise TraitError( "signature_scheme must start with 'hmac-', got %r" % new) hash_name = new.split('-', 1)[1] try: self.digest_mod = getattr(hashlib, hash_name) except AttributeError: raise TraitError("hashlib has no such attribute: %s" % hash_name) self._new_auth() digest_mod = Any() def _digest_mod_default(self): return hashlib.sha256 auth = Instance(hmac.HMAC) def _new_auth(self): if self.key: self.auth = hmac.HMAC(self.key, digestmod=self.digest_mod) else: self.auth = None digest_history = Set() digest_history_size = Integer( 2**16, config=True, help="""The maximum number of digests to remember. The digest history will be culled when it exceeds this value. """) keyfile = Unicode('', config=True, help="""path to file containing execution key.""") def _keyfile_changed(self, name, old, new): with open(new, 'rb') as f: self.key = f.read().strip() # for protecting against sends from forks pid = Integer() # serialization traits: pack = Any(default_packer) # the actual packer function def _pack_changed(self, name, old, new): if not callable(new): raise TypeError("packer must be callable, not %s" % type(new)) unpack = Any(default_unpacker) # the actual packer function def _unpack_changed(self, name, old, new): # unpacker is not checked - it is assumed to be if not callable(new): raise TypeError("unpacker must be callable, not %s" % type(new)) # thresholds: copy_threshold = Integer( 2**16, config=True, help= "Threshold (in bytes) beyond which a buffer should be sent without copying." ) buffer_threshold = Integer( MAX_BYTES, config=True, help= "Threshold (in bytes) beyond which an object's buffer should be extracted to avoid pickling." ) item_threshold = Integer( MAX_ITEMS, config=True, help= """The maximum number of items for a container to be introspected for custom serialization. Containers larger than this are pickled outright. """) def __init__(self, **kwargs): """create a Session object Parameters ---------- debug : bool whether to trigger extra debugging statements packer/unpacker : str : 'json', 'pickle' or import_string importstrings for methods to serialize message parts. If just 'json' or 'pickle', predefined JSON and pickle packers will be used. Otherwise, the entire importstring must be used. The functions must accept at least valid JSON input, and output *bytes*. For example, to use msgpack: packer = 'msgpack.packb', unpacker='msgpack.unpackb' pack/unpack : callables You can also set the pack/unpack callables for serialization directly. session : unicode (must be ascii) the ID of this Session object. The default is to generate a new UUID. bsession : bytes The session as bytes username : unicode username added to message headers. The default is to ask the OS. key : bytes The key used to initialize an HMAC signature. If unset, messages will not be signed or checked. signature_scheme : str The message digest scheme. Currently must be of the form 'hmac-HASH', where 'HASH' is a hashing function available in Python's hashlib. The default is 'hmac-sha256'. This is ignored if 'key' is empty. keyfile : filepath The file containing a key. If this is set, `key` will be initialized to the contents of the file. """ super(Session, self).__init__(**kwargs) self._check_packers() self.none = self.pack({}) # ensure self._session_default() if necessary, so bsession is defined: self.session self.pid = os.getpid() @property def msg_id(self): """always return new uuid""" return str(uuid.uuid4()) def _check_packers(self): """check packers for datetime support.""" pack = self.pack unpack = self.unpack # check simple serialization msg = dict(a=[1, 'hi']) try: packed = pack(msg) except Exception as e: msg = "packer '{packer}' could not serialize a simple message: {e}{jsonmsg}" if self.packer == 'json': jsonmsg = "\nzmq.utils.jsonapi.jsonmod = %s" % jsonapi.jsonmod else: jsonmsg = "" raise ValueError( msg.format(packer=self.packer, e=e, jsonmsg=jsonmsg)) # ensure packed message is bytes if not isinstance(packed, bytes): raise ValueError("message packed to %r, but bytes are required" % type(packed)) # check that unpack is pack's inverse try: unpacked = unpack(packed) assert unpacked == msg except Exception as e: msg = "unpacker '{unpacker}' could not handle output from packer '{packer}': {e}{jsonmsg}" if self.packer == 'json': jsonmsg = "\nzmq.utils.jsonapi.jsonmod = %s" % jsonapi.jsonmod else: jsonmsg = "" raise ValueError( msg.format(packer=self.packer, unpacker=self.unpacker, e=e, jsonmsg=jsonmsg)) # check datetime support msg = dict(t=datetime.now()) try: unpacked = unpack(pack(msg)) if isinstance(unpacked['t'], datetime): raise ValueError("Shouldn't deserialize to datetime") except Exception: self.pack = lambda o: pack(squash_dates(o)) self.unpack = lambda s: unpack(s) def msg_header(self, msg_type): return msg_header(self.msg_id, msg_type, self.username, self.session) def msg(self, msg_type, content=None, parent=None, header=None, metadata=None): """Return the nested message dict. This format is different from what is sent over the wire. The serialize/deserialize methods converts this nested message dict to the wire format, which is a list of message parts. """ msg = {} header = self.msg_header(msg_type) if header is None else header msg['header'] = header msg['msg_id'] = header['msg_id'] msg['msg_type'] = header['msg_type'] msg['parent_header'] = {} if parent is None else extract_header(parent) msg['content'] = {} if content is None else content msg['metadata'] = self.metadata.copy() if metadata is not None: msg['metadata'].update(metadata) return msg def sign(self, msg_list): """Sign a message with HMAC digest. If no auth, return b''. Parameters ---------- msg_list : list The [p_header,p_parent,p_content] part of the message list. """ if self.auth is None: return b'' h = self.auth.copy() for m in msg_list: h.update(m) return str_to_bytes(h.hexdigest()) def serialize(self, msg, ident=None): """Serialize the message components to bytes. This is roughly the inverse of deserialize. The serialize/deserialize methods work with full message lists, whereas pack/unpack work with the individual message parts in the message list. Parameters ---------- msg : dict or Message The next message dict as returned by the self.msg method. Returns ------- msg_list : list The list of bytes objects to be sent with the format:: [ident1, ident2, ..., DELIM, HMAC, p_header, p_parent, p_metadata, p_content, buffer1, buffer2, ...] In this list, the ``p_*`` entities are the packed or serialized versions, so if JSON is used, these are utf8 encoded JSON strings. """ content = msg.get('content', {}) if content is None: content = self.none elif isinstance(content, dict): content = self.pack(content) elif isinstance(content, bytes): # content is already packed, as in a relayed message pass elif isinstance(content, unicode_type): # should be bytes, but JSON often spits out unicode content = content.encode('utf8') else: raise TypeError("Content incorrect type: %s" % type(content)) real_message = [ self.pack(msg['header']), self.pack(msg['parent_header']), self.pack(msg['metadata']), content, ] to_send = [] if isinstance(ident, list): # accept list of idents to_send.extend(ident) elif ident is not None: to_send.append(ident) to_send.append(DELIM) signature = self.sign(real_message) to_send.append(signature) to_send.extend(real_message) return to_send def send(self, stream, msg_or_type, content=None, parent=None, ident=None, buffers=None, track=False, header=None, metadata=None): """Build and send a message via stream or socket. The message format used by this function internally is as follows: [ident1,ident2,...,DELIM,HMAC,p_header,p_parent,p_content, buffer1,buffer2,...] The serialize/deserialize methods convert the nested message dict into this format. Parameters ---------- stream : zmq.Socket or ZMQStream The socket-like object used to send the data. msg_or_type : str or Message/dict Normally, msg_or_type will be a msg_type unless a message is being sent more than once. If a header is supplied, this can be set to None and the msg_type will be pulled from the header. content : dict or None The content of the message (ignored if msg_or_type is a message). header : dict or None The header dict for the message (ignored if msg_to_type is a message). parent : Message or dict or None The parent or parent header describing the parent of this message (ignored if msg_or_type is a message). ident : bytes or list of bytes The zmq.IDENTITY routing path. metadata : dict or None The metadata describing the message buffers : list or None The already-serialized buffers to be appended to the message. track : bool Whether to track. Only for use with Sockets, because ZMQStream objects cannot track messages. Returns ------- msg : dict The constructed message. """ if not isinstance(stream, zmq.Socket): # ZMQStreams and dummy sockets do not support tracking. track = False if isinstance(msg_or_type, (Message, dict)): # We got a Message or message dict, not a msg_type so don't # build a new Message. msg = msg_or_type buffers = buffers or msg.get('buffers', []) else: msg = self.msg(msg_or_type, content=content, parent=parent, header=header, metadata=metadata) if not os.getpid() == self.pid: io.rprint("WARNING: attempted to send message from fork") io.rprint(msg) return buffers = [] if buffers is None else buffers if self.adapt_version: msg = adapt(msg, self.adapt_version) to_send = self.serialize(msg, ident) to_send.extend(buffers) longest = max([len(s) for s in to_send]) copy = (longest < self.copy_threshold) if buffers and track and not copy: # only really track when we are doing zero-copy buffers tracker = stream.send_multipart(to_send, copy=False, track=True) else: # use dummy tracker, which will be done immediately tracker = DONE stream.send_multipart(to_send, copy=copy) if self.debug: pprint.pprint(msg) pprint.pprint(to_send) pprint.pprint(buffers) msg['tracker'] = tracker return msg def send_raw(self, stream, msg_list, flags=0, copy=True, ident=None): """Send a raw message via ident path. This method is used to send a already serialized message. Parameters ---------- stream : ZMQStream or Socket The ZMQ stream or socket to use for sending the message. msg_list : list The serialized list of messages to send. This only includes the [p_header,p_parent,p_metadata,p_content,buffer1,buffer2,...] portion of the message. ident : ident or list A single ident or a list of idents to use in sending. """ to_send = [] if isinstance(ident, bytes): ident = [ident] if ident is not None: to_send.extend(ident) to_send.append(DELIM) to_send.append(self.sign(msg_list)) to_send.extend(msg_list) stream.send_multipart(to_send, flags, copy=copy) def recv(self, socket, mode=zmq.NOBLOCK, content=True, copy=True): """Receive and unpack a message. Parameters ---------- socket : ZMQStream or Socket The socket or stream to use in receiving. Returns ------- [idents], msg [idents] is a list of idents and msg is a nested message dict of same format as self.msg returns. """ if isinstance(socket, ZMQStream): socket = socket.socket try: msg_list = socket.recv_multipart(mode, copy=copy) except zmq.ZMQError as e: if e.errno == zmq.EAGAIN: # We can convert EAGAIN to None as we know in this case # recv_multipart won't return None. return None, None else: raise # split multipart message into identity list and message dict # invalid large messages can cause very expensive string comparisons idents, msg_list = self.feed_identities(msg_list, copy) try: return idents, self.deserialize(msg_list, content=content, copy=copy) except Exception as e: # TODO: handle it raise e def feed_identities(self, msg_list, copy=True): """Split the identities from the rest of the message. Feed until DELIM is reached, then return the prefix as idents and remainder as msg_list. This is easily broken by setting an IDENT to DELIM, but that would be silly. Parameters ---------- msg_list : a list of Message or bytes objects The message to be split. copy : bool flag determining whether the arguments are bytes or Messages Returns ------- (idents, msg_list) : two lists idents will always be a list of bytes, each of which is a ZMQ identity. msg_list will be a list of bytes or zmq.Messages of the form [HMAC,p_header,p_parent,p_content,buffer1,buffer2,...] and should be unpackable/unserializable via self.deserialize at this point. """ if copy: idx = msg_list.index(DELIM) return msg_list[:idx], msg_list[idx + 1:] else: failed = True for idx, m in enumerate(msg_list): if m.bytes == DELIM: failed = False break if failed: raise ValueError("DELIM not in msg_list") idents, msg_list = msg_list[:idx], msg_list[idx + 1:] return [m.bytes for m in idents], msg_list def _add_digest(self, signature): """add a digest to history to protect against replay attacks""" if self.digest_history_size == 0: # no history, never add digests return self.digest_history.add(signature) if len(self.digest_history) > self.digest_history_size: # threshold reached, cull 10% self._cull_digest_history() def _cull_digest_history(self): """cull the digest history Removes a randomly selected 10% of the digest history """ current = len(self.digest_history) n_to_cull = max(int(current // 10), current - self.digest_history_size) if n_to_cull >= current: self.digest_history = set() return to_cull = random.sample(self.digest_history, n_to_cull) self.digest_history.difference_update(to_cull) def deserialize(self, msg_list, content=True, copy=True): """Unserialize a msg_list to a nested message dict. This is roughly the inverse of serialize. The serialize/deserialize methods work with full message lists, whereas pack/unpack work with the individual message parts in the message list. Parameters ---------- msg_list : list of bytes or Message objects The list of message parts of the form [HMAC,p_header,p_parent, p_metadata,p_content,buffer1,buffer2,...]. content : bool (True) Whether to unpack the content dict (True), or leave it packed (False). copy : bool (True) Whether to return the bytes (True), or the non-copying Message object in each place (False). Returns ------- msg : dict The nested message dict with top-level keys [header, parent_header, content, buffers]. """ minlen = 5 message = {} if not copy: for i in range(minlen): msg_list[i] = msg_list[i].bytes if self.auth is not None: signature = msg_list[0] if not signature: raise ValueError("Unsigned Message") if signature in self.digest_history: raise ValueError("Duplicate Signature: %r" % signature) self._add_digest(signature) check = self.sign(msg_list[1:5]) if not compare_digest(signature, check): raise ValueError("Invalid Signature: %r" % signature) if not len(msg_list) >= minlen: raise TypeError( "malformed message, must have at least %i elements" % minlen) header = self.unpack(msg_list[1]) message['header'] = extract_dates(header) message['msg_id'] = header['msg_id'] message['msg_type'] = header['msg_type'] message['parent_header'] = extract_dates(self.unpack(msg_list[2])) message['metadata'] = self.unpack(msg_list[3]) if content: message['content'] = self.unpack(msg_list[4]) else: message['content'] = msg_list[4] message['buffers'] = msg_list[5:] # adapt to the current version return adapt(message) def unserialize(self, *args, **kwargs): warnings.warn( "Session.unserialize is deprecated. Use Session.deserialize.", DeprecationWarning, ) return self.deserialize(*args, **kwargs)
class EngineFactory(RegistrationFactory): """IPython engine""" # configurables: out_stream_factory = Type('IPython.zmq.iostream.OutStream', config=True, help="""The OutStream for handling stdout/err. Typically 'IPython.zmq.iostream.OutStream'""") display_hook_factory = Type('IPython.zmq.displayhook.ZMQDisplayHook', config=True, help="""The class for handling displayhook. Typically 'IPython.zmq.displayhook.ZMQDisplayHook'""") location = Unicode( config=True, help="""The location (an IP address) of the controller. This is used for disambiguating URLs, to determine whether loopback should be used to connect or the public address.""") timeout = CFloat( 2, config=True, help="""The time (in seconds) to wait for the Controller to respond to registration requests before giving up.""") sshserver = Unicode( config=True, help= """The SSH server to use for tunneling connections to the Controller.""" ) sshkey = Unicode( config=True, help= """The SSH private key file to use when tunneling connections to the Controller.""" ) paramiko = Bool( sys.platform == 'win32', config=True, help="""Whether to use paramiko instead of openssh for tunnels.""") # not configurable: user_ns = Dict() id = Integer(allow_none=True) registrar = Instance('zmq.eventloop.zmqstream.ZMQStream') kernel = Instance(Kernel) bident = CBytes() ident = Unicode() def _ident_changed(self, name, old, new): self.bident = cast_bytes(new) using_ssh = Bool(False) def __init__(self, **kwargs): super(EngineFactory, self).__init__(**kwargs) self.ident = self.session.session def init_connector(self): """construct connection function, which handles tunnels.""" self.using_ssh = bool(self.sshkey or self.sshserver) if self.sshkey and not self.sshserver: # We are using ssh directly to the controller, tunneling localhost to localhost self.sshserver = self.url.split('://')[1].split(':')[0] if self.using_ssh: if tunnel.try_passwordless_ssh(self.sshserver, self.sshkey, self.paramiko): password = False else: password = getpass("SSH Password for %s: " % self.sshserver) else: password = False def connect(s, url): url = disambiguate_url(url, self.location) if self.using_ssh: self.log.debug("Tunneling connection to %s via %s" % (url, self.sshserver)) return tunnel.tunnel_connection( s, url, self.sshserver, keyfile=self.sshkey, paramiko=self.paramiko, password=password, ) else: return s.connect(url) def maybe_tunnel(url): """like connect, but don't complete the connection (for use by heartbeat)""" url = disambiguate_url(url, self.location) if self.using_ssh: self.log.debug("Tunneling connection to %s via %s" % (url, self.sshserver)) url, tunnelobj = tunnel.open_tunnel( url, self.sshserver, keyfile=self.sshkey, paramiko=self.paramiko, password=password, ) return url return connect, maybe_tunnel def register(self): """send the registration_request""" self.log.info("Registering with controller at %s" % self.url) ctx = self.context connect, maybe_tunnel = self.init_connector() reg = ctx.socket(zmq.DEALER) reg.setsockopt(zmq.IDENTITY, self.bident) connect(reg, self.url) self.registrar = zmqstream.ZMQStream(reg, self.loop) content = dict(queue=self.ident, heartbeat=self.ident, control=self.ident) self.registrar.on_recv( lambda msg: self.complete_registration(msg, connect, maybe_tunnel)) # print (self.session.key) self.session.send(self.registrar, "registration_request", content=content) def complete_registration(self, msg, connect, maybe_tunnel): # print msg self._abort_dc.stop() ctx = self.context loop = self.loop identity = self.bident idents, msg = self.session.feed_identities(msg) msg = Message(self.session.unserialize(msg)) if msg.content.status == 'ok': self.id = int(msg.content.id) # launch heartbeat hb_addrs = msg.content.heartbeat # possibly forward hb ports with tunnels hb_addrs = [maybe_tunnel(addr) for addr in hb_addrs] heart = Heart(*map(str, hb_addrs), heart_id=identity) heart.start() # create Shell Streams (MUX, Task, etc.): queue_addr = msg.content.mux shell_addrs = [str(queue_addr)] task_addr = msg.content.task if task_addr: shell_addrs.append(str(task_addr)) # Uncomment this to go back to two-socket model # shell_streams = [] # for addr in shell_addrs: # stream = zmqstream.ZMQStream(ctx.socket(zmq.ROUTER), loop) # stream.setsockopt(zmq.IDENTITY, identity) # stream.connect(disambiguate_url(addr, self.location)) # shell_streams.append(stream) # Now use only one shell stream for mux and tasks stream = zmqstream.ZMQStream(ctx.socket(zmq.ROUTER), loop) stream.setsockopt(zmq.IDENTITY, identity) shell_streams = [stream] for addr in shell_addrs: connect(stream, addr) # end single stream-socket # control stream: control_addr = str(msg.content.control) control_stream = zmqstream.ZMQStream(ctx.socket(zmq.ROUTER), loop) control_stream.setsockopt(zmq.IDENTITY, identity) connect(control_stream, control_addr) # create iopub stream: iopub_addr = msg.content.iopub iopub_socket = ctx.socket(zmq.PUB) iopub_socket.setsockopt(zmq.IDENTITY, identity) connect(iopub_socket, iopub_addr) # disable history: self.config.HistoryManager.hist_file = ':memory:' # Redirect input streams and set a display hook. if self.out_stream_factory: sys.stdout = self.out_stream_factory(self.session, iopub_socket, u'stdout') sys.stdout.topic = cast_bytes('engine.%i.stdout' % self.id) sys.stderr = self.out_stream_factory(self.session, iopub_socket, u'stderr') sys.stderr.topic = cast_bytes('engine.%i.stderr' % self.id) if self.display_hook_factory: sys.displayhook = self.display_hook_factory( self.session, iopub_socket) sys.displayhook.topic = cast_bytes('engine.%i.pyout' % self.id) self.kernel = Kernel(config=self.config, int_id=self.id, ident=self.ident, session=self.session, control_stream=control_stream, shell_streams=shell_streams, iopub_socket=iopub_socket, loop=loop, user_ns=self.user_ns, log=self.log) self.kernel.shell.display_pub.topic = cast_bytes( 'engine.%i.displaypub' % self.id) # FIXME: This is a hack until IPKernelApp and IPEngineApp can be fully merged app = IPKernelApp(config=self.config, shell=self.kernel.shell, kernel=self.kernel, log=self.log) app.init_profile_dir() app.init_code() self.kernel.start() else: self.log.fatal("Registration Failed: %s" % msg) raise Exception("Registration Failed: %s" % msg) self.log.info("Completed registration with id %i" % self.id) def abort(self): self.log.fatal("Registration timed out after %.1f seconds" % self.timeout) if self.url.startswith('127.'): self.log.fatal(""" If the controller and engines are not on the same machine, you will have to instruct the controller to listen on an external IP (in ipcontroller_config.py): c.HubFactory.ip='*' # for all interfaces, internal and external c.HubFactory.ip='192.168.1.101' # or any interface that the engines can see or tunnel connections via ssh. """) self.session.send(self.registrar, "unregistration_request", content=dict(id=self.id)) time.sleep(1) sys.exit(255) def start(self): dc = ioloop.DelayedCallback(self.register, 0, self.loop) dc.start() self._abort_dc = ioloop.DelayedCallback(self.abort, self.timeout * 1000, self.loop) self._abort_dc.start()
class TaskScheduler(SessionFactory): """Python TaskScheduler object. This is the simplest object that supports msg_id based DAG dependencies. *Only* task msg_ids are checked, not msg_ids of jobs submitted via the MUX queue. """ hwm = Integer(1, config=True, help="""specify the High Water Mark (HWM) for the downstream socket in the Task scheduler. This is the maximum number of allowed outstanding tasks on each engine. The default (1) means that only one task can be outstanding on each engine. Setting TaskScheduler.hwm=0 means there is no limit, and the engines continue to be assigned tasks while they are working, effectively hiding network latency behind computation, but can result in an imbalance of work when submitting many heterogenous tasks all at once. Any positive value greater than one is a compromise between the two. """) scheme_name = Enum( ('leastload', 'pure', 'lru', 'plainrandom', 'weighted', 'twobin'), 'leastload', config=True, allow_none=False, help="""select the task scheduler scheme [default: Python LRU] Options are: 'pure', 'lru', 'plainrandom', 'weighted', 'twobin','leastload'""" ) def _scheme_name_changed(self, old, new): self.log.debug("Using scheme %r" % new) self.scheme = globals()[new] # input arguments: scheme = Instance(FunctionType) # function for determining the destination def _scheme_default(self): return leastload client_stream = Instance(zmqstream.ZMQStream) # client-facing stream engine_stream = Instance(zmqstream.ZMQStream) # engine-facing stream notifier_stream = Instance(zmqstream.ZMQStream) # hub-facing sub stream mon_stream = Instance(zmqstream.ZMQStream) # hub-facing pub stream query_stream = Instance(zmqstream.ZMQStream) # hub-facing DEALER stream # internals: graph = Dict() # dict by msg_id of [ msg_ids that depend on key ] retries = Dict() # dict by msg_id of retries remaining (non-neg ints) # waiting = List() # list of msg_ids ready to run, but haven't due to HWM depending = Dict() # dict by msg_id of Jobs pending = Dict() # dict by engine_uuid of submitted tasks completed = Dict() # dict by engine_uuid of completed tasks failed = Dict() # dict by engine_uuid of failed tasks destinations = Dict( ) # dict by msg_id of engine_uuids where jobs ran (reverse of completed+failed) clients = Dict() # dict by msg_id for who submitted the task targets = List() # list of target IDENTs loads = List() # list of engine loads # full = Set() # set of IDENTs that have HWM outstanding tasks all_completed = Set() # set of all completed tasks all_failed = Set() # set of all failed tasks all_done = Set() # set of all finished tasks=union(completed,failed) all_ids = Set() # set of all submitted task IDs auditor = Instance('zmq.eventloop.ioloop.PeriodicCallback') ident = CBytes() # ZMQ identity. This should just be self.session.session # but ensure Bytes def _ident_default(self): return self.session.bsession def start(self): self.query_stream.on_recv(self.dispatch_query_reply) self.session.send(self.query_stream, "connection_request", {}) self.engine_stream.on_recv(self.dispatch_result, copy=False) self.client_stream.on_recv(self.dispatch_submission, copy=False) self._notification_handlers = dict( registration_notification=self._register_engine, unregistration_notification=self._unregister_engine) self.notifier_stream.on_recv(self.dispatch_notification) self.auditor = ioloop.PeriodicCallback(self.audit_timeouts, 2e3, self.loop) # 1 Hz self.auditor.start() self.log.info("Scheduler started [%s]" % self.scheme_name) def resume_receiving(self): """Resume accepting jobs.""" self.client_stream.on_recv(self.dispatch_submission, copy=False) def stop_receiving(self): """Stop accepting jobs while there are no engines. Leave them in the ZMQ queue.""" self.client_stream.on_recv(None) #----------------------------------------------------------------------- # [Un]Registration Handling #----------------------------------------------------------------------- def dispatch_query_reply(self, msg): """handle reply to our initial connection request""" try: idents, msg = self.session.feed_identities(msg) except ValueError: self.log.warn("task::Invalid Message: %r", msg) return try: msg = self.session.unserialize(msg) except ValueError: self.log.warn("task::Unauthorized message from: %r" % idents) return content = msg['content'] for uuid in content.get('engines', {}).values(): self._register_engine(cast_bytes(uuid)) @util.log_errors def dispatch_notification(self, msg): """dispatch register/unregister events.""" try: idents, msg = self.session.feed_identities(msg) except ValueError: self.log.warn("task::Invalid Message: %r", msg) return try: msg = self.session.unserialize(msg) except ValueError: self.log.warn("task::Unauthorized message from: %r" % idents) return msg_type = msg['header']['msg_type'] handler = self._notification_handlers.get(msg_type, None) if handler is None: self.log.error("Unhandled message type: %r" % msg_type) else: try: handler(cast_bytes(msg['content']['uuid'])) except Exception: self.log.error("task::Invalid notification msg: %r", msg, exc_info=True) def _register_engine(self, uid): """New engine with ident `uid` became available.""" # head of the line: self.targets.insert(0, uid) self.loads.insert(0, 0) # initialize sets self.completed[uid] = set() self.failed[uid] = set() self.pending[uid] = {} # rescan the graph: self.update_graph(None) def _unregister_engine(self, uid): """Existing engine with ident `uid` became unavailable.""" if len(self.targets) == 1: # this was our only engine pass # handle any potentially finished tasks: self.engine_stream.flush() # don't pop destinations, because they might be used later # map(self.destinations.pop, self.completed.pop(uid)) # map(self.destinations.pop, self.failed.pop(uid)) # prevent this engine from receiving work idx = self.targets.index(uid) self.targets.pop(idx) self.loads.pop(idx) # wait 5 seconds before cleaning up pending jobs, since the results might # still be incoming if self.pending[uid]: dc = ioloop.DelayedCallback( lambda: self.handle_stranded_tasks(uid), 5000, self.loop) dc.start() else: self.completed.pop(uid) self.failed.pop(uid) def handle_stranded_tasks(self, engine): """Deal with jobs resident in an engine that died.""" lost = self.pending[engine] for msg_id in lost.keys(): if msg_id not in self.pending[engine]: # prevent double-handling of messages continue raw_msg = lost[msg_id].raw_msg idents, msg = self.session.feed_identities(raw_msg, copy=False) parent = self.session.unpack(msg[1].bytes) idents = [engine, idents[0]] # build fake error reply try: raise error.EngineError( "Engine %r died while running task %r" % (engine, msg_id)) except: content = error.wrap_exception() # build fake metadata md = dict( status=u'error', engine=engine, date=datetime.now(), ) msg = self.session.msg('apply_reply', content, parent=parent, metadata=md) raw_reply = map(zmq.Message, self.session.serialize(msg, ident=idents)) # and dispatch it self.dispatch_result(raw_reply) # finally scrub completed/failed lists self.completed.pop(engine) self.failed.pop(engine) #----------------------------------------------------------------------- # Job Submission #----------------------------------------------------------------------- @util.log_errors def dispatch_submission(self, raw_msg): """Dispatch job submission to appropriate handlers.""" # ensure targets up to date: self.notifier_stream.flush() try: idents, msg = self.session.feed_identities(raw_msg, copy=False) msg = self.session.unserialize(msg, content=False, copy=False) except Exception: self.log.error("task::Invaid task msg: %r" % raw_msg, exc_info=True) return # send to monitor self.mon_stream.send_multipart([b'intask'] + raw_msg, copy=False) header = msg['header'] md = msg['metadata'] msg_id = header['msg_id'] self.all_ids.add(msg_id) # get targets as a set of bytes objects # from a list of unicode objects targets = md.get('targets', []) targets = map(cast_bytes, targets) targets = set(targets) retries = md.get('retries', 0) self.retries[msg_id] = retries # time dependencies after = md.get('after', None) if after: after = Dependency(after) if after.all: if after.success: after = Dependency( after.difference(self.all_completed), success=after.success, failure=after.failure, all=after.all, ) if after.failure: after = Dependency( after.difference(self.all_failed), success=after.success, failure=after.failure, all=after.all, ) if after.check(self.all_completed, self.all_failed): # recast as empty set, if `after` already met, # to prevent unnecessary set comparisons after = MET else: after = MET # location dependencies follow = Dependency(md.get('follow', [])) # turn timeouts into datetime objects: timeout = md.get('timeout', None) if timeout: # cast to float, because jsonlib returns floats as decimal.Decimal, # which timedelta does not accept timeout = datetime.now() + timedelta(0, float(timeout), 0) job = Job( msg_id=msg_id, raw_msg=raw_msg, idents=idents, msg=msg, header=header, targets=targets, after=after, follow=follow, timeout=timeout, metadata=md, ) # validate and reduce dependencies: for dep in after, follow: if not dep: # empty dependency continue # check valid: if msg_id in dep or dep.difference(self.all_ids): self.depending[msg_id] = job return self.fail_unreachable(msg_id, error.InvalidDependency) # check if unreachable: if dep.unreachable(self.all_completed, self.all_failed): self.depending[msg_id] = job return self.fail_unreachable(msg_id) if after.check(self.all_completed, self.all_failed): # time deps already met, try to run if not self.maybe_run(job): # can't run yet if msg_id not in self.all_failed: # could have failed as unreachable self.save_unmet(job) else: self.save_unmet(job) def audit_timeouts(self): """Audit all waiting tasks for expired timeouts.""" now = datetime.now() for msg_id in self.depending.keys(): # must recheck, in case one failure cascaded to another: if msg_id in self.depending: job = self.depending[msg_id] if job.timeout and job.timeout < now: self.fail_unreachable(msg_id, error.TaskTimeout) def fail_unreachable(self, msg_id, why=error.ImpossibleDependency): """a task has become unreachable, send a reply with an ImpossibleDependency error.""" if msg_id not in self.depending: self.log.error("msg %r already failed!", msg_id) return job = self.depending.pop(msg_id) for mid in job.dependents: if mid in self.graph: self.graph[mid].remove(msg_id) try: raise why() except: content = error.wrap_exception() self.all_done.add(msg_id) self.all_failed.add(msg_id) msg = self.session.send(self.client_stream, 'apply_reply', content, parent=job.header, ident=job.idents) self.session.send(self.mon_stream, msg, ident=[b'outtask'] + job.idents) self.update_graph(msg_id, success=False) def maybe_run(self, job): """check location dependencies, and run if they are met.""" msg_id = job.msg_id self.log.debug("Attempting to assign task %s", msg_id) if not self.targets: # no engines, definitely can't run return False if job.follow or job.targets or job.blacklist or self.hwm: # we need a can_run filter def can_run(idx): # check hwm if self.hwm and self.loads[idx] == self.hwm: return False target = self.targets[idx] # check blacklist if target in job.blacklist: return False # check targets if job.targets and target not in job.targets: return False # check follow return job.follow.check(self.completed[target], self.failed[target]) indices = filter(can_run, range(len(self.targets))) if not indices: # couldn't run if job.follow.all: # check follow for impossibility dests = set() relevant = set() if job.follow.success: relevant = self.all_completed if job.follow.failure: relevant = relevant.union(self.all_failed) for m in job.follow.intersection(relevant): dests.add(self.destinations[m]) if len(dests) > 1: self.depending[msg_id] = job self.fail_unreachable(msg_id) return False if job.targets: # check blacklist+targets for impossibility job.targets.difference_update(job.blacklist) if not job.targets or not job.targets.intersection( self.targets): self.depending[msg_id] = job self.fail_unreachable(msg_id) return False return False else: indices = None self.submit_task(job, indices) return True def save_unmet(self, job): """Save a message for later submission when its dependencies are met.""" msg_id = job.msg_id self.depending[msg_id] = job # track the ids in follow or after, but not those already finished for dep_id in job.after.union(job.follow).difference(self.all_done): if dep_id not in self.graph: self.graph[dep_id] = set() self.graph[dep_id].add(msg_id) def submit_task(self, job, indices=None): """Submit a task to any of a subset of our targets.""" if indices: loads = [self.loads[i] for i in indices] else: loads = self.loads idx = self.scheme(loads) if indices: idx = indices[idx] target = self.targets[idx] # print (target, map(str, msg[:3])) # send job to the engine self.engine_stream.send(target, flags=zmq.SNDMORE, copy=False) self.engine_stream.send_multipart(job.raw_msg, copy=False) # update load self.add_job(idx) self.pending[target][job.msg_id] = job # notify Hub content = dict(msg_id=job.msg_id, engine_id=target.decode('ascii')) self.session.send(self.mon_stream, 'task_destination', content=content, ident=[b'tracktask', self.ident]) #----------------------------------------------------------------------- # Result Handling #----------------------------------------------------------------------- @util.log_errors def dispatch_result(self, raw_msg): """dispatch method for result replies""" try: idents, msg = self.session.feed_identities(raw_msg, copy=False) msg = self.session.unserialize(msg, content=False, copy=False) engine = idents[0] try: idx = self.targets.index(engine) except ValueError: pass # skip load-update for dead engines else: self.finish_job(idx) except Exception: self.log.error("task::Invaid result: %r", raw_msg, exc_info=True) return md = msg['metadata'] parent = msg['parent_header'] if md.get('dependencies_met', True): success = (md['status'] == 'ok') msg_id = parent['msg_id'] retries = self.retries[msg_id] if not success and retries > 0: # failed self.retries[msg_id] = retries - 1 self.handle_unmet_dependency(idents, parent) else: del self.retries[msg_id] # relay to client and update graph self.handle_result(idents, parent, raw_msg, success) # send to Hub monitor self.mon_stream.send_multipart([b'outtask'] + raw_msg, copy=False) else: self.handle_unmet_dependency(idents, parent) def handle_result(self, idents, parent, raw_msg, success=True): """handle a real task result, either success or failure""" # first, relay result to client engine = idents[0] client = idents[1] # swap_ids for ROUTER-ROUTER mirror raw_msg[:2] = [client, engine] # print (map(str, raw_msg[:4])) self.client_stream.send_multipart(raw_msg, copy=False) # now, update our data structures msg_id = parent['msg_id'] self.pending[engine].pop(msg_id) if success: self.completed[engine].add(msg_id) self.all_completed.add(msg_id) else: self.failed[engine].add(msg_id) self.all_failed.add(msg_id) self.all_done.add(msg_id) self.destinations[msg_id] = engine self.update_graph(msg_id, success) def handle_unmet_dependency(self, idents, parent): """handle an unmet dependency""" engine = idents[0] msg_id = parent['msg_id'] job = self.pending[engine].pop(msg_id) job.blacklist.add(engine) if job.blacklist == job.targets: self.depending[msg_id] = job self.fail_unreachable(msg_id) elif not self.maybe_run(job): # resubmit failed if msg_id not in self.all_failed: # put it back in our dependency tree self.save_unmet(job) if self.hwm: try: idx = self.targets.index(engine) except ValueError: pass # skip load-update for dead engines else: if self.loads[idx] == self.hwm - 1: self.update_graph(None) def update_graph(self, dep_id=None, success=True): """dep_id just finished. Update our dependency graph and submit any jobs that just became runable. Called with dep_id=None to update entire graph for hwm, but without finishing a task. """ # print ("\n\n***********") # pprint (dep_id) # pprint (self.graph) # pprint (self.depending) # pprint (self.all_completed) # pprint (self.all_failed) # print ("\n\n***********\n\n") # update any jobs that depended on the dependency jobs = self.graph.pop(dep_id, []) # recheck *all* jobs if # a) we have HWM and an engine just become no longer full # or b) dep_id was given as None if dep_id is None or self.hwm and any( [load == self.hwm - 1 for load in self.loads]): jobs = self.depending.keys() for msg_id in sorted( jobs, key=lambda msg_id: self.depending[msg_id].timestamp): job = self.depending[msg_id] if job.after.unreachable(self.all_completed, self.all_failed)\ or job.follow.unreachable(self.all_completed, self.all_failed): self.fail_unreachable(msg_id) elif job.after.check(self.all_completed, self.all_failed): # time deps met, maybe run if self.maybe_run(job): self.depending.pop(msg_id) for mid in job.dependents: if mid in self.graph: self.graph[mid].remove(msg_id) #---------------------------------------------------------------------- # methods to be overridden by subclasses #---------------------------------------------------------------------- def add_job(self, idx): """Called after self.targets[idx] just got the job with header. Override with subclasses. The default ordering is simple LRU. The default loads are the number of outstanding jobs.""" self.loads[idx] += 1 for lis in (self.targets, self.loads): lis.append(lis.pop(idx)) def finish_job(self, idx): """Called after self.targets[idx] just finished a job. Override with subclasses.""" self.loads[idx] -= 1
class EngineFactory(RegistrationFactory): """IPython engine""" # configurables: out_stream_factory=Type('IPython.kernel.zmq.iostream.OutStream', config=True, help="""The OutStream for handling stdout/err. Typically 'IPython.kernel.zmq.iostream.OutStream'""") display_hook_factory=Type('IPython.kernel.zmq.displayhook.ZMQDisplayHook', config=True, help="""The class for handling displayhook. Typically 'IPython.kernel.zmq.displayhook.ZMQDisplayHook'""") location=Unicode(config=True, help="""The location (an IP address) of the controller. This is used for disambiguating URLs, to determine whether loopback should be used to connect or the public address.""") timeout=Float(5.0, config=True, help="""The time (in seconds) to wait for the Controller to respond to registration requests before giving up.""") max_heartbeat_misses=Integer(50, config=True, help="""The maximum number of times a check for the heartbeat ping of a controller can be missed before shutting down the engine. If set to 0, the check is disabled.""") sshserver=Unicode(config=True, help="""The SSH server to use for tunneling connections to the Controller.""") sshkey=Unicode(config=True, help="""The SSH private key file to use when tunneling connections to the Controller.""") paramiko=Bool(sys.platform == 'win32', config=True, help="""Whether to use paramiko instead of openssh for tunnels.""") @property def tunnel_mod(self): from zmq.ssh import tunnel return tunnel # not configurable: connection_info = Dict() user_ns = Dict() id = Integer(allow_none=True) registrar = Instance('zmq.eventloop.zmqstream.ZMQStream') kernel = Instance(Kernel) hb_check_period=Integer() # States for the heartbeat monitoring # Initial values for monitored and pinged must satisfy "monitored > pinged == False" so that # during the first check no "missed" ping is reported. Must be floats for Python 3 compatibility. _hb_last_pinged = 0.0 _hb_last_monitored = 0.0 _hb_missed_beats = 0 # The zmq Stream which receives the pings from the Heart _hb_listener = None bident = CBytes() ident = Unicode() def _ident_changed(self, name, old, new): self.bident = cast_bytes(new) using_ssh=Bool(False) def __init__(self, **kwargs): super(EngineFactory, self).__init__(**kwargs) self.ident = self.session.session def init_connector(self): """construct connection function, which handles tunnels.""" self.using_ssh = bool(self.sshkey or self.sshserver) if self.sshkey and not self.sshserver: # We are using ssh directly to the controller, tunneling localhost to localhost self.sshserver = self.url.split('://')[1].split(':')[0] if self.using_ssh: if self.tunnel_mod.try_passwordless_ssh(self.sshserver, self.sshkey, self.paramiko): password=False else: password = getpass("SSH Password for %s: "%self.sshserver) else: password = False def connect(s, url): url = disambiguate_url(url, self.location) if self.using_ssh: self.log.debug("Tunneling connection to %s via %s", url, self.sshserver) return self.tunnel_mod.tunnel_connection(s, url, self.sshserver, keyfile=self.sshkey, paramiko=self.paramiko, password=password, ) else: return s.connect(url) def maybe_tunnel(url): """like connect, but don't complete the connection (for use by heartbeat)""" url = disambiguate_url(url, self.location) if self.using_ssh: self.log.debug("Tunneling connection to %s via %s", url, self.sshserver) url, tunnelobj = self.tunnel_mod.open_tunnel(url, self.sshserver, keyfile=self.sshkey, paramiko=self.paramiko, password=password, ) return str(url) return connect, maybe_tunnel def register(self): """send the registration_request""" self.log.info("Registering with controller at %s"%self.url) ctx = self.context connect,maybe_tunnel = self.init_connector() reg = ctx.socket(zmq.DEALER) reg.setsockopt(zmq.IDENTITY, self.bident) connect(reg, self.url) self.registrar = zmqstream.ZMQStream(reg, self.loop) content = dict(uuid=self.ident) self.registrar.on_recv(lambda msg: self.complete_registration(msg, connect, maybe_tunnel)) # print (self.session.key) self.session.send(self.registrar, "registration_request", content=content) def _report_ping(self, msg): """Callback for when the heartmonitor.Heart receives a ping""" #self.log.debug("Received a ping: %s", msg) self._hb_last_pinged = time.time() def complete_registration(self, msg, connect, maybe_tunnel): # print msg self.loop.remove_timeout(self._abort_timeout) ctx = self.context loop = self.loop identity = self.bident idents,msg = self.session.feed_identities(msg) msg = self.session.deserialize(msg) content = msg['content'] info = self.connection_info def url(key): """get zmq url for given channel""" return str(info["interface"] + ":%i" % info[key]) if content['status'] == 'ok': self.id = int(content['id']) # launch heartbeat # possibly forward hb ports with tunnels hb_ping = maybe_tunnel(url('hb_ping')) hb_pong = maybe_tunnel(url('hb_pong')) hb_monitor = None if self.max_heartbeat_misses > 0: # Add a monitor socket which will record the last time a ping was seen mon = self.context.socket(zmq.SUB) mport = mon.bind_to_random_port('tcp://%s' % localhost()) mon.setsockopt(zmq.SUBSCRIBE, b"") self._hb_listener = zmqstream.ZMQStream(mon, self.loop) self._hb_listener.on_recv(self._report_ping) hb_monitor = "tcp://%s:%i" % (localhost(), mport) heart = Heart(hb_ping, hb_pong, hb_monitor , heart_id=identity) heart.start() # create Shell Connections (MUX, Task, etc.): shell_addrs = url('mux'), url('task') # Use only one shell stream for mux and tasks stream = zmqstream.ZMQStream(ctx.socket(zmq.ROUTER), loop) stream.setsockopt(zmq.IDENTITY, identity) shell_streams = [stream] for addr in shell_addrs: connect(stream, addr) # control stream: control_addr = url('control') control_stream = zmqstream.ZMQStream(ctx.socket(zmq.ROUTER), loop) control_stream.setsockopt(zmq.IDENTITY, identity) connect(control_stream, control_addr) # create iopub stream: iopub_addr = url('iopub') iopub_socket = ctx.socket(zmq.PUB) iopub_socket.setsockopt(zmq.IDENTITY, identity) connect(iopub_socket, iopub_addr) # disable history: self.config.HistoryManager.hist_file = ':memory:' # Redirect input streams and set a display hook. if self.out_stream_factory: sys.stdout = self.out_stream_factory(self.session, iopub_socket, u'stdout') sys.stdout.topic = cast_bytes('engine.%i.stdout' % self.id) sys.stderr = self.out_stream_factory(self.session, iopub_socket, u'stderr') sys.stderr.topic = cast_bytes('engine.%i.stderr' % self.id) if self.display_hook_factory: sys.displayhook = self.display_hook_factory(self.session, iopub_socket) sys.displayhook.topic = cast_bytes('engine.%i.execute_result' % self.id) self.kernel = Kernel(parent=self, int_id=self.id, ident=self.ident, session=self.session, control_stream=control_stream, shell_streams=shell_streams, iopub_socket=iopub_socket, loop=loop, user_ns=self.user_ns, log=self.log) self.kernel.shell.display_pub.topic = cast_bytes('engine.%i.displaypub' % self.id) # periodically check the heartbeat pings of the controller # Should be started here and not in "start()" so that the right period can be taken # from the hubs HeartBeatMonitor.period if self.max_heartbeat_misses > 0: # Use a slightly bigger check period than the hub signal period to not warn unnecessary self.hb_check_period = int(content['hb_period'])+10 self.log.info("Starting to monitor the heartbeat signal from the hub every %i ms." , self.hb_check_period) self._hb_reporter = ioloop.PeriodicCallback(self._hb_monitor, self.hb_check_period, self.loop) self._hb_reporter.start() else: self.log.info("Monitoring of the heartbeat signal from the hub is not enabled.") # FIXME: This is a hack until IPKernelApp and IPEngineApp can be fully merged app = IPKernelApp(parent=self, shell=self.kernel.shell, kernel=self.kernel, log=self.log) app.init_profile_dir() app.init_code() self.kernel.start() else: self.log.fatal("Registration Failed: %s"%msg) raise Exception("Registration Failed: %s"%msg) self.log.info("Completed registration with id %i"%self.id) def abort(self): self.log.fatal("Registration timed out after %.1f seconds"%self.timeout) if self.url.startswith('127.'): self.log.fatal(""" If the controller and engines are not on the same machine, you will have to instruct the controller to listen on an external IP (in ipcontroller_config.py): c.HubFactory.ip='*' # for all interfaces, internal and external c.HubFactory.ip='192.168.1.101' # or any interface that the engines can see or tunnel connections via ssh. """) self.session.send(self.registrar, "unregistration_request", content=dict(id=self.id)) time.sleep(1) sys.exit(255) def _hb_monitor(self): """Callback to monitor the heartbeat from the controller""" self._hb_listener.flush() if self._hb_last_monitored > self._hb_last_pinged: self._hb_missed_beats += 1 self.log.warn("No heartbeat in the last %s ms (%s time(s) in a row).", self.hb_check_period, self._hb_missed_beats) else: #self.log.debug("Heartbeat received (after missing %s beats).", self._hb_missed_beats) self._hb_missed_beats = 0 if self._hb_missed_beats >= self.max_heartbeat_misses: self.log.fatal("Maximum number of heartbeats misses reached (%s times %s ms), shutting down.", self.max_heartbeat_misses, self.hb_check_period) self.session.send(self.registrar, "unregistration_request", content=dict(id=self.id)) self.loop.stop() self._hb_last_monitored = time.time() def start(self): loop = self.loop def _start(): self.register() self._abort_timeout = loop.add_timeout(loop.time() + self.timeout, self.abort) self.loop.add_callback(_start)
class Kernel(SessionFactory): #--------------------------------------------------------------------------- # Kernel interface #--------------------------------------------------------------------------- # kwargs: exec_lines = List(Unicode, config=True, help="List of lines to execute") # identities: int_id = Int(-1) bident = CBytes() ident = Unicode() def _ident_changed(self, name, old, new): self.bident = asbytes(new) user_ns = Dict(config=True, help="""Set the user's namespace of the Kernel""") control_stream = Instance(zmqstream.ZMQStream) task_stream = Instance(zmqstream.ZMQStream) iopub_stream = Instance(zmqstream.ZMQStream) client = Instance('IPython.parallel.Client') # internals shell_streams = List() compiler = Instance(CommandCompiler, (), {}) completer = Instance(KernelCompleter) aborted = Set() shell_handlers = Dict() control_handlers = Dict() def _set_prefix(self): self.prefix = "engine.%s" % self.int_id def _connect_completer(self): self.completer = KernelCompleter(self.user_ns) def __init__(self, **kwargs): super(Kernel, self).__init__(**kwargs) self._set_prefix() self._connect_completer() self.on_trait_change(self._set_prefix, 'id') self.on_trait_change(self._connect_completer, 'user_ns') # Build dict of handlers for message types for msg_type in [ 'execute_request', 'complete_request', 'apply_request', 'clear_request' ]: self.shell_handlers[msg_type] = getattr(self, msg_type) for msg_type in ['shutdown_request', 'abort_request' ] + self.shell_handlers.keys(): self.control_handlers[msg_type] = getattr(self, msg_type) self._initial_exec_lines() def _wrap_exception(self, method=None): e_info = dict(engine_uuid=self.ident, engine_id=self.int_id, method=method) content = wrap_exception(e_info) return content def _initial_exec_lines(self): s = _Passer() content = dict(silent=True, user_variable=[], user_expressions=[]) for line in self.exec_lines: self.log.debug("executing initialization: %s" % line) content.update({'code': line}) msg = self.session.msg('execute_request', content) self.execute_request(s, [], msg) #-------------------- control handlers ----------------------------- def abort_queues(self): for stream in self.shell_streams: if stream: self.abort_queue(stream) def abort_queue(self, stream): while True: idents, msg = self.session.recv(stream, zmq.NOBLOCK, content=True) if msg is None: return self.log.info("Aborting:") self.log.info(str(msg)) msg_type = msg['header']['msg_type'] reply_type = msg_type.split('_')[0] + '_reply' # reply_msg = self.session.msg(reply_type, {'status' : 'aborted'}, msg) # self.reply_socket.send(ident,zmq.SNDMORE) # self.reply_socket.send_json(reply_msg) reply_msg = self.session.send(stream, reply_type, content={'status': 'aborted'}, parent=msg, ident=idents) self.log.debug(str(reply_msg)) # We need to wait a bit for requests to come in. This can probably # be set shorter for true asynchronous clients. time.sleep(0.05) def abort_request(self, stream, ident, parent): """abort a specifig msg by id""" msg_ids = parent['content'].get('msg_ids', None) if isinstance(msg_ids, basestring): msg_ids = [msg_ids] if not msg_ids: self.abort_queues() for mid in msg_ids: self.aborted.add(str(mid)) content = dict(status='ok') reply_msg = self.session.send(stream, 'abort_reply', content=content, parent=parent, ident=ident) self.log.debug(str(reply_msg)) def shutdown_request(self, stream, ident, parent): """kill ourself. This should really be handled in an external process""" try: self.abort_queues() except: content = self._wrap_exception('shutdown') else: content = dict(parent['content']) content['status'] = 'ok' msg = self.session.send(stream, 'shutdown_reply', content=content, parent=parent, ident=ident) self.log.debug(str(msg)) dc = ioloop.DelayedCallback(lambda: sys.exit(0), 1000, self.loop) dc.start() def dispatch_control(self, msg): idents, msg = self.session.feed_identities(msg, copy=False) try: msg = self.session.unserialize(msg, content=True, copy=False) except: self.log.error("Invalid Message", exc_info=True) return else: self.log.debug("Control received, %s", msg) header = msg['header'] msg_id = header['msg_id'] msg_type = header['msg_type'] handler = self.control_handlers.get(msg_type, None) if handler is None: self.log.error("UNKNOWN CONTROL MESSAGE TYPE: %r" % msg_type) else: handler(self.control_stream, idents, msg) #-------------------- queue helpers ------------------------------ def check_dependencies(self, dependencies): if not dependencies: return True if len(dependencies) == 2 and dependencies[0] in 'any all'.split(): anyorall = dependencies[0] dependencies = dependencies[1] else: anyorall = 'all' results = self.client.get_results(dependencies, status_only=True) if results['status'] != 'ok': return False if anyorall == 'any': if not results['completed']: return False else: if results['pending']: return False return True def check_aborted(self, msg_id): return msg_id in self.aborted #-------------------- queue handlers ----------------------------- def clear_request(self, stream, idents, parent): """Clear our namespace.""" self.user_ns = {} msg = self.session.send(stream, 'clear_reply', ident=idents, parent=parent, content=dict(status='ok')) self._initial_exec_lines() def execute_request(self, stream, ident, parent): self.log.debug('execute request %s' % parent) try: code = parent[u'content'][u'code'] except: self.log.error("Got bad msg: %s" % parent, exc_info=True) return self.session.send(self.iopub_stream, u'pyin', {u'code': code}, parent=parent, ident=asbytes('%s.pyin' % self.prefix)) started = datetime.now() try: comp_code = self.compiler(code, '<zmq-kernel>') # allow for not overriding displayhook if hasattr(sys.displayhook, 'set_parent'): sys.displayhook.set_parent(parent) sys.stdout.set_parent(parent) sys.stderr.set_parent(parent) exec comp_code in self.user_ns, self.user_ns except: exc_content = self._wrap_exception('execute') # exc_msg = self.session.msg(u'pyerr', exc_content, parent) self.session.send(self.iopub_stream, u'pyerr', exc_content, parent=parent, ident=asbytes('%s.pyerr' % self.prefix)) reply_content = exc_content else: reply_content = {'status': 'ok'} reply_msg = self.session.send(stream, u'execute_reply', reply_content, parent=parent, ident=ident, subheader=dict(started=started)) self.log.debug(str(reply_msg)) if reply_msg['content']['status'] == u'error': self.abort_queues() def complete_request(self, stream, ident, parent): matches = {'matches': self.complete(parent), 'status': 'ok'} completion_msg = self.session.send(stream, 'complete_reply', matches, parent, ident) # print >> sys.__stdout__, completion_msg def complete(self, msg): return self.completer.complete(msg.content.line, msg.content.text) def apply_request(self, stream, ident, parent): # flush previous reply, so this request won't block it stream.flush(zmq.POLLOUT) try: content = parent[u'content'] bufs = parent[u'buffers'] msg_id = parent['header']['msg_id'] # bound = parent['header'].get('bound', False) except: self.log.error("Got bad msg: %s" % parent, exc_info=True) return # pyin_msg = self.session.msg(u'pyin',{u'code':code}, parent=parent) # self.iopub_stream.send(pyin_msg) # self.session.send(self.iopub_stream, u'pyin', {u'code':code},parent=parent) sub = { 'dependencies_met': True, 'engine': self.ident, 'started': datetime.now() } try: # allow for not overriding displayhook if hasattr(sys.displayhook, 'set_parent'): sys.displayhook.set_parent(parent) sys.stdout.set_parent(parent) sys.stderr.set_parent(parent) # exec "f(*args,**kwargs)" in self.user_ns, self.user_ns working = self.user_ns # suffix = prefix = "_" + str(msg_id).replace("-", "") + "_" f, args, kwargs = unpack_apply_message(bufs, working, copy=False) # if bound: # bound_ns = Namespace(working) # args = [bound_ns]+list(args) fname = getattr(f, '__name__', 'f') fname = prefix + "f" argname = prefix + "args" kwargname = prefix + "kwargs" resultname = prefix + "result" ns = {fname: f, argname: args, kwargname: kwargs, resultname: None} # print ns working.update(ns) code = "%s=%s(*%s,**%s)" % (resultname, fname, argname, kwargname) try: exec code in working, working result = working.get(resultname) finally: for key in ns.iterkeys(): working.pop(key) # if bound: # working.update(bound_ns) packed_result, buf = serialize_object(result) result_buf = [packed_result] + buf except: exc_content = self._wrap_exception('apply') # exc_msg = self.session.msg(u'pyerr', exc_content, parent) self.session.send(self.iopub_stream, u'pyerr', exc_content, parent=parent, ident=asbytes('%s.pyerr' % self.prefix)) reply_content = exc_content result_buf = [] if exc_content['ename'] == 'UnmetDependency': sub['dependencies_met'] = False else: reply_content = {'status': 'ok'} # put 'ok'/'error' status in header, for scheduler introspection: sub['status'] = reply_content['status'] reply_msg = self.session.send(stream, u'apply_reply', reply_content, parent=parent, ident=ident, buffers=result_buf, subheader=sub) # flush i/o # should this be before reply_msg is sent, like in the single-kernel code, # or should nothing get in the way of real results? sys.stdout.flush() sys.stderr.flush() def dispatch_queue(self, stream, msg): self.control_stream.flush() idents, msg = self.session.feed_identities(msg, copy=False) try: msg = self.session.unserialize(msg, content=True, copy=False) except: self.log.error("Invalid Message", exc_info=True) return else: self.log.debug("Message received, %s", msg) header = msg['header'] msg_id = header['msg_id'] msg_type = msg['header']['msg_type'] if self.check_aborted(msg_id): self.aborted.remove(msg_id) # is it safe to assume a msg_id will not be resubmitted? reply_type = msg_type.split('_')[0] + '_reply' status = {'status': 'aborted'} reply_msg = self.session.send(stream, reply_type, subheader=status, content=status, parent=msg, ident=idents) return handler = self.shell_handlers.get(msg_type, None) if handler is None: self.log.error("UNKNOWN MESSAGE TYPE: %r" % msg_type) else: handler(stream, idents, msg) def start(self): #### stream mode: if self.control_stream: self.control_stream.on_recv(self.dispatch_control, copy=False) self.control_stream.on_err(printer) def make_dispatcher(stream): def dispatcher(msg): return self.dispatch_queue(stream, msg) return dispatcher for s in self.shell_streams: s.on_recv(make_dispatcher(s), copy=False) s.on_err(printer) if self.iopub_stream: self.iopub_stream.on_err(printer)
class AmberSimulator(Device): name = 'AMBER' path = 'msmaccelerator.simulate.amber_simulation.AmberSimulator' short_description = 'Run a single round of dynamics with AMBER' long_description = '''This device will connect to the msmaccelerator server, request the initial conditions with which to start a simulation, and propagate dynamics''' mdin = FilePath(config=True, exists=True, isfile=True, help="""AMBER .in file controlling the production run. If no production is desired, do not set this parameter.""") workdir = CBytes(config=True, help="""Directory to work in. If not set, we'll requirest a temporary directory from the OS and clean it up when we're finished. This option is useful for debugging.""") executable = Enum(['pmemd', 'pmemd.cuda', 'pmemd.cuda.MPI'], config=True, default_value='pmemd', help="Which AMBER executable to use?") precommand = Unicode(u'', config=True, help="Something to run before the command, like mpirun") prmtop = FilePath(config=True, exists=True, isfile=True, help="""Parameter/topology file for the system""") amber_home = FilePath(exists=True, isdir=True, help='Home directory for AMBER installation') def _amber_home_default(self): if 'AMBERHOME' not in os.environ: raise KeyError("You need to set the AMBERHOME environment variable") return os.environ['AMBERHOME'] aliases = dict(mdin='AmberSimulator.mdin', precommand='AmberSimulator.precommand', prmtop='AmberSimulator.prmtop', zmq_port='Device.zmq_port', zmq_url='Device.zmq_url', executable='AmberSimulator.executable') def start(self): super(AmberSimulator, self).start() def error(self, msg): self.log.error(msg) self.exit(1) def on_startup_message(self, msg): """This method is called when the device receives its startup message from the server. """ assert msg.header.msg_type in ['simulate'] # only allowed RPC return getattr(self, msg.header.msg_type)(msg.header, msg.content) def simulate(self, header, content): """Run the simulation in subprocesses to invoke the AMBER binaries""" if content.starting_state.protocol == 'localfs': if not content.starting_state.path.endswith('.inpcrd'): raise ValueError('starting state must have inpcrd extension. ' 'did you start server in amber mode? ' 'starting_state.path=%s' % content.starting_state.path) else: raise NotImplementedError('Only localfs transport is currently ' 'supported.') template = '{precommand} {binary} -O -i {mdin} -o {mdout} -p {prmtop} -c {inpcrd} -r {restart} -x {traj}' # RUNNING PRODUCTION with cd_context('amber_workdir', logger=self.log): base = splitext(basename(self.mdin))[0] binary = join(self.amber_home, 'bin', self.executable) mdout = base + '.out' restart = base + '.restart' traj = base + '.nc' cmd = template.format(binary=binary, mdin=relpath(self.mdin), mdout=mdout, prmtop=relpath(self.prmtop), inpcrd=relpath(content.starting_state.path), restart=restart, precommand=self.precommand, traj=relpath(content.output.path)).split() self.log.info('Executing Command: %s' % cmd) subprocess.check_output(cmd) self.send_recv(msg_type='simulation_done', content={ 'status': 'success', 'output': { 'protocol': 'localfs', 'path': content.output.path } })
class EngineFactory(RegistrationFactory): """IPython engine""" # configurables: out_stream_factory = Type('IPython.zmq.iostream.OutStream', config=True, help="""The OutStream for handling stdout/err. Typically 'IPython.zmq.iostream.OutStream'""") display_hook_factory = Type('IPython.zmq.displayhook.ZMQDisplayHook', config=True, help="""The class for handling displayhook. Typically 'IPython.zmq.displayhook.ZMQDisplayHook'""") location = Unicode( config=True, help="""The location (an IP address) of the controller. This is used for disambiguating URLs, to determine whether loopback should be used to connect or the public address.""") timeout = CFloat( 2, config=True, help="""The time (in seconds) to wait for the Controller to respond to registration requests before giving up.""") # not configurable: user_ns = Dict() id = Int(allow_none=True) registrar = Instance('zmq.eventloop.zmqstream.ZMQStream') kernel = Instance(Kernel) bident = CBytes() ident = Unicode() def _ident_changed(self, name, old, new): self.bident = asbytes(new) def __init__(self, **kwargs): super(EngineFactory, self).__init__(**kwargs) self.ident = self.session.session ctx = self.context reg = ctx.socket(zmq.XREQ) reg.setsockopt(zmq.IDENTITY, self.bident) reg.connect(self.url) self.registrar = zmqstream.ZMQStream(reg, self.loop) def register(self): """send the registration_request""" self.log.info("Registering with controller at %s" % self.url) content = dict(queue=self.ident, heartbeat=self.ident, control=self.ident) self.registrar.on_recv(self.complete_registration) # print (self.session.key) self.session.send(self.registrar, "registration_request", content=content) def complete_registration(self, msg): # print msg self._abort_dc.stop() ctx = self.context loop = self.loop identity = self.bident idents, msg = self.session.feed_identities(msg) msg = Message(self.session.unpack_message(msg)) if msg.content.status == 'ok': self.id = int(msg.content.id) # create Shell Streams (MUX, Task, etc.): queue_addr = msg.content.mux shell_addrs = [str(queue_addr)] task_addr = msg.content.task if task_addr: shell_addrs.append(str(task_addr)) # Uncomment this to go back to two-socket model # shell_streams = [] # for addr in shell_addrs: # stream = zmqstream.ZMQStream(ctx.socket(zmq.XREP), loop) # stream.setsockopt(zmq.IDENTITY, identity) # stream.connect(disambiguate_url(addr, self.location)) # shell_streams.append(stream) # Now use only one shell stream for mux and tasks stream = zmqstream.ZMQStream(ctx.socket(zmq.XREP), loop) stream.setsockopt(zmq.IDENTITY, identity) shell_streams = [stream] for addr in shell_addrs: stream.connect(disambiguate_url(addr, self.location)) # end single stream-socket # control stream: control_addr = str(msg.content.control) control_stream = zmqstream.ZMQStream(ctx.socket(zmq.XREP), loop) control_stream.setsockopt(zmq.IDENTITY, identity) control_stream.connect( disambiguate_url(control_addr, self.location)) # create iopub stream: iopub_addr = msg.content.iopub iopub_stream = zmqstream.ZMQStream(ctx.socket(zmq.PUB), loop) iopub_stream.setsockopt(zmq.IDENTITY, identity) iopub_stream.connect(disambiguate_url(iopub_addr, self.location)) # launch heartbeat hb_addrs = msg.content.heartbeat # print (hb_addrs) # # Redirect input streams and set a display hook. if self.out_stream_factory: sys.stdout = self.out_stream_factory(self.session, iopub_stream, u'stdout') sys.stdout.topic = 'engine.%i.stdout' % self.id sys.stderr = self.out_stream_factory(self.session, iopub_stream, u'stderr') sys.stderr.topic = 'engine.%i.stderr' % self.id if self.display_hook_factory: sys.displayhook = self.display_hook_factory( self.session, iopub_stream) sys.displayhook.topic = 'engine.%i.pyout' % self.id self.kernel = Kernel(config=self.config, int_id=self.id, ident=self.ident, session=self.session, control_stream=control_stream, shell_streams=shell_streams, iopub_stream=iopub_stream, loop=loop, user_ns=self.user_ns, log=self.log) self.kernel.start() hb_addrs = [ disambiguate_url(addr, self.location) for addr in hb_addrs ] heart = Heart(*map(str, hb_addrs), heart_id=identity) heart.start() else: self.log.fatal("Registration Failed: %s" % msg) raise Exception("Registration Failed: %s" % msg) self.log.info("Completed registration with id %i" % self.id) def abort(self): self.log.fatal("Registration timed out after %.1f seconds" % self.timeout) self.session.send(self.registrar, "unregistration_request", content=dict(id=self.id)) time.sleep(1) sys.exit(255) def start(self): dc = ioloop.DelayedCallback(self.register, 0, self.loop) dc.start() self._abort_dc = ioloop.DelayedCallback(self.abort, self.timeout * 1000, self.loop) self._abort_dc.start()