addftool 0.1.7__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/sync.py

@@ -35,6 +35,8 @@ def add_args(parser):
     parser.add_argument("--md5_verify", action='store_true', default=False,
                         help="whether to verify the md5 of the file after sync, default is False.")
     parser.add_argument("--port", help="the port for torchrun, default is 29501", type=int, default=29501)
+    parser.add_argument("--torchrun_alias", type=str, default="torchrun",
+                        help="the alias of torchrun, default is torchrun. If you use torchrun, please set it to torchrun.")
     # distributed downloader from blob
     parser.add_argument("--donwload_nodes", help="download nodes, default is node-0", type=int, default=1)
     parser.add_argument("folder", help="the folder need to sync", type=str)
@@ -104,18 +106,22 @@ def sync_main(args):
 
     # divide the files into chunks for each node by file size
     sorted_files = sorted(file_size_list.items(), key=lambda x: x[1], reverse=True)
+    zero_files = []
 
     groups = [[] for _ in range(args.donwload_nodes)]
     for i, (file_name, file_size) in enumerate(sorted_files):
+        if file_size == 0:
+            zero_files.append(file_name)
+            continue
         groups[i % args.donwload_nodes].append(file_name)
     
     # create a temp folder to save the downloaded files
-    temp_config_dir = tempfile.mktemp()
-    os.makedirs(temp_config_dir, exist_ok=True)
-    print(f"Temp config dir: {temp_config_dir}")
+    local_temp_config_dir = tempfile.mktemp()
+    os.makedirs(local_temp_config_dir, exist_ok=True)
+    print(f"Temp config dir: {local_temp_config_dir}")
     
     for i, group in enumerate(groups):
-        group_file_path = os.path.join(temp_config_dir, f"node_{i}.txt")
+        group_file_path = os.path.join(local_temp_config_dir, f"node_{i}.txt")
         total_size = 0
         with open(group_file_path, "w") as f:
             for file_name in group:
@@ -123,6 +129,12 @@ def sync_main(args):
                 total_size += file_size_list[file_name]
         print(f"Node-{i} will download {len(group)} files, total size: {total_size} bytes")
     
+    with open(os.path.join(local_temp_config_dir, "zero_files.txt"), "w") as f:
+        for file_name in zero_files:
+            f.write(file_name + "\n")
+    
+    print(f"Detect {len(zero_files)} files with size 0 bytes, they will be special handled.")
+
     with open(args.hostfile, "r") as f:
         host_list = []
         for line in f:
@@ -135,24 +147,31 @@ def sync_main(args):
     
     print(f"Find {len(host_list)} hosts in hostfile: {args.hostfile}")
     connection_list = []
+
+    # avoid the temp_config_dir to be the same as remote_temp_config_dir
+    remote_temp_config_dir = tempfile.mktemp()
+    while remote_temp_config_dir == local_temp_config_dir:
+        remote_temp_config_dir = tempfile.mktemp()
+
     master_addr = get_ip_via_ssh(host_list[0])
     for i, host in enumerate(host_list):
         # copy this .py file to the remote host
         put_commands = []
-        put_commands.append((__file__, os.path.join(temp_config_dir, "sync.py")))
+        put_commands.append((__file__, os.path.join(remote_temp_config_dir, "sync.py")))
         if i < args.donwload_nodes:
-            local_group_file = os.path.join(temp_config_dir, f"node_{i}.txt")
-            put_commands.append((local_group_file, os.path.join(temp_config_dir, f"node_{i}.txt")))
+            local_group_file = os.path.join(local_temp_config_dir, f"node_{i}.txt")
+            put_commands.append((local_group_file, os.path.join(remote_temp_config_dir, f"node_{i}.txt")))
+        put_commands.append((os.path.join(local_temp_config_dir, "zero_files.txt"), os.path.join(remote_temp_config_dir, "zero_files.txt")))
 
         commnads = f"export SAS_TOKEN=\"{sas_token}\""
-        commnads += f" && torchrun --nproc_per_node=1 --nnodes={len(host_list)} --node_rank={i} --master_addr={master_addr} --master_port={args.port}"
-        commnads += f" {temp_config_dir}/sync.py {args.folder} --tool {args.tool} --from_blob_url {args.from_blob_url}"
+        commnads += f" && {args.torchrun_alias} --nproc_per_node=1 --nnodes={len(host_list)} --node_rank={i} --master_addr={master_addr} --master_port={args.port}"
+        commnads += f" {remote_temp_config_dir}/sync.py {args.folder} --tool {args.tool} --from_blob_url {args.from_blob_url}"
         if args.md5_verify:
             commnads += " --md5_verify"
         if i < args.donwload_nodes:
