<commit_before>"""Pilot functionality"""

from __future__ import absolute_import, division, print_function

import os
import sys
import time
import logging
import tempfile
import shutil
from functools import partial
from multiprocessing import Process
from collections import namedtuple
from datetime import timedelta
from glob import glob
import signal

from iceprod.core import to_file, constants
from iceprod.core import exe_json
from iceprod.core.util import (Task_Resources, Task_Resource_Overusage,
                               get_task_resources)
from iceprod.core.dataclasses import Number, String
import iceprod.core.logger

logger = logging.getLogger('pilot')
logger.setLevel(logging.DEBUG)

from tornado import gen
from tornado.ioloop import IOLoop
from tornado.locks import Condition

try:
    import psutil
except ImportError:
    psutil = None

try:
    from setproctitle import setproctitle
except ImportError:
    def setproctitle(name):
        pass

Task = namedtuple('Task', ['p','process','resources','tmpdir'])

def process_wrapper(func, title):
    """
    Set process title. Set log file, stdout, stderr. Then go on to call func.

    Args:
        func (callable): The function that we really want to run
        title (str): The new process title
    """
    try:
        setproctitle(title)
    except:
        pass

    iceprod.core.logger.new_file(constants['stdlog'])
    stdout = partial(to_file,sys.stdout,constants['stdout'])
    stderr = partial(to_file,sys.stderr,constants['stderr'])
    with stdout(), stderr():
        func()

class Pilot(object):
    """
    A pilot task runner.

    The pilot allows multiple tasks to run in sequence or parallel.
    It keeps track of resource usage, killing anything that goes over
    requested amounts.

    Args:
        config (dict): the configuration dictionary
        runner (callable): the task/config runner
    """
    def __init__(self, config, runner, pilot_id, debug=False, run_timeout=60):
        self.config = config
        self.runner = runner
        self.pilot_id = pilot_id
        self.debug = debug
        self.run_timeout = timedelta(seconds=run_timeout)

        try:
            setproctitle('iceprod2_pilot({})'.format(pilot_id))
        except:
            pass

        # set up resources for pilot
        self.lock = Condition()
        self.resources = get_task_resources()
        if 'resources' in config['options']:
            self.resources.update(config['options']['resources'])

        self.start_time = time.time()

        # set up monitor
        self.ioloop = IOLoop.current()
        self.tasks = {}
        if psutil:
            self.ioloop.add_callback(self.monitor)
        else:
            logger.warn('no psutil. not checking resource usage')

        # set up signal handler
        def handler(signum, frame):
            logger.critical('termination signal received')
            self.ioloop.add_callback_from_signal(self.term_handler)
        signal.signal(signal.SIGTERM, handler)

        self.ioloop.run_sync(self.run)

        if self.debug:
            # append out, err, log
            for dirs in glob('tmp*'):
                for filename in (constants['stdout'], constants['stderr'],
                                 constants['stdlog']):
                    if os.path.exists(os.path.join(dirs,filename)):
                        with open(filename,'a') as f:
                            print('', file=f)
                            print('----',dirs,'----', file=f)
                            with open(os.path.join(dirs,filename)) as f2:
                                print(f2.read(), file=f)

    def term_handler(self):
        """Handle a SIGTERM gracefully"""
        logger.info('checking resources after SIGTERM')
        start_time = time.time()
        for task_id in list(self.tasks):
            task = self.tasks[task_id]
            reason = 'pilot SIGTERM'
            if psutil:
                used_resources = {r:0 for r in task.resources}
                if not task.p.is_alive():
                    logger.info('%r not alive', task_id)
                else:
                    try:
                        processes = [task.process]+task.process.children()
                        for p in processes:
                            used_resources['cpu'] += p.cpu_percent()/100.0
                            used_resources['memory'] += p.memory_info().rss/1000000000.0
                        used_resources['time'] = (start_time - task.process.create_time())/3600.0
                        used_resources['disk'] = du(task.tmpdir)/1000000000.0
                        logger.debug('task %r used %r',task_id,used_resources)
                    except:
                        logger.warn('error getting resources', exc_info=True)

                    for r in used_resources:
                        if used_resources[r]-task.resources[r] > 0:
                            reason = 'Resource overuse for {}: {}'.format(r,used_resources[r])
                            break
                    logger.warn('kill %r for %s', task_id, reason)
                    processes.reverse() # kill children first
                    for p in processes:
                        try:
                            p.terminate()
                        except:
                            logger.warn('error terminating process for %r',
                                        task_id, exc_info=True)

                    def on_terminate(proc):
                        logger.info("process %r terminated with exit code %r",
                                    proc, proc.returncode)
                    try:
                        gone, alive = psutil.wait_procs(processes, timeout=0.1,
                                                        callback=on_terminate)
                        for p in alive:
                            p.kill()
                    except:
                        logger.warn('failed to kill processes for %r', task_id,
                                    exc_info=True)

            # clean up task
            self.clean_task(task)
            del self.tasks[task_id]
            exe_json.task_kill(task_id, resources=used_resources,
                               reason=reason)

        # stop the pilot
        exe_json.update_pilot(self.pilot_id, tasks='')
        self.ioloop.stop()

    @gen.coroutine
    def monitor(self):
        """Monitor the tasks, killing any that go over resource limits"""
        try:
            sleep_time = 0.1 # check every X seconds
            disk_sleep_time = 180
            disk_start_time = time.time()
            while True:
                logger.info('pilot monitor - checking resource usage')
                start_time = time.time()
                for task_id in list(self.tasks):
                    task = self.tasks[task_id]
                    if not task.p.is_alive():
                        logger.info('not alive, so notify')
                        self.lock.notify()
                        continue

                    used_resources = {r:0 for r in task.resources}
                    try:
                        processes = [task.process]+task.process.children()
                        for p in processes:
                            with p.oneshot():
                                used_resources['cpu'] += p.cpu_percent()/100.0
                                used_resources['memory'] += p.memory_info().rss/1000000000.0
                                if not used_resources['time']:
                                    used_resources['time'] = (start_time - task.process.create_time())/3600.0
                        if start_time - disk_start_time > disk_sleep_time:
                            disk_start_time = start_time
                            used_resources['disk'] = du(task.tmpdir)/1000000000.0
                        logger.debug('task %r used %r',task_id,used_resources)
                    except:
                        logger.warn('error getting resources', exc_info=True)
                        continue

                    kill = False
                    reason = ''
                    for r in used_resources:
                        overusage = used_resources[r]-task.resources[r]
                        if overusage > 0:
                            overusage_percent = used_resources[r]*1.0/task.resources[r]
                            if (r in Task_Resource_Overusage and
                                (overusage_percent < Task_Resource_Overusage[r]['ignore'] or
                                 (overusage < self.resources[r] and
                                  overusage_percent < Task_Resource_Overusage[r]['allowed']
                                ))):
                                logger.info('managable overusage of %s for %r',
                                            r, task_id)
                                continue
                            kill = True
                            reason = 'Resource overuse for {}: {}'.format(r,used_resources[r])
                            break
                    if kill:
                        logger.warn('kill %r for going over resources %r',
                                    task_id, used_resources)
                        processes.reverse() # kill children first
                        for p in processes:
                            try:
                                p.terminate()
                            except:
                                logger.warn('error terminating process for %r',
                                            task_id, exc_info=True)

                        def on_terminate(proc):
                            logger.info("process %r terminated with exit code %r",
                                        proc, proc.returncode)
                        try:
                            gone, alive = psutil.wait_procs(processes, timeout=0.1,
                                                            callback=on_terminate)
                            for p in alive:
                                p.kill()
                        except:
                            logger.warn('failed to kill processes for %r', task_id,
                                        exc_info=True)
                        self.clean_task(task)
                        del self.tasks[task_id]
                        exe_json.task_kill(task_id, resources=used_resources,
                                           reason=reason)
                        # try to queue another task
                        logger.info('killed, so notify')
                        self.lock.notify()

                duration = time.time()-start_time
                logger.debug('sleep_time %d, duration %d',sleep_time,duration)
                if duration < sleep_time:
                    yield self.lock.wait(timeout=timedelta(seconds=sleep_time-duration))
        except Exception:
            logger.error('pilot monitor died', exc_info=True)
            raise

    @gen.coroutine
    def run(self):
        """Run the pilot"""
        errors = 0
        tasks_running = 0
        running = True
        while running:
            while running:
                try:
                    task_config = exe_json.downloadtask(self.config['options']['gridspec'],
                                                        resources=self.resources)
                except Exception:
                    errors += 1
                    if errors > 5:
                        running = False
                    logger.error('cannot download task. current error count is %d',
                                 errors, exc_info=True)
                    continue
                logger.info('task config: %r', task_config)

                if task_config is None:
                    logger.info('no task available')
                    if not self.tasks:
                        running = False
                    break
                else:
                    try:
                        task_id = task_config['options']['task_id']
                    except Exception:
                        errors += 1
                        if errors > 5:
                            running = False
                        logger.warn('error getting task_id from config')
                        continue
                    try:
                        self.create_task(task_config)
                    except Exception:
                        errors += 1
                        if errors > 5:
                            running = False
                        logger.warn('error creating task %s', task_id,
                                    exc_info=True)
                        exe_json.task_kill(task_id, reason='failed to create task')
                    else:
                        tasks_running += 1
                        for r in self.resources:
                            if r in ('cpu','memory','disk', 'gpu'):
                                self.resources[r] -= task_config['options']['resources'][r]
                            elif isinstance(self.resources[r], String):
                                self.resources[r] = self.resources[r].split(',')
                        exe_json.update_pilot(self.pilot_id, tasks=','.join(self.tasks))

                if not any(self.resources.values()):
                    logger.info('no resources left, so wait for tasks to finish')
                    break

            # wait until we can queue more tasks
            while running or self.tasks:
                logger.info('wait while tasks are running. timeout=%r',self.run_timeout)
                ret = yield self.lock.wait(timeout=self.run_timeout)
                logger.info('yield returned %r',ret)
                # check if any processes have died
                for task_id in list(self.tasks):
                    if not self.tasks[task_id].p.is_alive():
                        self.clean_task(self.tasks[task_id])
                        del self.tasks[task_id]
                if len(self.tasks) < tasks_running:
                    logger.info('%d tasks removed', tasks_running-len(self.tasks))
                    tasks_running = len(self.tasks)
                    exe_json.update_pilot(self.pilot_id, tasks=','.join(self.tasks))
                    if running:
                        break

        if errors >= 5:
            logger.critical('too many errors when running tasks')
        else:
            logger.warn('cleanly stopping pilot')

    def create_task(self, config):
        """
        Create a new Task and start running it

        Args:
            config (dict): The task config
        """
        task_id = config['options']['task_id']

        # add grid-specific config
        for k in self.config['options']:
            if k == 'resources':
                pass
            elif k not in config['options']:
                config['options'][k] = self.config['options'][k]

        # add resources
        if 'resources' not in config['options']:
            config['options']['resources'] = {}
        for r in self.resources:
            if r in config['options']['resources']:
                if config['options']['resources'][r] > self.resources[r]:
                    raise Exception('requested {} of {}, but pilot has {}'.format(
                        config['options']['resources'][r], r, self.resources[r]))
            else:
                config['options']['resources'][r] = self.resources[r]

        # run task in tmp dir
        main_dir = os.getcwd()
        try:
            tmpdir = tempfile.mkdtemp(dir=main_dir)

            # symlink important files
            if 'ssl' in config['options']:
                for f in config['options']['ssl']:
                    os.symlink(os.path.join(main_dir,config['options']['ssl'][f]),
                               os.path.join(tmpdir,config['options']['ssl'][f]))

            # start the task
            os.chdir(tmpdir)
            p = Process(target=partial(process_wrapper, partial(self.runner, config),
                                       'iceprod_task_{}'.format(task_id)))
            p.start()
            if psutil:
                ps = psutil.Process(p.pid)
                ps.nice(psutil.Process().nice()+1)
            else:
                ps = None
            r = config['options']['resources']
            self.tasks[task_id] = Task(p=p, process=ps, resources=r,
                                       tmpdir=tmpdir)
        except:
            logger.error('error creating task', exc_info=True)
        finally:
            os.chdir(main_dir)

    def clean_task(self, task):
        """
        Clean up a Task

        Args:
            task (Task): the Task
        """
        logger.info('cleaning task at %s', task.tmpdir)
        try:
            # return resources to pilot
            for r in task.resources:
                if r in ('cpu','memory','disk', 'gpu'):
                    self.resources[r] += task.resources[r]
            if not self.debug:
                # clean tmpdir
                shutil.rmtree(task.tmpdir)
        except Exception:
            logger.warn('error cleaning up task', exc_info=True)

