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

addftool/blob.py

@@ -101,12 +101,13 @@ def install_main(args):
 def mount_main(args):
     sas_token = get_token(args, info=True)
 
+    cache_gate = 0 if args.no_cache else 1
     template = {
         'logging': {'type': 'silent', 'level': 'log_off'}, 
         'components': ['libfuse', 'file_cache', 'attr_cache', 'azstorage'], 
-        'libfuse': {'attribute-expiration-sec': 120, 'entry-expiration-sec': 120, 'negative-entry-expiration-sec': 240}, 
-        'file_cache': {'path': '', 'timeout-sec': 120, 'max-size-mb': 1024 * int(args.file_cache_size)}, 
-        'attr_cache': {'timeout-sec': 7200}, 
+        'libfuse': {'attribute-expiration-sec': 120 * cache_gate, 'entry-expiration-sec': 120 * cache_gate, 'negative-entry-expiration-sec': 240 * cache_gate}, 
+        'file_cache': {'path': '', 'timeout-sec': 120 * cache_gate, 'max-size-mb': 1024 * int(args.file_cache_size)}, 
+        'attr_cache': {'timeout-sec': 7200 * cache_gate}, 
         'azstorage': {
             'type': 'block', 
             'endpoint': '', 
@@ -145,6 +146,8 @@ def mount_main(args):
     print("Create config file: ", temp_config)
 
     command = f"blobfuse2 mount {args.mount} --config-file={temp_config}"
+    if args.no_cache:
+        command += " -o direct_io"
     if args.allow_other:
         command += " --allow-other"
         # to avoid "Error: fusermount3: option allow_other only allowed if 'user_allow_other' is set in /etc/fuse.conf"
@@ -235,6 +238,7 @@ def add_args(parser):
     mount_parser.add_argument("-o", "--allow-other", help="allow other", action="store_true", default=False)
     mount_parser.add_argument("--file-cache-size", help="file cache size (GB)", default=256)
     mount_parser.add_argument("--sudo", help="sudo", action="store_true")
+    mount_parser.add_argument("--no_cache", help="no cache", action="store_true")
 
     token_parser = subparsers.add_parser('token', help='Token help')
     add_api(token_parser)

addftool/broadcast_folder.py

@@ -3,13 +3,20 @@ import time
 import fnmatch
 import subprocess
 import hashlib
+import sys
+import warnings
+try:
+    from addftool.ssh import get_client, get_ssh_config, handle_hosts_outputs
+    from pssh.clients import ParallelSSHClient
+except ImportError:
+    pass
 from concurrent.futures import ThreadPoolExecutor
-
-from fabric import Connection, ThreadingGroup
+import gevent
 
 try:
     import torch
     import torch.distributed as dist
+    import torch.multiprocessing as mp
     from torch.distributed import init_process_group, destroy_process_group
     _torch_is_available = True
 except ImportError:
@@ -29,9 +36,8 @@ def add_args(parser):
 
     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("--port", help="the port for torch, default is 29501", type=int, default=29501)
+    parser.add_argument("--python_alias", type=str, default="python")
     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.")
 
@@ -55,28 +61,13 @@ def add_args(parser):
                         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.")
+
+    parser.add_argument("--worker_args", type=str, default="")
     
     # 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"
@@ -150,20 +141,40 @@ def broadcast_folder_main(args):
                 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])
+    
+    user_ssh_config, systen_ssh_config = get_ssh_config()
+    not_successed_hosts = host_list[:]
+    for try_count in range(3):
+        client = get_client(not_successed_hosts, user_ssh_configs=user_ssh_config, system_ssh_configs=systen_ssh_config)
+        mkdir_cmds = client.run_command(f"mkdir -p {remote_temp_config_dir}", stop_on_errors=False)
+        client.join(mkdir_cmds, timeout=10)
+        assert isinstance(client, ParallelSSHClient), "Failed to create ParallelSSHClient"
+        cmds = client.scp_send(__file__, os.path.join(remote_temp_config_dir, "broadcast_folder.py"))
+        gevent.joinall(cmds, raise_error=False, timeout=10)
+        resend_hosts = []
+        for host, cmd in zip(not_successed_hosts, cmds):
+            if cmd.exception:
+                print(f"Failed to copy file to {host}: {cmd.exception}")
+                resend_hosts.append(host)
+        
+        if len(resend_hosts) == 0:
+            break
+        not_successed_hosts = resend_hosts
+    
+    if len(resend_hosts) > 0:
+        print(f"Failed to copy file to {len(resend_hosts)} hosts: {resend_hosts}")
+        sys.exit(1)
 
+    host_commands = []
     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 OMP_NUM_THREADS=1 "
+        commnads = "PYTHONUNBUFFERED=1 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} "
+        commnads += f"{args.python_alias} {remote_temp_config_dir}/broadcast_folder.py {args.folder} --tool {args.tool} --transfer_ranks_per_node {args.transfer_ranks_per_node} "
+        commnads += f" --worker_args {master_addr}_{args.port}_{len(host_list)}_{args.transfer_ranks_per_node}_{i} "
         if args.contain_md5_files:
             commnads += " --contain_md5_files"
         if args.include_string:
