因为组里面有服务器GPU资源，有些时候想要查看各个服务器的GPU资源使用情况，如果是用ssh连接，再使用nvidia-smi命令查看太过麻烦了，便想要写一个python脚本，只需要在本机上运行代码即可查看目标服务器GPU使用情况。

代码思路如下：

graph TD
Step1:使用paramiko库远程连接服务器 --> Step2:使用命令行nvidia-smi获取相关信息 -->Step3:将获得的输出流转化为DataFrame进行打印

Step1:使用paramiko库远程连接服务器

在这里主要是使用了paramiko库，因为连接方式使用的是密钥文件的方式所以传入的是key_file而不是password。如果是密码连接的方式可以自行对下面的代码进行修改。

python
def ssh_command_with_key(hostname, port, username, key_file, command):
    # 创建SSH客户端
    client = paramiko.SSHClient()
    # 自动添加服务器的主机名和密钥到本地HostKeys对象，不推荐在生产环境中使用
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

    try:
        # 使用私钥文件
        mykey = paramiko.RSAKey.from_private_key_file(key_file)

        # 连接到远程服务器
        client.connect(hostname, port, username, pkey=mykey)

        # 执行命令
        stdin, stdout, stderr = client.exec_command(command)

        # 读取命令的标准输出和标准错误
        output = stdout.read()
        error = stderr.read()

        return output, error

    finally:
        # 关闭连接
        client.close()

Step2:使用命令行nvidia-smi获取相关信息

通过上面的连接服务器代码，能让我们运行command指令，这里主要涉及到下面两个命令，分别是显卡信息查询和任务信息查询，即分别查询显卡状态以及当前正在运行的任务状态。同时添加--format=csv参数，方便后续读入DataFrame。

shell
nvidia-smi --query-gpu=timestamp,name,memory.used,memory.total --format=csv

shell
nvidia-smi --query-compute-apps=gpu_name,pid,process_name,used_memory --format=csv

Step3:将获得的输出流转化为DataFrame进行打印

将获取到的输入流读入为DataFrame格式，方便查看。

python
        output, error = ssh_command_with_key(hostname, port, username, key_file, command)
        csv_file = io.StringIO(output.decode())
        df = pd.read_csv(csv_file)

完整代码

python
import io
import paramiko
import pandas as pd
from datetime import datetime

# 设置显示选项
pd.set_option('display.max_rows', 100)
pd.set_option('display.max_columns', 5)
pd.set_option('display.width', 1000)


def ssh_command_with_key(hostname, port, username, key_file, command):
    # 创建SSH客户端
    client = paramiko.SSHClient()
    # 自动添加服务器的主机名和密钥到本地HostKeys对象，不推荐在生产环境中使用
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

    try:
        # 使用私钥文件
        mykey = paramiko.RSAKey.from_private_key_file(key_file)

        # 连接到远程服务器
        client.connect(hostname, port, username, pkey=mykey)

        # 执行命令
        stdin, stdout, stderr = client.exec_command(command)

        # 读取命令的标准输出和标准错误
        output = stdout.read()
        error = stderr.read()

        return output, error

    finally:
        # 关闭连接
        client.close()

def sever_gpu_status(hostname, port, username, key_file, commands):
    # time_stamp print
    now = datetime.now()
    print("TIME_STAMP:", now.strftime("%Y-%m-%d %H:%M:%S"))
    print('-' * 120)
    for command in commands:
        output, error = ssh_command_with_key(hostname, port, username, key_file, command)
        csv_file = io.StringIO(output.decode())
        df = pd.read_csv(csv_file)
        print(df)
        print('-' * 120)

# 使用示例
hostname = ''
port = xx
username = ''
key_file = r''
# commands = ['nvidia-smi --query-compute-apps=gpu_name,pid,process_name,used_memory --format=csv']
commands = ['nvidia-smi --query-gpu=timestamp,name,memory.used,memory.total --format=csv',
            'nvidia-smi --query-compute-apps=gpu_name,pid,process_name,used_memory --format=csv']
msg = sever_gpu_status(hostname, port, username, key_file, commands)

代码运行结果如下：