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

addftool/addf_portal.py

@@ -4,6 +4,8 @@ from addftool.sync import add_sync_args, sync_main
 from addftool.deploy import add_deploy_args, deploy_main
 from addftool.broadcast_folder import add_broadcast_folder_args, broadcast_folder_main
 
+from addftool.blob import add_blob_args, blob_main
+
 
 def get_args():
     parser = argparse.ArgumentParser(description="Addf's tool")
@@ -13,6 +15,7 @@ def get_args():
     add_sync_args(subparsers)
     add_deploy_args(subparsers)
     add_broadcast_folder_args(subparsers)
+    add_blob_args(subparsers)
 
     return parser.parse_args()
 
@@ -27,6 +30,8 @@ def main():
         deploy_main(args)
     elif args.command == "broadcast-folder":
         broadcast_folder_main(args)
+    elif args.command == "blob":
+        blob_main(args)
     else:
         print("Unknown command: ", args.command)
 

addftool/blob.py

@@ -199,13 +199,29 @@ def get_sas_token(api_url, name, container, key, info=False):
         return f"Error: {response.status_code}"
 
 
-def main():
-    # exmaple usage: addfblob install
-    # exmaple usage: addfblob mount -k <key> -a <api> -b <buffer> -m <mount_dir>
-    # exmaple usage: addfblob token -k <key> -a <api>
-    parser = argparse.ArgumentParser(description="Addf's tool")
+def parse_blob_account_and_container_from_url(blob_url):
+    """
+    Parse the blob URL to extract account name and container name.
+    """
+    if not blob_url.startswith("https://"):
+        raise ValueError(f"Invalid blob URL: {blob_url}")
+
+    parts = blob_url.split("/")
+    account_name = parts[2].split(".")[0]
+    container_name = parts[3].strip()
+
+    return account_name, container_name
+
+
+def get_sas_token_for_blob_url(api_url, blob_url, key, info=False):
+    an, cn = parse_blob_account_and_container_from_url(blob_url)
+    if info:
+        print(f"Parse blob URL: {blob_url}, account name: {an}, container name: {cn}")
+    return get_sas_token(api_url, an, cn, key, info=info)
+
 
-    subparsers = parser.add_subparsers(dest='command', help='Sub-command help')
+def add_args(parser):
+    subparsers = parser.add_subparsers(dest='blob_command', help='Sub-command help')
     install_parser = subparsers.add_parser('install', help='Install help')
     # install_parser.add_argument("-o", "--output_script", help="output script", default=None)
     install_parser.add_argument("--packages", help="packages", default="fuse3 blobfuse2 azcopy")
@@ -223,19 +239,33 @@ def main():
     token_parser = subparsers.add_parser('token', help='Token help')
     add_api(token_parser)
 
-    args = parser.parse_args()
+def add_blob_args(subparsers):
+    deploy_parser = subparsers.add_parser('blob', help='Blob help')
+    add_args(deploy_parser)
+
 
+def blob_main(args):
     # check os is linux/unix and current user, set --sudo if current user is not root
     if os.name == 'posix' and os.getuid() != 0:
         args.sudo = True
 
-    if args.command == 'install':
+    if args.blob_command == 'install':
         install_main(args)
-    elif args.command == 'mount':
+    elif args.blob_command == 'mount':
         mount_main(args)
-    elif args.command == 'token':
+    elif args.blob_command == 'token':
         print(get_token(args, info=False))
 
 
+def main():
+    # exmaple usage: addfblob install
+    # exmaple usage: addfblob mount -k <key> -a <api> -b <buffer> -m <mount_dir>
+    # exmaple usage: addfblob token -k <key> -a <api>
+    parser = argparse.ArgumentParser(description="Addf's tool")
+    add_args(parser)
+    args = parser.parse_args()
+    blob_main(args)
+
+
 if __name__ == "__main__":
     main()

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/deploy/ssh_server.py

