From c6d998c4dfaa747edd031eac5c71630cf7465c3a Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Sat, 11 Jul 2020 01:40:50 +0300
Subject: [PATCH] Add terminate process and rmtree utilities

---
 trains_agent/helper/base.py    | 19 ++++++++++++++++++-
 trains_agent/helper/process.py | 24 ++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/trains_agent/helper/base.py b/trains_agent/helper/base.py
index a99e389..6d022f7 100644
--- a/trains_agent/helper/base.py
+++ b/trains_agent/helper/base.py
@@ -173,13 +173,30 @@ def normalize_path(*paths):
 
 
 def safe_remove_file(filename, error_message=None):
+    # noinspection PyBroadException
     try:
-        os.remove(filename)
+        if filename:
+            os.remove(filename)
     except Exception:
         if error_message:
             print(error_message)
 
 
+def safe_remove_tree(filename):
+    if not filename:
+        return
+    # noinspection PyBroadException
+    try:
+        shutil.rmtree(filename, ignore_errors=True)
+    except Exception:
+        pass
+    # noinspection PyBroadException
+    try:
+        os.remove(filename)
+    except Exception:
+        pass
+
+
 def get_python_path(script_dir, entry_point, package_api):
     try:
         python_path_sep = ';' if is_windows_platform() else ':'
diff --git a/trains_agent/helper/process.py b/trains_agent/helper/process.py
index 9676819..c5b364b 100644
--- a/trains_agent/helper/process.py
+++ b/trains_agent/helper/process.py
@@ -11,6 +11,7 @@ from copy import deepcopy
 from distutils.spawn import find_executable
 from itertools import chain, repeat, islice
 from os.path import devnull
+from time import sleep
 from typing import Union, Text, Sequence, Any, TypeVar, Callable
 
 import psutil
@@ -41,6 +42,29 @@ def get_bash_output(cmd, strip=False, stderr=subprocess.STDOUT, stdin=False):
     return output if not strip or not output else output.strip()
 
 
+def terminate_process(pid, timeout=10.):
+    # noinspection PyBroadException
+    try:
+        proc = psutil.Process(pid)
+        proc.terminate()
+        cnt = 0
+        while proc.is_running() and cnt < timeout:
+            sleep(1.)
+            cnt += 1
+        proc.terminate()
+        cnt = 0
+        while proc.is_running() and cnt < timeout:
+            sleep(1.)
+            cnt += 1
+        proc.kill()
+    except Exception:
+        pass
+    # noinspection PyBroadException
+    try:
+        return not psutil.Process(pid).is_running()
+    except Exception:
+        return True
+
 def kill_all_child_processes(pid=None):
     # get current process if pid not provided
     include_parent = True