addftool 0.1.2__py3-none-any.whl → 0.1.3__py3-none-any.whl

addftool/addf_portal.py

@@ -0,0 +1,23 @@
+import argparse
+from addftool.process import add_killer_args, killer_main
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description="Addf's tool")
+
+    subparsers = parser.add_subparsers(dest='command', help='Sub-command help')
+    add_killer_args(subparsers)
+
+    return parser.parse_args()
+
+
+def main():
+    args = get_args()
+    if args.command == "kill":
+        killer_main(args)
+    else:
+        print("Unknown command: ", args.command)
+
+
+if __name__ == "__main__":
+    main()

addftool/process/__init__.py

@@ -0,0 +1,87 @@
+import psutil
+import time
+
+from .utils import get_processes, get_process_using_rocm, get_process_using_cuda
+
+
+def add_killer_args(subparsers):
+    process_killer_parser = subparsers.add_parser('kill', help='process kill')
+
+    process_killer_parser.add_argument("--prefix", help="prefix of command", default=None)
+    process_killer_parser.add_argument("--contains", help="contains of command", default=None)
+    process_killer_parser.add_argument("--timeout", help="timeout of command", default=5, type=int)
+    process_killer_parser.add_argument("--try_count", help="try count of command", default=3, type=int)
+
+    process_killer_parser.add_argument("--rocm", help="kill process using rocm", action='store_true', default=False)
+    process_killer_parser.add_argument("--cuda", help="kill process using cuda", action='store_true', default=False)
+
+    process_killer_parser.add_argument("-v", "--view", help="view process", action='store_true', default=False)
+
+
+def kill_process(processes, timeout=5, try_count=3):
+    for process in processes:
+        try:
+            proc = psutil.Process(process['pid'])
+            print(f"Killing process {proc.pid} - {proc.cmdline()}")
+            proc.terminate()
+        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
+            print(f"Process {process['pid']} not found or access denied")
+            continue
+    
+    time.sleep(timeout)
+    for _ in range(try_count):
+        all_terminated = True
+        for process in processes:
+            try:
+                proc = psutil.Process(process['pid'])
+                print("Process still running: ", proc.pid)
+                proc.kill()
+                all_terminated = False
+            except (psutil.NoSuchProcess):
+                continue
+            except (psutil.AccessDenied):
+                print(f"Process {process['pid']} access denied")
+                continue
+            except (psutil.ZombieProcess):
+                print(f"Process {process['pid']} is a zombie process")
+                continue
+            except Exception as e:
+                print(f"Error killing process {process['pid']}: {e}")
+                continue
+        if all_terminated:
+            print("All processes terminated")
+            break
+    
+    if not all_terminated:
+        print("Some processes are still running. Please check manually.")
+        exit(1)
+
+
+def find_and_kill_process(prefix=None, contains=None, use_rocm=False, use_cuda=False, timeout=5, try_count=3, only_view=False):
+    do_not_do_anything = prefix is None and contains is None
+    if use_rocm:
+        processes = get_process_using_rocm()
+    elif use_cuda:
+        processes = get_process_using_cuda()
+    elif do_not_do_anything and not only_view:
+        print("Use top or htop to find the process you want to kill")
+        return
+    else:
+        processes = get_processes(command_prefix=prefix, prefix_contains=contains)
+    
+    if only_view:
+        print(f"Found {len(processes)} processes")
+        for process in processes:
+            print(f"PID: {process['pid']}, Command: {' '.join(process['command'])}")
+        return
+
+    if len(processes) > 0:
+        print(f"Found {len(processes)} processes to kill")
+        kill_process(processes, timeout, try_count)
+
+
+def killer_main(args):
+    find_and_kill_process(
+        args.prefix, args.contains, use_rocm=args.rocm, use_cuda=args.cuda, 
+        timeout=args.timeout, try_count=args.try_count, only_view=args.view,
+    )

addftool/process/utils.py

