#!/usr/bin/env python3 # -*- coding: utf-8 -*- ''' Large bipartite graph stress test. ''' import logging import multiprocessing import time import atexit from common import connection_argument_parser, assert_equal, \ OutputData, execute_till_success, \ batch, render, SessionCache def parse_args(): ''' Parses user arguments :return: parsed arguments ''' parser = connection_argument_parser() parser.add_argument('--worker-count', type=int, default=multiprocessing.cpu_count(), help='Number of concurrent workers.') parser.add_argument("--logging", default="INFO", choices=["INFO", "DEBUG", "WARNING", "ERROR"], help="Logging level") parser.add_argument('--u-count', type=int, default=100, help='Size of U set in the bipartite graph.') parser.add_argument('--v-count', type=int, default=100, help='Size of V set in the bipartite graph.') parser.add_argument('--vertex-batch-size', type=int, default=100, help="Create vertices in batches of this size.") parser.add_argument('--edge-batching', action='store_true', help='Create edges in batches.') parser.add_argument('--edge-batch-size', type=int, default=100, help='Number of edges in a batch when edges ' 'are created in batches.') return parser.parse_args() log = logging.getLogger(__name__) args = parse_args() output_data = OutputData() atexit.register(SessionCache.cleanup) def create_u_v_edges(u): ''' Creates nodes and checks that all nodes were created. create edges from one vertex in U set to all vertex of V set :param worker_id: worker id :return: tuple (worker_id, create execution time, time unit) ''' start_time = time.time() session = SessionCache.argument_session(args) no_failures = 0 match_u = 'MATCH (u:U {id: %d})' % u if args.edge_batching: # TODO: try to randomize execution, the execution time should # be smaller, add randomize flag for v_id_batch in batch(range(args.v_count), args.edge_batch_size): match_v = render(" MATCH (v{0}:V {{id: {0}}})", v_id_batch) create_u = render(" CREATE (u)-[:R]->(v{0})", v_id_batch) query = match_u + "".join(match_v) + "".join(create_u) no_failures += execute_till_success(session, query)[1] else: no_failures += execute_till_success( session, match_u + ' MATCH (v:V) CREATE (u)-[:R]->(v)')[1] end_time = time.time() return u, end_time - start_time, "s", no_failures def traverse_from_u_worker(u): ''' Traverses edges starting from an element of U set. Traversed labels are: :U -> :V -> :U. ''' session = SessionCache.argument_session(args) start_time = time.time() assert_equal( args.u_count * args.v_count - args.v_count, # cypher morphism session.run("MATCH (u1:U {id: %s})-[e1]->(v:V)<-[e2]-(u2:U) " "RETURN count(v) AS cnt" % u).data()[0]['cnt'], "Number of traversed edges started " "from U(id:%s) is wrong!. " % u + "Expected: %s Actual: %s") end_time = time.time() return u, end_time - start_time, 's' def traverse_from_v_worker(v): ''' Traverses edges starting from an element of V set. Traversed labels are: :V -> :U -> :V. ''' session = SessionCache.argument_session(args) start_time = time.time() assert_equal( args.u_count * args.v_count - args.u_count, # cypher morphism session.run("MATCH (v1:V {id: %s})<-[e1]-(u:U)-[e2]->(v2:V) " "RETURN count(u) AS cnt" % v).data()[0]['cnt'], "Number of traversed edges started " "from V(id:%s) is wrong!. " % v + "Expected: %s Actual: %s") end_time = time.time() return v, end_time - start_time, 's' def execution_handler(): ''' Initializes client processes, database and starts the execution. ''' # instance cleanup session = SessionCache.argument_session(args) start_time = time.time() # clean existing database session.run('MATCH (n) DETACH DELETE n').consume() cleanup_end_time = time.time() output_data.add_measurement("cleanup_time", cleanup_end_time - start_time) # create U vertices for b in batch(render('CREATE (:U {{id: {}}})', range(args.u_count)), args.vertex_batch_size): session.run(" ".join(b)).consume() # create V vertices for b in batch(render('CREATE (:V {{id: {}}})', range(args.v_count)), args.vertex_batch_size): session.run(" ".join(b)).consume() vertices_create_end_time = time.time() output_data.add_measurement( 'vertices_create_time', vertices_create_end_time - cleanup_end_time) # concurrent create execution & tests with multiprocessing.Pool(args.worker_count) as p: create_edges_start_time = time.time() for worker_id, create_time, time_unit, no_failures in \ p.map(create_u_v_edges, [i for i in range(args.u_count)]): log.info('Worker ID: %s; Create time: %s%s Failures: %s' % (worker_id, create_time, time_unit, no_failures)) create_edges_end_time = time.time() output_data.add_measurement( 'edges_create_time', create_edges_end_time - create_edges_start_time) # check total number of edges assert_equal( args.v_count * args.u_count, session.run( 'MATCH ()-[r]->() ' 'RETURN count(r) AS cnt').data()[0]['cnt'], "Total number of edges isn't correct! Expected: %s Actual: %s") # check traversals starting from all elements of U traverse_from_u_start_time = time.time() for u, traverse_u_time, time_unit in \ p.map(traverse_from_u_worker, [i for i in range(args.u_count)]): log.info("U {id: %s} %s%s" % (u, traverse_u_time, time_unit)) traverse_from_u_end_time = time.time() output_data.add_measurement( 'traverse_from_u_time', traverse_from_u_end_time - traverse_from_u_start_time) # check traversals starting from all elements of V traverse_from_v_start_time = time.time() for v, traverse_v_time, time_unit in \ p.map(traverse_from_v_worker, [i for i in range(args.v_count)]): log.info("V {id: %s} %s%s" % (v, traverse_v_time, time_unit)) traverse_from_v_end_time = time.time() output_data.add_measurement( 'traverse_from_v_time', traverse_from_v_end_time - traverse_from_v_start_time) # check total number of vertices assert_equal( args.v_count + args.u_count, session.run('MATCH (n) RETURN count(n) AS cnt').data()[0]['cnt'], "Total number of vertices isn't correct! Expected: %s Actual: %s") # check total number of edges assert_equal( args.v_count * args.u_count, session.run( 'MATCH ()-[r]->() RETURN count(r) AS cnt').data()[0]['cnt'], "Total number of edges isn't correct! Expected: %s Actual: %s") end_time = time.time() output_data.add_measurement("total_execution_time", end_time - start_time) if __name__ == '__main__': logging.basicConfig(level=args.logging) if args.logging != "DEBUG": logging.getLogger("neo4j").setLevel(logging.WARNING) output_data.add_status("stress_test_name", "bipartite") output_data.add_status("number_of_vertices", args.u_count + args.v_count) output_data.add_status("number_of_edges", args.u_count * args.v_count) output_data.add_status("vertex_batch_size", args.vertex_batch_size) output_data.add_status("edge_batching", args.edge_batching) output_data.add_status("edge_batch_size", args.edge_batch_size) execution_handler() if args.logging in ["DEBUG", "INFO"]: output_data.dump()