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

addftool/deploy/azure.py

@@ -11,6 +11,7 @@ def deploy_azure(packages):
     command_prefix = "sudo " if need_sudo() else ""
     command = "dpkg -i /tmp/packages-microsoft-prod.deb"
     execute_command(command_prefix + command)
+    execute_command(command_prefix + "apt-get update")
 
     install_packages(packages)
 

addftool/sync.py

@@ -0,0 +1,341 @@
+import os
+import sys
+import time
+import subprocess
+import hashlib
+import tempfile
+from multiprocessing import Queue, Process
+
+try:
+    from fabric import Connection, ThreadingGroup
+except ImportError:
+    Connection = object
+
+try:
+    import torch
+    import torch.distributed as dist
+    from torch.distributed import init_process_group, destroy_process_group
+    _torch_is_available = True
+except ImportError:
+    _torch_is_available = False
+
+
+def add_sync_args(subparsers):
+    deploy_parser = subparsers.add_parser('sync', help='download and sync folder from master node to other nodes')
+    add_args(deploy_parser)
+
+
+def add_args(parser):
+    parser.add_argument("--from_blob_url", help="download from blob url to master node before sync", type=str, default="")
+    parser.add_argument("--tool", help="tool name", type=str, default="torch_nccl", choices=["torch_nccl"])
+    parser.add_argument("--hostfile", help="host file, sync file from node-0 to others", type=str, default="")
+
+    parser.add_argument("--download_index_file", type=str, default="", 
+                        help="the file to save the download index, should be generated by master node.")
+    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)
+
+
+class ConnectionWithCommand(Connection):
+    def __init__(self, host, temp_config_dir, puts, command):
+        super().__init__(host)
+        self.command = command
+        self.puts = puts
+        self.temp_config_dir = temp_config_dir
+
+    def run(self, command, **kwargs):
+        super().run(f"mkdir -p {self.temp_config_dir}", **kwargs)
+        for src, dest in self.puts:
+            self.put(src, remote=dest)
+        super().run(self.command, **kwargs)
+        if command:
+            super().run(command, **kwargs)
+
+
+def get_ip_via_ssh(hostname):
+    if hostname == "localhost":
+        return "127.0.0.1"
+    try:
+        cmd = ["ssh", hostname, "hostname -I | awk '{print $1}'"]
+        result = subprocess.run(cmd, capture_output=True, text=True, timeout=5)
+        
+        if result.returncode == 0:
+            ip = result.stdout.strip()
+            return ip
+        else:
+            print(f"SSH {hostname} failed: {result.stderr}")
+            return None
+    except Exception as e:
+        print(f"Error executing SSH command on {hostname}: {e}")
+        return None
+
+
+def sync_main(args):
+    sas_token = os.environ.get("SAS_TOKEN")
+    if not sas_token:
+        raise ValueError("SAS_TOKEN environment variable is not set.")
+    
+    try:
+        list_operation = subprocess.run(
+            ["azcopy", "list", args.from_blob_url + sas_token, "--machine-readable"], 
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True
+        )
+        if list_operation.returncode != 0:
+            raise RuntimeError(f"Failed to list blob: {list_operation.stderr}")
+    except Exception as e:
+        raise RuntimeError(f"Error while listing blob: {e}")
+    
+    file_size_list = {}
+    for line in list_operation.stdout.splitlines():
+        # print(line)
+        parts = line.split("; Content Length:")
+        if len(parts) != 2:
+            print(f"INFO: {line}")
+            continue
+        file_name = parts[0].strip()
+        file_size = int(parts[1])
+        file_size_list[file_name] = file_size
+
+    # 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
+    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(local_temp_config_dir, f"node_{i}.txt")
+        total_size = 0
+        with open(group_file_path, "w") as f:
+            for file_name in group:
+                f.write(file_name + "\n")
+                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:
+            line = line.strip()
+            if line and not line.startswith("#"):
+                host_list.append(line)
+    
+    if len(host_list) < len(groups):
+        raise ValueError(f"Number of hosts in hostfile {len(host_list)} is less than number of download nodes {len(groups)}")
+    
+    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(remote_temp_config_dir, "sync.py")))
+        if i < args.donwload_nodes:
+            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" && {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 {remote_temp_config_dir}/node_{i}.txt"
+        
+        connection_list.append(ConnectionWithCommand(host, remote_temp_config_dir, put_commands, commnads))
+    
+    group = ThreadingGroup.from_connections(connection_list)
+    group.run('echo "Hello"', hide=False)
+
+
+def download_files_from_blob(queue, blob_url, sas_token, folder, download_files, node_rank):
+    # This function should implement the logic to download files from blob storage
+    # using the provided blob_url and sas_token. The downloaded files should be
+    # saved in the specified folder.
+    if not blob_url.endswith("/"):
+        blob_url += "/"
+    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)
+        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(
+                    ["azcopy", "copy", blob_url + file_name + sas_token, file_path],
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE,
+                    text=True
+                )
+                if download_status.returncode != 0:
+                    raise RuntimeError(f"Failed to download {file_name}: {download_status.stderr}")
+                print(f"Rank {node_rank}: Downloaded {file_name} successfully, from {blob_url} to {file_path}")
+                queue.put(file_path)
+                break
+            except Exception as e:
+                print(f"Rank {node_rank}: Download failed: {e}")
+
+
+def sync_file_from_rank(rank, file_path, from_rank, md5_verify=False):
+    if rank == from_rank:
+        with open(file_path, "rb") as f:
+            data = f.read()
+        num_bytes = len(data)
+        if md5_verify:
+            md5 = hashlib.md5()
+            md5.update(data)
+            md5_value = md5.hexdigest()
+        else:
+            md5_value = ""
+        obj_list = [file_path, num_bytes, md5_value]
+        dist.broadcast_object_list(obj_list, src=from_rank)
+        tensor = torch.frombuffer(data, dtype=torch.uint8)
+        tensor = tensor.cuda()
+    else:
+        obj_list = [0, "", ""]
+        dist.broadcast_object_list(obj_list, src=from_rank)
+        file_path, num_bytes, md5_value = obj_list
+        tensor = torch.empty(num_bytes, dtype=torch.uint8, device='cuda')
+    
+    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:
+            md5 = hashlib.md5()
+            md5.update(tensor.cpu().numpy())
+            md5_value_recv = md5.hexdigest()
+            if md5_value_recv != md5_value:
+                raise ValueError(f"MD5 mismatch for file {file_path}: {md5_value_recv} != {md5_value}")
+            else:
+                print(f"Node-{rank} verified file {file_path} with MD5: {md5_value_recv}")
+
+
+def sync_worker(args):
+    assert args.tool in ["torch_nccl"], f"tool {args.tool} is not supported"
+    if not _torch_is_available:
+        raise ImportError("Torch is not available. Please install torch to use this feature.")
+    start_time = time.time()
+
+    init_process_group(backend='nccl')
+    node_rank = int(os.environ['RANK'])
+    world_size = int(os.environ['WORLD_SIZE'])
+
+    print(f"rank {node_rank} start sync worker, args = {args}, nccl init time: {time.time() - start_time:.2f}s")
+
+    if world_size < 2:
+        raise ValueError("World size must be at least 2 for distributed download.")
+    
+    download_queue = Queue()
+
+    download_files = []
+    transfered_files = set()
+    if args.download_index_file:
+        with open(args.download_index_file, "r") as f:
+            for line in f:
+                download_files.append(line.strip())
+
+    download_process = Process(
+        target=download_files_from_blob, 
+        args=(download_queue, args.from_blob_url, os.environ["SAS_TOKEN"], args.folder, download_files, node_rank), 
+    )
+    download_process.start()
+
+    last_download = None
+
+    while True:
+        if len(download_files) == len(transfered_files):
+            status_code = world_size + 1
+        elif last_download is not None:
+            status_code = node_rank
+        else:
+            try:
+                last_download = download_queue.get(timeout=1)
+                status_code = node_rank
+            except Exception as e:
+                status_code = world_size
+        
+        global_status_code = torch.tensor(status_code).cuda()
+        dist.all_reduce(global_status_code, op=dist.ReduceOp.MIN)
+        global_status_code = global_status_code.item()
+
+        if global_status_code == world_size + 1:
+            print(f"Node-{node_rank} finished downloading all files, time taken: {time.time() - start_time:.2f}s")
+            break
+        elif global_status_code == world_size:
+            if node_rank == 0:
+                print(f"All nodes is waiting for other nodes to finish downloading...")
+            time.sleep(1)
+        elif global_status_code == node_rank:
+            print(f"Node-{node_rank} is downloaded {last_download}, prepare to broadcast it..., time taken: {time.time() - start_time:.2f}s")
+            sync_file_from_rank(node_rank, last_download, node_rank, md5_verify=args.md5_verify)
+            transfered_files.add(last_download)
+            last_download = None
+        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")
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Addf's tool")
+    add_args(parser)
+    args = parser.parse_args()
+    if args.hostfile:
+        sync_main(args)
+    else:
+        sync_worker(args)

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

