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

addftool/addf_portal.py

@@ -1,5 +1,7 @@
 import argparse
 from addftool.process import add_killer_args, killer_main
+from addftool.sync import add_sync_args, sync_main
+from addftool.deploy import add_deploy_args, deploy_main
 
 
 def get_args():
@@ -7,6 +9,8 @@ def get_args():
 
     subparsers = parser.add_subparsers(dest='command', help='Sub-command help')
     add_killer_args(subparsers)
+    add_sync_args(subparsers)
+    add_deploy_args(subparsers)
 
     return parser.parse_args()
 
@@ -15,6 +19,10 @@ def main():
     args = get_args()
     if args.command == "kill":
         killer_main(args)
+    elif args.command == "sync":
+        sync_main(args)
+    elif args.command == "deploy":
+        deploy_main(args)
     else:
         print("Unknown command: ", args.command)
 

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/process/__init__.py

@@ -52,6 +52,7 @@ def kill_process(processes, timeout=5, try_count=3):
         if all_terminated:
             print("All processes terminated")
             break
+        time.sleep(timeout)
     
     if not all_terminated:
         print("Some processes are still running. Please check manually.")

addftool/sync.py

@@ -0,0 +1,301 @@
+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)
+    # 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)
+
+    groups = [[] for _ in range(args.donwload_nodes)]
+    for i, (file_name, file_size) in enumerate(sorted_files):
+        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}")
+    
+    for i, group in enumerate(groups):
+        group_file_path = os.path.join(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(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 = []
+    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")))
+        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")))
+
+        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}"
+        if args.md5_verify:
+            commnads += " --md5_verify"
+        if i < args.donwload_nodes:
+            commnads += f" --download_index_file {temp_config_dir}/node_{i}.txt"
+        
+        connection_list.append(ConnectionWithCommand(host, 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)
+        os.makedirs(os.path.dirname(file_path), 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:
+        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)
+
+    dist.barrier()
+    download_process.join()
+    destroy_process_group()
+    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/util.py

@@ -1,12 +1,75 @@
+import os
 import subprocess
 
 
-def execute_command(command, to_file=None):
+def execute_command(command, to_file=None, only_stdout=True, hide=False):
     if to_file is not None:
         to_file.write(command + "\n")
         return None
     else:
-        print("Execute command: ", command)
-        result = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE).stdout.read().decode()
-        print("Result: ", result)
-        return result
+        if not hide:
+            print("Execute command: ", command)
+        result = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE)
+        result.wait()
+        print(f"Return code: {result.returncode}")
+        if result.stdout is not None:
+            stdout = result.stdout.read().decode()
+            print(f"Stdout: {stdout}")
+        else:
+            stdout = None
+        if only_stdout:
+            if not hide and stdout is not None:
+                print(stdout)
+            return stdout
+        if result.stderr is not None:
+            stderr = result.stderr.read().decode()
+            print(f"Stderr: {stderr}")
+        else:
+            stderr = None
+
+        return {'stdout': stdout, 'stderr': stderr, 'returncode': result.returncode}
+
+
+def need_sudo():
+    return os.name == 'posix' and os.getuid() != 0
+
+
+def is_running_in_docker():
+    return os.path.exists('/.dockerenv') or \
+           any('docker' in line for line in open('/proc/self/cgroup', 'r')) if os.path.exists('/proc/self/cgroup') else False or \
+           os.environ.get('container') == 'docker' or \
+           os.environ.get('DOCKER') == 'true' or \
+           os.environ.get('DOCKER_CONTAINER') == 'yes'
+
+
+def get_ubuntu_version():  
+    with open("/etc/os-release") as f:
+        for line in f:
+            if line.startswith("VERSION_ID="):
+                version = line.split("=")[1].strip().strip('"')
+                return version
+    return "22.04"
+
+
+def check_package_installed(package):
+    command = f"dpkg -l | grep {package}"
+    result = execute_command(command)
+    if result is not None and package in result:
+        return True
+    return False
+
+
+def install_packages(package_list):
+    to_install = []
+    for package in package_list:
+        if check_package_installed(package):
+            print(f"{package} is already installed")
+            continue
+        to_install.append(package)
+    
+    if len(to_install) > 0:
+        packages = " ".join(to_install)
+        command = f"apt-get install -y {packages}"
+        if need_sudo():
+            command = "sudo " + command
+        execute_command(command)

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

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

addftool-0.1.7.dist-info/RECORD

@@ -0,0 +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/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/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/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)

addftool-0.1.5.dist-info/RECORD

@@ -1,17 +0,0 @@
-addftool/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-addftool/addf_portal.py,sha256=U52TdNcwWOEvv_C5r-guWYxn3ntzwYI2eBzJIE7IdcY,493
-addftool/blob.py,sha256=NZOItDyFUIdV1tfhJZJJBEzGy296CE5NCictTzP4OPc,8282
-addftool/tool.py,sha256=EuKQ2t2InN7yB-_oYLcdsA7vRqzRGTunwIxplUSqEG0,2054
-addftool/util.py,sha256=Q3A68vJDxgfeNiEFmk54HuMuworVndocXpSbVpvGMfc,362
-addftool/deploy/__init__.py,sha256=tpyoTh3SqAQojPizsJDvQohu1Pcb3-w-DP5sO4-5lBM,1220
-addftool/deploy/azure.py,sha256=UQR1hOEYUtsm2fbWBczsnEB_mh7yUuN2NDv3sgMMsac,1246
-addftool/deploy/ssh_server.py,sha256=f2T8fgwACVljPfdcimMywUjsFnLCWRde7iWPAILpRz8,5463
-addftool/process/__init__.py,sha256=OB-cZXP1jK7l8uN8nKfhg_bCX6Slz6DeeBpEodt-IK4,3515
-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.5.dist-info/METADATA,sha256=8fCVZ1r4rS685bprhRuQE5Le7MV3gu8dbNRkvFfo05w,148
-addftool-0.1.5.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-addftool-0.1.5.dist-info/entry_points.txt,sha256=9lkmuWMInwUAtev8w8poNkNd7iML9Bjd5CBCFVxg2b8,111
-addftool-0.1.5.dist-info/top_level.txt,sha256=jqj56-plrBbyzY0tIxB6wPzjAA8kte4hUlajyyQygN4,9
-addftool-0.1.5.dist-info/RECORD,,