pip install psutil
./safe_run.py bash long_script.sh
After finished. all related process will be killed. only tested in linux.
#!/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() |