Merge pull request #175 from epikhinm/master

Added taskset affinity for phantom and tankcore

docs/configuration.rst

@@ -219,6 +219,7 @@ Advanced options:
 * **method_prefix** - object's type, that has a functionality to create test requests. Default: method_stream
 * **source_log_prefix** - prefix, added to class name that reads source data. Default: empty
 * **method_options** - additional options for method objects. It is used for Elliptics etc. Default: empty
+* **affinity** - set a phantom's CPU affinity. Example: '0-3' enabling first 4 cores, '0,1,2,16,17,18' enabling 6 cores. Default: empty
 
 TLS/SSL additional options (ssl=1 is required):
 
@@ -681,6 +682,8 @@ Options
 * **artifacts_base_dir** - base directory for artifacts storing. Temporary artifacts files are stored here. Default: current directory
 * **artifacts_dir** - directory where to keep artifacts after test. Default: directory in ``artifacts_base_dir`` named in  Date/Time format.
 * **flush_config_to** - dump configuration options after each tank step (`yandex.tank steps. sorry, russian only <http://clubs.ya.ru/yandex-tank/replies.xml?item_no=6>`_) to that file 
+* **taskset_path** - path to taskset command. Default: taskset
+* **affinity** - set a yandex-tank's (python process and load generator process) CPU affinity. Example: '0-3' enabling first 4 cores, '0,1,2,16,17,18' enabling 6 cores. Default: empty
 
 Tips&Tricks
 ^^^^^^^^^^^

yandextank/core/tankcore.py

@@ -207,10 +207,13 @@ def __init__(self):
         self.lock_file = None
         self.flush_config_to = None
         self.lock_dir = None
+        self.taskset_path = None
+        self.taskset_affinity = None
         self.set_option(self.SECTION, self.PID_OPTION, str(os.getpid()))
 
     def get_available_options(self):
-        return ["artifacts_base_dir", "artifacts_dir", "flush_config_to"]
+        return ["artifacts_base_dir", "artifacts_dir", "flush_config_to",
+                "taskset_path", "affinity"]
 
     def load_configs(self, configs):
         """ Tells core to load configs set into options storage """
@@ -242,6 +245,8 @@ def load_plugins(self):
             self.SECTION, "artifacts_base_dir", self.artifacts_base_dir)
         self.artifacts_base_dir = os.path.expanduser(base_dir)
         self.artifacts_dir = self.get_option(self.SECTION, "artifacts_dir", "")
+        self.taskset_path = self.get_option(self.SECTION, 'taskset_path', 'taskset')
+        self.taskset_affinity = self.get_option(self.SECTION, 'affinity', '')
 
         options = self.config.get_options(self.SECTION, self.PLUGIN_PREFIX)
         for (plugin_name, plugin_path) in options:
@@ -263,6 +268,8 @@ def plugins_configure(self):
             os.chmod(self.artifacts_base_dir, 0755)
 
         self.log.info("Configuring plugins...")
+        if self.taskset_affinity != '':
+            self.taskset(os.getpid(), self.taskset_path, self.taskset_affinity)
         for plugin in self.plugins:
             self.log.debug("Configuring %s", plugin)
             plugin.configure()
@@ -359,6 +366,17 @@ def plugins_post_process(self, retcode):
 
         return retcode
 
+    def taskset(self, pid, path, affinity):
+        if affinity != '':
+            args = "%s -pc %s %s" % (path, affinity, pid)
+            retcode, stdout, stderr = execute(
+                args, shell=True, poll_period=0.1, catch_out=True)
+            self.log.debug('taskset stdout: %s', stdout)
+            if retcode != 0:
+                raise KeyError(stderr)
+            else:
+                self.log.info("Enabled taskset for pid %s with affinity %s", str(pid), affinity)
+
     def __collect_artifacts(self):
         self.log.debug("Collecting artifacts")
         if not self.artifacts_dir:

yandextank/plugins/Phantom/plugin.py

@@ -8,6 +8,7 @@
 import time
 import datetime
 import string
+import multiprocessing as mp
 
 from yandextank.plugins import ConsoleScreen
 from yandextank.plugins.Aggregator import \
@@ -41,6 +42,9 @@ def __init__(self, core):
         self.phantom_start_time = time.time()
         self.buffered_seconds = "2"
 
+        self.taskset_affinity = None
+        self.cpu_count = mp.cpu_count()
+
         self.phantom = None
         self.cached_info = None
         self.phantom_stderr = None
@@ -52,7 +56,8 @@ def get_key():
         return __file__
 
     def get_available_options(self):
-        opts = ["eta_file", "phantom_path", "buffered_seconds", "exclude_markers", ]
+        opts = ["eta_file", "phantom_path", "buffered_seconds", "exclude_markers"]
+        opts += 'affinity'
         opts += [PhantomConfig.OPTION_PHOUT, self.OPTION_CONFIG]
         opts += PhantomConfig.get_available_options()
         return opts
@@ -68,6 +73,7 @@ def configure(self):
             self.get_option("buffered_seconds", self.buffered_seconds))
         self.exclude_markers = set(filter((lambda marker: marker != ''),
                                         self.get_option('exclude_markers', []).split(' ')))
+        self.taskset_affinity = self.get_option('affinity', '')
 
         try:
             autostop = self.core.get_plugin_of_type(AutostopPlugin)
@@ -148,6 +154,10 @@ def start_test(self):
             args = [self.phantom_path, 'run', self.config]
             self.log.debug(
                 "Starting %s with arguments: %s", self.phantom_path, args)
+            if self.taskset_affinity != '':
+                args = [self.core.taskset_path, '-c', self.taskset_affinity] + args
+                self.log.debug(
+                    'Enabling taskset for phantom with affinity: %s, cores count: %d', self.taskset_affinity, self.cpu_count)
             self.phantom_start_time = time.time()
             phantom_stderr_file = self.core.mkstemp(
                 ".log", "phantom_stdout_stderr_")