addftool 0.1.4__py3-none-any.whl → 0.1.6__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/__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

@@ -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/__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/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.4.dist-info → addftool-0.1.6.dist-info}/METADATA

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

addftool-0.1.6.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/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/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.4.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=OB-cZXP1jK7l8uN8nKfhg_bCX6Slz6DeeBpEodt-IK4,3515
-addftool/process/utils.py,sha256=me4HqMz5OgRcQMUJmVhKdTJh4SW5BB-pd_lq7g8-UwE,2252
-addftool-0.1.4.dist-info/METADATA,sha256=ohQCdQdVevHRzMRu7ByPhBGxqPlxmEKnDxNmfcy_WII,148
-addftool-0.1.4.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-addftool-0.1.4.dist-info/entry_points.txt,sha256=9lkmuWMInwUAtev8w8poNkNd7iML9Bjd5CBCFVxg2b8,111
-addftool-0.1.4.dist-info/top_level.txt,sha256=jqj56-plrBbyzY0tIxB6wPzjAA8kte4hUlajyyQygN4,9
-addftool-0.1.4.dist-info/RECORD,,