def du(path):
    """
    Perform a "du" on a path, getting the disk usage.

    Args:
        path (str): The path to analyze

    Returns:
        int: bytes used
    """
    total = 0
    for root,dirs,files in os.walk(path):
        for d in list(dirs):
            p = os.path.join(root,d)
            if os.path.islink(p):
                dirs.remove(d)
        for f in files:
            p = os.path.join(root,f)
            if not os.path.islink(p):
                total += os.path.getsize(p)
    return total
<commit_msg>more logging. run monitor every 0.2 seconds<commit_after>"""Pilot functionality"""

from __future__ import absolute_import, division, print_function

import os
import sys
import time
import logging
import tempfile
import shutil
from functools import partial
from multiprocessing import Process
from collections import namedtuple
from datetime import timedelta
from glob import glob
import signal

from iceprod.core import to_file, constants
from iceprod.core import exe_json
from iceprod.core.util import (Task_Resources, Task_Resource_Overusage,
                               get_task_resources)
from iceprod.core.dataclasses import Number, String
import iceprod.core.logger

logger = logging.getLogger('pilot')
logger.setLevel(logging.DEBUG)

from tornado import gen
from tornado.ioloop import IOLoop
from tornado.locks import Condition

try:
    import psutil
except ImportError:
    psutil = None

try:
    from setproctitle import setproctitle
except ImportError:
    def setproctitle(name):
        pass

Task = namedtuple('Task', ['p','process','resources','tmpdir'])

def process_wrapper(func, title):
    """
    Set process title. Set log file, stdout, stderr. Then go on to call func.

    Args:
        func (callable): The function that we really want to run
        title (str): The new process title
    """
    try:
        setproctitle(title)
    except:
        pass

    iceprod.core.logger.new_file(constants['stdlog'])
    stdout = partial(to_file,sys.stdout,constants['stdout'])
    stderr = partial(to_file,sys.stderr,constants['stderr'])
    with stdout(), stderr():
        func()

class Pilot(object):
    """
    A pilot task runner.

    The pilot allows multiple tasks to run in sequence or parallel.
    It keeps track of resource usage, killing anything that goes over
    requested amounts.

    Args:
        config (dict): the configuration dictionary
        runner (callable): the task/config runner
    """
    def __init__(self, config, runner, pilot_id, debug=False, run_timeout=60):
        self.config = config
        self.runner = runner
        self.pilot_id = pilot_id
        self.debug = debug
        self.run_timeout = timedelta(seconds=run_timeout)

        try:
            setproctitle('iceprod2_pilot({})'.format(pilot_id))
        except:
            pass

        # set up resources for pilot
        self.lock = Condition()
        self.resources = get_task_resources()
        if 'resources' in config['options']:
            self.resources.update(config['options']['resources'])

        self.start_time = time.time()

        # set up monitor
        self.ioloop = IOLoop.current()
        self.tasks = {}
        if psutil:
            self.ioloop.add_callback(self.monitor)
        else:
            logger.warn('no psutil. not checking resource usage')

        # set up signal handler
        def handler(signum, frame):
            logger.critical('termination signal received')
            self.ioloop.add_callback_from_signal(self.term_handler)
        signal.signal(signal.SIGTERM, handler)

        self.ioloop.run_sync(self.run)

        if self.debug:
            # append out, err, log
            for dirs in glob('tmp*'):
                for filename in (constants['stdout'], constants['stderr'],
                                 constants['stdlog']):
                    if os.path.exists(os.path.join(dirs,filename)):
                        with open(filename,'a') as f:
                            print('', file=f)
                            print('----',dirs,'----', file=f)
                            with open(os.path.join(dirs,filename)) as f2:
                                print(f2.read(), file=f)

    def term_handler(self):
        """Handle a SIGTERM gracefully"""
        logger.info('checking resources after SIGTERM')
        start_time = time.time()
        for task_id in list(self.tasks):
            task = self.tasks[task_id]
            reason = 'pilot SIGTERM'
            if psutil:
                used_resources = {r:0 for r in task.resources}
                if not task.p.is_alive():
                    logger.info('%r not alive', task_id)
                else:
                    try:
                        processes = [task.process]+task.process.children()
                        for p in processes:
                            with p.oneshot():
                                used_resources['cpu'] += p.cpu_percent()/100.0
                                used_resources['memory'] += p.memory_info().rss/1000000000.0
                                if not used_resources['time']:
                                    used_resources['time'] = (start_time - task.process.create_time())/3600.0
                        used_resources['disk'] = du(task.tmpdir)/1000000000.0
                        logger.debug('task %r used %r',task_id,used_resources)
                    except:
                        logger.warn('error getting resources', exc_info=True)

                    for r in used_resources:
                        if used_resources[r]-task.resources[r] > 0:
                            reason = 'Resource overuse for {}: {}'.format(r,used_resources[r])
                            break
                    logger.warn('kill %r for %s', task_id, reason)
                    processes.reverse() # kill children first
                    for p in processes:
                        try:
                            p.terminate()
                        except:
                            logger.warn('error terminating process for %r',
                                        task_id, exc_info=True)

                    def on_terminate(proc):
                        logger.info("process %r terminated with exit code %r",
                                    proc, proc.returncode)
                    try:
                        gone, alive = psutil.wait_procs(processes, timeout=0.1,
                                                        callback=on_terminate)
                        for p in alive:
                            p.kill()
                    except:
                        logger.warn('failed to kill processes for %r', task_id,
                                    exc_info=True)

            # clean up task
            self.clean_task(task)
            del self.tasks[task_id]
            exe_json.task_kill(task_id, resources=used_resources,
                               reason=reason)

        # stop the pilot
        exe_json.update_pilot(self.pilot_id, tasks='')
        self.ioloop.stop()

    @gen.coroutine
    def monitor(self):
        """Monitor the tasks, killing any that go over resource limits"""
        try:
            sleep_time = 0.2 # check every X seconds
            disk_sleep_time = 180
            disk_start_time = time.time()
            while True:
                logger.info('pilot monitor - checking resource usage')
                start_time = time.time()
                for task_id in list(self.tasks):
                    task = self.tasks[task_id]
                    if not task.p.is_alive():
                        logger.info('not alive, so notify')
                        self.lock.notify()
                        continue

                    used_resources = {r:0 for r in task.resources}
                    try:
                        processes = [task.process]+task.process.children()
                        for p in processes:
                            with p.oneshot():
                                used_resources['cpu'] += p.cpu_percent()/100.0
                                used_resources['memory'] += p.memory_info().rss/1000000000.0
                                if not used_resources['time']:
                                    used_resources['time'] = (start_time - task.process.create_time())/3600.0
                        if start_time - disk_start_time > disk_sleep_time:
                            disk_start_time = start_time
                            used_resources['disk'] = du(task.tmpdir)/1000000000.0
                        logger.debug('task %r used %r',task_id,used_resources)
                    except:
                        logger.warn('error getting resources', exc_info=True)
                        continue

                    kill = False
                    reason = ''
                    for r in used_resources:
                        overusage = used_resources[r]-task.resources[r]
                        if overusage > 0:
                            overusage_percent = used_resources[r]*1.0/task.resources[r]
                            if (r in Task_Resource_Overusage and
                                (overusage_percent < Task_Resource_Overusage[r]['ignore'] or
                                 (overusage < self.resources[r] and
                                  overusage_percent < Task_Resource_Overusage[r]['allowed']
                                ))):
                                logger.info('managable overusage of %s for %r',
                                            r, task_id)
                                continue
                            kill = True
                            reason = 'Resource overuse for {}: {}'.format(r,used_resources[r])
                            break
                    if kill:
                        logger.warn('kill %r for going over resources %r',
                                    task_id, used_resources)
                        processes.reverse() # kill children first
                        for p in processes:
                            try:
                                p.terminate()
                            except:
                                logger.warn('error terminating process for %r',
                                            task_id, exc_info=True)

                        def on_terminate(proc):
                            logger.info("process %r terminated with exit code %r",
                                        proc, proc.returncode)
                        try:
                            gone, alive = psutil.wait_procs(processes, timeout=0.1,
                                                            callback=on_terminate)
                            for p in alive:
                                p.kill()
                        except:
                            logger.warn('failed to kill processes for %r', task_id,
                                        exc_info=True)
                        self.clean_task(task)
                        del self.tasks[task_id]
                        exe_json.task_kill(task_id, resources=used_resources,
                                           reason=reason)
                        # try to queue another task
                        logger.info('killed, so notify')
                        self.lock.notify()

                duration = time.time()-start_time
                logger.debug('sleep_time %d, duration %d',sleep_time,duration)
                if duration < sleep_time:
                    yield self.lock.wait(timeout=timedelta(seconds=sleep_time-duration))
        except Exception:
            logger.error('pilot monitor died', exc_info=True)
            raise
        logger.info('pilot monitor exiting')

    @gen.coroutine
    def run(self):
        """Run the pilot"""
        errors = 0
        tasks_running = 0
        running = True
        while running:
            while running:
                try:
                    task_config = exe_json.downloadtask(self.config['options']['gridspec'],
                                                        resources=self.resources)
                except Exception:
                    errors += 1
                    if errors > 5:
                        running = False
                    logger.error('cannot download task. current error count is %d',
                                 errors, exc_info=True)
                    continue
                logger.info('task config: %r', task_config)

                if task_config is None:
                    logger.info('no task available')
                    if not self.tasks:
                        running = False
                    break
                else:
                    try:
                        task_id = task_config['options']['task_id']
                    except Exception:
                        errors += 1
                        if errors > 5:
                            running = False
                        logger.warn('error getting task_id from config')
                        continue
                    try:
                        self.create_task(task_config)
                    except Exception:
                        errors += 1
                        if errors > 5:
                            running = False
                        logger.warn('error creating task %s', task_id,
                                    exc_info=True)
                        exe_json.task_kill(task_id, reason='failed to create task')
                    else:
                        tasks_running += 1
                        for r in self.resources:
                            if r in ('cpu','memory','disk', 'gpu'):
                                self.resources[r] -= task_config['options']['resources'][r]
                            elif isinstance(self.resources[r], String):
                                self.resources[r] = self.resources[r].split(',')
                        exe_json.update_pilot(self.pilot_id, tasks=','.join(self.tasks))

                if not any(self.resources.values()):
                    logger.info('no resources left, so wait for tasks to finish')
                    break

            # wait until we can queue more tasks
            while running or self.tasks:
                logger.info('wait while tasks are running. timeout=%r',self.run_timeout)
                ret = yield self.lock.wait(timeout=self.run_timeout)
                logger.debug('yield returned %r',ret)
                # check if any processes have died
                for task_id in list(self.tasks):
                    if not self.tasks[task_id].p.is_alive():
                        self.clean_task(self.tasks[task_id])
                        del self.tasks[task_id]
                if len(self.tasks) < tasks_running:
                    logger.info('%d tasks removed', tasks_running-len(self.tasks))
                    tasks_running = len(self.tasks)
                    exe_json.update_pilot(self.pilot_id, tasks=','.join(self.tasks))
                    if running:
                        break

        if errors >= 5:
            logger.critical('too many errors when running tasks')
        else:
            logger.warn('cleanly stopping pilot')

    def create_task(self, config):
        """
        Create a new Task and start running it

        Args:
            config (dict): The task config
        """
        task_id = config['options']['task_id']

        # add grid-specific config
        for k in self.config['options']:
            if k == 'resources':
                pass
            elif k not in config['options']:
                config['options'][k] = self.config['options'][k]

        # add resources
        if 'resources' not in config['options']:
            config['options']['resources'] = {}
        for r in self.resources:
            if r in config['options']['resources']:
                if config['options']['resources'][r] > self.resources[r]:
                    raise Exception('requested {} of {}, but pilot has {}'.format(
                        config['options']['resources'][r], r, self.resources[r]))
            else:
                config['options']['resources'][r] = self.resources[r]

        # run task in tmp dir
        main_dir = os.getcwd()
        try:
            tmpdir = tempfile.mkdtemp(dir=main_dir)

            # symlink important files
            if 'ssl' in config['options']:
                for f in config['options']['ssl']:
                    os.symlink(os.path.join(main_dir,config['options']['ssl'][f]),
                               os.path.join(tmpdir,config['options']['ssl'][f]))

            # start the task
            os.chdir(tmpdir)
            p = Process(target=partial(process_wrapper, partial(self.runner, config),
                                       'iceprod_task_{}'.format(task_id)))
            p.start()
            if psutil:
                ps = psutil.Process(p.pid)
                ps.nice(psutil.Process().nice()+1)
            else:
                ps = None
            r = config['options']['resources']
            self.tasks[task_id] = Task(p=p, process=ps, resources=r,
                                       tmpdir=tmpdir)
        except:
            logger.error('error creating task', exc_info=True)
        finally:
            os.chdir(main_dir)

    def clean_task(self, task):
        """
        Clean up a Task

        Args:
            task (Task): the Task
        """
        logger.info('cleaning task at %s', task.tmpdir)
        try:
            # return resources to pilot
            for r in task.resources:
                if r in ('cpu','memory','disk', 'gpu'):
                    self.resources[r] += task.resources[r]
            if not self.debug:
                # clean tmpdir
                shutil.rmtree(task.tmpdir)
        except Exception:
            logger.warn('error cleaning up task', exc_info=True)

