Python下基于paramiko打造自动配置ubuntu22.04

import paramiko
import time

def execute_command(ssh_client, command):
    stdin, stdout, stderr = ssh_client.exec_command(command)
    exit_status = stdout.channel.recv_exit_status()
    output = stdout.read().decode("utf-8")
    return exit_status, output

def wait_for_reboot(ssh_client, hostname, port, username, password):
    while True:
        time.sleep(600)  # 等待一段时间，避免频繁尝试连接
        try:
            ssh_client.connect(hostname, port, username, password)
            ssh_client.close()
        except paramiko.ssh_exception.SSHException:
            print("主机尚未重启完成，继续等待...")
        else:
            print("主机重启完成")
            break

def run_as_root(ssh_client, sudo_password, commands):
    for command in commands:
        print(f"正在执行命令：{command}")
        # 使用sudo切换到root权限执行命令
        full_command = f"echo '{sudo_password}' | sudo -S {command}"
        exit_status, command_output = execute_command(ssh_client, full_command)
        print("命令执行结果：")
        print(command_output)
        print(f"退出状态：{exit_status}")

if __name__ == "__main__":
    hostname = "192.168.154.136"
    port = 22
    username = "lwp"
    password = "ABCabc@123"
    sudo_password = "ABCabc@123"  # sudo密码

    try:
        # 创建SSH客户端实例
        ssh_client = paramiko.SSHClient()
        ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        # 连接到Linux主机
        ssh_client.connect(hostname, port, username, password)

        # 切换到root权限并执行命令
        root_commands = [
            # 在此添加要在root权限下执行的命令
            "apt update -y",
            "apt install vim -y",
            "apt install build-essential -y",
            "apt install stress -y",
            "apt install htop -y",
            "apt install lm-sensors -y",

            # 安装nohup命令
            "apt-get install coreutils",
            "mkdir DRIVER",
            "mkdir GPU",
                        
            # 下载Nvidia 530.41驱动
            "wget https://us.download.nvidia.com/XFree86/Linux-x86_64/530.41.03/NVIDIA-Linux-x86_64-530.41.03.run -P DRIVER/",
            "cd DRIVER/",
            "chmod +x ./DRIVER/NVIDIA-Linux-x86_64-530.41.03.run",
            "CD ../",

            # 下载CUDA12.0驱动并等待8分钟
            "wget --show-progress --no-check-certificate -c -b https://developer.download.nvidia.com/compute/cuda/12.0.0/local_installers/cuda_12.0.0_525.60.13_linux.run -P DRIVER/",
            #"echo '等待8分钟...' && sleep 480",  # 等待8分钟（8分钟 = 480秒）

            
            # 编辑/etc/profile文件并在末尾添加两行
            "tee -a /etc/profile <<< 'export PATH=/usr/local/cuda/bin:$PATH'",
            "tee -a /etc/profile <<< 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH'",
                
            # 编辑/etc/modprobe.d/blacklist.conf文件并在末尾添加两行
            "tee -a /etc/modprobe.d/blacklist.conf <<< 'blacklist nouveau'",
            "tee -a /etc/modprobe.d/blacklist.conf <<< 'options nouveau modeset=0'",
        
            # 执行sudo update-initramfs -u命令更新内核并返回执行结果
            "update-initramfs -u",
            # ... 其他命令 ...
        ]
        run_as_root(ssh_client, sudo_password, root_commands)

        run_as_root(ssh_client, sudo_password, root_commands)

        # 执行重启命令
        reboot_command = "shutdown -r now"
        print(f"正在以root权限执行重启命令：{reboot_command}")
        exit_status, command_output = execute_command(ssh_client, reboot_command)
        print("命令执行结果：")
        print(command_output)
        print(f"退出状态：{exit_status}")

        # 等待主机重启完成
        wait_for_reboot(ssh_client, hostname, port, username, password)

        # 重新连接到主机
        ssh_client.connect(hostname, port, username, password)

        # 继续以root权限执行剩余的命令
        remaining_root_commands = [
            # 在此添加在root权限下执行的剩余命令
            "lsmod | grep nouveau",
            # "",

            # 安装Nvidia驱动
            "touch ./DRIVER/nvidia_install.log",
            "nohup sh ./DRIVER/NVIDIA-Linux-x86_64-530.41.03.run --no-opengl-files > ./DRIVER/nvidia_install.log 2>&1 &",
            "echo '等待NVIDIA驱动安装完成...' && sleep 300",
            "nvidia-smi",

            # 安装CUDA12.0
            "chmod +x ./DRIVER/cuda_12.0.0_525.60.13_linux.run",
            "touch ./DRIVER/cuda_install.log",
            "nohup ./DRIVER/cuda_12.0.0_525.60.13_linux.run --silent --driver --no-opengl-libs --no-drm --no-opengl-flavor > ./DRIVER/cuda_install.log 2>&1 &",
            "echo '等待CUDA驱动安装完成...' && sleep 300",
            "source /etc/profile",
            "nvcc -V",
           
            # 下载gpu-burn并解压
            "wget http://wili.cc/blog/entries/gpu-burn/gpu_burn-1.1.tar.gz -P GPU/",
            "tar -zxvf ./GPU/gpu_burn-1.1.tar.gz -C GPU/",

            # 编辑Makefile
            "sed -i 's/-arch=compute_30/-arch=compute_60/' ./GPU/Makefile",
            "cd GPU/",
            "make",

            
            # ...
        ]
        run_as_root(ssh_client, sudo_password, remaining_root_commands)

    except paramiko.AuthenticationException:
        print("认证失败，请检查您的凭据。")
    except paramiko.SSHException as ssh_ex:
        print(f"SSH连接错误：{ssh_ex}")
    except Exception as ex:
        print(f"错误：{ex}")
    finally:
        # 关闭SSH连接
        ssh_client.close()