Allow setting nice level and proc affinity w/psutil

bin/zk-stats-daemon

@@ -11,6 +11,7 @@ if os.getenv("ZKTRAFFIC_SOURCE") is not None:
   sys.path.insert(0, ".")
 
 from zktraffic.endpoints.stats_server import StatsServer
+from zktraffic.base.process import ProcessOptions
 
 from twitter.common import app, log
 from twitter.common.http import HttpServer
@@ -57,6 +58,12 @@ app.add_option("--niceness",
                type=int,
                default=0,
                help="set the niceness")
+app.add_option("--set-cpu-affinity",
+               dest="cpu_affinity",
+               metavar="CPU#[,CPU#]",
+               type=str,
+               default=None,
+               help="A comma-separated list of CPU cores to pin this process to")
 app.add_option("--max-queued-requests",
                type=int,
                default=400000,
@@ -88,8 +95,13 @@ def main(_, opts):
 
   signal.signal(signal.SIGINT, signal.SIG_DFL)
 
+  process = ProcessOptions()
+
   if opts.niceness >= 0:
-    os.nice(opts.niceness)
+    process.set_niceness(opts.niceness)
+
+  if opts.cpu_affinity:
+    process.set_cpu_affinity(opts.cpu_affinity)
 
   server = Server()
   server.mount_routes(stats)

requirements.txt

@@ -1,4 +1,5 @@
 dpkt-fix
 nose
-twitter.common.log
+psutil>=2.1.0
 scapy==2.2.0-dev
+twitter.common.log

setup.py

@@ -63,18 +63,20 @@ def get_version():
       install_requires=[
           'ansicolors',
           'dpkt-fix',
+          'psutil>=2.1.0',
+          'scapy==2.2.0-dev',
           'twitter.common.app',
           'twitter.common.collections',
           'twitter.common.exceptions',
           'twitter.common.http',
           'twitter.common.log',
-          'scapy==2.2.0-dev',
       ],
       tests_require=[
           'dpkt-fix',
           'nose',
-          'twitter.common.log',
+          'psutil>=2.1.0',
           'scapy==2.2.0-dev',
+          'twitter.common.log',
       ],
       extras_require={
           'test': [

zktraffic/base/process.py

@@ -0,0 +1,78 @@
+# ==================================================================================================
+# Copyright 2014 Twitter, Inc.
+# --------------------------------------------------------------------------------------------------
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this work except in compliance with the License.
+# You may obtain a copy of the License in the LICENSE file, or at:
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==================================================================================================
+
+
+""" Process helpers used to improve zktraffic operation """
+
+from twitter.common import log
+
+import psutil
+from psutil import AccessDenied, NoSuchProcess
+
+
+class ProcessOptions(object):
+
+    def __init__(self):
+        self.process = psutil.Process()
+
+    def set_cpu_affinity(self, cpu_affinity_csv):
+        """
+        Set CPU affinity for this process
+        :param cpu_affinity_csv: A comma-separated string representing CPU cores
+        """
+        try:
+            cpu_list = self.parse_cpu_affinity(cpu_affinity_csv)
+            self.process.cpu_affinity(cpu_list)
+        except (OSError, ValueError) as e:
+            log.warn('unable to set cpu affinity: {}, on process: {}'.format(cpu_affinity_csv, e))
+        except AttributeError:
+            log.warn('cpu affinity is not available on your platform')
+
+    def get_cpu_affinity(self):
+        """
+        Get CPU affinity of this process
+        :return: a list() of CPU cores this processes is pinned to
+        """
+        return self.process.cpu_affinity()
+
+    def set_niceness(self, nice_level):
+        """
+        Set the nice level of this process
+        :param nice_level: the nice level to set
+        """
+        try:
+            # TODO (phobos182): double check that psutil does not allow negative nice values
+            if not 0 <= nice_level <= 20:
+                raise ValueError('nice level must be between 0 and 20')
+            self.process.nice(nice_level)
+        except(EnvironmentError, ValueError, AccessDenied, NoSuchProcess) as e:
+            log.warn('unable to set nice level on process: {}'.format(e))
+
+    def get_niceness(self):
+        """
+        Get nice level of this process
+        :return: an int() representing the nice level of this process
+        """
+        return self.process.nice()
+
+    @staticmethod
+    def parse_cpu_affinity(cpu_affinity_csv):
+        """
+        Static method to parse a csv string string into a list of integers
+        :param cpu_affinity_csv: a CSV of cpu cores to parse
+        :return: a list() of integers representing CPU cores
+        """
+        return [int(_) for _ in cpu_affinity_csv.split(',')]

zktraffic/tests/test_process_options.py

@@ -0,0 +1,61 @@
+# ==================================================================================================
+# Copyright 2014 Twitter, Inc.
+# --------------------------------------------------------------------------------------------------
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this work except in compliance with the License.
+# You may obtain a copy of the License in the LICENSE file, or at:
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==================================================================================================
+
+
+import mock
+import os
+
+from zktraffic.base.process import ProcessOptions
+
+import psutil
+
+
+proc = ProcessOptions()
+
+
+def test_niceness():
+  """
+  Test CPU niceness calls
+  """
+  proc.set_niceness(15)
+  assert proc.get_niceness() == 15
+
+
+def test_cpu_affinity_parsing():
+    """
+    Test CPU affinity list parsing
+    """
+    assert ProcessOptions.parse_cpu_affinity('0,1,2,3,4,5') == [0, 1, 2, 3, 4, 5]
+    assert ProcessOptions.parse_cpu_affinity('2') == [2]
+
+
+def test_cpu_affinity():
+  """
+  Test CPU affinity setting
+  """
+  def mock_cpu_affinity_handler(self, *args, **kwargs):
+      return [0, 1]
+
+  # if running in TravisCI, mock the cpu_affinity() method call
+  if os.environ.get('TRAVIS'):
+      with mock.patch.object(psutil.Process, 'cpu_affinity', create=True, new=mock_cpu_affinity_handler):
+          proc.set_cpu_affinity('0,1')
+          assert proc.get_cpu_affinity() == [0, 1]
+  else:
+    proc.set_cpu_affinity('0,1')
+    assert proc.get_cpu_affinity() == [0, 1]
+    proc.set_cpu_affinity('1')
+    assert proc.get_cpu_affinity() == [1]