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

addftool/deploy/__init__.py

@@ -0,0 +1,35 @@
+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
+
+
+def add_deploy_args(subparsers: argparse._SubParsersAction):
+    """
+    Add deploy arguments to the parser.
+    Args:
+        subparsers (argparse._SubParsersAction): The subparsers object to add the arguments to.
+    """
+    deploy_parsers = subparsers.add_parser('deploy', help='deploy/install software, script, tools, etc.')
+    deploy_subparsers = deploy_parsers.add_subparsers(dest='deploy_type', help='Deploy options')
+
+    ssh_server_parser = deploy_subparsers.add_parser('ssh-server', help='SSH options')
+    add_deploy_ssh_server_args(ssh_server_parser)
+
+    azure_parser = deploy_subparsers.add_parser('azure', help='Azure options')
+    add_deploy_azure_args(azure_parser)
+
+
+def deploy_main(args):
+    """
+    Main function for deploy.
+    Args:
+        args (argparse.Namespace): The parsed arguments.
+    """
+    if args.deploy_type == 'ssh-server':
+        deploy_ssh_server_main(args)
+    elif args.deploy_type == 'azure':
+        deploy_azure_main(args)
+    else:
+        print("Unknown deploy type")
+        print("Deploy type: ", args.deploy_type)

addftool/deploy/azure.py

@@ -0,0 +1,39 @@
+from addftool.util import need_sudo, execute_command, get_ubuntu_version, install_packages
+
+
+def deploy_azure(packages):
+    ubuntu_version = get_ubuntu_version()
+    print("Get ubuntu version: ", ubuntu_version)
+
+    command = f"wget https://packages.microsoft.com/config/ubuntu/{ubuntu_version}/packages-microsoft-prod.deb -O /tmp/packages-microsoft-prod.deb"
+    print("Install packages-microsoft-prod.deb")
+    execute_command(command)
+    command_prefix = "sudo " if need_sudo() else ""
+    command = "dpkg -i /tmp/packages-microsoft-prod.deb"
+    execute_command(command_prefix + command)
+
+    install_packages(packages)
+
+
+def add_deploy_azure_args(parser):
+    """
+    Add Azure deployment arguments to the parser.
+    Args:
+        parser (argparse.ArgumentParser): The argument parser to add the arguments to.
+    """
+    parser.add_argument("packages", help="packages", default=['fuse3', 'blobfuse2', 'azcopy'], nargs="*")
+
+def deploy_azure_main(args):
+    # print(args.packages)
+    deploy_azure(args.packages)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Azure deployment arguments")
+    add_deploy_azure_args(parser)
+
+    args = parser.parse_args()
+
+    deploy_azure_main(args)

addftool/deploy/ssh_server.py