@@ -174,11 +185,20 @@ def broadcast_folder_main(args):
             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))
+        host_commands.append(commnads)
+        print(f"Run command on {host}: {commnads}")
+
+    client = get_client(host_list, user_ssh_configs=user_ssh_config, system_ssh_configs=systen_ssh_config)
+    if True:
+        cmds = client.run_command("%s", host_args=host_commands, stop_on_errors=False)
+        handle_hosts_outputs(cmds)
+        for host, cmd in zip(host_list, cmds):
+            if cmd.exception or cmd.exit_code != 0:
+                print(f"Failed to run command on {host}: {cmd.exception}, EXIT CODE: {cmd.exit_code}")
+                sys.exit(1)
+            print(f"Command on {host} finished with exit code {cmd.exit_code}")
     
-    group = ThreadingGroup.from_connections(connection_list)
-    group.run('echo "Hello"', hide=False)
+    print(f"All nodes finished broadcasting files")
 
 
 def download_files_from_blob(queue, blob_url, sas_token, folder, download_files, node_rank):
@@ -259,16 +279,14 @@ def broadcast_file_from_rank(rank, file_path, from_rank, device, file_size, max_
                 raise ValueError(f"MD5 verification failed for {file_path}: {file_md5} != {src_md5_str}")
 
 
-def broadcast_folder_worker(args):
+def broadcast_folder_worker(local_rank, node_rank, world_size, master_addr, master_port, 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'])
+    global_rank = local_rank + node_rank * args.transfer_ranks_per_node
+    init_process_group(backend='nccl', init_method=f"tcp://{master_addr}:{master_port}", rank=global_rank, world_size=world_size)
     num_nodes = world_size // args.transfer_ranks_per_node
     worker_rank = local_rank
 
@@ -291,11 +309,15 @@ def broadcast_folder_worker(args):
     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")
-    
+    if global_rank == 0:
+        print(f"rank {global_rank} start broadcast worker, args = {args}, nccl init time: {time.time() - start_time:.2f}s")
+    else:
+        print(f"rank {global_rank} nccl init time: {time.time() - start_time:.2f}s")
+
     file_size_dict = {}
     
     if global_rank == 0:
+        warnings.filterwarnings("ignore", category=UserWarning, message="The given buffer is not writable, and PyTorch does not support non-writable tensors.")
         # 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()]
@@ -397,6 +419,23 @@ def broadcast_folder_worker(args):
     print(f"Rank {global_rank} finished broadcasting all files, time taken: {time.time() - start_time:.2f}s")
 
 
+def broadcast_node_main(args):
+    if not _torch_is_available:
+        raise ImportError("Torch is not available. Please install torch to use this feature.")
+    parts = args.worker_args.split("_")
+    master_addr = parts[0]
+    master_port = int(parts[1])
+    num_nodes = int(parts[2])
+    num_ranks_per_node = int(parts[3])
+    node_rank = int(parts[4])
+    world_size = num_nodes * num_ranks_per_node
+    
+    mp.spawn(
+        broadcast_folder_worker, nprocs=num_ranks_per_node, join=True, 
+        args=(node_rank, world_size, master_addr, master_port, args), 
+    )
+
+
 if __name__ == "__main__":
     import argparse
 
@@ -406,4 +445,4 @@ if __name__ == "__main__":
     if args.hostfile:
         broadcast_folder_main(args)
     else:
-        broadcast_folder_worker(args)
+        broadcast_node_main(args)

addftool/ssh/__init__.py

@@ -0,0 +1,128 @@
+import re
+import os
+from pssh.clients import ParallelSSHClient
+from pssh.config import HostConfig
+import gevent
+
+
+def get_host_config(hostname, configs):
+    """Get configuration for a specific host from all configs."""
+    host_config = {}
+    
+    # Check for exact hostname match
+    if hostname in configs:
+        host_config.update(configs[hostname])
+    
+    # Check for wildcard matches
+    for pattern, config in configs.items():
+        if '*' in pattern or '?' in pattern:
+            # Convert SSH glob pattern to regex pattern
+            regex_pattern = pattern.replace('.', '\\.').replace('*', '.*').replace('?', '.')
+            if re.match(f"^{regex_pattern}$", hostname):
+                host_config.update(config)
+    
+    return host_config
+
+
+def parse_ssh_config_file(file_path):
+    """Parse SSH config file into a dictionary of host configurations."""
+    host_configs = {}
+    current_host = None
+    
+    if not os.path.exists(file_path):
+        return host_configs
+
+    try:
+        with open(file_path, 'r') as f:
+            for line in f:
+                line = line.strip()
+                if not line or line.startswith('#'):
+                    continue
+                    
+                if line.lower().startswith('host ') and not line.lower().startswith('host *'):
+                    hosts = line.split()[1:]
+                    for host in hosts:
+                        current_host = host
+                        if current_host not in host_configs:
+                            host_configs[current_host] = {}
+                elif current_host and ' ' in line:
+                    key, value = line.split(None, 1)
+                    host_configs[current_host][key.lower()] = value
+    except Exception as e:
+        print(f"Error reading {file_path}: {e}")
+    
+    return host_configs
+
+
+def get_ssh_config():
+    user_config_path = os.path.expanduser("~/.ssh/config")
+    system_config_path = "/etc/ssh/ssh_config"
+    
+    user_configs = parse_ssh_config_file(user_config_path)
+    system_configs = parse_ssh_config_file(system_config_path)
+    return user_configs, system_configs
+
+
+def get_client(hosts, user_ssh_configs=None, system_ssh_configs=None):
+    if user_ssh_configs is None or system_ssh_configs is None:
+        user_ssh_configs, system_ssh_configs = get_ssh_config()
+
+    to_connect = []
+    for hostname in hosts:
+        host_ssh_config = get_host_config(hostname, user_ssh_configs)
+        if not host_ssh_config:
+            host_ssh_config = get_host_config(hostname, system_ssh_configs)
+        
+        if host_ssh_config:
+            host_config = HostConfig(
+                user=host_ssh_config.get('user', os.getenv('USER', 'root')),
+                port=int(host_ssh_config.get('port', 22)),
+                private_key=host_ssh_config.get('identityfile', None),
+            )
+            host_name = host_ssh_config.get('hostname', hostname)
+            to_connect.append((host_name, host_config))
+        else:
+            print(f"No config found for host {hostname}")
+
+    # Create a ParallelSSHClient with the list of hosts
+    client = ParallelSSHClient([host[0] for host in to_connect], host_config=[host[1] for host in to_connect])
+    return client
+
+
+def handle_stream(host, stream_in, stream_name, print_call=None):
+    if print_call is None:
+        print_call = print
+    try:
+        if stream_in:
+            for line in stream_in:
+                prefix = " ERROR" if stream_name == "stderr" else ""
+                print_call(f"[{host}]{prefix}: {line}")
+                gevent.sleep(0)  # 让出控制权
+    except Exception as e:
+        print(f"[{host}] {stream_name} Exception: {e}")
+
+
+def _stdout_log(line):
+    """Log the output line."""
+    print(line)
+
+
+def _stderr_log(line):
+    """Log the error line."""
+    print("ERROR: " + line)
+
+
+def handle_hosts_outputs(hosts_outputs, out_log=None, err_log=None):
+    """Handle the outputs from the SSH command execution."""
+    if out_log is None:
+        out_log = _stdout_log
+    if err_log is None:
+        err_log = _stderr_log
+    jobs = []
+    for output in hosts_outputs:
+        host_name = output.host
+        if output:
+            jobs.append(gevent.spawn(handle_stream, host_name, output.stdout, "stdout", out_log))
+            jobs.append(gevent.spawn(handle_stream, host_name, output.stderr, "stderr", err_log))
+
+    gevent.joinall(jobs, raise_error=False)

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

@@ -1,8 +1,10 @@
 Metadata-Version: 2.4
 Name: addftool
-Version: 0.2.5
+Version: 0.2.7
 Requires-Dist: cryptography
 Requires-Dist: requests
 Requires-Dist: PyYAML
 Requires-Dist: psutil
 Requires-Dist: fabric
+Requires-Dist: gevent
+Requires-Dist: parallel-ssh

{addftool-0.2.5.dist-info → addftool-0.2.7.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=X9tvvMT7cCDbfKqE7hUGaFbBOHIecIgD8xcT34Bqb_8,17708
+addftool/blob.py,sha256=y1HZaDBUNeXicVytvwpRXwufvvrgxR33ruBlYpxnSa4,9453
+addftool/broadcast_folder.py,sha256=GQBuSL8Ch537V_fSBHesWyqT3KRYry68pbYOKy2bDj4,19619
 addftool/sync.py,sha256=ZpYxbM8uiPFrV7ODmOaM7asVPCWaxBixA-arVc-1kfs,14045
 addftool/tool.py,sha256=FmxRY3-pP0_Z0zCUAngjmEMmPUruMftg_iUlB1t2TnQ,2001
 addftool/util.py,sha256=zlNLu8Be8cGIpNRqBw8_0q7nFxWlsJ9cToN62ohjdXE,2335
@@ -11,9 +11,9 @@ addftool/deploy/ssh_server.py,sha256=7glpJJNROskpqPkeYrTc2MbVzRendUZLv-ZgPs6HCq8
 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.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/ssh/__init__.py,sha256=h5_rCO0A6q2Yw9vFguQZZp_ApAJsT1dcnKnbKKZ0cDM,4409
+addftool-0.2.7.dist-info/METADATA,sha256=7wElFYgZp3OX387bAsiQTwzfrc4pHH8dlK6vPnDULWU,220
+addftool-0.2.7.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
+addftool-0.2.7.dist-info/entry_points.txt,sha256=9lkmuWMInwUAtev8w8poNkNd7iML9Bjd5CBCFVxg2b8,111
+addftool-0.2.7.dist-info/top_level.txt,sha256=jqj56-plrBbyzY0tIxB6wPzjAA8kte4hUlajyyQygN4,9
+addftool-0.2.7.dist-info/RECORD,,

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

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