addftool 0.1.8__py3-none-any.whl → 0.1.9__py3-none-any.whl

addftool/process/__init__.py

@@ -11,6 +11,8 @@ def add_killer_args(subparsers):
     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("--contain_arg", help="args of command", default="", type=str)
+
     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)
 
@@ -59,7 +61,7 @@ def kill_process(processes, timeout=5, try_count=3):
         exit(1)
 
 
-def find_and_kill_process(command="", contains=False, use_rocm=False, use_cuda=False, timeout=5, try_count=3, only_view=False):
+def find_and_kill_process(command="", contains=False, contain_arg="", use_rocm=False, use_cuda=False, timeout=5, try_count=3, only_view=False):
     do_not_do_anything = command is None or len(command) == 0
     if use_rocm:
         processes = get_process_using_rocm()
@@ -69,7 +71,7 @@ def find_and_kill_process(command="", contains=False, use_rocm=False, use_cuda=F
         print("Use top or htop to find the process you want to kill")
         return
     else:
-        processes = get_processes(command=command, contains=contains)
+        processes = get_processes(command=command, contains=contains, contain_arg=contain_arg)
     
     if only_view:
         print(f"Found {len(processes)} processes")
@@ -83,7 +85,8 @@ def find_and_kill_process(command="", contains=False, use_rocm=False, use_cuda=F
 
 
 def killer_main(args):
+    print(args)
     find_and_kill_process(
-        args.name, args.contains, use_rocm=args.rocm, use_cuda=args.cuda, 
+        args.name, args.contains, args.contain_arg, use_rocm=args.rocm, use_cuda=args.cuda, 
         timeout=args.timeout, try_count=args.try_count, only_view=args.view,
     )

addftool/process/utils.py

@@ -2,7 +2,7 @@ import psutil
 import subprocess
 
 
-def get_processes(command="", contains=False, pids=None):
+def get_processes(command="", contains=False, pids=None, contain_arg=""):
     """获取进程的PID和命令"""
     
     processes = []
@@ -21,6 +21,17 @@ def get_processes(command="", contains=False, pids=None):
                     continue
             if len(cmdline) > 1 and cmdline[1].endswith('addf'):
                 continue
+            if len(contain_arg) > 0:
+                if len(cmdline) <= 1:
+                    continue
+                flag = False
+                for arg in cmdline[1:]:
+                    if contain_arg in arg:
+                        flag = True
+                        break
+                if not flag:
+                    continue
+
             processes.append({'pid': pid, 'command': cmdline})
         except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
             pass

{addftool-0.1.8.dist-info → addftool-0.1.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: addftool
-Version: 0.1.8
+Version: 0.1.9
 Requires-Dist: cryptography
 Requires-Dist: requests
 Requires-Dist: PyYAML

{addftool-0.1.8.dist-info → addftool-0.1.9.dist-info}/RECORD

@@ -7,11 +7,11 @@ addftool/util.py,sha256=zlNLu8Be8cGIpNRqBw8_0q7nFxWlsJ9cToN62ohjdXE,2335
 addftool/deploy/__init__.py,sha256=tpyoTh3SqAQojPizsJDvQohu1Pcb3-w-DP5sO4-5lBM,1220
 addftool/deploy/azure.py,sha256=_o_9Eh8cVwLDAqvfyRYBtQRHs_Gul-nCs2ZXttwO1bk,1301
 addftool/deploy/ssh_server.py,sha256=f2T8fgwACVljPfdcimMywUjsFnLCWRde7iWPAILpRz8,5463
-addftool/process/__init__.py,sha256=gPdGsjMEET6crzOz4Iw5cmf6RR1toXGovydRXv8Uagk,3543
-addftool/process/utils.py,sha256=me4HqMz5OgRcQMUJmVhKdTJh4SW5BB-pd_lq7g8-UwE,2252
+addftool/process/__init__.py,sha256=Dze8OrcyjQlAbPrjE_h8bMi8W4b3OJyZOjTucPrkJvM,3721
+addftool/process/utils.py,sha256=JldxnwanLJOgxaPgmCJh7SeBRaaj5rFxWWxh1hpsvbA,2609
 addftool/ssh/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-addftool-0.1.8.dist-info/METADATA,sha256=f50lOq51j55hNh2hnk6SdAni0E7MXHec81sBCOHZ_ro,170
-addftool-0.1.8.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-addftool-0.1.8.dist-info/entry_points.txt,sha256=9lkmuWMInwUAtev8w8poNkNd7iML9Bjd5CBCFVxg2b8,111
-addftool-0.1.8.dist-info/top_level.txt,sha256=jqj56-plrBbyzY0tIxB6wPzjAA8kte4hUlajyyQygN4,9
-addftool-0.1.8.dist-info/RECORD,,
+addftool-0.1.9.dist-info/METADATA,sha256=VEfXNesNw0S0lZl2_hYrJhrUyRN9SpP1iT1aIhtgKbU,170
+addftool-0.1.9.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+addftool-0.1.9.dist-info/entry_points.txt,sha256=9lkmuWMInwUAtev8w8poNkNd7iML9Bjd5CBCFVxg2b8,111
+addftool-0.1.9.dist-info/top_level.txt,sha256=jqj56-plrBbyzY0tIxB6wPzjAA8kte4hUlajyyQygN4,9
+addftool-0.1.9.dist-info/RECORD,,