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 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_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_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 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 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 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)