def du(path):
    """
    Perform a "du" on a path, getting the disk usage.

    Args:
        path (str): The path to analyze

    Returns:
        int: bytes used
    """
    logger.info('du of %s', path)
    total = 0
    for root,dirs,files in os.walk(path):
        for d in list(dirs):
            p = os.path.join(root,d)
            if os.path.islink(p):
                dirs.remove(d)
        for f in files:
            p = os.path.join(root,f)
            if not os.path.islink(p):
                total += os.path.getsize(p)
    logger.info('du of %s finished: %r', path, total)
    return total
<|endoftext|>


<commit_before>#!/usr/bin/python
# sendandreceivearguments.py
# Author: Adrien Emery
# Make sure the you have the SendAndReceiveArguments example loaded onto the Arduino
from __future__ import print_function
from future import standard_library
standard_library.install_aliases()
from builtins import object
import random
import sys
import serial
import time
import threading

from globals import MyGlobals
from cmdmessenger import CmdMessenger
from serial.tools import list_ports

class MessengerController(object):

    def __init__(self):

        # make sure this baudrate matches the baudrate on the Arduino
        self.running = False
        self.baud = 9600
        self.temperature = {}
        self.temperature["Hot"]  = 100
        self.temperature["Mash"] = 100
        self.temperature["Boil"] = 100

        self.valve_status = []
        self.commands = ['acknowledge',
                         'error',
                         'ping',
                         'SetPin',
                         'PwmPin',
                         'ReadTemperature',
                         'DumpInWater',
                         'DumpInWater_reached',
                         'Resistor'
                         ]

        if (MyGlobals.args['hardware_disconnected'] is False):
            try:
                # try to open the first available usb port
                self.port_name = self.list_usb_ports()[0][0]
                self.serial_port = serial.Serial(self.port_name, self.baud, timeout=0, rtscts=True)
            except (serial.SerialException, IndexError):
                raise SystemExit('Could not open serial port.')
            else:
                time.sleep(2)
                self.messenger = CmdMessenger(self.serial_port)

            # send a command that the arduino will acknowledge
            self.messenger.send_cmd(self.commands.index('acknowledge'))
            # Wait until the arduino sends and acknowledgement back
            status = self.messenger.wait_for_ack(ackid=self.commands.index('acknowledge'), timeout=5)
            if (status is None):
                print('Arduino timeout')
                sys.exit(255)

            # attach callbacks
            self.messenger.attach(func=self.on_error, msgid=self.commands.index('error'))
            self.messenger.attach(func=self.on_read_temperature,
                                  msgid=self.commands.index('ReadTemperature'))
            self.messenger.attach(func=self.on_dump_in_reached,
                                  msgid=self.commands.index('DumpInWater_reached'))

            thread = threading.Thread(target=self.run, args=())
            thread.daemon = True                            # Daemonize thread
            thread.start()                                  # Start the execution

    def list_usb_ports(self):
        """ Use the grep generator to get a list of all USB ports.
        """
        ports = [port for port in list_ports.grep('USB')]
        return ports

    def on_error(self, received_command, *args, **kwargs):
        """Callback function to handle errors
        """
        print(('Error:', args[0][0]))

    def on_read_temperature(self, received_command, *args, **kwargs):
        """ Callback on temperature """

        self.temperature["Hot"]  = int(args[0][0])/100
        self.temperature["Mash"] = int(args[0][1])/100
        self.temperature["Boil"] = int(args[0][2])/100

    def on_dump_in_reached(self, received_command, *args, **kwargs):
        """ Callback on dump_in reached """

        self.valve_status[args[0][0]] = False

    def stop(self):
        self.running = False

    def ping(self):
        """ Send a ping (blocking) """
        self.messenger.send_cmd(self.commands.index('ping'))

        self.messenger.wait_for_ack(ackid=self.commands.index('acknowledge'))

    def set_pin(self, pin, value):
        """ Set a boolean value to a pin (blocking) """
        self.messenger.send_cmd(self.commands.index('SetPin'), pin, value)

        self.messenger.wait_for_ack(ackid=self.commands.index('acknowledge'))

    def set_pwm_pin(self, pin, value):
        """ Set a pwm value to a pin (blocking)
            value is supposed to be between 0-1
        """
        self.messenger.send_cmd(self.commands.index('PwmPin'), pin, int(value*255))

        self.messenger.wait_for_ack(ackid=self.commands.index('acknowledge'))

    def set_resistors(self, mash, boil, hot):
        print(mash)
        self.messenger.send_cmd(self.commands.index('Resistor'), hot, mash, boil, 100 - mash - boil - hot)

    def dump_in_water(self, valve, value):
        """ Set a pwm value to a pin (blocking)
            value is supposed to be between 0-65 liters
        """
        self.messenger.send_cmd(self.commands.index('DumpInWater'), valve, value)
        self.valve_status[valve] = True
        self.messenger.wait_for_ack(ackid=self.commands.index('acknowledge'))

    def run(self):
        """Main loop to send and receive data from the Arduino
        """
        while (True):
            self.messenger.feed_in_data()
            time.sleep(0.01)  # please do not hurt my core

