def main(): arguments = docopt.docopt(__doc__) logger.info(arguments) context = zmq.Context() metadata = json.loads(arguments['--metadata']) url = arguments['<url>'] if arguments['--request'] == 'SUB': sub = context.socket(zmq.SUB) sub.connect(url) sub.setsockopt(zmq.SUBSCRIBE, '') while True: arr, metadata = recv_array(sub) logger.info("metadata: %s\n array: %s", metadata, arr) elif arguments['--request'] == 'REQ': req = context.socket(zmq.REQ) req.connect(url) value = None send_array(req, value, metadata=metadata) # wait for reply arr, metadata = recv_array(req) logger.info("metadata: %s\n array: %s", metadata, arr) elif arguments['--request'] == 'PUSH': req = context.socket(zmq.PUSH) req.connect(url) value = None send_array(req, value, metadata=metadata) logger.info("metadata: %s", metadata)
def test_missing_scalar(arr, req, rep): """send an array with missing data as a scalar""" arr_masked = np.ma.masked_less(arr, 2) # test if it works if we use a numpy scalar arr_masked.fill_value = np.int32(9999) mmi.send_array(req, arr_masked) received, metadata = mmi.recv_array(rep) numpy.testing.assert_array_equal(arr_masked, received)
def test_metadata_only(self): """send a message with only metadata""" req = ctx.socket(zmq.REQ) req.connect('tcp://localhost:9002') rep = ctx.socket(zmq.REP) rep.bind('tcp://*:9002') mmi.send_array(req, A=None) _, metadata = mmi.recv_array(rep) self.assertTrue('timestamp' in metadata)
def test_sndrcv(self): A = np.array([1,2,3]) req = ctx.socket(zmq.REQ) req.connect('tcp://localhost:9002') rep = ctx.socket(zmq.REP) rep.bind('tcp://*:9002') mmi.send_array(req, A) B, metadata = mmi.recv_array(rep) numpy.testing.assert_array_equal(A, B)
def test_missing(self): """send an array with missing data""" A = np.array([1, 2, 3, 4]) A = np.ma.masked_less(A, 2) req = ctx.socket(zmq.REQ) req.connect('tcp://localhost:9002') rep = ctx.socket(zmq.REP) rep.bind('tcp://*:9002') mmi.send_array(req, A) B, metadata = mmi.recv_array(rep) numpy.testing.assert_array_equal(A, B)
def run(self): """run the model""" model = self.model configfile = self.configfile interval = self.interval sockets = self.sockets model.initialize(configfile) if model.state == 'pause': logger.info( "model initialized and started in pause mode, waiting for requests" ) else: logger.info("model started and initialized, running") if self.tracker: self.register() atexit.register(self.unregister) self.process_incoming() # Keep on counting indefinitely counter = itertools.count() logger.info("Entering timeloop...") for i in counter: while model.state == "pause": # keep waiting for messages when paused # process_incoming should set model.state to play self.process_incoming() else: # otherwise process messages once and continue self.process_incoming() if model.state == "quit": break # lookup dt or use -1 (default) dt = model.get_time_step() or -1 model.update(dt) # check counter, if not a multiple of interval, skip this step if i % interval: continue for key in self.output_vars: value = model.get_var(key) metadata = {'name': key, 'iteration': i} # 4ms for 1M doubles logger.debug("sending {}".format(metadata)) if 'pub' in sockets: send_array(sockets['pub'], value, metadata=metadata) logger.info("Finalizing...") model.finalize()
def finalize(self): """ Finalize the module """ method = "finalize" A = None metadata = {method : -1} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket)
def update(self, dt): """ Advance the module with timestep dt """ method = "update" A = None metadata = {method : dt} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket)
def initialize(self, configfile=None): """ Initialize the module """ method = "initialize" A = None metadata = {method : configfile} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket)
def set_var(self, name, var): """ Set the variable name with the values of var """ method = "set_var" A = var metadata = {method : name} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket)
def test_missing_scalar(self): """send an array with missing data as a scalar""" A = np.array([1, 2, 3, 4]) A = np.ma.masked_less(A, 2) # test if it works if we use a numpy scalar A.fill_value = np.int32(9999) req = ctx.socket(zmq.REQ) req.connect('tcp://localhost:9002') rep = ctx.socket(zmq.REP) rep.bind('tcp://*:9002') mmi.send_array(req, A) B, metadata = mmi.recv_array(rep) numpy.testing.assert_array_equal(A, B)
def get_var_name(self, i): """ Return variable name """ method = "get_var_name" A = None metadata = {method : i} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket) return metadata[method]
def set_var_slice(self, name, start, count, var): """ Set the variable name with the values of var """ method = "set_var_slice" A = var metadata = {method: name, "start": start, "count": count} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags)
def initialize(self, configfile=None): """ Initialize the module """ method = "initialize" A = None metadata = {method: configfile} send_array(self.socket, A, metadata) A, metadata = recv_array( self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags)
def get_start_time(self): """ Return start time """ method = "get_start_time" A = None metadata = {method : -1} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket) return metadata[method]
def get_current_time(self): """ Return current time of simulation """ method = "get_current_time" A = None metadata = {method : -1} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket) return metadata[method]
def get_var(self, name): """ Return an nd array from model library """ method = "get_var" A = None metadata = {method : name} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket) return A
def get_var_count(self): """ Return number of variables """ method = "get_var_count" A = None metadata = {method : -1} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket) return metadata[method]
def get_var_rank(self, name): """ Return variable rank """ method = "get_var_rank" A = None metadata = {method : name} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket) return metadata[method]
def get_var_shape(self, name): """ Return variable shape """ method = "get_var_shape" A = None metadata = {method : rank} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket) return metadata[method]
def update(self, dt): """ Advance the module with timestep dt """ method = "update" A = None metadata = {method: dt} send_array(self.socket, A, metadata) A, metadata = recv_array( self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags)
def set_var(self, name, var): """ Set the variable name with the values of var """ method = "set_var" A = var metadata = {method: name} send_array(self.socket, A, metadata) A, metadata = recv_array( self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags)
def finalize(self): """ Finalize the module """ method = "finalize" A = None metadata = {method: -1} send_array(self.socket, A, metadata) A, metadata = recv_array( self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags)
def update(self, dt): """ Advance the module with timestep dt """ method = "update" A = None metadata = {method: dt} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags)
def set_var(self, name, var): """ Set the variable name with the values of var """ method = "set_var" A = var metadata = {method: name} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags)
def set_current_time(self, t): """ Set current time of simulation """ method = "set_current_time" A = None metadata = {method: t} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags)
def finalize(self): """ Finalize the module """ method = "finalize" A = None metadata = {method: -1} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags)
def initialize(self, configfile=None): """ Initialize the module """ method = "initialize" A = None metadata = {method: configfile} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags)
def get_var_count(self): """ Return number of variables """ method = "get_var_count" A = None metadata = {method: -1} send_array(self.socket, A, metadata) A, metadata = recv_array( self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags) return metadata[method]
def get_var_type(self, name): """ Return variable name """ method = "get_var_type" A = None metadata = {method: name} send_array(self.socket, A, metadata) A, metadata = recv_array( self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags) return metadata[method]
def remote(self, action): """ Function specific for MMI, not BMI. action is one of: "play", "stop", "pause", "rewind", "quit" """ method = "remote" A = None metadata = {method: action} send_array(self.socket, A, metadata) A, metadata = recv_array( self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags) return metadata[method]
def get_current_time(self): """ Return current time of simulation """ method = "get_current_time" A = None metadata = {method: -1} send_array(self.socket, A, metadata) A, metadata = recv_array( self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags) return metadata[method]
def get_start_time(self): """ Return start time """ method = "get_start_time" A = None metadata = {method: -1} send_array(self.socket, A, metadata) A, metadata = recv_array( self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags) return metadata[method]
def get_var(self, name): """ Return an nd array from model library """ method = "get_var" A = None metadata = {method: name} send_array(self.socket, A, metadata) A, metadata = recv_array( self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags) return A
def get_var_shape(self, name): """ Return variable shape """ method = "get_var_shape" A = None metadata = {method: name} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags) return metadata[method]
def get_start_time(self): """ Return start time """ method = "get_start_time" A = None metadata = {method: -1} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags) return metadata[method]
def remote(self, action): """ Function specific for MMI, not BMI. action is one of: "play", "stop", "pause", "rewind", "quit" """ method = "remote" A = None metadata = {method: action} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags) return metadata[method]
def get_var_count(self): """ Return number of variables """ method = "get_var_count" A = None metadata = {method: -1} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags) return metadata[method]
def get_var(self, name): """ Return an nd array from model library """ method = "get_var" A = None metadata = {method: name} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags) return A
def get_time_step(self): """ Return time step of simulation """ method = "get_time_step" A = None metadata = {method: -1} send_array(self.socket, A, metadata) A, metadata = recv_array(self.socket, poll=self.poll, poll_timeout=self.poll_timeout, flags=self.zmq_flags) return metadata[method]
def runner( arguments, wrapper_class, wrapper_kwargs={}, extra_metadata={}): """ MMI Runner """ # make a socket that replies to message with the grid # if we are running mpi we want to know the rank if arguments['--mpi']: import mpi4py.MPI comm = mpi4py.MPI.COMM_WORLD rank = comm.Get_rank() size = comm.Get_size() else: # or we assume it's 0 rank = 0 size = 1 if arguments["--port"] == "random": # ports will be filled in ports = {} else: port = int(arguments['--port']) ports = { "REQ": port + 0, "PUSH": port + 1, "SUB": port + 2 } # if we want to communicate with separate domains # we have to setup a socket for each of them if arguments['--mpi'] == 'all': # use a socket for each rank rank for port in ports: ports[port] += (rank * 3) # now we can create the zmq sockets and poller sockets = {} if rank == 0 or arguments['--mpi'] == 'all': # Create sockets sockets = create_sockets(ports) logger.debug("ZMQ MMI Ports:") for key, value in ports.items(): logger.debug("%s = %d" % (key, value)) # not so verbose # bmi.wrapper.logger.setLevel(logging.WARN) wrapper = wrapper_class( engine=arguments['<engine>'], configfile=arguments['<configfile>'], **wrapper_kwargs) # for replying to grid requests with wrapper as model: model.state = "play" if arguments["--pause"]: model.state = "pause" logger.info("model initialized and started in pause mode, waiting for requests ...") else: logger.info("model started and initialized, running ...") if arguments["--track"]: server = arguments["--track"] metadata = {} # update connection information from external service # You might want to disable this if you have some sort of sense of privacy try: metadata["ifconfig"] = requests.get("http://ipinfo.io/json").json() except: pass # except requests.exceptions.ConnectionError: # logger.exception("Could not read ip info from ipinfo.io") # pass # except simplejson.scanner.JSONDecodeError: # logger.exception("Could not parse ip info from ipinfo.io") # pass # node metadata["node"] = platform.node() metadata.update({ "engine": arguments['<engine>'], "configfile": arguments['<configfile>'], "ports": ports, "rank": rank, "size": size }) metadata_filename = os.path.join( os.path.dirname(arguments['<configfile>']), "metadata.json" ) if os.path.isfile(metadata_filename): metadata.update(json.load(open(metadata_filename))) metadata.update(extra_metadata) register(server, metadata) if arguments["--track"]: atexit.register(unregister, server, metadata) # Start a reply process in the background, with variables available # after initialization, sent all at once as py_obj data = { var: model.get_var(var) for var in arguments['-g'] } process_incoming(model, sockets, data) # Keep on counting indefinitely counter = itertools.count() for i in counter: while model.state == "pause": # keep waiting for messages when paused process_incoming(model, sockets, data) else: # otherwise process messages once and continue process_incoming(model, sockets, data) logger.debug("i %s", i) # paused ... model.update(-1) # check counter if arguments.get('--interval') and (i % int(arguments['--interval'])): continue for key in arguments['-o']: value = model.get_var(key) metadata = {'name': key, 'iteration': i} # 4ms for 1M doubles logger.debug("sending {}".format(metadata)) if 'pub' in sockets: send_array(sockets['pub'], value, metadata=metadata)
var: subgrid.get_nd(var, sliced=True) for var in arguments.globalvariables } # add the quad_grid for easy plotting data["quad_grid"] = python_subgrid.plotting.make_quad_grid(subgrid) process_incoming(subgrid, poller, rep, pull, data) # Keep on counting indefinitely counter = itertools.count() for i in counter: # Any requests? while paused: process_incoming(subgrid, poller, rep, pull, data) # Calculate subgrid.update(-1) # check counter if (i % arguments.interval): continue for key in arguments.outputvariables: value = subgrid.get_nd(key, sliced=True) metadata = {'name': key, 'iteration': i} # 4ms for 1M doubles logger.info("sending {}".format(metadata)) send_array(pub, value, metadata=metadata)
def process_incoming(model, sockets, data): """ process incoming messages data is a dict with several arrays """ # Check for new messages if not sockets: return # unpack sockets poller = sockets['poller'] rep = sockets['rep'] # pull = sockets['pull'] pub = sockets['pub'] items = poller.poll(10) for sock, n in items: for i in range(n): A, metadata = recv_array(sock) logger.debug("got metadata: %s", metadata) var = None # bmi actions if "update" in metadata: dt = float(metadata["update"]) logger.debug("updating with dt %s", dt) model.update(dt) metadata["dt"] = dt elif "get_var" in metadata: name = metadata["get_var"] # temporary implementation if metadata.get("copy", False): var = model.get_var(name).copy() else: var = model.get_var(name) logger.debug("sending variable %s with shape %s", name, var.shape) metadata['name'] = name # assert socket is req socket elif "get_var_count" in metadata: # temporary implementation n = model.get_var_count() metadata['get_var_count'] = n # assert socket is req socket elif "get_var_rank" in metadata: # temporary implementation var_name = metadata['get_var_rank'] n = model.get_var_rank(var_name) metadata['get_var_rank'] = n # assert socket is req socket elif "get_var_shape" in metadata: # temporary implementation var_name = metadata['get_var_shape'] n = model.get_var_shape(var_name) metadata['get_var_shape'] = tuple([int(item) for item in n]) # assert socket is req socket elif "get_var_type" in metadata: # temporary implementation var_name = metadata['get_var_type'] n = model.get_var_type(var_name) metadata['get_var_type'] = n # assert socket is req socket elif "get_var_name" in metadata: i = int(metadata["get_var_name"]) name = model.get_var_name(i) metadata['get_var_name'] = name # assert socket is req socket elif "set_var" in metadata: name = metadata["set_var"] logger.debug("setting variable %s", name) model.set_var(name, A) metadata["name"] = name # !? elif "set_var_slice" in metadata: name = metadata["set_var_slice"] logger.debug("setting variable %s", name) start = metadata["start"] count = metadata["count"] model.set_var_slice(name, start, count, A) metadata["name"] = name # !? elif "set_var_index" in metadata: name = metadata["set_var_index"] logger.debug("setting variable %s using index index", name) index = metadata["index"] # TODO: test if this is fast enough. # Otherwise move to BMI ++ but that is # a bit of a burden on implementers var = model.get_var(name).copy() var.flat[index] = A model.set_var(name, var) metadata["name"] = name # !? elif "get_current_time" in metadata: metadata["get_current_time"] t = model.get_current_time() metadata['get_current_time'] = t elif "get_time_step" in metadata: metadata["get_time_step"] dt = model.get_time_step() metadata['get_time_step'] = dt elif "get_end_time" in metadata: metadata["get_end_time"] t = model.get_end_time() metadata['get_end_time'] = t elif "get_start_time" in metadata: metadata["get_start_time"] t = model.get_start_time() metadata['get_start_time'] = t # assert socket is req socket # custom actions elif "remote" in metadata: assert metadata["remote"] in {"play", "stop", "pause", "rewind"} model.state = metadata["remote"] elif "operator" in metadata: # TODO: support same operators as MPI_ops here...., # TODO: reduce before apply # TODO: assert pull socket pass # S = tuple(slice(*x) for x in action['slice']) # print(repr(arr[S])) # if action['operator'] == 'setitem': # arr[S] = data # elif action['operator'] == 'add': # arr[S] += data elif "initialize" in metadata: config_file = metadata["initialize"] model.initialize(config_file) elif "finalize" in metadata: model.finalize() else: logger.warn("got unknown message {} from socket {}".format(str(metadata), sock)) if sock.socket_type == zmq.REP: # reply send_array(rep, var, metadata=metadata) # any getter requested through the pull socket? elif any(x.startswith("get_") for x in metadata) and sock.socket_type == zmq.PULL: # return through the pub socket send_array(pub, var, metadata=metadata)