def __init__(self, addrs, ctx, identity, name, version=0):
Node.__init__(self, identity, addrs['graph'], addrs['msg'], ctx=ctx)
ZmqHandler.__init__(self, addrs['results'], ctx=ctx)
self.comm = GraphCommHandler(name, addrs['comm'], ctx=ctx)
self._name = name
self.version = version
self.exceptions = {}
def graph_comm(request, ipc_dir, event_loop):
ctxs = []
name = "graph"
addr = "ipc://%s/graphcomm_async" % ipc_dir
# check the requested parameters
if isinstance(request.param, dict):
conf = request.param
use_async = False
else:
use_async, conf = request.param
if use_async:
ctx = zmq.Context()
manager = FakeManager(ctx, addr, conf)
ctxs.append(ctx)
comm = AsyncGraphCommHandler(name, addr)
ctxs.append(comm._ctx)
else:
comm = GraphCommHandler(name, addr)
ctx = comm._ctx
manager = FakeManager(ctx, addr, conf)
ctxs.append(ctx)
yield comm, conf
comm._sock.send_string('test_exit')
manager.thread.join()
# clean up all the zmq stuff
for ctx in ctxs:
ctx.destroy()
def graph_comm_simple(request, ipc_dir):
ctxs = []
name = "graph"
addr = "ipc://%s/graphcomm_simple" % ipc_dir
comm = GraphCommHandler(name, addr, request.param)
ctxs.append(comm._ctx)
yield comm, addr
# clean up all the zmq stuff
for ctx in ctxs:
ctx.destroy()
def flowchart(request, workerjson, broker, ipc_dir, graphmgr_addr,
qevent_loop):
try:
from pytest_cov.embed import cleanup_on_sigterm
cleanup_on_sigterm()
except ImportError:
pass
parser = build_parser()
args = parser.parse_args([
"-n", "1", '-i',
str(ipc_dir), '--headless',
'%s://%s' % (request.param, workerjson)
])
queue = mp.Queue()
ami = mp.Process(name='ami', target=run_ami, args=(args, queue))
ami.start()
try:
# wait for ami to be fully up before updating the sources
with GraphCommHandler(graphmgr_addr.name, graphmgr_addr.comm) as comm:
while not comm.sources:
time.sleep(0.1)
with Flowchart(broker_addr=broker.broker_sub_addr,
graphmgr_addr=graphmgr_addr,
checkpoint_addr=broker.checkpoint_pub_addr) as fc:
qevent_loop.run_until_complete(fc.updateSources(init=True))
yield (fc, broker)
except Exception as e:
# let the fixture exit 'gracefully' if it fails
print("error setting up flowchart fixture:", e)
yield None
finally:
queue.put(None)
ami.join(2)
# if ami still hasn't exitted then kill it
if ami.is_alive():
ami.terminate()
ami.join()
if ami.exitcode == 0 or ami.exitcode == -signal.SIGTERM:
return 0
else:
print('AMI exited with non-zero status code: %d' % ami.exitcode)
return 1
def start_ami(request, workerjson):
try:
from pytest_cov.embed import cleanup_on_sigterm
cleanup_on_sigterm()
except ImportError:
pass
parser = build_parser()
args = parser.parse_args([
"-n", "1", '--headless', '--tcp',
'%s://%s' % (request.param, workerjson)
])
queue = mp.Queue()
ami = mp.Process(name='ami', target=run_ami, args=(args, queue))
ami.start()
try:
host = "127.0.0.1"
comm_addr = "tcp://%s:%d" % (host, BasePort + Ports.Comm)
with GraphCommHandler(args.graph_name, comm_addr) as comm_handler:
yield comm_handler
except Exception as e:
# let the fixture exit 'gracefully' if it fails
print(e)
yield None
finally:
queue.put(None)
ami.join(1)
# if ami still hasn't exitted then kill it
if ami.is_alive():
ami.terminate()
ami.join(1)
if ami.exitcode == 0 or ami.exitcode == -signal.SIGTERM:
return 0
else:
print('AMI exited with non-zero status code: %d' % ami.exitcode)
return 1
def run_ami(args, queue=None):
xtcdir = None
ipcdir = None
owns_ipcdir = True
flags = {}
if queue is None:
queue = mp.Queue()
if args.ipc:
ipcdir = tempfile.mkdtemp()
owns_ipcdir = True
elif args.ipc_dir is not None:
ipcdir = args.ipc_dir
owns_ipcdir = False
if ipcdir is None:
host = "127.0.0.1"
comm_addr = "tcp://%s:%d" % (host, args.port)
graph_addr = "tcp://%s:%d" % (host, args.port+1)
collector_addr = "tcp://%s:%d" % (host, args.port+2)
globalcol_addr = "tcp://%s:%d" % (host, args.port+3)
results_addr = "tcp://%s:%d" % (host, args.port+4)
export_addr = "tcp://%s:%d" % (host, args.port+5)
msg_addr = "tcp://%s:%d" % (host, args.port+6)
info_addr = "tcp://%s:%d" % (host, args.port+7)
view_addr = "tcp://%s:%d" % (host, args.port+8)
profile_addr = "tcp://%s:%d" % (host, args.port+9)
else:
collector_addr = "ipc://%s/node_collector" % ipcdir
globalcol_addr = "ipc://%s/collector" % ipcdir
graph_addr = "ipc://%s/graph" % ipcdir
comm_addr = "ipc://%s/comm" % ipcdir
results_addr = "ipc://%s/results" % ipcdir
export_addr = "ipc://%s/export" % ipcdir
msg_addr = "ipc://%s/message" % ipcdir
info_addr = "ipc://%s/info" % ipcdir
view_addr = "ipc://%s/view" % ipcdir
profile_addr = "ipc://%s/profile" % ipcdir
procs = []
client_proc = None
log_handlers = [logging.StreamHandler()]
if args.headless or args.console:
console_fmt = logging.Formatter(LogConfig.BasicFormat)
log_handlers[0].setFormatter(console_fmt)
if args.log_file is not None:
log_handlers.append(logging.FileHandler(args.log_file))
log_level = getattr(logging, args.log_level.upper(), logging.INFO)
logging.basicConfig(format=LogConfig.FullFormat, level=log_level, handlers=log_handlers)
try:
for flag in args.flags:
try:
key, value = flag.split('=')
flags[key] = value
except ValueError:
logger.exception("Problem parsing data source flag %s", flag)
if args.source is not None:
src_url_match = re.match('(?P<prot>.*)://(?P<body>.*)', args.source)
if src_url_match:
src_cfg = src_url_match.groups()
else:
logger.critical("Invalid data source config string: %s", args.source)
return 1
else:
src_cfg = None
for i in range(args.num_workers):
proc = mp.Process(
name='worker%03d-n0' % i,
target=functools.partial(_sys_exit, run_worker),
args=(i, args.num_workers, args.heartbeat, src_cfg,
collector_addr, graph_addr, msg_addr, export_addr, flags, args.prometheus_dir, args.hutch)
)
proc.daemon = True
proc.start()
procs.append(proc)
collector_proc = mp.Process(
name='nodecol-n0',
target=functools.partial(_sys_exit, run_node_collector),
args=(0, args.num_workers, collector_addr, globalcol_addr, graph_addr, msg_addr,
args.prometheus_dir, args.hutch)
)
collector_proc.daemon = True
collector_proc.start()
procs.append(collector_proc)
globalcol_proc = mp.Process(
name='globalcol',
target=functools.partial(_sys_exit, run_global_collector),
args=(0, 1, globalcol_addr, results_addr, graph_addr, msg_addr,
args.prometheus_dir, args.hutch)
)
globalcol_proc.daemon = True
globalcol_proc.start()
procs.append(globalcol_proc)
manager_proc = mp.Process(
name='manager',
target=functools.partial(_sys_exit, run_manager),
args=(args.num_workers, 1, results_addr, graph_addr, comm_addr, msg_addr, info_addr, export_addr,
view_addr, profile_addr, args.prometheus_dir, args.hutch)
)
manager_proc.daemon = True
manager_proc.start()
procs.append(manager_proc)
if args.export:
if run_export is None:
logger.critical("Export module is not available: p4p needs to be installed to use the export feature!")
return 1
export_proc = mp.Process(
name='export',
target=functools.partial(_sys_exit, run_export),
args=(args.export, comm_addr, export_addr, args.aggregate)
)
export_proc.daemon = True
export_proc.start()
procs.append(export_proc)
if not (args.console or args.headless):
client_proc = mp.Process(
name='client',
target=run_client,
args=(args.graph_name, comm_addr, info_addr, view_addr, profile_addr, args.load, args.gui_mode,
args.prometheus_dir, args.hutch)
)
client_proc.daemon = False
client_proc.start()
procs.append(client_proc)
# register a signal handler for cleanup on sigterm
signal.signal(signal.SIGTERM, functools.partial(_sig_handler, procs))
if args.console:
run_console(args.graph_name, comm_addr, args.load)
elif args.headless:
if args.load:
comm_handler = GraphCommHandler(args.graph_name, comm_addr)
comm_handler.load(args.load)
while queue.empty():
pass
else:
client_proc.join()
except KeyboardInterrupt:
logger.info("Worker killed by user...")
finally:
failed_proc = cleanup(procs)
# cleanup ipc directories
if owns_ipcdir and ipcdir is not None and os.path.exists(ipcdir):
shutil.rmtree(ipcdir)
if xtcdir is not None and os.path.exists(xtcdir):
shutil.rmtree(xtcdir)
# return a non-zero status code if any workerss died
if client_proc is not None and client_proc.exitcode != 0:
return client_proc.exitcode
elif failed_proc:
return 1
def flowchart_hdf(request, tmp_path, qtbot, broker, ipc_dir, graphmgr_addr,
qevent_loop):
try:
from pytest_cov.embed import cleanup_on_sigterm
cleanup_on_sigterm()
except ImportError:
pass
cfg = {
"interval": 0.01,
"init_time": 0.1,
"bound": 100,
"repeat": True,
"files": [os.path.join('tests/graphs', request.param[0])]
}
fname = os.path.join(tmp_path, 'worker.json')
with open(fname, 'w') as fd:
json.dump(cfg, fd)
parser = build_parser()
args = parser.parse_args(
["-n", "1", '-i',
str(ipc_dir), '--headless',
'hdf5://%s' % fname])
queue = mp.Queue()
ami = mp.Process(name='ami', target=run_ami, args=(args, queue))
ami.start()
try:
# wait for ami to be fully up before updating the sources
comm = GraphCommHandler(graphmgr_addr.name, graphmgr_addr.comm)
while not comm.sources:
time.sleep(0.1)
with Flowchart(broker_addr=broker.broker_sub_addr,
graphmgr_addr=graphmgr_addr,
checkpoint_addr=broker.checkpoint_pub_addr) as fc:
qevent_loop.run_until_complete(fc.updateSources(init=True))
qtbot.addWidget(fc.widget())
fc.loadFile(os.path.join('tests/graphs', request.param[1]))
yield (fc, broker, comm)
except Exception as e:
# let the fixture exit 'gracefully' if it fails
print("error setting up flowchart fixture:", e)
yield None
finally:
queue.put(None)
ami.join(2)
# if ami still hasn't exitted then kill it
if ami.is_alive():
ami.terminate()
ami.join()
if ami.exitcode == 0 or ami.exitcode == -signal.SIGTERM:
return 0
else:
print('AMI exited with non-zero status code: %d' % ami.exitcode)
return 1
class ResultsInjector(Node, ZmqHandler):
def __init__(self, addrs, ctx, identity, name, version=0):
Node.__init__(self, identity, addrs['graph'], addrs['msg'], ctx=ctx)
ZmqHandler.__init__(self, addrs['results'], ctx=ctx)
self.comm = GraphCommHandler(name, addrs['comm'], ctx=ctx)
self._name = name
self.version = version
self.exceptions = {}
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, traceback):
self.close()
def close(self):
self.graph_comm.close()
self.node_msg_comm.close()
self.collector.close()
self.comm.close()
@property
def name(self):
return self._name
@property
def wait_counter(self):
hb = self.comm.heartbeat
return hb if hb is not None else -1
def recv_graph(self, name, version, args, payload):
self.store_graph(name, version, payload)
def recv_graph_add(self, name, version, args, payload):
self.store_graph(name, version, payload)
def recv_graph_del(self, name, version, args, payload):
self.store_graph(name, version, payload)
def recv_graph_purge(self, name, version, args, payload):
if name in self.graphs:
del self.graphs[name]
def recv_graph_exception(self, name, version, exception):
if name not in self.exceptions:
self.exceptions[name] = {}
self.exceptions[name][version] = exception
def store_graph(self, name, version, payload):
if name not in self.graphs:
self.graphs[name] = {}
self.graphs[name][version] = payload
def mark(self):
count = self.wait_counter + 1
self.collector_message(self.node, Heartbeat(count, 0), self.name, self.version, {})
return count
def partition(self, payload, wait=False):
payload = {src: at.dumps(typ) for src, typ in payload.items()}
self.message(MsgTypes.Transition, self.node, Transition(Transitions.Configure, payload))
if wait:
self.wait_for(self.mark())
else:
return self.mark
def data(self, hb, payload, wait=False):
self.collector_message(self.node, Heartbeat(hb, 0), self.name, self.version, payload)
if wait:
self.wait_for(hb)
else:
return hb
def wait_for(self, counter):
while self.wait_counter < counter:
time.sleep(0.01)
def wait_graph(self, timeout=None):
if timeout is None:
self.graph_comm.recv()
return True
else:
failed = True
start = time.time()
while time.time() - start < timeout:
try:
self.graph_comm.recv(False)
failed = False
break
except zmq.Again:
pass
return not failed