addftool 0.1.9__py3-none-any.whl → 0.2.0__py3-none-any.whl

addftool/addf_portal.py

@@ -1,7 +1,5 @@
 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():
@@ -9,8 +7,6 @@ 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()
 
@@ -19,10 +15,6 @@ 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/broadcast_folder.py

@@ -0,0 +1,283 @@
+import os
+import time
+import subprocess
+import hashlib
+
+from fabric import Connection, ThreadingGroup
+
+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_broadcast_folder_args(subparsers):
+    deploy_parser = subparsers.add_parser('broadcast-folder', help='broadcast folder from master node to other nodes')
+    add_args(deploy_parser)
+
+
+def add_args(parser):
+    parser.add_argument("--tool", help="tool name", type=str, default="torch_nccl", choices=["torch_nccl"])
+    parser.add_argument("--hostfile", help="host file, broadcast file from node-0 to others", type=str, default="")
+
+    parser.add_argument("--download_timeout", help="download timeout, default is 30 min", type=int, default=60 * 30)
+
+    parser.add_argument("--md5_verify", action='store_true', default=False,
+                        help="whether to verify the md5 of the file after broadcast, 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.")
+    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("--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, " \
+                        "If empty, will transfer the files from the node-0's local folder.")
+    
+    # distributed downloader from blob
+    parser.add_argument("folder", help="the folder need to broadcast", 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)
+        print(f"Running command on {self.original_host}: {self.command}")
+        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 broadcast_folder_main(args):
+    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)
+    
+    print(f"Find {len(host_list)} hosts in hostfile: {args.hostfile}")
+    connection_list = []
+
+    remote_temp_config_dir = "/tmp/broadcast_temp_config_dir"
+
+    master_addr = get_ip_via_ssh(host_list[0])
+
+    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 "
+        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.from_blob_url and i == 0:
+            commnads += f" --from_blob_url {args.from_blob_url}"
+        if args.md5_verify:
+            commnads += " --md5_verify"
+        
+        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):
+    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 broadcast_file_from_rank(rank, file_path, from_rank, device, file_size, max_chunk_size=250 * 1024 * 1024, md5_verify=False, group=None):
+    if file_size == 0:
+        if rank != from_rank:
+            with open(file_path, "wb") as f:
+                f.write(b"")
+        return
+
+    if group is None:
+        group = dist.group.WORLD
+    pinned_cpu_tensor = torch.empty(min(max_chunk_size, file_size), dtype=torch.uint8).pin_memory()
+
+    for offset in range(0, file_size, max_chunk_size):
+        read_size = min(max_chunk_size, file_size - offset)
+        if rank == from_rank:
+            with open(file_path, "rb") as f:
+                f.seek(offset)
+                data = f.read(read_size)
+            tensor = torch.frombuffer(data, dtype=torch.uint8).contiguous().pin_memory()
+            tensor = tensor.to(device)
+            torch.cuda.synchronize()
+            dist.broadcast(tensor, src=from_rank, group=group, async_op=True)
+        else:
+            tensor = torch.empty(read_size, dtype=torch.uint8, device=device)
+            dist.broadcast(tensor, src=from_rank, group=group)
+            torch.cuda.synchronize()
+
+            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, "ab" if offset > 0 else "wb") as f:
+                pinned_cpu_tensor[:read_size].copy_(tensor)
+                np_array = pinned_cpu_tensor.numpy()[:read_size]
+                np_array.tofile(f)
+    
+    if md5_verify:
+        if rank == from_rank:
+            with open(file_path, "rb") as f:
+                file_md5 = hashlib.md5(f.read()).hexdigest()
+            md5_tensor = torch.frombuffer(file_md5.encode('utf-8'), dtype=torch.uint8).to(device)
+            dist.broadcast(md5_tensor, src=from_rank, group=group)
+        else:
+            with open(file_path, "rb") as f:
+                file_md5 = hashlib.md5(f.read()).hexdigest()
+            src_md5_tensor = torch.empty(32, dtype=torch.uint8, device=device)
+            dist.broadcast(src_md5_tensor, src=from_rank, group=group)
+
+            src_md5_str = src_md5_tensor.cpu().numpy().tobytes().decode('utf-8')
+            if file_md5 != src_md5_str:
+                raise ValueError(f"MD5 verification failed for {file_path}: {file_md5} != {src_md5_str}")
+
+
+def broadcast_folder_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')
+    global_rank = int(os.environ['RANK'])
+    local_rank = int(os.environ['LOCAL_RANK'])
+    world_size = int(os.environ['WORLD_SIZE'])
+    num_nodes = world_size // args.transfer_ranks_per_node
+    worker_rank = local_rank
+
+    num_gpus = torch.cuda.device_count()
+    torch.cuda.set_device(worker_rank % num_gpus)
+
+    device = torch.device(f"cuda:{worker_rank % num_gpus}")
+    
+    workers_groups = []
+    for i in range(args.transfer_ranks_per_node):
+        worker_ranks = []
+        for j in range(num_nodes):
+            worker_ranks.append(j * args.transfer_ranks_per_node + i)
+        
+        if global_rank == 0:
+            print(f"worker group {i} ranks: {worker_ranks}")
+        
+        workers_groups.append(dist.new_group(worker_ranks))
+
+    if global_rank == 0:
+        print(f"Init {len(workers_groups)} worker groups")
+
+    print(f"rank {global_rank} start broadcast worker, args = {args}, nccl init time: {time.time() - start_time:.2f}s")
+    
+    file_size_dict = {}
+    
+    if global_rank == 0:
+        if args.from_blob_url:
+            raise NotImplementedError("Downloading files from blob storage is not implemented yet.")
+
+        if not os.path.exists(args.folder):
+            raise ValueError(f"Folder {args.folder} does not exist.")
+        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
+        
+        # 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]        
+        obj_list = [file_list, file_size_list]
+        dist.broadcast_object_list(obj_list, src=0)
+    else:
+        obj_list = [None, None]
+        dist.broadcast_object_list(obj_list, src=0)
+        file_list, file_size_list = obj_list
+    
+    print(f"Rank {global_rank}: {len(file_list)} files, total size: {sum(file_size_list) / (1024 * 1024):.2f} MB, time taken: {time.time() - start_time:.2f}s")
+    
+    worker_g = workers_groups[worker_rank]
+    from_rank = global_rank % args.transfer_ranks_per_node
+    for i in range(len(file_list)):
+        if i % args.transfer_ranks_per_node == worker_rank:
+            file_path = file_list[i]
+            file_size = file_size_list[i]
+            broadcast_file_from_rank(
+                global_rank, file_path, from_rank, device,
+                file_size, md5_verify=args.md5_verify, group=worker_g,
+            )
+            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")
+    
+    dist.barrier()
+    for i in range(len(workers_groups)):
+        if i != worker_rank:
+            dist.destroy_process_group(workers_groups[i])
+    destroy_process_group()
+
+    print(f"Rank {global_rank} finished broadcasting 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:
+        broadcast_folder_main(args)
+    else:
+        broadcast_folder_worker(args)

addftool/deploy/__init__.py

@@ -2,6 +2,7 @@ import argparse
 
 from .ssh_server import add_deploy_ssh_server_args, deploy_ssh_server_main
 from .azure import add_deploy_azure_args, deploy_azure_main
+from .vscode_server import add_deploy_vscode_server_args, deploy_vscode_server_main
 
 
 def add_deploy_args(subparsers: argparse._SubParsersAction):
@@ -19,6 +20,9 @@ def add_deploy_args(subparsers: argparse._SubParsersAction):
     azure_parser = deploy_subparsers.add_parser('azure', help='Azure options')
     add_deploy_azure_args(azure_parser)
 
+    vscode_server_parser = deploy_subparsers.add_parser('vscode-server', help='VS Code options')
+    add_deploy_vscode_server_args(vscode_server_parser)
+
 
 def deploy_main(args):
     """
@@ -30,6 +34,8 @@ def deploy_main(args):
         deploy_ssh_server_main(args)
     elif args.deploy_type == 'azure':
         deploy_azure_main(args)
+    elif args.deploy_type == 'vscode-server':
+        deploy_vscode_server_main(args)
     else:
         print("Unknown deploy type")
         print("Deploy type: ", args.deploy_type)

addftool/deploy/vscode_server.py

@@ -0,0 +1,74 @@
+from addftool.util import need_sudo, execute_command, install_packages
+
+
+def deploy_vscode_server(install_tunnel=True):
+    """
+    Deploy Visual Studio Code and optionally start VS Code Tunnel.
+    
+    Args:
+        install_tunnel (bool): Whether to start VS Code Tunnel after installation
+    """
+    command_prefix = "sudo " if need_sudo() else ""
+    
+    # Install prerequisites
+    print("Installing prerequisites...")
+    execute_command(command_prefix + "apt-get install -y wget gpg")
+    
+    # Add Microsoft's GPG key
+    print("Adding Microsoft's GPG key...")
+    execute_command("wget -qO- https://packages.microsoft.com/keys/microsoft.asc | gpg --dearmor > packages.microsoft.gpg")
+    execute_command(command_prefix + "install -D -o root -g root -m 644 packages.microsoft.gpg /etc/apt/keyrings/packages.microsoft.gpg")
+    
+    # Add VS Code repository
+    print("Adding VS Code repository...")
+    execute_command(command_prefix + "sh -c 'echo \"deb [arch=amd64,arm64,armhf signed-by=/etc/apt/keyrings/packages.microsoft.gpg] "
+                   "https://packages.microsoft.com/repos/code stable main\" > /etc/apt/sources.list.d/vscode.list'")
+    
+    # Clean up
+    execute_command("rm -f packages.microsoft.gpg")
+    
+    # Install VS Code
+    print("Installing VS Code...")
+    execute_command(command_prefix + "apt install -y apt-transport-https")
+    execute_command(command_prefix + "apt update")
+    execute_command(command_prefix + "apt install -y code")
+    
+    # Start VS Code Tunnel if requested
+    if install_tunnel:
+        print("Starting VS Code Tunnel...")
+        execute_command("code tunnel --accept-server-license-terms")
+    
+    print("VS Code installation completed successfully!")
+
+
+def add_deploy_vscode_server_args(parser):
+    """
+    Add VS Code deployment arguments to the parser.
+    
+    Args:
+        parser (argparse.ArgumentParser): The argument parser to add the arguments to.
+    """
+    parser.add_argument("--no-tunnel", dest="install_tunnel", action="store_false", 
+                        help="Skip starting VS Code Tunnel after installation")
+    parser.set_defaults(install_tunnel=True)
+
+
+def deploy_vscode_server_main(args):
+    """
+    Main function for VS Code deployment.
+    
+    Args:
+        args: Parsed command-line arguments
+    """
+    deploy_vscode_server(install_tunnel=args.install_tunnel)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="VS Code deployment arguments")
+    add_deploy_vscode_server_args(parser)
+
+    args = parser.parse_args()
+
+    deploy_vscode_server_main(args)

addftool/util.py

@@ -1,75 +1,12 @@
-import os
 import subprocess
 
 
-def execute_command(command, to_file=None, only_stdout=True, hide=False):
+def execute_command(command, to_file=None):
     if to_file is not None:
         to_file.write(command + "\n")
         return None
     else:
-        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)
+        print("Execute command: ", command)
+        result = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE).stdout.read().decode()
+        print("Result: ", result)
+        return result

{addftool-0.1.9.dist-info → addftool-0.2.0.dist-info}/METADATA

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

{addftool-0.1.9.dist-info → addftool-0.2.0.dist-info}/RECORD

@@ -1,17 +1,19 @@
 addftool/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-addftool/addf_portal.py,sha256=w2LgsoutfnrKhtrQAXouUMwLqnsp5ALlsBYUWg8n9NM,781
+addftool/addf_portal.py,sha256=U52TdNcwWOEvv_C5r-guWYxn3ntzwYI2eBzJIE7IdcY,493
 addftool/blob.py,sha256=NZOItDyFUIdV1tfhJZJJBEzGy296CE5NCictTzP4OPc,8282
+addftool/broadcast_folder.py,sha256=bEOr-8Q14DhZFS658hab4U-9HvAW3EGluEriilIdEXQ,11976
 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/util.py,sha256=Q3A68vJDxgfeNiEFmk54HuMuworVndocXpSbVpvGMfc,362
+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/vscode_server.py,sha256=9KoWQdK3ESgKtdACEHTng2OuQKlIZjh_R_qJNT-XOv8,2609
 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.1.9.dist-info/METADATA,sha256=VEfXNesNw0S0lZl2_hYrJhrUyRN9SpP1iT1aIhtgKbU,170
-addftool-0.1.9.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-addftool-0.1.9.dist-info/entry_points.txt,sha256=9lkmuWMInwUAtev8w8poNkNd7iML9Bjd5CBCFVxg2b8,111
-addftool-0.1.9.dist-info/top_level.txt,sha256=jqj56-plrBbyzY0tIxB6wPzjAA8kte4hUlajyyQygN4,9
-addftool-0.1.9.dist-info/RECORD,,
+addftool-0.2.0.dist-info/METADATA,sha256=Ja75t6IT9XdIMjCUsZKLsx-QJJ9Yo2UfYtpVUC1AYPw,170
+addftool-0.2.0.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+addftool-0.2.0.dist-info/entry_points.txt,sha256=9lkmuWMInwUAtev8w8poNkNd7iML9Bjd5CBCFVxg2b8,111
+addftool-0.2.0.dist-info/top_level.txt,sha256=jqj56-plrBbyzY0tIxB6wPzjAA8kte4hUlajyyQygN4,9
+addftool-0.2.0.dist-info/RECORD,,