atlas-config

#!/usr/bin/env python

import sys, errno, os, platform, time, glob, subprocess, signal, argparse

parser = argparse.ArgumentParser(
    description='(Re-)Build ATLAS (http://math-atlas.sourceforge.net) '
    'according to the SAGE_ATLAS_ARCH environment variable')

parser.add_argument('--unthrottle', nargs=1, type=int,
                    help='switch CPU throttling off until PID finishes',
                    metavar='PID')

parser.add_argument('--archdef', action='store_const', const=True,
                    help='build archdef tarball and save it to the current directory')


def pid_exists(pid):
    """
    Check whether pid exists in the current process table.
    """
    if pid < 0:
        return False
    try:
        os.kill(pid, 0)
    except OSError, e:
        return e.errno == errno.EPERM
    else:
        return True


def unthrottle_posix(pid):
    cpus = list(glob.glob('/sys/devices/system/cpu/cpu*/cpufreq/scaling_governor'))
    scaling_governor = []
    for cpu in cpus:
        with open(cpu, 'r') as f:
            scaling_governor.append(f.readline())
    for cpu in cpus:
        with open(cpu, 'w') as f:
            f.write('performance')

    def signal_handler(signum, frame):
        print 'Signal hanadler called with signal', signum
        for cpu, governor in zip(cpus, scaling_governor):
            with open(cpu, 'w') as f:
                f.write(governor)
        print 'Reverted throttling to the previous behaviour.'
        sys.exit(0)

    signal.signal(signal.SIGPIPE, signal_handler)
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    try:
        while pid_exists(pid):
            time.sleep(1)
    except KeyboardInterrupt:
        pass

    signal_handler(None, None)


def unthrottle(pid):
    if os.name == 'posix':
        unthrottle_posix(pid)
    else:
        print 'I don\'t know how to unthrottle your system ('+platform.system()+')'
        sys.exit(3)



def is_throttled_posix():
    cpus = list(glob.glob('/sys/devices/system/cpu/cpu*/cpufreq/scaling_governor'))
    for cpu in cpus:
        with open(cpu, 'r') as f:
            if f.readline().strip() != 'performance':
                return True
    return False


def is_throttled():
    if os.name == 'posix':
        return is_throttled_posix()
    else:
        print 'I don\'t know how to unthrottle your system ('+platform.system()+')'
        sys.exit(3)


def check_root():
    if os.name == 'posix':
        if os.getuid() == 0:
            return
        print '\nError: You need to be root to (un)throttle the CPU.\n'
        sys.exit(1)
    else:
        print 'I don\'t know how to probe administrator rights on your system ('+platform.system()+')'
        sys.exit(3)


def check_nonroot():
    if os.name == 'posix':
        if os.getuid() > 0:
            return
        print '\nError: You are crazy to run this as root, exiting.\n'
        sys.exit(2)
    else:
        print 'I don\'t know how to probe administrator rights on your system ('+platform.system()+')'
        sys.exit(3)


command = None

def wait_for_command():
    global command
    if command is None:
        return
    else:
        command.send_signal(signal.SIGTERM)

import atexit
atexit.register(wait_for_command)


def unthrottle_self():
    if os.name == 'posix':
        global command
        print 'Running sudo atlas-config --unthrottle to turn CPU throttling off.'
        command = subprocess.Popen(['sudo', os.path.abspath( __file__ ),
                                    '--unthrottle', str(os.getpid())])
    else:
        print 'I don\'t know how to unthrottle your system ('+platform.system()+')'
        sys.exit(2)

    print 'Waiting for CPU throttling to be turned off...'
    while is_throttled():
        time.sleep(1)



if __name__ == '__main__':
    args = parser.parse_args()

    if args.unthrottle:
        check_root()
        unthrottle(args.unthrottle[0])

    check_nonroot()
    unthrottle_self()
    print
    if os.environ.get('SAGE_FAT_BINARY', 'no') == 'yes':
        print 'Building ATLAS with SAGE_FAT_BINARY (generic archdefs).'
    elif os.environ.has_key('SAGE_ATLAS_ARCH'):
        print 'Building ATLAS with SAGE_ATLAS_ARCH =', os.environ['SAGE_ATLAS_ARCH']
    else:
        print 'Building ATLAS without specifying architecture.'
    print 'This may take many hours during which you should leave the computer otherwise'
    print 'idle to obtain accurate timings.'
    if args.archdef:
        os.environ['SAGE_ATLAS_SAVE_ARCHDEF'] = os.getcwd()
        print 'The resulting <archdef>.tar.bz2 will be saved in '+os.getcwd()
    print
    print 'You have 5 seconds to interrupt...'
    time.sleep(1)
    print 'You have 4 seconds to interrupt...'
    time.sleep(1)
    print 'You have 3 seconds to interrupt...'
    time.sleep(1)
    print 'You have 2 seconds to interrupt...'
    time.sleep(1)
    print 'You have 1 second to interrupt...'
    time.sleep(1)

    sys.stdout.flush()
    sys.stderr.flush()
    os.system('sage -f atlas')