if __name__ == '__main__':
    msg = MessengerController()

    try:
        print('Press Ctrl+C to exit...')
        msg.run()
    except KeyboardInterrupt:
        #msg.stop()
        msg.set_pwm_pin(5, 0.5)
        print('Exiting...')
        time.sleep(100)
<commit_msg>stupid print<commit_after>#!/usr/bin/python
# sendandreceivearguments.py
# Author: Adrien Emery
# Make sure the you have the SendAndReceiveArguments example loaded onto the Arduino
from __future__ import print_function
from future import standard_library
standard_library.install_aliases()
from builtins import object
import random
import sys
import serial
import time
import threading

from globals import MyGlobals
from cmdmessenger import CmdMessenger
from serial.tools import list_ports

class MessengerController(object):

    def __init__(self):

        # make sure this baudrate matches the baudrate on the Arduino
        self.running = False
        self.baud = 9600
        self.temperature = {}
        self.temperature["Hot"]  = 100
        self.temperature["Mash"] = 100
        self.temperature["Boil"] = 100

        self.valve_status = []
        self.commands = ['acknowledge',
                         'error',
                         'ping',
                         'SetPin',
                         'PwmPin',
                         'ReadTemperature',
                         'DumpInWater',
                         'DumpInWater_reached',
                         'Resistor'
                         ]

        if (MyGlobals.args['hardware_disconnected'] is False):
            try:
                # try to open the first available usb port
                self.port_name = self.list_usb_ports()[0][0]
                self.serial_port = serial.Serial(self.port_name, self.baud, timeout=0, rtscts=True)
            except (serial.SerialException, IndexError):
                raise SystemExit('Could not open serial port.')
            else:
                time.sleep(2)
                self.messenger = CmdMessenger(self.serial_port)

            # send a command that the arduino will acknowledge
            self.messenger.send_cmd(self.commands.index('acknowledge'))
            # Wait until the arduino sends and acknowledgement back
            status = self.messenger.wait_for_ack(ackid=self.commands.index('acknowledge'), timeout=5)
            if (status is None):
                print('Arduino timeout')
                sys.exit(255)

            # attach callbacks
            self.messenger.attach(func=self.on_error, msgid=self.commands.index('error'))
            self.messenger.attach(func=self.on_read_temperature,
                                  msgid=self.commands.index('ReadTemperature'))
            self.messenger.attach(func=self.on_dump_in_reached,
                                  msgid=self.commands.index('DumpInWater_reached'))

            thread = threading.Thread(target=self.run, args=())
            thread.daemon = True                            # Daemonize thread
            thread.start()                                  # Start the execution

    def list_usb_ports(self):
        """ Use the grep generator to get a list of all USB ports.
        """
        ports = [port for port in list_ports.grep('USB')]
        return ports

    def on_error(self, received_command, *args, **kwargs):
        """Callback function to handle errors
        """
        print(('Error:', args[0][0]))

    def on_read_temperature(self, received_command, *args, **kwargs):
        """ Callback on temperature """

        self.temperature["Hot"]  = int(args[0][0])/100
        self.temperature["Mash"] = int(args[0][1])/100
        self.temperature["Boil"] = int(args[0][2])/100

    def on_dump_in_reached(self, received_command, *args, **kwargs):
        """ Callback on dump_in reached """

        self.valve_status[args[0][0]] = False

    def stop(self):
        self.running = False

    def ping(self):
        """ Send a ping (blocking) """
        self.messenger.send_cmd(self.commands.index('ping'))

        self.messenger.wait_for_ack(ackid=self.commands.index('acknowledge'))

    def set_pin(self, pin, value):
        """ Set a boolean value to a pin (blocking) """
        self.messenger.send_cmd(self.commands.index('SetPin'), pin, value)

        self.messenger.wait_for_ack(ackid=self.commands.index('acknowledge'))

    def set_pwm_pin(self, pin, value):
        """ Set a pwm value to a pin (blocking)
            value is supposed to be between 0-1
        """
        self.messenger.send_cmd(self.commands.index('PwmPin'), pin, int(value*255))

        self.messenger.wait_for_ack(ackid=self.commands.index('acknowledge'))

    def set_resistors(self, mash, boil, hot):
        self.messenger.send_cmd(self.commands.index('Resistor'), hot, mash, boil, 100 - mash - boil - hot)

    def dump_in_water(self, valve, value):
        """ Set a pwm value to a pin (blocking)
            value is supposed to be between 0-65 liters
        """
        self.messenger.send_cmd(self.commands.index('DumpInWater'), valve, value)
        self.valve_status[valve] = True
        self.messenger.wait_for_ack(ackid=self.commands.index('acknowledge'))

    def run(self):
        """Main loop to send and receive data from the Arduino
        """
        while (True):
            self.messenger.feed_in_data()
            time.sleep(0.01)  # please do not hurt my core

if __name__ == '__main__':
    msg = MessengerController()

    try:
        print('Press Ctrl+C to exit...')
        msg.run()
    except KeyboardInterrupt:
        #msg.stop()
        msg.set_pwm_pin(5, 0.5)
        print('Exiting...')
        time.sleep(100)
<|endoftext|>


<commit_before># Copyright 2015 The Bazel Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License 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.
"""Java AppEngine support for Bazel.

For now, it only support bundling a WebApp and running locally.

To create a WebApp for Google AppEngine, add the rules:
appengine_war(
  name = "MyWebApp",
  # Jars to use for the classpath in the webapp.
  jars = ["//java/com/google/examples/mywebapp:java"],
  # data to put in the webapp, the directory structure of the data set
  # will be maintained.
  data = ["//java/com/google/examples/mywebapp:data"],
  # Data's root path, it will be considered as the root of the data files.
  # If unspecified, the path to the current package will be used. The path is
  # relative to current package or, relative to the workspace root if starting
  # with a leading slash.
  data_path = "/java/com/google/examples/mywebapp",
)

To test locally:
bazel run :MyWebApp

To deploy on Google app engine:
bazel run :MyWebApp.deploy

You can also make directly a single target for it with:

java_war(
  name = "MyWebApp",
  srcs = glob(["**/*.java"]),
  resources = ["..."],
  data = ["..."],
  data_path = "...",
)

Resources will be put in the classpath whereas data will be bundled at the root
of the war file. This is strictly equivalent to (it is actually a convenience
macros that translate to that):

java_library(
  name = "libMyWebApp",
  srcs = glob(["**/*.java"]),
  resources = ["..."],
)

appengine_war(
  name = "MyWebApp",
  jars = [":libMyWebApp"],
  data = ["..."],
  data_path = "...",
)

Finally, the appengine macro also create a .deploy target that will try to use the
AppEngine SDK to upload your application to AppEngine. It takes an optional argument: the
APP_ID. If not specified, it uses the default APP_ID provided in the application
web.xml.
"""

jar_filetype = FileType([".jar"])

def _add_file(in_file, output, path = None):
  output_path = output
  input_path = in_file.path
  if path and in_file.short_path.startswith(path):
    output_path += in_file.short_path[len(path):]
  return [
      "mkdir -p $(dirname %s)" % output_path,
      "test -L %s || ln -s $(pwd)/%s %s" % (output_path, input_path, output_path)
      ]

def _make_war(zipper, input_dir, output):
  return [
      "(root=$(pwd);" +
      ("cd %s &&" % input_dir) +
      ("${root}/%s Cc ${root}/%s $(find .))" % (zipper.path, output.path))
      ]

def _common_substring(str1, str2):
  i = 0
  res = ""
  for c in str1:
    if str2[i] != c:
      return res
    res += c
    i += 1
  return res

def _short_path_dirname(path):
  sp = path.short_path
  return sp[0:len(sp)-len(path.basename)-1]

def _war_impl(ctxt):
  """Implementation of the rule that creates
     - the war
     - the script to deploy
  """
  zipper = ctxt.file._zipper

  data_path = ctxt.attr.data_path
  if not data_path:
    data_path = _short_path_dirname(ctxt.outputs.war)
  elif data_path[0] == "/":
    data_path = data_path[1:]
  else:  # relative path
    data_path = _short_path_dirname(ctxt.outputs.war) + "/" + data_path

  war = ctxt.outputs.war
  build_output = war.path + ".build_output"
  cmd = [
      "set -e;rm -rf " + build_output,
      "mkdir -p " + build_output
      ]

  inputs = [zipper]
  cmd += ["mkdir -p %s/WEB-INF/lib" % build_output]

  transitive_deps = depset()
  for jar in ctxt.attr.jars:
    if hasattr(jar, "java"):  # java_library, java_import
      transitive_deps += jar.java.transitive_runtime_deps
    elif hasattr(jar, "files"):  # a jar file
      transitive_deps += jar.files

  for dep in transitive_deps:
    cmd += _add_file(dep, build_output + "/WEB-INF/lib")
    inputs.append(dep)

  for jar in ctxt.files._appengine_deps:
    cmd += _add_file(jar, build_output + "/WEB-INF/lib")
    inputs.append(jar)

  inputs += ctxt.files.data
  for res in ctxt.files.data:
    # Add the data file
    cmd += _add_file(res, build_output, path = data_path)

  cmd += _make_war(zipper, build_output, war)

  ctxt.action(
      inputs = inputs,
      outputs = [war],
      mnemonic="WAR",
      command="\n".join(cmd),
      use_default_shell_env=True)

  executable = ctxt.outputs.executable
  appengine_sdk = None
  for f in ctxt.files._appengine_sdk:
    if not appengine_sdk:
      appengine_sdk = f.short_path
    elif not f.path.startswith(appengine_sdk):
      appengine_sdk = _common_substring(appengine_sdk, f.short_path)
  if not appengine_sdk:
    fail("could not find appengine files",
         attr = str(ctxt.attr._appengine_sdk.label))

  classpath = ["${JAVA_RUNFILES}/%s" % jar.short_path for jar in transitive_deps]
  classpath += [
      "${JAVA_RUNFILES}/%s" % jar.short_path
      for jar in ctxt.files._appengine_deps
  ]

  substitutions = {
      "%{workspace_name}" : ctxt.workspace_name,
      "%{zipper}": ctxt.file._zipper.short_path,
      "%{war}": ctxt.outputs.war.short_path,
      "%{java}": ctxt.file._java.short_path,
      "%{appengine_sdk}": appengine_sdk,
      "%{classpath}":  (":".join(classpath)),
      "%{data_path}": data_path
  }

  ctxt.template_action(
      output = executable,
      template = ctxt.file._runner_template,
      substitutions = substitutions,
      executable = True)
  ctxt.template_action(
      output = ctxt.outputs.deploy_sh,
      template = ctxt.file._deploy_template,
      substitutions = substitutions)

  runfiles = ctxt.runfiles(files = [war, executable]
                           + list(transitive_deps)
                           + inputs
                           + ctxt.files._appengine_sdk
                           + [ctxt.file._java, ctxt.file._zipper])
  return struct(runfiles = runfiles)

