Pass assigned cpus to harness by flags

Summary:
Since we have different kind of workers in Apollo we should pass
assigned cpus to harness from apollo generate script and not define them
in harness or in benchmarks.

Reviewers: mferencevic

Reviewed By: mferencevic

Subscribers: pullbot

Differential Revision: https://phabricator.memgraph.io/D916

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.

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)

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

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)

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