@@ -1,7 +1,8 @@
 Metadata-Version: 2.4
 Name: addftool
-Version: 0.1.6
+Version: 0.1.8
 Requires-Dist: cryptography
 Requires-Dist: requests
 Requires-Dist: PyYAML
 Requires-Dist: psutil
+Requires-Dist: fabric

{addftool-0.1.6.dist-info → addftool-0.1.8.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=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=UQR1hOEYUtsm2fbWBczsnEB_mh7yUuN2NDv3sgMMsac,1246
+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/ssh/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-addftool/sync/__init__.py,sha256=wOqFCOA51rFUttBjOO44W3Fc66mhX5ir2R89lsO6gR0,1702
-addftool-0.1.6.dist-info/METADATA,sha256=SQSnBFDxmD836XQAgZDtqvpzT_Z1LewNKTRCfdMC5ag,148
-addftool-0.1.6.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-addftool-0.1.6.dist-info/entry_points.txt,sha256=9lkmuWMInwUAtev8w8poNkNd7iML9Bjd5CBCFVxg2b8,111
-addftool-0.1.6.dist-info/top_level.txt,sha256=jqj56-plrBbyzY0tIxB6wPzjAA8kte4hUlajyyQygN4,9
-addftool-0.1.6.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,,

addftool/sync/__init__.py

@@ -1,42 +0,0 @@
-import os
-
-
-def add_sync_args(subparsers):
-    process_killer_parser = subparsers.add_parser('sync', help='download and sync folder from master node to other nodes')
-
-    process_killer_parser.add_argument("--from_blob_url", help="download from blob url to master node before sync", type=str, default="")
-    process_killer_parser.add_argument("--sas_token", help="sas token for blob url", type=str, default="")
-    process_killer_parser.add_argument("--tool", help="tool name", type=str, default="torch_nccl", choices=["torch_nccl", "rsync"])
-    process_killer_parser.add_argument("--hostfile", help="host file, sync file from node-0 to others", type=str, default="")
-
-    # distributed downloader from blob
-    process_killer_parser.add_argument("--donwload_nodes", help="download nodes, default is node-0", type=int, default=1)
-
-    process_killer_parser.add_argument("folder", nargs='?', help="the folder need to sync", type=str, default="")
-
-
-def sync_main(args):
-    print(args)
-    exit(0)
-    if args.source == "" or args.target == "":
-        print("Please provide source and target folder")
-        return
-
-    # check if source is a folder
-    if not os.path.isdir(args.source):
-        print(f"Source {args.source} is not a folder")
-        return
-
-    # check if target is a folder
-    if not os.path.isdir(args.target):
-        print(f"Target {args.target} is not a folder")
-        return
-
-    # check if source and target are the same
-    if os.path.abspath(args.source) == os.path.abspath(args.target):
-        print(f"Source and target are the same")
-        return
-
-    # sync source to target
-    command = f"rsync -avz --delete {args.source} {args.target}"
-    os.system(command)