Skip to content

Commit

Permalink
Merge pull request #3 from phobos182/master
Browse files Browse the repository at this point in the history
Allow setting nice level and proc affinity w/ psutil
  • Loading branch information
rgs1 committed Oct 7, 2014
2 parents 1812920 + 6c56362 commit 93f57d3
Show file tree
Hide file tree
Showing 5 changed files with 158 additions and 4 deletions.
14 changes: 13 additions & 1 deletion bin/zk-stats-daemon
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand Down
3 changes: 2 additions & 1 deletion requirements.txt
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
dpkt-fix
nose
twitter.common.log
psutil>=2.1.0
scapy==2.2.0-dev
twitter.common.log
6 changes: 4 additions & 2 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -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': [
Expand Down
78 changes: 78 additions & 0 deletions zktraffic/base/process.py
Original file line number Diff line number Diff line change
@@ -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(',')]
61 changes: 61 additions & 0 deletions zktraffic/tests/test_process_options.py
Original file line number Diff line number Diff line change
@@ -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]

0 comments on commit 93f57d3

Please sign in to comment.