Stupid detached head...

1 files changed, 355 insertions, 0 deletions

diff --git a/aprio.py b/aprio.py
new file mode 100755
index 0000000..b0a986c
--- /dev/null
+++ b/aprio.py
@@ -0,0 +1,355 @@
+#!/usr/bin/env python
+#
+# aprio is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# aprio is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with aprio.  If not, see <http://www.gnu.org/licenses/>.
+""" aprio stands for "Automatic Priority"
+"""
+
+import os
+import logging, logging.config
+import math
+import time
+from datetime import timedelta
+
+
+try:
+    import psutil
+except ImportError:
+    print("psutil module not found!")
+    exit(1)
+
+try:
+    import daemon
+except ImportError:
+    print("daemon module not found!")
+    exit(1)
+
+
+try:
+    import argparse
+except ImportError:
+    print("argparse module not found!")
+    exit(1)
+
+
+class Transpire(object):
+    """Supply common time constants
+    """
+    def __init__(self):
+        self.second = timedelta(seconds=1).total_seconds()
+        self.minute = timedelta(minutes=1).total_seconds()
+        self.hour = timedelta(hours=1).total_seconds()
+        self.day = timedelta(days=1).total_seconds()
+        self.week = timedelta(weeks=1).total_seconds()
+        self.month = timedelta(weeks=4).total_seconds()
+
+
+ELAPSED = Transpire()
+CONFIG = {}
+CONFIG['LOAD_THRESHOLD'] = psutil.cpu_count() / 2
+if CONFIG['LOAD_THRESHOLD'] < 1:
+    CONFIG['LOAD_THRESHOLD'] = psutil.cpu_count()
+CONFIG['CPU_THRESHOLD'] = 50.0
+CONFIG['CPUTIME_THRESHOLD'] = ELAPSED.second
+CONFIG['TIME_SCALE'] = 'week'
+CONFIG['TIME_FACTOR'] = 1
+CONFIG['POLL'] = 3
+CONFIG['TEST_MODE'] = False
+CONFIG['VERBOSE'] = False
+CONFIG['QUITE'] = False
+
+
+def renice(proc, nice_value=0):
+    """Change the process priority of a psutil.Process object
+    """
+    logger = logging.getLogger(__name__)
+    nice_current = 255
+    try:
+        pid = proc.pid
+        nice_previous = proc.get_nice()
+        
+        if nice_previous < 0:
+            return
+
+        if nice_value <= nice_previous:
+            return
+
+        if not CONFIG['TEST_MODE']:
+            proc.set_nice(nice_value)
+
+        if not CONFIG['TEST_MODE']:
+            nice_current = proc.get_nice()
+        else:
+            nice_current = nice_value
+
+        logger.info("{0}:Priority modified ({1} -> {2})"
+            .format(pid, nice_previous, nice_current))
+    except psutil.AccessDenied:
+        logger.warning("{0}:{1}:Permission denied setting nice to {2}"
+            .format(pid, proc.username(), nice_value))
+    except psutil.NoSuchProcess:
+        return
+    return nice_current
+
+
+def convert_nice(proc, **kwargs):
+    """Analyzes a process' total kernel time, or the time since the process
+    began.  If the time meets or exceeds a defined threshold, an
+    appropriate nice value will be applied to the process.
+
+    Keyword Arguments:
+    model -- (default 'relative')
+        'kernel' = Total CPU time accumulated
+        'relative' = Total time elapsed since process started
+    nice_min -- (default 0)
+    nice_max -- (default 20)
+    time_scale -- (default 1 month)
+    """
+    logger = logging.getLogger(__name__)
+    model = 'relative'
+    nice_min = 0
+    nice_max = 20
+    time_scale = ELAPSED.month
+    
+    if kwargs.has_key('model'):
+        model = kwargs['model']
+    
+    if kwargs.has_key('nice_min'):
+        nice_min = kwargs['nice_min']
+        
+    if kwargs.has_key('nice_max'):
+        nice_max = kwargs['nice_max']
+        
+    if kwargs.has_key('time_scale'):
+        time_scale = kwargs['time_scale']
+
+    if model == 'kernel':
+        time_user, time_system = proc.cpu_times()  
+        total_time = time_user + time_system
+    elif model == 'relative':
+        total_time = time.time() - proc.create_time()
+    else:
+        raise ValueError('"{0}" is not a valid time model'.format(model))
+        
+    logger.debug('Time model "{0}"'.format(model))
+    nice = 0
+    try:
+        nice = nice_min - total_time * (-nice_max / (time_scale + 1))
+        if nice > nice_max:
+            nice = nice_max
+    except ZeroDivisionError:
+        return int(nice_max)
+    finally:
+        return int(math.floor(nice))
+    
+    return nice 
+
+
+def filter_processes(cpu_threshold=CONFIG['CPU_THRESHOLD'],
+                    cputime_threshold=CONFIG['CPUTIME_THRESHOLD'],
+                    **kwargs):
+    """Yield a filtered process list matching system usage criteria.
+    
+    Keyword ARGUMENTS:
+    cpu_threshold -- must exceed CPU% (default 50.0)
+    cputime_threshold -- must exceed CPU TIME in seconds (default 1.0)
+    user -- yield processes owned by a particular account
+    """
+    logger = logging.getLogger(__name__)
+    user = ""
+    if kwargs.has_key('user'):
+        user = kwargs['user']
+
+    for proc in psutil.process_iter():
+        try:
+            pid = proc.pid
+            username = proc.username()
+            user_time, system_time = proc.cpu_times()
+            cputime_total = user_time + system_time
+            uid, euid, _ = proc.uids()
+
+            if uid == 0 or euid == 0:
+                continue
+
+
+            cpu = proc.get_cpu_percent(interval=0.05)
+            if cpu > cpu_threshold:
+                logger.debug("{0}:cpu_threshold ({1}% > {2}%)"
+                    .format(pid, cpu, cpu_threshold))
+                if cputime_total > cputime_threshold:
+                    logger.debug("{0}:cputime_threshold ({1} > {2})"
+                        .format(pid, cputime_total, cputime_threshold))
+                    if user:
+                        if user != username:
+                            continue
+                    yield proc
+
+        except psutil.NoSuchProcess as ex:
+            logger.debug("{0}:disappeared".format(ex.pid))
+
+
+def main(args):
+    """ Poll system for bad processes
+    """
+    logger = logging.getLogger(__name__)
+    CONFIG['POLL'] = args.poll
+    CONFIG['VERBOSE'] = args.verbose
+    CONFIG['QUIET'] = args.quiet
+    CONFIG['TEST_MODE'] = args.test
+    CONFIG['CPU_THRESHOLD'] = args.cpu_threshold
+    CONFIG['CPUTIME_THRESHOLD'] = args.cputime_threshold
+    CONFIG['LOAD_THRESHOLD'] = args.load_threshold
+    CONFIG['DAEMON'] = args.daemon
+    user = args.user
+    
+    
+    load_sleep = False
+    load_warn = False
+
+    while(True):
+        load = os.getloadavg()
+        load = sum(load) / len(load)
+        if load < CONFIG['LOAD_THRESHOLD']:
+            load_sleep = True
+            load_warn = False
+            if load_sleep:
+                logger.debug("load_threshold nominal ({0} < {1})"
+                    .format(load, CONFIG['LOAD_THRESHOLD']))
+            load_sleep = False
+            time.sleep(CONFIG['POLL'])
+            continue
+        else:
+            load_warn = True
+
+        if load_warn:
+            logger.debug("load_threshold exceeded ({0} > {1})"
+                .format(load, CONFIG['LOAD_THRESHOLD']))
+
+        for bad in filter_processes(CONFIG['CPU_THRESHOLD'],
+                                    CONFIG['CPUTIME_THRESHOLD'],
+                                    user=user):
+            try:
+                nice = convert_nice(bad, model='kernel')
+
+                if not nice:
+                    nice = convert_nice(bad)
+
+                if nice != 0:
+                    renice(bad, nice)
+
+            except psutil.NoSuchProcess:
+                continue
+        time.sleep(CONFIG['POLL'])
+
+
+if __name__ == "__main__":
+
+    PARSER = argparse.ArgumentParser()
+    PARSER.add_argument('--daemon',
+        '-d',
+        action='store_true',
+        help="Fork into background")
+
+    PARSER.add_argument('--logfile',
+        '-L',
+        action='store',
+        default="",
+        type=str,
+        help="Log output to filename")
+
+    PARSER.add_argument('--user',
+        '-u',
+        default="",
+        type=str,
+        help='Limit to specific user')    
+
+    PARSER.add_argument('--cpu-threshold',
+        '-c',
+        default=CONFIG['CPU_THRESHOLD'],
+        type=float,
+        help='Trigger after n%%')
+
+    PARSER.add_argument('--cputime-threshold',
+        '-t',
+        default=CONFIG['CPUTIME_THRESHOLD'],
+        type=float,
+        help='Trigger after n seconds')
+
+    PARSER.add_argument('--load-threshold',
+        '-l',
+        default=CONFIG['LOAD_THRESHOLD'],
+        type=float,
+        help='Trigger after n load average')
+
+    PARSER.add_argument('--time-scale',
+        default=CONFIG['TIME_SCALE'],
+        type=str,
+        help='second, day, week, month, or year')
+    
+    PARSER.add_argument('--time-factor',
+        default=CONFIG['TIME_FACTOR'],
+        type=float,
+        help='NICE_SCALE = (TIME_SCALE * TIME_FACTOR)')    
+    
+    PARSER.add_argument('--poll',
+        '-p',
+        default=CONFIG['POLL'],
+        type=float,
+        help='Wait n seconds between polling processes')
+
+    PARSER.add_argument('--test',
+        '-T',
+        action='store_true',
+        default=False,
+        help='Do not modify processes; report only.')
+
+    PARSER.add_argument('--verbose',
+        '-v',
+        action='store_true',
+        default=False,
+        help='Verbose output')
+
+    PARSER.add_argument('--quiet',
+        '-q',
+        action='store_true',
+        default=False,
+        help='Suppress output')
+    
+    ARGUMENTS = PARSER.parse_args()
+
+    FORMAT = "%(levelname)s:%(asctime)s:%(funcName)s:%(message)s"
+    if ARGUMENTS.logfile:
+        logging.basicConfig(filename=os.path.abspath(ARGUMENTS.logfile),
+            format=FORMAT)
+    else:
+        logging.basicConfig(format=FORMAT) 
+
+    logging.basicConfig(level=logging.INFO)
+    LOGGER = logging.getLogger(__name__)
+    LOGGER.setLevel(logging.INFO)
+
+     
+    if ARGUMENTS.verbose:
+        LOGGER.setLevel(logging.DEBUG)
+
+    if ARGUMENTS.test:
+        LOGGER.debug('Test mode (processes will not be modified)')
+    
+    if ARGUMENTS.daemon:
+        LOGGER.debug('Daemon mode')
+        with daemon.DaemonContext():
+            main(ARGUMENTS)
+    else:
+        LOGGER.debug('Foreground mode')
+        main(ARGUMENTS)
+