diff --git a/tests/macro_benchmark/clients.py b/tests/macro_benchmark/clients.py index f7a869a39..ef9b70555 100644 --- a/tests/macro_benchmark/clients.py +++ b/tests/macro_benchmark/clients.py @@ -3,26 +3,23 @@ import os import time import json import tempfile -from common import get_absolute_path, WALL_TIME, CPU_TIME, MAX_MEMORY +from common import get_absolute_path, WALL_TIME, CPU_TIME, MAX_MEMORY, set_cpus log = logging.getLogger(__name__) try: import jail - APOLLO = True except: import jail_faker as jail - APOLLO = False # This could be a function, not a class, but we want to reuse jail process since # we can instantiate only 8 of them. class QueryClient: - def __init__(self, args, cpus=None): + def __init__(self, args): self.log = logging.getLogger("QueryClient") self.client = jail.get_process() - if cpus: - self.client.set_cpus(cpus) + set_cpus("client-cpu-ids", self.client, args) def __call__(self, queries, database, num_client_workers): self.log.debug("execute('%s')", str(queries)) @@ -79,11 +76,10 @@ class QueryClient: class LongRunningClient: - def __init__(self, args, cpus=None): + def __init__(self, args): self.log = logging.getLogger("LongRunningClient") self.client = jail.get_process() - if cpus: - self.client.set_cpus(cpus) + set_cpus("client-cpu-ids", self.client, args) # TODO: This is quite similar to __call__ method of QueryClient. Remove # duplication. diff --git a/tests/macro_benchmark/common.py b/tests/macro_benchmark/common.py index b70b43a57..5c7fc27c7 100644 --- a/tests/macro_benchmark/common.py +++ b/tests/macro_benchmark/common.py @@ -1,4 +1,13 @@ import os +from argparse import ArgumentParser + +try: + import jail + APOLLO = True +except: + import jail_faker as jail + APOLLO = False + WALL_TIME = "wall_time" CPU_TIME = "cpu_time" @@ -6,6 +15,7 @@ MAX_MEMORY = "max_memory" DIR_PATH = os.path.dirname(os.path.realpath(__file__)) + def get_absolute_path(path, base=""): if base == "build": extra = "../../build" @@ -18,3 +28,21 @@ def get_absolute_path(path, base=""): else: extra = "" return os.path.normpath(os.path.join(DIR_PATH, extra, path)) + + +# Assign process to cpus passed by flag_name flag from args. If process is +# running on Apollo that flag is obligatory, otherwise is ignored. flag_name +# should not contain leading dashed. +def set_cpus(flag_name, process, args): + argp = ArgumentParser() + # named, optional arguments + argp.add_argument("--" + flag_name, nargs="+", type=int, help="cpus that " + "will be used by process. Obligatory on Apollo, ignored " + "otherwise.") + args, _ = argp.parse_known_args(args) + attr_flag_name = flag_name.replace("-", "_") + cpus = getattr(args, attr_flag_name) + assert not APOLLO or cpus, \ + "flag --{} is obligatory on Apollo".format(flag_name) + if cpus: + process.set_cpus(cpus) diff --git a/tests/macro_benchmark/databases.py b/tests/macro_benchmark/databases.py index 2801f069c..0b7ba523a 100644 --- a/tests/macro_benchmark/databases.py +++ b/tests/macro_benchmark/databases.py @@ -6,14 +6,12 @@ from collections import defaultdict import tempfile import shutil import time -from common import get_absolute_path +from common import get_absolute_path, set_cpus try: import jail - APOLLO = True except: import jail_faker as jail - APOLLO = False def wait_for_server(port, delay=1.0): @@ -27,7 +25,7 @@ class Memgraph: """ Knows how to start and stop memgraph. """ - def __init__(self, args, config, num_workers, cpus=None): + def __init__(self, args, config, num_workers): self.log = logging.getLogger("MemgraphRunner") argp = ArgumentParser("MemgraphArgumentParser") argp.add_argument("--runner-bin", @@ -39,8 +37,7 @@ class Memgraph: self.config = config self.num_workers = num_workers self.database_bin = jail.get_process() - if cpus: - self.database_bin.set_cpus(cpus) + set_cpus("database-cpu-ids", self.database_bin, args) def start(self): self.log.info("start") @@ -70,7 +67,7 @@ class Neo: """ Knows how to start and stop neo4j. """ - def __init__(self, args, config, cpus=None): + def __init__(self, args, config): self.log = logging.getLogger("NeoRunner") argp = ArgumentParser("NeoArgumentParser") argp.add_argument("--runner-bin", default=get_absolute_path( @@ -83,8 +80,7 @@ class Neo: self.args, _ = argp.parse_known_args(args) self.config = config self.database_bin = jail.get_process() - if cpus: - self.database_bin.set_cpus(cpus) + set_cpus("database-cpu-ids", self.database_bin, args) def start(self): self.log.info("start") @@ -141,7 +137,7 @@ class Postgres: self.args, _ = argp.parse_known_args(args) self.username = "macro_benchmark" self.database_bin = jail.get_process() - self.database_bin.set_cpus(cpus) + set_cpus("database-cpu-ids", self.database_bin, args) def start(self): self.log.info("start") diff --git a/tests/macro_benchmark/query_suite.py b/tests/macro_benchmark/query_suite.py index 7c05fc6a3..f39e2c515 100644 --- a/tests/macro_benchmark/query_suite.py +++ b/tests/macro_benchmark/query_suite.py @@ -175,7 +175,7 @@ class _QueryRunner: def __init__(self, args, database): self.log = logging.getLogger("_HarnessClientRunner") self.database = database - self.query_client = QueryClient(args, [2, 3]) + self.query_client = QueryClient(args) def start(self): self.database.start() @@ -200,7 +200,7 @@ class MemgraphRunner(_QueryRunner): argp.add_argument("--num-workers", help="Number of workers") self.args, remaining_args = argp.parse_known_args(args) database = Memgraph(remaining_args, self.args.runner_config, - self.args.num_workers, [1]) + self.args.num_workers) super(MemgraphRunner, self).__init__(remaining_args, database) @@ -214,5 +214,5 @@ class NeoRunner(_QueryRunner): default=get_absolute_path("config/neo4j.conf"), help="Path to neo config file") self.args, remaining_args = argp.parse_known_args(args) - database = Neo(remaining_args, self.args.runner_config, [1]) + database = Neo(remaining_args, self.args.runner_config) super(NeoRunner, self).__init__(remaining_args, database) diff --git a/tools/apollo/generate b/tools/apollo/generate index 58cc9e6d3..c7e47a7d2 100755 --- a/tools/apollo/generate +++ b/tools/apollo/generate @@ -187,7 +187,7 @@ binary_release_link_path = os.path.join(BUILD_RELEASE_DIR, "memgraph") MACRO_BENCHMARK_ARGS = ( "QuerySuite MemgraphRunner " "--groups aggregation 1000_create unwind_create dense_expand match " - "--no-strict") + "--no-strict --database-cpu-ids 1 --client-cpu-ids 2") macro_bench_path = os.path.join(BASE_DIR, "tests", "macro_benchmark") harness_client_binaries = os.path.join(BUILD_RELEASE_DIR, "tests", "macro_benchmark")