@@ -67,7 +67,10 @@ def configure_ssh_on_ubuntu(port, username, ssh_public_key="", password=""):
             execute_command(command_prefix + "apt-get install -y openssh-server")
 
         print("Modifying SSH configuration...")
+        # if "#Port 22" in sshd_config
         execute_command(command_prefix + f'sed -i "s/#Port 22/Port {port}/" /etc/ssh/sshd_config')
+        # if Port xxx in sshd_config
+        execute_command(command_prefix + f'sed -i "s/Port [0-9]*/Port {port}/" /etc/ssh/sshd_config')
         execute_command(command_prefix + 'sed -i "s/#PermitRootLogin prohibit-password/PermitRootLogin yes/" /etc/ssh/sshd_config')
 
         if not os.path.exists("/etc/ssh/ssh_host_rsa_key"):

addftool/tool.py

@@ -71,9 +71,5 @@ def main():
             print("No such init option")
 
 
-def pull():
-    execute_command("git pull origin")
-
-
 if __name__ == "__main__":
     main()

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

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

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

@@ -1,19 +1,19 @@
 addftool/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-addftool/addf_portal.py,sha256=pw6ivnwMfmK4WC-1Gs5cMndXKQRbMm6VvDyxXJDNrm4,991
-addftool/blob.py,sha256=NZOItDyFUIdV1tfhJZJJBEzGy296CE5NCictTzP4OPc,8282
-addftool/broadcast_folder.py,sha256=bEOr-8Q14DhZFS658hab4U-9HvAW3EGluEriilIdEXQ,11976
+addftool/addf_portal.py,sha256=6XjwGs5m2mRVDWVvCPOiqn1NlxDcGBTQ9Kr_0g5RsJc,1130
+addftool/blob.py,sha256=YT3eZrC9mdYtntDlwEI5MUHFw98wKtBC_pHEdtqvsv4,9206
+addftool/broadcast_folder.py,sha256=X9tvvMT7cCDbfKqE7hUGaFbBOHIecIgD8xcT34Bqb_8,17708
 addftool/sync.py,sha256=ZpYxbM8uiPFrV7ODmOaM7asVPCWaxBixA-arVc-1kfs,14045
-addftool/tool.py,sha256=EuKQ2t2InN7yB-_oYLcdsA7vRqzRGTunwIxplUSqEG0,2054
+addftool/tool.py,sha256=FmxRY3-pP0_Z0zCUAngjmEMmPUruMftg_iUlB1t2TnQ,2001
 addftool/util.py,sha256=zlNLu8Be8cGIpNRqBw8_0q7nFxWlsJ9cToN62ohjdXE,2335
 addftool/deploy/__init__.py,sha256=UL8b-Idt7lStlMiOm8oTZ65fdzYz99Fgzq2Gaw8WsZc,1544
 addftool/deploy/azure.py,sha256=_o_9Eh8cVwLDAqvfyRYBtQRHs_Gul-nCs2ZXttwO1bk,1301
-addftool/deploy/ssh_server.py,sha256=f2T8fgwACVljPfdcimMywUjsFnLCWRde7iWPAILpRz8,5463
+addftool/deploy/ssh_server.py,sha256=7glpJJNROskpqPkeYrTc2MbVzRendUZLv-ZgPs6HCq8,5641
 addftool/deploy/vscode_server.py,sha256=tLtSvlcK2fEOaw6udWt8dNELVhwv9F59hF5DJJ-1Nak,2666
 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.3.dist-info/METADATA,sha256=LjfdP_1jJ0U2rBMd0DXdoDSmzfXtOpfrf17XX_LceT8,170
-addftool-0.2.3.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-addftool-0.2.3.dist-info/entry_points.txt,sha256=9lkmuWMInwUAtev8w8poNkNd7iML9Bjd5CBCFVxg2b8,111
-addftool-0.2.3.dist-info/top_level.txt,sha256=jqj56-plrBbyzY0tIxB6wPzjAA8kte4hUlajyyQygN4,9
-addftool-0.2.3.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.3.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