appengine_war_base = rule(
    _war_impl,
    attrs = {
        "_java": attr.label(
            default = Label("@bazel_tools//tools/jdk:java"),
            single_file = True,
        ),
        "_zipper": attr.label(
            default = Label("@bazel_tools//tools/zip:zipper"),
            single_file = True,
        ),
        "_runner_template": attr.label(
            default = Label("//appengine:runner_template"),
            single_file = True,
        ),
        "_deploy_template": attr.label(
            default = Label("//appengine:deploy_template"),
            single_file = True,
        ),
        "_appengine_sdk": attr.label(
            default = Label("@com_google_appengine_java//:sdk"),
        ),
        "_appengine_deps": attr.label_list(
            default = [Label("@com_google_appengine_java//:api")],
        ),
        "jars": attr.label_list(
            allow_files = jar_filetype,
            mandatory = True,
        ),
        "data": attr.label_list(allow_files = True),
        "data_path": attr.string(),
    },
    executable = True,
    outputs = {
        "war": "%{name}.war",
        "deploy_sh": "%{name}_deploy.sh",
    },
)

def java_war(name, data=[], data_path=None, **kwargs):
  """Convenience macro to call appengine_war with Java sources rather than jar.
  """
  native.java_library(name = "lib%s" % name, **kwargs)
  appengine_war(name = name,
                jars = ["lib%s" % name],
                data=data,
                data_path=data_path)

def appengine_war(name, jars, data, data_path, testonly = 0):
  """Convenience macro that builds the war and offers an executable
     target to deploy on Google app engine.
  """
  appengine_war_base(
      name = name,
      jars = jars,
      data = data,
      data_path = data_path,
      testonly = testonly,
  )
  # Create the executable rule to deploy
  native.sh_binary(
      name = "%s.deploy" % name,
      srcs = ["%s_deploy.sh" % name],
      data = [name],
      testonly = testonly,
  )

APPENGINE_VERSION = "1.9.54"

APPENGINE_DIR = "appengine-java-sdk-" + APPENGINE_VERSION

APPENGINE_BUILD_FILE = """
# BUILD file to use the Java AppEngine SDK with a remote repository.
java_import(
    name = "jars",
    jars = glob(["lib/**/*.jar"]),
    visibility = ["//visibility:public"],
)

java_import(
    name = "user",
    jars = glob(["lib/user/*.jar"]),
    visibility = ["//visibility:public"],
)

java_import(
    name = "api",
    jars = [
        "lib/agent/appengine-agent.jar",
        "lib/appengine-tools-api.jar",
        "lib/impl/appengine-api.jar",
    ],
    visibility = ["//visibility:public"],
    neverlink = 1,
)

filegroup(
    name = "sdk",
    srcs = glob(["**"]),
    visibility = ["//visibility:public"],
)
"""

def appengine_repositories():
  native.new_http_archive(
    name = "com_google_appengine_java",
    url = "http://central.maven.org/maven2/com/google/appengine/appengine-java-sdk/%s/%s.zip" % (APPENGINE_VERSION, APPENGINE_DIR),
    sha256 = "79311d74886a555e135afb47395197ce5c69feb5e5557707dfd7bcddbad45e89",
    build_file_content = APPENGINE_BUILD_FILE,
    strip_prefix = APPENGINE_DIR,
  )

  native.maven_jar(
      name = "javax_servlet_api",
      artifact = "javax.servlet:servlet-api:2.5",
  )
