addftool 0.2.4__py3-none-any.whl → 0.2.5__py3-none-any.whl

addftool/broadcast_folder.py

@@ -1,7 +1,9 @@
 import os
 import time
+import fnmatch
 import subprocess
 import hashlib
+from concurrent.futures import ThreadPoolExecutor
 
 from fabric import Connection, ThreadingGroup
 
@@ -33,6 +35,22 @@ def add_args(parser):
     parser.add_argument("--transfer_ranks_per_node", type=int, default=8,
                         help="the number of ranks per node to transfer the files, default is 8.")
 
+    parser.add_argument("--contain_md5_files", action='store_true', default=False,
+                        help="whether to contain the md5 files in the folder, default is False. " \
+                        "If True, the md5 files will be transferred to the other nodes and verified. " \
+                        "If False, the md5 files will be ignored.")
+
+    parser.add_argument("--include-string", type=str, default="",
+                        help="the string to include the files, default is empty. " \
+                        "Such as *.py, *.yaml, *.json, \"*.pt;*.pth\" etc. " \
+                        "Only node-0 will include the files from the folder, " \
+                        "If empty, will transfer all the files from the node-0's local folder.")
+    parser.add_argument("--exclude-string", type=str, default="",
+                        help="the string to exclude the files, default is empty. " \
+                        "Such as *.py, *.yaml, *.json, \"*.pt;*.pth\" etc. " \
+                        "Only node-0 will exclude the files from the folder, " \
+                        "If empty, will transfer all the files from the node-0's local folder.")
+
     parser.add_argument("--from_blob_url", type=str, default="",
                         help="the blob url to download from, default is empty. " \
                         "Only node-0 will download the files from the blob url, " \
@@ -77,6 +95,52 @@ def get_ip_via_ssh(hostname):
         return None
 
 
+def parallel_check_md5(file_list, expected_md5s):
+    """
+    Parallel check MD5 checksums for the given files.
+    
+    Args:
+        file_list: List of file paths to check
+        md5_dir: Directory containing MD5 files
+        
+    Returns:
+        True if all MD5 checksums match, False otherwise
+    """
+    def calculate_md5(file_path):
+        # Call md5sum and capture output
+        result = subprocess.run(["md5sum", file_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+        if result.returncode != 0:
+            print(f"Failed to calculate MD5 for {file_path}: {result.stderr}")
+            return file_path, None
+        
+        # md5sum output format: "<md5_hash>  <file_path>"
+        md5_hash = result.stdout.strip().split()[0]
+        return file_path, md5_hash
+    
+    # Calculate MD5 checksums in parallel
+    with ThreadPoolExecutor(max_workers=8) as executor:
+        results = list(executor.map(calculate_md5, file_list))
+    
+    # Check if all MD5s match
+    all_match = True
+    for file_path, actual_md5 in results:
+        if actual_md5 is None:
+            all_match = False
+            continue
+        
+        if file_path not in expected_md5s:
+            print(f"No expected MD5 for {file_path}")
+            all_match = False
+            continue
+            
+        expected_md5 = expected_md5s[file_path]
+        if actual_md5 != expected_md5:
+            print(f"MD5 mismatch for {file_path}: expected {expected_md5}, got {actual_md5}")
+            all_match = False
+    
+    return all_match
+
+
 def broadcast_folder_main(args):
     with open(args.hostfile, "r") as f:
         host_list = []
@@ -95,11 +159,17 @@ def broadcast_folder_main(args):
     for i, host in enumerate(host_list):
         put_commands = []
         put_commands.append((__file__, os.path.join(remote_temp_config_dir, "broadcast.py")))
-        commnads = "NCCL_IB_DISABLE=0 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 "
+        commnads = "NCCL_IB_DISABLE=0 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 OMP_NUM_THREADS=1 "
         if os.environ.get("SAS_TOKEN") is not None and i == 0:
             commnads += f"SAS_TOKEN=\"{os.environ['SAS_TOKEN']}\" "
         commnads += f"{args.torchrun_alias} --nproc_per_node={args.transfer_ranks_per_node} --nnodes={len(host_list)} --node_rank={i} --master_addr={master_addr} --master_port={args.port}"
         commnads += f" {remote_temp_config_dir}/broadcast.py {args.folder} --tool {args.tool} --transfer_ranks_per_node {args.transfer_ranks_per_node} "
+        if args.contain_md5_files:
+            commnads += " --contain_md5_files"
+        if args.include_string:
+            commnads += f" --include-string \"{args.include_string}\""
+        if args.exclude_string:
+            commnads += f" --exclude-string \"{args.exclude_string}\""
         if args.from_blob_url and i == 0:
             commnads += f" --from_blob_url {args.from_blob_url}"
         if args.md5_verify:
@@ -226,20 +296,53 @@ def broadcast_folder_worker(args):
     file_size_dict = {}
     
     if global_rank == 0:
-        if args.from_blob_url:
-            raise NotImplementedError("Downloading files from blob storage is not implemented yet.")
-
+        # Parse include and exclude patterns
+        include_patterns = [p.strip() for p in args.include_string.split(";") if p.strip()]
+        exclude_patterns = [p.strip() for p in args.exclude_string.split(";") if p.strip()]
+        
+        print(f"Include patterns: {include_patterns}")
+        print(f"Exclude patterns: {exclude_patterns}")
+        
         if not os.path.exists(args.folder):
             raise ValueError(f"Folder {args.folder} does not exist.")
+        
+        # Gather and filter files in a single pass
+        file_size_dict = {}
         for root, dirs, files in os.walk(args.folder):
             for file in files:
                 file_path = os.path.join(root, file)
-                file_size = os.path.getsize(file_path)
-                file_size_dict[file_path] = file_size
+                file_name = os.path.basename(file_path)
+                
+                # Skip md5 files if not containing them
+                if file_name.endswith(".md5") and not args.contain_md5_files:
+                    continue
+                    
+                # Apply include filters first (if any)
+                included = not include_patterns  # Include by default if no include patterns
+                if include_patterns:
+                    for pattern in include_patterns:
+                        if fnmatch.fnmatch(file_name, pattern):
+                            included = True
+                            break
+                
+                # Then apply exclude filters
+                if included and exclude_patterns:
+                    for pattern in exclude_patterns:
+                        if fnmatch.fnmatch(file_name, pattern):
+                            included = False
+                            break
+                
+                # Add to file dict if passes both filters
+                if included:
+                    file_size_dict[file_path] = os.path.getsize(file_path)
+        
+        print(f"After filtering: {len(file_size_dict)} files selected for transfer")
+        if len(include_patterns) > 0 or len(exclude_patterns) > 0:
+            print(f"Files selected for transfer: {file_size_dict.keys()}")
         
         # sort the file list by size
         file_list = sorted(file_size_dict.keys(), key=lambda x: file_size_dict[x], reverse=True)
-        file_size_list = [file_size_dict[file] for file in file_list]        
+        file_size_list = [file_size_dict[file] for file in file_list]
         obj_list = [file_list, file_size_list]
         dist.broadcast_object_list(obj_list, src=0)
     else:
@@ -251,6 +354,7 @@ def broadcast_folder_worker(args):
     
     worker_g = workers_groups[worker_rank]
     from_rank = global_rank % args.transfer_ranks_per_node
+    broadcast_file_list = []
     for i in range(len(file_list)):
         if i % args.transfer_ranks_per_node == worker_rank:
             file_path = file_list[i]
@@ -261,6 +365,7 @@ def broadcast_folder_worker(args):
             )
             if global_rank == from_rank:
                 print(f"Group {global_rank} finished broadcasting {file_path}, size: {file_size / (1024 * 1024):.2f} MB, time taken: {time.time() - start_time:.2f}s")
+            broadcast_file_list.append(file_path)
     
     dist.barrier()
     for i in range(len(workers_groups)):
@@ -268,6 +373,27 @@ def broadcast_folder_worker(args):
             dist.destroy_process_group(workers_groups[i])
     destroy_process_group()
 
+    if args.contain_md5_files and global_rank % args.transfer_ranks_per_node == 0:
+        to_verify_files = []
+        excepted_md5s = {}
+        for file_path in file_list:
+            if not file_path.endswith(".md5"):
+                md5_file_path = file_path + ".md5"
+                if os.path.exists(md5_file_path):
+                    with open(md5_file_path, "r") as f:
+                        md5_hash = f.read().strip()
+                    excepted_md5s[file_path] = md5_hash
+                    to_verify_files.append(file_path)
+                else:
+                    print(f"MD5 file {md5_file_path} not found, skipping verification.")
+        
+        # Verify MD5 checksums
+        if not parallel_check_md5(to_verify_files, excepted_md5s):
+            print(f"MD5 verification failed for some files, please check the logs.")
+            raise ValueError("MD5 verification failed.")
+        else:
+            print(f"Rank-{global_rank}: MD5 verification passed for all files.")
+
     print(f"Rank {global_rank} finished broadcasting all files, time taken: {time.time() - start_time:.2f}s")
 
 

{addftool-0.2.4.dist-info → addftool-0.2.5.dist-info}/METADATA

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

{addftool-0.2.4.dist-info → addftool-0.2.5.dist-info}/RECORD

@@ -1,7 +1,7 @@
 addftool/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 addftool/addf_portal.py,sha256=6XjwGs5m2mRVDWVvCPOiqn1NlxDcGBTQ9Kr_0g5RsJc,1130
 addftool/blob.py,sha256=YT3eZrC9mdYtntDlwEI5MUHFw98wKtBC_pHEdtqvsv4,9206
-addftool/broadcast_folder.py,sha256=bEOr-8Q14DhZFS658hab4U-9HvAW3EGluEriilIdEXQ,11976
+addftool/broadcast_folder.py,sha256=X9tvvMT7cCDbfKqE7hUGaFbBOHIecIgD8xcT34Bqb_8,17708
 addftool/sync.py,sha256=ZpYxbM8uiPFrV7ODmOaM7asVPCWaxBixA-arVc-1kfs,14045
 addftool/tool.py,sha256=FmxRY3-pP0_Z0zCUAngjmEMmPUruMftg_iUlB1t2TnQ,2001
 addftool/util.py,sha256=zlNLu8Be8cGIpNRqBw8_0q7nFxWlsJ9cToN62ohjdXE,2335
@@ -12,8 +12,8 @@ addftool/deploy/vscode_server.py,sha256=tLtSvlcK2fEOaw6udWt8dNELVhwv9F59hF5DJJ-1
 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.2.4.dist-info/METADATA,sha256=Mxypih4dkYUmHFyyVBqbZLVf4CYJqapkfHMJ2zCMVh0,170
-addftool-0.2.4.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-addftool-0.2.4.dist-info/entry_points.txt,sha256=9lkmuWMInwUAtev8w8poNkNd7iML9Bjd5CBCFVxg2b8,111
-addftool-0.2.4.dist-info/top_level.txt,sha256=jqj56-plrBbyzY0tIxB6wPzjAA8kte4hUlajyyQygN4,9
-addftool-0.2.4.dist-info/RECORD,,
+addftool-0.2.5.dist-info/METADATA,sha256=e0fYxOCuq9cps5H7TZS-5wxSN3sqrs9e6g2acvJcMSw,170
+addftool-0.2.5.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
+addftool-0.2.5.dist-info/entry_points.txt,sha256=9lkmuWMInwUAtev8w8poNkNd7iML9Bjd5CBCFVxg2b8,111
+addftool-0.2.5.dist-info/top_level.txt,sha256=jqj56-plrBbyzY0tIxB6wPzjAA8kte4hUlajyyQygN4,9
+addftool-0.2.5.dist-info/RECORD,,

{addftool-0.2.4.dist-info → addftool-0.2.5.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.0)
+Generator: setuptools (79.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any