-            commnads += f" --download_index_file {temp_config_dir}/node_{i}.txt"
+            commnads += f" --download_index_file {remote_temp_config_dir}/node_{i}.txt"
         
-        connection_list.append(ConnectionWithCommand(host, temp_config_dir, put_commands, commnads))
+        connection_list.append(ConnectionWithCommand(host, remote_temp_config_dir, put_commands, commnads))
     
     group = ThreadingGroup.from_connections(connection_list)
     group.run('echo "Hello"', hide=False)
@@ -167,7 +186,9 @@ def download_files_from_blob(queue, blob_url, sas_token, folder, download_files,
     print(f"Node-{node_rank} start downloading {len(download_files)} files from {blob_url} to {folder}")
     for file_name in download_files:
         file_path = os.path.join(folder, file_name)
-        os.makedirs(os.path.dirname(file_path), exist_ok=True)
+        file_dir = os.path.dirname(file_path)
+        if not os.path.exists(file_dir):
+            os.makedirs(file_dir, exist_ok=True)
         for try_count in range(3):
             try:
                 download_status = subprocess.run(
@@ -208,6 +229,9 @@ def sync_file_from_rank(rank, file_path, from_rank, md5_verify=False):
     
     dist.broadcast(tensor, src=from_rank)
     if rank != from_rank:
+        file_dir = os.path.dirname(file_path)
+        if not os.path.exists(file_dir):
+            os.makedirs(file_dir, exist_ok=True)
         with open(file_path, "wb") as f:
             tensor.cpu().numpy().tofile(f)
         if md5_verify:
@@ -283,9 +307,25 @@ def sync_worker(args):
         else:
             sync_file_from_rank(node_rank, "", global_status_code, md5_verify=args.md5_verify)
 
+    print(f"Node-{node_rank} finished downloading files, time taken: {time.time() - start_time:.2f}s")
     dist.barrier()
     download_process.join()
     destroy_process_group()
+
+    # current directory
+    zero_file = os.path.join(__file__, "zero_files.txt")
+    if os.path.exists(zero_file):
+        with open(zero_file, "r") as f:
+            zero_files = [line.strip() for line in f]
+        for zero_file_name in zero_files:
+            zero_file_path = os.path.join(args.folder, zero_file_name)
+            zero_file_dir = os.path.dirname(zero_file_path)
+            if not os.path.exists(zero_file_dir):
+                os.makedirs(zero_file_dir, exist_ok=True)
+            with open(zero_file_path, "wb") as f:
+                f.write(b"")
+        print(f"Node-{node_rank} handled {len(zero_files)} files with size 0 bytes, time taken: {time.time() - start_time:.2f}s")
+
     print(f"Node-{node_rank} finished syncing all files, time taken: {time.time() - start_time:.2f}s")
 
 

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

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

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

@@ -1,17 +1,17 @@
 addftool/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 addftool/addf_portal.py,sha256=w2LgsoutfnrKhtrQAXouUMwLqnsp5ALlsBYUWg8n9NM,781
 addftool/blob.py,sha256=NZOItDyFUIdV1tfhJZJJBEzGy296CE5NCictTzP4OPc,8282
-addftool/sync.py,sha256=ZFxDeHLYcoTftQUKqVeXJqpkHnqwPg1SvrVUm0kt9OI,12019
+addftool/sync.py,sha256=ZpYxbM8uiPFrV7ODmOaM7asVPCWaxBixA-arVc-1kfs,14045
 addftool/tool.py,sha256=EuKQ2t2InN7yB-_oYLcdsA7vRqzRGTunwIxplUSqEG0,2054
 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.7.dist-info/METADATA,sha256=m-g6ENcdOOB7ufkIibSl5L8tYflejPWzlqorWwT2j1A,170
-addftool-0.1.7.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-addftool-0.1.7.dist-info/entry_points.txt,sha256=9lkmuWMInwUAtev8w8poNkNd7iML9Bjd5CBCFVxg2b8,111
-addftool-0.1.7.dist-info/top_level.txt,sha256=jqj56-plrBbyzY0tIxB6wPzjAA8kte4hUlajyyQygN4,9
-addftool-0.1.7.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,,