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

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.8.dist-info}/METADATA

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

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

@@ -1,7 +1,7 @@
 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
@@ -10,8 +10,8 @@ addftool/deploy/ssh_server.py,sha256=f2T8fgwACVljPfdcimMywUjsFnLCWRde7iWPAILpRz8
 addftool/process/__init__.py,sha256=gPdGsjMEET6crzOz4Iw5cmf6RR1toXGovydRXv8Uagk,3543
 addftool/process/utils.py,sha256=me4HqMz5OgRcQMUJmVhKdTJh4SW5BB-pd_lq7g8-UwE,2252
 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.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,,