@@ -0,0 +1,137 @@
+import subprocess
+import argparse
+import os
+from addftool.util import need_sudo, execute_command, is_running_in_docker
+
+
+def convert_ssh2_to_openssh(ssh2_public_key_content):
+    lines = ssh2_public_key_content.strip().split('\n')
+    encoded_key = ""
+    start_found = False
+    for line in lines:
+        line = line.strip()
+        if line == "---- BEGIN SSH2 PUBLIC KEY ----":
+            start_found = True
+            continue
+        elif line == "---- END SSH2 PUBLIC KEY ----":
+            break
+        elif start_found and line.startswith("Comment:"):
+            comment = line.split(":", 1)[1].strip().strip('"')
+        elif start_found:
+            encoded_key += line
+
+    if encoded_key:
+        # Assuming it's an RSA key based on the comment
+        openssh_key = f"ssh-rsa {encoded_key} \"{comment}\""
+        return openssh_key
+    else:
+        return None
+
+
+def parse_key_from_file(file_path):
+    """
+    解析 SSH 公钥文件，支持 SSH2 和 OpenSSH 格式
+    :param file_path: 公钥文件路径
+    :return: 公钥内容
+    """
+    with open(file_path, "r") as f:
+        content = f.read()
+        if "---- BEGIN SSH2 PUBLIC KEY ----" in content:
+            return convert_ssh2_to_openssh(content)
+        else:
+            return content.strip()
+
+
+def configure_ssh_on_ubuntu(port, username, ssh_public_key="", password=""):
+    if not is_running_in_docker():
+        print("This script is not running in a Docker container. Exiting...")
+        return
+    print(f"deploy ssh server on port: {port}, username: {username}, ssh_public_key: {ssh_public_key}, password: {password is not None}")
+    if not(len(ssh_public_key) > 0 or len(password) > 0):
+        print("ssh_public_key or password must be provided.")
+        return 
+    command_prefix = "sudo " if need_sudo() else ""
+    try:
+        need_install_ssh = True
+        try:
+            # execute_command(["which", "sshd"])
+            response = execute_command("which sshd", only_stdout=False)
+            if response["returncode"] == 0:
+                print("SSH server is already installed.")
+                need_install_ssh = False
+        except subprocess.CalledProcessError:
+            pass
+        if need_install_ssh:
+            print("SSH server is not installed. Installing...")
+            execute_command(command_prefix + "apt-get update")
+            execute_command(command_prefix + "apt-get install -y openssh-server")
+
+        print("Modifying SSH configuration...")
+        execute_command(command_prefix + f'sed -i "s/#Port 22/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"):
+            print("SSH host keys not found. Generating new keys...")
+            execute_command(command_prefix + 'ssh-keygen -t rsa -f /etc/ssh/ssh_host_rsa_key -N ""')
+
+        print(f'Creating user {username} and configuring SSH key...')
+        try:
+            # run_command(["id", "-u", username])
+            execute_command(command_prefix + f"id -u {username}")
+        except subprocess.CalledProcessError:
+            print(f"User {username} does not exist. Creating user...")
+            execute_command(command_prefix + f"useradd -m -s /bin/bash {username}")
+
+        ssh_dir = f"/home/{username}/.ssh" if username != "root" else "/root/.ssh"
+        authorized_keys_file = f"{ssh_dir}/authorized_keys"
+
+
+        execute_command(command_prefix + f"mkdir -p {ssh_dir}")
+        execute_command(command_prefix + f"chmod 700 {ssh_dir}")
+
+        if ssh_public_key:
+            key_content = parse_key_from_file(ssh_public_key)
+            execute_command(command_prefix + f"bash -c 'echo \"{key_content}\" >> {authorized_keys_file}'", hide=True)
+        execute_command(command_prefix + f"chmod 600 {authorized_keys_file}")
+        execute_command(command_prefix + f"chown -R {username}:{username} {ssh_dir}")
+
+        if password:
+            execute_command(command_prefix + f"echo '{username}:{password}' | chpasswd", hide=True)
+
+        print("Restarting SSH service...")
+        # run_command(["sudo", "systemctl", "restart", "ssh"])
+        execute_command("/etc/init.d/ssh restart")
+
+        print(f"ssh is configured! You can connect using the following command: ssh -p {port} {username}@<host_ip>")
+
+    except subprocess.CalledProcessError as e:
+        print(f"Command called process error: {e}")
+    except FileNotFoundError as e:
+        print(f"File not found error: {e}")
+    except Exception as e:
+        print(f"An error occurred: {e}")
+
+def add_deploy_ssh_server_args(parser):
+    """添加 SSH 配置参数"""
+    parser.add_argument("--port", type=int, required=True, help="SSH server port")
+    parser.add_argument("--username", required=True, help="SSH username")
+    parser.add_argument("--password", help="SSH password", type=str, default="")
+    parser.add_argument("--ssh-public-key", help="SSH public key", type=str, default="")
+    
+
+
+def deploy_ssh_server_main(args):
+    """
+    部署 SSH 服务器
+    :param args: 命令行参数
+    """
+    configure_ssh_on_ubuntu(args.port, args.username, args.ssh_public_key, args.password)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="在 Ubuntu 系统中配置 SSH")
+    add_deploy_ssh_server_args(parser)
+
+    args = parser.parse_args()
+
+    deploy_ssh_server_main(args)

addftool/process/__init__.py

@@ -7,8 +7,7 @@ from .utils import get_processes, get_process_using_rocm, get_process_using_cuda
 def add_killer_args(subparsers):
     process_killer_parser = subparsers.add_parser('kill', help='process kill')
 
