run command and kill by env

README.md

Install

pip install psutil

Usage

./safe_run.py bash long_script.sh

After finished. all related process will be killed. only tested in linux.

safe_run.py

#!/usr/bin/env python
# coding: utf-8

import os
import sys
import time
import signal

import psutil


CRED = '\033[31m'
CRST = '\033[0m'

SIGNALS_TO_NAMES_DICT = dict((getattr(signal, n), n) \
    for n in dir(signal) if n.startswith('SIG') and '_' not in n )

class Command(object):
    ENV_NAME = 'COMMAND_ID'

    def __init__(self, *args):
        self._args = args
        self._proc = None
        self._uid = str(time.time())

    def start(self):
        env = os.environ.copy()
        env[self.ENV_NAME] = self._uid
        self._proc = psutil.Popen(list(self._args), stdout=sys.stdout, stderr=sys.stderr, env=env)
        return self._proc

    @property
    def _processes(self):
        procs = []
        for proc in psutil.process_iter():
            try:
                env = proc.environ()
                if env.get(self.ENV_NAME) == self._uid:
                    procs.append(proc)
            except psutil.AccessDenied:
                pass
            except psutil.NoSuchProcess:
                pass
        return procs

    def stop_signal(self, signal=signal.SIGTERM):
        procs = self._processes
        if not procs:
            return
        print >>sys.stderr, CRED+'Send signal {0} to all related process(num={1}).'.format(
                SIGNALS_TO_NAMES_DICT[signal], len(procs))+CRST
        for p in procs:
            try:
                p.send_signal(signal)
            except psutil.NoSuchProcess:
                pass

    def stop(self, timeout=2.0):
        self.stop_signal(signal=signal.SIGTERM)
        deadline = time.time() + timeout # wait for 3s
        for p in self._processes:
            now = time.time()
            if now >= deadline:
                break
            try:
                p.wait(timeout=deadline-now)
            except psutil.TimeoutExpired:
                break

        self.stop_signal(signal=signal.SIGKILL)


def main():
    if len(sys.argv) == 1:
        sys.exit('Usage: %s [-t timeout] <command> [args...]' % sys.argv[0])

    timeout = None
    args = sys.argv[1:]
    if sys.argv[1] == '-t':
        timeout = float(sys.argv[2])
        args = sys.argv[3:]

    c = Command(*args)
    try:
        c.start().wait(timeout)
    except KeyboardInterrupt:
        print >>sys.stderr, CRED+'Signal Interrupt catched. try to stop ...'+CRST
    except psutil.TimeoutExpired:
        print >>sys.stderr, (CRED+'Command run timeout(%.1f)'+CRST) % timeout
    finally:
        c.stop()


if __name__ == '__main__':
    main()