@@ -0,0 +1,70 @@
+import psutil
+import subprocess
+
+
+def get_processes(command_prefix=None, prefix_contains=None, pids=None):
+    """获取进程的PID和命令"""
+    
+    processes = []
+    for proc in psutil.process_iter(['pid', 'cmdline']):
+        try:
+            pid = proc.info['pid']
+            cmdline = proc.info['cmdline']
+            if cmdline is None or len(cmdline) == 0:
+                continue
+            if pids is not None and pid not in pids:
+                continue
+            if command_prefix and not cmdline[0] == command_prefix:
+                continue
+            if prefix_contains and prefix_contains not in cmdline[0]:
+                continue
+            if len(cmdline) > 1 and cmdline[1].endswith('addf'):
+                continue
+            processes.append({'pid': pid, 'command': cmdline})
+        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
+            pass
+            
+    return processes
+
+
+def get_process_using_rocm():
+    result = subprocess.run(['rocm-smi', '--showpidgpus'], capture_output=True, text=True)
+    if result.returncode != 0:
+        return False
+
+    gpu_pids = []
+    output_lines = result.stdout.strip().split('\n')
+    for line in output_lines:
+        if "PID" in line and ":" in line:
+            # Extract PID from lines like "GPU[0-7]: PID: 12345"
+            try:
+                pid_part = line.split("PID:")[1].strip()
+                pid = int(pid_part.split()[0])
+                gpu_pids.append(pid)
+            except (IndexError, ValueError):
+                continue
+    
+    # Remove duplicates
+    gpu_pids = list(set(gpu_pids))
+    return get_processes(pids=gpu_pids)
+
+
+def get_process_using_cuda():
+    result = subprocess.run(['nvidia-smi', '--query-compute-apps=pid', '--format=csv,noheader'], capture_output=True, text=True)
+    if result.returncode != 0:
+        return False
+
+    gpu_pids = []
+    output_lines = result.stdout.strip().split('\n')
+    for line in output_lines:
+        try:
+            pid = int(line.strip())
+            gpu_pids.append(pid)
+        except ValueError:
+            continue
+    
+    # Remove duplicates
+    gpu_pids = list(set(gpu_pids))
+
+    # get process id and cmdline
+    return get_processes(pids=gpu_pids)

{addftool-0.1.2.dist-info → addftool-0.1.3.dist-info}/METADATA

@@ -1,6 +1,7 @@
 Metadata-Version: 2.4
 Name: addftool
-Version: 0.1.2
+Version: 0.1.3
 Requires-Dist: cryptography
 Requires-Dist: requests
 Requires-Dist: PyYAML
+Requires-Dist: psutil

addftool-0.1.3.dist-info/RECORD

@@ -0,0 +1,12 @@
+addftool/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+addftool/addf_portal.py,sha256=U52TdNcwWOEvv_C5r-guWYxn3ntzwYI2eBzJIE7IdcY,493
+addftool/blob.py,sha256=NZOItDyFUIdV1tfhJZJJBEzGy296CE5NCictTzP4OPc,8282
+addftool/tool.py,sha256=EuKQ2t2InN7yB-_oYLcdsA7vRqzRGTunwIxplUSqEG0,2054
+addftool/util.py,sha256=Q3A68vJDxgfeNiEFmk54HuMuworVndocXpSbVpvGMfc,362
+addftool/process/__init__.py,sha256=EFJ90RcF04M4imGm5NLgb9rmpwxBslZUL_4LKA52uv8,3487
+addftool/process/utils.py,sha256=kS1bYYiWymB1_SObqQLfGiVbZdCJqd57SO6DrwN5H1s,2246
+addftool-0.1.3.dist-info/METADATA,sha256=YwSmBYP8v9T_PnVr-3WM-V2KkCcPg17IwMKllfp-pc0,148
+addftool-0.1.3.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+addftool-0.1.3.dist-info/entry_points.txt,sha256=9lkmuWMInwUAtev8w8poNkNd7iML9Bjd5CBCFVxg2b8,111
+addftool-0.1.3.dist-info/top_level.txt,sha256=jqj56-plrBbyzY0tIxB6wPzjAA8kte4hUlajyyQygN4,9
+addftool-0.1.3.dist-info/RECORD,,

{addftool-0.1.2.dist-info → addftool-0.1.3.dist-info}/entry_points.txt

@@ -1,3 +1,4 @@
 [console_scripts]
+addf = addftool.addf_portal:main
 addfblob = addftool.blob:main
 addftool = addftool.tool:main

addftool-0.1.2.dist-info/RECORD

@@ -1,9 +0,0 @@
-addftool/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-addftool/blob.py,sha256=NZOItDyFUIdV1tfhJZJJBEzGy296CE5NCictTzP4OPc,8282
-addftool/tool.py,sha256=EuKQ2t2InN7yB-_oYLcdsA7vRqzRGTunwIxplUSqEG0,2054
-addftool/util.py,sha256=Q3A68vJDxgfeNiEFmk54HuMuworVndocXpSbVpvGMfc,362
-addftool-0.1.2.dist-info/METADATA,sha256=gmakYu_DywA1K1z3yEoLN0h43nBUFbtNy0JmKcEbTgw,126
-addftool-0.1.2.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-addftool-0.1.2.dist-info/entry_points.txt,sha256=h5TlQy4AQw-J55HqI7FRkIVRUfw-x3pStFWCoYtoCTM,78
-addftool-0.1.2.dist-info/top_level.txt,sha256=jqj56-plrBbyzY0tIxB6wPzjAA8kte4hUlajyyQygN4,9
-addftool-0.1.2.dist-info/RECORD,,