-    process_killer_parser.add_argument("--prefix", help="prefix of command", default=None)
-    process_killer_parser.add_argument("--contains", help="contains of command", default=None)
+    process_killer_parser.add_argument("-c", "--contains", help="contains of command", action='store_true', default=False)
     process_killer_parser.add_argument("--timeout", help="timeout of command", default=5, type=int)
     process_killer_parser.add_argument("--try_count", help="try count of command", default=3, type=int)
 
@@ -17,6 +16,8 @@ def add_killer_args(subparsers):
 
     process_killer_parser.add_argument("-v", "--view", help="view process", action='store_true', default=False)
 
+    process_killer_parser.add_argument("name", nargs='?', help="name of process", type=str, default="")
+
 
 def kill_process(processes, timeout=5, try_count=3):
     for process in processes:
@@ -57,8 +58,8 @@ def kill_process(processes, timeout=5, try_count=3):
         exit(1)
 
 
-def find_and_kill_process(prefix=None, contains=None, use_rocm=False, use_cuda=False, timeout=5, try_count=3, only_view=False):
-    do_not_do_anything = prefix is None and contains is None
+def find_and_kill_process(command="", contains=False, use_rocm=False, use_cuda=False, timeout=5, try_count=3, only_view=False):
+    do_not_do_anything = command is None or len(command) == 0
     if use_rocm:
         processes = get_process_using_rocm()
     elif use_cuda:
@@ -67,7 +68,7 @@ def find_and_kill_process(prefix=None, contains=None, use_rocm=False, use_cuda=F
         print("Use top or htop to find the process you want to kill")
         return
     else:
-        processes = get_processes(command_prefix=prefix, prefix_contains=contains)
+        processes = get_processes(command=command, contains=contains)
     
     if only_view:
         print(f"Found {len(processes)} processes")
@@ -82,6 +83,6 @@ def find_and_kill_process(prefix=None, contains=None, use_rocm=False, use_cuda=F
 
 def killer_main(args):
     find_and_kill_process(
-        args.prefix, args.contains, use_rocm=args.rocm, use_cuda=args.cuda, 
+        args.name, args.contains, use_rocm=args.rocm, use_cuda=args.cuda, 
         timeout=args.timeout, try_count=args.try_count, only_view=args.view,
     )

addftool/process/utils.py

@@ -2,7 +2,7 @@ import psutil
 import subprocess
 
 
-def get_processes(command_prefix=None, prefix_contains=None, pids=None):
+def get_processes(command="", contains=False, pids=None):
     """获取进程的PID和命令"""
     
     processes = []
@@ -14,10 +14,11 @@ def get_processes(command_prefix=None, prefix_contains=None, pids=None):
                 continue
             if pids is not None and pid not in pids:
                 continue
-            if command_prefix and not cmdline[0] == command_prefix:
-                continue
-            if prefix_contains and prefix_contains not in cmdline[0]:
-                continue
+            if len(command) > 0:
+                if contains and command not in cmdline[0]:
+                    continue
+                if not contains and command != cmdline[0]:
+                    continue
             if len(cmdline) > 1 and cmdline[1].endswith('addf'):
                 continue
             processes.append({'pid': pid, 'command': cmdline})

addftool/sync/__init__.py

@@ -0,0 +1,42 @@
+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.3.dist-info → addftool-0.1.5.dist-info}/METADATA

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

addftool-0.1.5.dist-info/RECORD

@@ -0,0 +1,17 @@
+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,,

addftool-0.1.3.dist-info/RECORD

@@ -1,12 +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/process/__init__.py,sha256=EFJ90RcF04M4imGm5NLgb9rmpwxBslZUL_4LKA52uv8,3487
-addftool/process/utils.py,sha256=kS1bYYiWymB1_SObqQLfGiVbZdCJqd57SO6DrwN5H1s,2246
-addftool-0.1.3.dist-info/METADATA,sha256=YwSmBYP8v9T_PnVr-3WM-V2KkCcPg17IwMKllfp-pc0,148
-addftool-0.1.3.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-addftool-0.1.3.dist-info/entry_points.txt,sha256=9lkmuWMInwUAtev8w8poNkNd7iML9Bjd5CBCFVxg2b8,111
-addftool-0.1.3.dist-info/top_level.txt,sha256=jqj56-plrBbyzY0tIxB6wPzjAA8kte4hUlajyyQygN4,9
-addftool-0.1.3.dist-info/RECORD,,