<commit_msg>update appengine library to 1.9.57 (#57)<commit_after># Copyright 2015 The Bazel Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License 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.
"""Java AppEngine support for Bazel.

For now, it only support bundling a WebApp and running locally.

To create a WebApp for Google AppEngine, add the rules:
appengine_war(
  name = "MyWebApp",
  # Jars to use for the classpath in the webapp.
  jars = ["//java/com/google/examples/mywebapp:java"],
  # data to put in the webapp, the directory structure of the data set
  # will be maintained.
  data = ["//java/com/google/examples/mywebapp:data"],
  # Data's root path, it will be considered as the root of the data files.
  # If unspecified, the path to the current package will be used. The path is
  # relative to current package or, relative to the workspace root if starting
  # with a leading slash.
  data_path = "/java/com/google/examples/mywebapp",
)

To test locally:
bazel run :MyWebApp

To deploy on Google app engine:
bazel run :MyWebApp.deploy

You can also make directly a single target for it with:

java_war(
  name = "MyWebApp",
  srcs = glob(["**/*.java"]),
  resources = ["..."],
  data = ["..."],
  data_path = "...",
)

Resources will be put in the classpath whereas data will be bundled at the root
of the war file. This is strictly equivalent to (it is actually a convenience
macros that translate to that):

java_library(
  name = "libMyWebApp",
  srcs = glob(["**/*.java"]),
  resources = ["..."],
)

appengine_war(
  name = "MyWebApp",
  jars = [":libMyWebApp"],
  data = ["..."],
  data_path = "...",
)

Finally, the appengine macro also create a .deploy target that will try to use the
AppEngine SDK to upload your application to AppEngine. It takes an optional argument: the
APP_ID. If not specified, it uses the default APP_ID provided in the application
web.xml.
"""

jar_filetype = FileType([".jar"])

def _add_file(in_file, output, path = None):
  output_path = output
  input_path = in_file.path
  if path and in_file.short_path.startswith(path):
    output_path += in_file.short_path[len(path):]
  return [
      "mkdir -p $(dirname %s)" % output_path,
      "test -L %s || ln -s $(pwd)/%s %s" % (output_path, input_path, output_path)
      ]

def _make_war(zipper, input_dir, output):
  return [
      "(root=$(pwd);" +
      ("cd %s &&" % input_dir) +
      ("${root}/%s Cc ${root}/%s $(find .))" % (zipper.path, output.path))
      ]

def _common_substring(str1, str2):
  i = 0
  res = ""
  for c in str1:
    if str2[i] != c:
      return res
    res += c
    i += 1
  return res

def _short_path_dirname(path):
  sp = path.short_path
  return sp[0:len(sp)-len(path.basename)-1]

def _war_impl(ctxt):
  """Implementation of the rule that creates
     - the war
     - the script to deploy
  """
  zipper = ctxt.file._zipper

  data_path = ctxt.attr.data_path
  if not data_path:
    data_path = _short_path_dirname(ctxt.outputs.war)
  elif data_path[0] == "/":
    data_path = data_path[1:]
  else:  # relative path
    data_path = _short_path_dirname(ctxt.outputs.war) + "/" + data_path

  war = ctxt.outputs.war
  build_output = war.path + ".build_output"
  cmd = [
      "set -e;rm -rf " + build_output,
      "mkdir -p " + build_output
      ]

  inputs = [zipper]
  cmd += ["mkdir -p %s/WEB-INF/lib" % build_output]

  transitive_deps = depset()
  for jar in ctxt.attr.jars:
    if hasattr(jar, "java"):  # java_library, java_import
      transitive_deps += jar.java.transitive_runtime_deps
    elif hasattr(jar, "files"):  # a jar file
      transitive_deps += jar.files

  for dep in transitive_deps:
    cmd += _add_file(dep, build_output + "/WEB-INF/lib")
    inputs.append(dep)

  for jar in ctxt.files._appengine_deps:
    cmd += _add_file(jar, build_output + "/WEB-INF/lib")
    inputs.append(jar)

  inputs += ctxt.files.data
  for res in ctxt.files.data:
    # Add the data file
    cmd += _add_file(res, build_output, path = data_path)

  cmd += _make_war(zipper, build_output, war)

  ctxt.action(
      inputs = inputs,
      outputs = [war],
      mnemonic="WAR",
      command="\n".join(cmd),
      use_default_shell_env=True)

  executable = ctxt.outputs.executable
  appengine_sdk = None
  for f in ctxt.files._appengine_sdk:
    if not appengine_sdk:
      appengine_sdk = f.short_path
    elif not f.path.startswith(appengine_sdk):
      appengine_sdk = _common_substring(appengine_sdk, f.short_path)
  if not appengine_sdk:
    fail("could not find appengine files",
         attr = str(ctxt.attr._appengine_sdk.label))

  classpath = ["${JAVA_RUNFILES}/%s" % jar.short_path for jar in transitive_deps]
  classpath += [
      "${JAVA_RUNFILES}/%s" % jar.short_path
      for jar in ctxt.files._appengine_deps
  ]

  substitutions = {
      "%{workspace_name}" : ctxt.workspace_name,
      "%{zipper}": ctxt.file._zipper.short_path,
      "%{war}": ctxt.outputs.war.short_path,
      "%{java}": ctxt.file._java.short_path,
      "%{appengine_sdk}": appengine_sdk,
      "%{classpath}":  (":".join(classpath)),
      "%{data_path}": data_path
  }

  ctxt.template_action(
      output = executable,
      template = ctxt.file._runner_template,
      substitutions = substitutions,
      executable = True)
  ctxt.template_action(
      output = ctxt.outputs.deploy_sh,
      template = ctxt.file._deploy_template,
      substitutions = substitutions)

  runfiles = ctxt.runfiles(files = [war, executable]
                           + list(transitive_deps)
                           + inputs
                           + ctxt.files._appengine_sdk
                           + [ctxt.file._java, ctxt.file._zipper])
  return struct(runfiles = runfiles)

appengine_war_base = rule(
    _war_impl,
    attrs = {
        "_java": attr.label(
            default = Label("@bazel_tools//tools/jdk:java"),
            single_file = True,
        ),
        "_zipper": attr.label(
            default = Label("@bazel_tools//tools/zip:zipper"),
            single_file = True,
        ),
        "_runner_template": attr.label(
            default = Label("//appengine:runner_template"),
            single_file = True,
        ),
        "_deploy_template": attr.label(
            default = Label("//appengine:deploy_template"),
            single_file = True,
        ),
        "_appengine_sdk": attr.label(
            default = Label("@com_google_appengine_java//:sdk"),
        ),
        "_appengine_deps": attr.label_list(
            default = [Label("@com_google_appengine_java//:api")],
        ),
        "jars": attr.label_list(
            allow_files = jar_filetype,
            mandatory = True,
        ),
        "data": attr.label_list(allow_files = True),
        "data_path": attr.string(),
    },
    executable = True,
    outputs = {
        "war": "%{name}.war",
        "deploy_sh": "%{name}_deploy.sh",
    },
)

def java_war(name, data=[], data_path=None, **kwargs):
  """Convenience macro to call appengine_war with Java sources rather than jar.
  """
  native.java_library(name = "lib%s" % name, **kwargs)
  appengine_war(name = name,
                jars = ["lib%s" % name],
                data=data,
                data_path=data_path)

def appengine_war(name, jars, data, data_path, testonly = 0):
  """Convenience macro that builds the war and offers an executable
     target to deploy on Google app engine.
  """
  appengine_war_base(
      name = name,
      jars = jars,
      data = data,
      data_path = data_path,
      testonly = testonly,
  )
  # Create the executable rule to deploy
  native.sh_binary(
      name = "%s.deploy" % name,
      srcs = ["%s_deploy.sh" % name],
      data = [name],
      testonly = testonly,
  )

APPENGINE_VERSION = "1.9.57"

APPENGINE_DIR = "appengine-java-sdk-" + APPENGINE_VERSION

APPENGINE_BUILD_FILE = """
# BUILD file to use the Java AppEngine SDK with a remote repository.
java_import(
    name = "jars",
    jars = glob(["lib/**/*.jar"]),
    visibility = ["//visibility:public"],
)

java_import(
    name = "user",
    jars = glob(["lib/user/*.jar"]),
    visibility = ["//visibility:public"],
)

java_import(
    name = "api",
    jars = [
        "lib/agent/appengine-agent.jar",
        "lib/appengine-tools-api.jar",
        "lib/impl/appengine-api.jar",
    ],
    visibility = ["//visibility:public"],
    neverlink = 1,
)

filegroup(
    name = "sdk",
    srcs = glob(["**"]),
    visibility = ["//visibility:public"],
)
"""

def appengine_repositories():
  native.new_http_archive(
    name = "com_google_appengine_java",
    url = "http://central.maven.org/maven2/com/google/appengine/appengine-java-sdk/%s/%s.zip" % (APPENGINE_VERSION, APPENGINE_DIR),
    sha256 = "63f89be498d1e7462c03fe2b66d7773a9e5dd04ffb71e3b59a5fa74a7b810997",
    build_file_content = APPENGINE_BUILD_FILE,
    strip_prefix = APPENGINE_DIR,
  )

  native.maven_jar(
      name = "javax_servlet_api",
      artifact = "javax.servlet:servlet-api:2.5",
  )
<|endoftext|>


<commit_before>
#--IMPORT_ALL_FROM_FUTURE--#

'''
@date 2014-11-16
@author Hong-she Liang <starofrainnight@gmail.com>
'''

# Import the global selenium unit, not our selenium .
global_selenium = __import__('selenium')
import types
import time

def set_attribute(self, name, value):
    value = value.replace(r"'", r"\'") # Replace all r"'" with r"\'"
    value = value.replace("\n", r"\n")
    value = value.replace("\r", r"\r")
    value = value.replace("\t", r"\t")
    script = "arguments[0].setAttribute('%s', '%s')"  % (name, value)
    self._parent.execute_script(script, self)

def force_focus(self):
    global_selenium.webdriver.ActionChains(self._parent).move_to_element(self).perform()

def force_click(self):
    self._parent.execute_script("arguments[0].click();", self);
        <commit_msg>Use javascript to force focus on an element, because the action chains seems take no effect!<commit_after>
#--IMPORT_ALL_FROM_FUTURE--#

'''
@date 2014-11-16
@author Hong-she Liang <starofrainnight@gmail.com>
'''

# Import the global selenium unit, not our selenium .
global_selenium = __import__('selenium')
import types
import time

def set_attribute(self, name, value):
    value = value.replace(r"'", r"\'") # Replace all r"'" with r"\'"
    value = value.replace("\n", r"\n")
    value = value.replace("\r", r"\r")
    value = value.replace("\t", r"\t")
    script = "arguments[0].setAttribute('%s', '%s');"  % (name, value)
    self._parent.execute_script(script, self)

def force_focus(self):
    self._parent.execute_script("arguments[0].focus();", self);

def force_click(self):
    self._parent.execute_script("arguments[0].click();", self);
        <|endoftext|>


<commit_before>"""

"""

import os
import numpy as np
from whaler.dataprep import IO

class Analysis():
    """
    """
    def __init__(self):
        self.loc = os.getcwd()
        self.structs = next(os.walk('.'))[1]
        self.logfile = IO('whaler.log', self.loc)

    def groundstates_all(self, outname="groundstates.csv"):
        """Compares the energies of each calculated spin state for a structure
        and writes the energy differences as a table."""

        results = [self.spinstates(struct) for struct in self.structs]
        columns = [] #turn list of rows into list of columns

        # write table as groundstates.out file.
        writer = IO(outname, self.loc)

        headers = np.array(['Structures', 'S', 'T', 'P', 'D', 'Q'])

        writer.tabulate_data(columns, headers, 'Structures')

    def spinstates(self, structure):
        """For a given structure, identifies all of the files optimizing
        geometries in different spin states. Verifies convergence, and then
        finds the final single-point energy for each file. Returns an array of
        energies of the various spin states.
        Possibilities: S T P D Q (for S = 0, 1, 2, 1/2, 3/2)
        """
        path = os.path.join(self.loc, structure)
        geologs = list(filter(
                        lambda file: self.isfiletype(file, path, "geo.log"),
                        os.listdir(path)
                        ))
        print(geologs)

    def isfiletype(self, file, path, suffix):
        """
        """
        return file.endswith(suffix) and self.isvalid(file, path)

    def isvalid(self, file, path):
        """
        """
        reader = IO(file, path)
        end = reader.tail(2)
        if 'ABORTING THE RUN\n' in end:
            self.logfile.appendline(file + ' aborted abnormally.')
            return False
        elif 'ORCA TERMINATED NORMALLY' in end[0]:
            return True
        else:
            self.logfile.appendline(file + ' has unknown structure.')
            return False

    def getcalctype(self, file):
        """Takes a chemical computation file and gives the calc type labels,
        based on the filename formulation: xxxxxxx_NSyyy.log, where x chars
        refer to the structure name, N is the iteration number, S is the spin
        state label, and yyy is the optimization type.
        """
        return ''<commit_msg>Filter geo.log files based on iteration, normal termination, and convergence.<commit_after>"""

"""

import os
import numpy as np
from whaler.dataprep import IO

class Analysis():
    """
    """
    def __init__(self):
        self.loc = os.getcwd()
        self.structs = next(os.walk('.'))[1]
        self.logfile = IO('whaler.log', self.loc)

    def groundstates_all(self, outname="groundstates.csv"):
        """Compares the energies of each calculated spin state for a structure
        and writes the energy differences as a table."""

        results = [self.spinstates(struct) for struct in self.structs]
        columns = [] #turn list of rows into list of columns

        # write table as groundstates.out file.
        writer = IO(outname, self.loc)

        headers = np.array(['Structures', 'S', 'T', 'P', 'D', 'Q'])

        writer.tabulate_data(columns, headers, 'Structures')

    def spinstates(self, structure):
        """For a given structure, identifies all of the files optimizing
        geometries in different spin states. Verifies convergence, and then
        finds the final single-point energy for each file. Returns an array of
        energies of the various spin states.
        Possibilities: S T P D Q (for S = 0, 1, 2, 1/2, 3/2)
        """
        path = os.path.join(self.loc, structure)
        files = os.listdir(path) # Starting file list.

        # Narrows it down to geo.log files.
        geologs = list(filter(
                        lambda file: file.endswith("geo.log"),
                        files
                        ))
        print(geologs)

        # Unpacks filetypes and checks that the highest iteration is examined.
        ftypes = {file:self.getcalctype(file) for file in geologs}

        print(ftypes)

        iter, state, type = (zip(*ftypes.values()))

        curriter = max(iter)
        currfiles = {k:v for (k,v) in ftypes.items() if v[0] == curriter}

        print(currfiles)

        # Removes invalid files, marking the log.
        validlogs = {
            k:v for (k,v) in currfiles.items() if self.isvalid(k, path)}

        print(validlogs)

    def getcalctype(self, file):
        """Takes a chemical computation file and gives the calc type labels,
        based on the filename formulation: xxxxxxx_NSyyy.log, where x chars
        refer to the structure name, N is the iteration number, S is the spin
        state label, and yyy is the optimization type.
        """
        labels = file.split('_')[-1]
        iter = int(labels[0])
        state = labels[1]
        type = labels.split('.')[0][2:]
        return (iter, state, type)

    def isvalid(self, file, path):
        """
        """
        reader = IO(file, path)
        end = reader.tail(2)
        if 'ABORTING THE RUN\n' in end:
            self.logfile.appendline(file + ' aborted abnormally.')
            return False
        elif 'ORCA TERMINATED NORMALLY' in end[0]:
            return self.isconverged(file, path)
        else:
            self.logfile.appendline(file + ' has unknown structure.')
            return False

    def isconverged(self, file, path, chunk=100):
        """
        """
        reader = IO(file, path)
        tail = reader.tail(chunk)
        if chunk > 1000:
            self.logfile.appendline(file + ' has unknown structure.')
            return False
        elif 'WARNING!!!!!!!\n' in tail:
            self.logfile.appendline(file + ' has not converged.')
            return False
        elif '*** OPTIMIZATION RUN DONE ***' in ''.join(tail):
            return True
        else:
            self.converged(file, path, chunk+100)

        <|endoftext|>


<commit_before>"""Pilot functionality"""

from __future__ import absolute_import, division, print_function

import os
import sys
import time
import logging
import tempfile
import shutil
from functools import partial
from multiprocessing import Process
from collections import namedtuple
from datetime import timedelta
from glob import glob

from iceprod.core import to_file, constants
from iceprod.core import exe_json
from iceprod.core.util import Task_Resources, get_task_resources
from iceprod.core.dataclasses import Number, String
import iceprod.core.logger

logger = logging.getLogger('pilot')
logger.setLevel(logging.DEBUG)

from tornado import gen
from tornado.ioloop import IOLoop
from tornado.locks import Condition

try:
    import psutil
except ImportError:
    psutil = None

try:
    from setproctitle import setproctitle
except ImportError:
    def setproctitle(name):
        pass

Task = namedtuple('Task', ['p','process','resources','tmpdir'])

def process_wrapper(func, title):
    """
    Set process title. Set log file, stdout, stderr. Then go on to call func.

    Args:
        func (callable): The function that we really want to run
        title (str): The new process title
    """
    try:
        setproctitle(title)
    except Exception:
        pass

    iceprod.core.logger.new_file(constants['stdlog'])
    stdout = partial(to_file,sys.stdout,constants['stdout'])
    stderr = partial(to_file,sys.stderr,constants['stderr'])
    with stdout(), stderr():
        func()

class Pilot(object):
    """
    A pilot task runner.

    The pilot allows multiple tasks to run in sequence or parallel.
    It keeps track of resource usage, killing anything that goes over
    requested amounts.

    Args:
        config (dict): the configuration dictionary
        runner (callable): the task/config runner
    """
    def __init__(self, config, runner, pilot_id, debug=False, run_timeout=60):
        self.config = config
        self.runner = runner
        self.pilot_id = pilot_id
        self.debug = debug
        self.run_timeout = timedelta(seconds=run_timeout)

        # set up resources for pilot
        self.lock = Condition()
        self.resources = get_task_resources()
        if 'resources' in config['options']:
            self.resources.update(config['options']['resources'])

        self.start_time = time.time()

        # set up monitor
        self.ioloop = IOLoop.current()
        self.tasks = {}
        if psutil:
            self.ioloop.add_callback(self.monitor)
        else:
            logger.warn('no psutil. not checking resource usage')

        self.ioloop.run_sync(self.run)

        if self.debug:
            # append out, err, log
            for dirs in glob('tmp*'):
                for filename in (constants['stdout'], constants['stderr'],
                                 constants['stdlog']):
                    if os.path.exists(os.path.join(dirs,filename)):
                        with open(filename,'a') as f:
                            print('')
                            print('----',dirs,'----', file=f)
                            with open(os.path.join(dirs,filename)) as f2:
                                print(f2.read(), file=f)

    @gen.coroutine
    def monitor(self):
        """Monitor the tasks, killing any that go over resource limits"""
        try:
            sleep_time = 0.1 # check every X seconds
            disk_sleep_time = 180
            disk_start_time = time.time()
            while True:
                logger.info('pilot monitor - checking resource usage')
                start_time = time.time()
                for task_id in list(self.tasks):
                    task = self.tasks[task_id]
                    if not task.p.is_alive():
                        logger.info('not alive, so notify')
                        self.lock.notify()
                        continue

                    used_resources = {r:0 for r in task.resources}
                    processes = [task.process]+task.process.children()
                    for p in processes:
                        used_resources['cpu'] += p.cpu_percent()/100.0
                        used_resources['memory'] += p.memory_info().rss/1000000000.0
                    used_resources['time'] = (start_time - task.process.create_time())/3600.0
                    if start_time - disk_start_time > disk_sleep_time:
                        disk_start_time = start_time
                        used_resources['disk'] = du(task.tmpdir)/1000000000.0
                    logger.debug('task %r used %r',task_id,used_resources)

                    kill = False
                    reason = ''
                    for r in used_resources:
                        if used_resources[r] > task.resources[r]:
                            kill = True
                            reason = 'Resource overuse for {}: {}'.format(r,used_resources[r])
                            break
                    if kill:
                        logger.warn('kill %r for going over resources %r',
                                    task_id, used_resources)
                        processes.reverse() # kill children first
                        for p in processes:
                            p.terminate()

                        def on_terminate(proc):
                            logger.info("process %r terminated with exit code %r",
                                        proc, proc.returncode)
                        gone, alive = psutil.wait_procs(processes, timeout=0.1,
                                                        callback=on_terminate)
                        for p in alive:
                            p.kill()
                        self.clean_task(task)
                        del self.tasks[task_id]
                        exe_json.task_kill(task_id, resources=used_resources,
                                           reason=reason)
                        # try to queue another task
                        logger.info('killed, so notify')
                        self.lock.notify()

                duration = time.time()-start_time
                logger.debug('sleep_time %d, duration %d',sleep_time,duration)
                if duration < sleep_time:
                    yield self.lock.wait(timeout=timedelta(seconds=sleep_time-duration))
        except Exception:
            logger.error('pilot monitor died', exc_info=True)
            raise

    @gen.coroutine
    def run(self):
        """Run the pilot"""
        errors = 0
        tasks_running = 0
        running = True
        while running:
            while running:
                try:
                    task_config = exe_json.downloadtask(self.config['options']['gridspec'],
                                                        resources=self.resources)
                except Exception:
                    errors += 1
                    if errors > 5:
                        running = False
                    logger.error('cannot download task. current error count is %d',
                                 errors, exc_info=True)
                    continue
                logger.info('task config: %r', task_config)

                if task_config is None:
                    logger.info('no task available')
                    if not self.tasks:
                        running = False
                    break
                else:
                    try:
                        task_id = task_config['options']['task_id']
                    except Exception:
                        errors += 1
                        if errors > 5:
                            running = False
                        logger.warn('error getting task_id from config')
                        continue
                    try:
                        self.create_task(task_config)
                    except Exception:
                        errors += 1
                        if errors > 5:
                            running = False
                        logger.warn('error creating task %s', task_id,
                                    exc_info=True)
                        exe_json.task_kill(task_id, reason='failed to create task')
                    else:
                        tasks_running += 1
                        for r in self.resources:
                            if r in ('cpu','memory','disk', 'gpu'):
                                self.resources[r] -= task_config['options']['resources'][r]
                            elif isinstance(self.resources[r], String):
                                self.resources[r] = self.resources[r].split(',')
                        exe_json.update_pilot(self.pilot_id, tasks=','.join(self.tasks))

            # wait until we can queue more tasks
            while running or self.tasks:
                logger.info('wait while tasks are running. timeout=%r',self.run_timeout)
                ret = yield self.lock.wait(timeout=self.run_timeout)
                logger.info('yield returned %r',ret)
                # check if any processes have died
                for task_id in list(self.tasks):
                    if not self.tasks[task_id].p.is_alive():
                        self.clean_task(self.tasks[task_id])
                        del self.tasks[task_id]
                if len(self.tasks) < tasks_running:
                    logger.info('%d tasks removed', tasks_running-len(self.tasks))
                    tasks_running = len(self.tasks)
                    exe_json.update_pilot(self.pilot_id, tasks=','.join(self.tasks))
                    if running:
                        break

        if errors >= 5:
            logger.critical('too many errors when running tasks')
        else:
            logger.warn('cleanly stopping pilot')

    def create_task(self, config):
        """
        Create a new Task and start running it

        Args:
            config (dict): The task config
        """
        task_id = config['options']['task_id']

        # add grid-specific config
        for k in self.config['options']:
            if k == 'resources':
                pass
            elif k not in config['options']:
                config['options'][k] = self.config['options'][k]

        # add resources
        if 'resources' not in config['options']:
            config['options']['resources'] = {}
        for r in self.resources:
            if r in config['options']['resources']:
                if config['options']['resources'][r] > self.resources[r]:
                    raise Exception('requested {} of {}, but pilot has {}'.format(
                        config['options']['resources'][r], r, self.resources[r]))
            else:
                config['options']['resources'][r] = self.resources[r]

        # run task in tmp dir
        main_dir = os.getcwd()
        try:
            tmpdir = tempfile.mkdtemp(dir=main_dir)

            # symlink important files
            if 'ssl' in config['options']:
                for f in config['options']['ssl']:
                    os.symlink(config['options']['ssl'][f],
                               os.path.join(tmpdir,config['options']['ssl'][f]))

            # start the task
            os.chdir(tmpdir)
            p = Process(target=partial(process_wrapper, partial(self.runner, config),
                                       'iceprod_task_{}'.format(task_id)))
            p.start()
            if psutil:
                ps = psutil.Process(p.pid)
                ps.nice(psutil.Process().nice()+1)
            else:
                ps = None
            r = config['options']['resources']
            self.tasks[task_id] = Task(p=p, process=ps, resources=r,
                                       tmpdir=tmpdir)
        finally:
            os.chdir(main_dir)

    def clean_task(self, task):
        """
        Clean up a Task

        Args:
            task (Task): the Task
        """
        logger.info('cleaning task at %s', task.tmpdir)
        try:
            # return resources to pilot
            for r in task.resources:
                if r in ('cpu','memory','disk', 'gpu'):
                    self.resources[r] += task.resources[r]
            if not self.debug:
                # clean tmpdir
                shutil.rmtree(task.tmpdir)
        except Exception:
            logger.warn('error cleaning up task', exc_info=True)

def du(path):
    """
    Perform a "du" on a path, getting the disk usage.

    Args:
        path (str): The path to analyze

    Returns:
        int: bytes used
    """
    total = 0
    for root,dirs,files in os.walk(path):
        for d in list(dirs):
            p = os.path.join(root,d)
            if os.path.islink(p):
                dirs.remove(d)
        for f in files:
            p = os.path.join(root,f)
            if not os.path.islink(p):
                total += os.path.getsize(p)
    return total
<commit_msg>try to fix the ssl symlink loop<commit_after>"""Pilot functionality"""

from __future__ import absolute_import, division, print_function

import os
import sys
import time
import logging
import tempfile
import shutil
from functools import partial
from multiprocessing import Process
from collections import namedtuple
from datetime import timedelta
from glob import glob

from iceprod.core import to_file, constants
from iceprod.core import exe_json
from iceprod.core.util import Task_Resources, get_task_resources
from iceprod.core.dataclasses import Number, String
import iceprod.core.logger

logger = logging.getLogger('pilot')
logger.setLevel(logging.DEBUG)

from tornado import gen
from tornado.ioloop import IOLoop
from tornado.locks import Condition

try:
    import psutil
except ImportError:
    psutil = None

try:
    from setproctitle import setproctitle
except ImportError:
    def setproctitle(name):
        pass

Task = namedtuple('Task', ['p','process','resources','tmpdir'])

def process_wrapper(func, title):
    """
    Set process title. Set log file, stdout, stderr. Then go on to call func.

    Args:
        func (callable): The function that we really want to run
        title (str): The new process title
    """
    try:
        setproctitle(title)
    except Exception:
        pass

    iceprod.core.logger.new_file(constants['stdlog'])
    stdout = partial(to_file,sys.stdout,constants['stdout'])
    stderr = partial(to_file,sys.stderr,constants['stderr'])
    with stdout(), stderr():
        func()

class Pilot(object):
    """
    A pilot task runner.

    The pilot allows multiple tasks to run in sequence or parallel.
    It keeps track of resource usage, killing anything that goes over
    requested amounts.

    Args:
        config (dict): the configuration dictionary
        runner (callable): the task/config runner
    """
    def __init__(self, config, runner, pilot_id, debug=False, run_timeout=60):
        self.config = config
        self.runner = runner
        self.pilot_id = pilot_id
        self.debug = debug
        self.run_timeout = timedelta(seconds=run_timeout)

        # set up resources for pilot
        self.lock = Condition()
        self.resources = get_task_resources()
        if 'resources' in config['options']:
            self.resources.update(config['options']['resources'])

        self.start_time = time.time()

        # set up monitor
        self.ioloop = IOLoop.current()
        self.tasks = {}
        if psutil:
            self.ioloop.add_callback(self.monitor)
        else:
            logger.warn('no psutil. not checking resource usage')

        self.ioloop.run_sync(self.run)

        if self.debug:
            # append out, err, log
            for dirs in glob('tmp*'):
                for filename in (constants['stdout'], constants['stderr'],
                                 constants['stdlog']):
                    if os.path.exists(os.path.join(dirs,filename)):
                        with open(filename,'a') as f:
                            print('')
                            print('----',dirs,'----', file=f)
                            with open(os.path.join(dirs,filename)) as f2:
                                print(f2.read(), file=f)

    @gen.coroutine
    def monitor(self):
        """Monitor the tasks, killing any that go over resource limits"""
        try:
            sleep_time = 0.1 # check every X seconds
            disk_sleep_time = 180
            disk_start_time = time.time()
            while True:
                logger.info('pilot monitor - checking resource usage')
                start_time = time.time()
                for task_id in list(self.tasks):
                    task = self.tasks[task_id]
                    if not task.p.is_alive():
                        logger.info('not alive, so notify')
                        self.lock.notify()
                        continue

                    used_resources = {r:0 for r in task.resources}
                    processes = [task.process]+task.process.children()
                    for p in processes:
                        used_resources['cpu'] += p.cpu_percent()/100.0
                        used_resources['memory'] += p.memory_info().rss/1000000000.0
                    used_resources['time'] = (start_time - task.process.create_time())/3600.0
                    if start_time - disk_start_time > disk_sleep_time:
                        disk_start_time = start_time
                        used_resources['disk'] = du(task.tmpdir)/1000000000.0
                    logger.debug('task %r used %r',task_id,used_resources)

                    kill = False
                    reason = ''
                    for r in used_resources:
                        if used_resources[r] > task.resources[r]:
                            kill = True
                            reason = 'Resource overuse for {}: {}'.format(r,used_resources[r])
                            break
                    if kill:
                        logger.warn('kill %r for going over resources %r',
                                    task_id, used_resources)
                        processes.reverse() # kill children first
                        for p in processes:
                            p.terminate()

                        def on_terminate(proc):
                            logger.info("process %r terminated with exit code %r",
                                        proc, proc.returncode)
                        gone, alive = psutil.wait_procs(processes, timeout=0.1,
                                                        callback=on_terminate)
                        for p in alive:
                            p.kill()
                        self.clean_task(task)
                        del self.tasks[task_id]
                        exe_json.task_kill(task_id, resources=used_resources,
                                           reason=reason)
                        # try to queue another task
                        logger.info('killed, so notify')
                        self.lock.notify()

                duration = time.time()-start_time
                logger.debug('sleep_time %d, duration %d',sleep_time,duration)
                if duration < sleep_time:
                    yield self.lock.wait(timeout=timedelta(seconds=sleep_time-duration))
        except Exception:
            logger.error('pilot monitor died', exc_info=True)
            raise

    @gen.coroutine
    def run(self):
        """Run the pilot"""
        errors = 0
        tasks_running = 0
        running = True
        while running:
            while running:
                try:
                    task_config = exe_json.downloadtask(self.config['options']['gridspec'],
                                                        resources=self.resources)
                except Exception:
                    errors += 1
                    if errors > 5:
                        running = False
                    logger.error('cannot download task. current error count is %d',
                                 errors, exc_info=True)
                    continue
                logger.info('task config: %r', task_config)

                if task_config is None:
                    logger.info('no task available')
                    if not self.tasks:
                        running = False
                    break
                else:
                    try:
                        task_id = task_config['options']['task_id']
                    except Exception:
                        errors += 1
                        if errors > 5:
                            running = False
                        logger.warn('error getting task_id from config')
                        continue
                    try:
                        self.create_task(task_config)
                    except Exception:
                        errors += 1
                        if errors > 5:
                            running = False
                        logger.warn('error creating task %s', task_id,
                                    exc_info=True)
                        exe_json.task_kill(task_id, reason='failed to create task')
                    else:
                        tasks_running += 1
                        for r in self.resources:
                            if r in ('cpu','memory','disk', 'gpu'):
                                self.resources[r] -= task_config['options']['resources'][r]
                            elif isinstance(self.resources[r], String):
                                self.resources[r] = self.resources[r].split(',')
                        exe_json.update_pilot(self.pilot_id, tasks=','.join(self.tasks))

            # wait until we can queue more tasks
            while running or self.tasks:
                logger.info('wait while tasks are running. timeout=%r',self.run_timeout)
                ret = yield self.lock.wait(timeout=self.run_timeout)
                logger.info('yield returned %r',ret)
                # check if any processes have died
                for task_id in list(self.tasks):
                    if not self.tasks[task_id].p.is_alive():
                        self.clean_task(self.tasks[task_id])
                        del self.tasks[task_id]
                if len(self.tasks) < tasks_running:
                    logger.info('%d tasks removed', tasks_running-len(self.tasks))
                    tasks_running = len(self.tasks)
                    exe_json.update_pilot(self.pilot_id, tasks=','.join(self.tasks))
                    if running:
                        break

        if errors >= 5:
            logger.critical('too many errors when running tasks')
        else:
            logger.warn('cleanly stopping pilot')

    def create_task(self, config):
        """
        Create a new Task and start running it

        Args:
            config (dict): The task config
        """
        task_id = config['options']['task_id']

        # add grid-specific config
        for k in self.config['options']:
            if k == 'resources':
                pass
            elif k not in config['options']:
                config['options'][k] = self.config['options'][k]

        # add resources
        if 'resources' not in config['options']:
            config['options']['resources'] = {}
        for r in self.resources:
            if r in config['options']['resources']:
                if config['options']['resources'][r] > self.resources[r]:
                    raise Exception('requested {} of {}, but pilot has {}'.format(
                        config['options']['resources'][r], r, self.resources[r]))
            else:
                config['options']['resources'][r] = self.resources[r]

        # run task in tmp dir
        main_dir = os.getcwd()
        try:
            tmpdir = tempfile.mkdtemp(dir=main_dir)

            # symlink important files
            if 'ssl' in config['options']:
                for f in config['options']['ssl']:
                    os.symlink(os.path.join(main_dir,config['options']['ssl'][f]),
                               os.path.join(tmpdir,config['options']['ssl'][f]))

            # start the task
            os.chdir(tmpdir)
            p = Process(target=partial(process_wrapper, partial(self.runner, config),
                                       'iceprod_task_{}'.format(task_id)))
            p.start()
            if psutil:
                ps = psutil.Process(p.pid)
                ps.nice(psutil.Process().nice()+1)
            else:
                ps = None
            r = config['options']['resources']
            self.tasks[task_id] = Task(p=p, process=ps, resources=r,
                                       tmpdir=tmpdir)
        finally:
            os.chdir(main_dir)

    def clean_task(self, task):
        """
        Clean up a Task

        Args:
            task (Task): the Task
        """
        logger.info('cleaning task at %s', task.tmpdir)
        try:
            # return resources to pilot
            for r in task.resources:
                if r in ('cpu','memory','disk', 'gpu'):
                    self.resources[r] += task.resources[r]
            if not self.debug:
                # clean tmpdir
                shutil.rmtree(task.tmpdir)
        except Exception:
            logger.warn('error cleaning up task', exc_info=True)

def du(path):
    """
    Perform a "du" on a path, getting the disk usage.

    Args:
        path (str): The path to analyze

    Returns:
        int: bytes used
    """
    total = 0
    for root,dirs,files in os.walk(path):
        for d in list(dirs):
            p = os.path.join(root,d)
            if os.path.islink(p):
                dirs.remove(d)
        for f in files:
            p = os.path.join(root,f)
            if not os.path.islink(p):
                total += os.path.getsize(p)
    return total
<|endoftext|>