89 行
3.2 KiB
Python
89 行
3.2 KiB
Python
import paramiko
|
|
import logging
|
|
|
|
# 配置日志
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger("GPUCollector")
|
|
|
|
class GPUCollector:
|
|
"""
|
|
负责通过 SSH 远程连接到服务器并采集 NVIDIA GPU 状态的类
|
|
"""
|
|
def __init__(self, server_config):
|
|
self.alias = server_config.get('alias', 'Unknown')
|
|
self.ip = server_config.get('ip')
|
|
self.port = server_config.get('port', 22)
|
|
self.username = server_config.get('username')
|
|
self.password = server_config.get('password')
|
|
self.timeout = 10
|
|
|
|
def fetch_gpu_data(self):
|
|
"""
|
|
执行远程 nvidia-smi 命令并解析结果
|
|
返回: List[Dict] 包含每张显卡的详细信息，失败则返回 None
|
|
"""
|
|
ssh = paramiko.SSHClient()
|
|
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
|
|
|
try:
|
|
# 建立连接
|
|
ssh.connect(
|
|
hostname=self.ip,
|
|
port=self.port,
|
|
username=self.username,
|
|
password=self.password,
|
|
timeout=self.timeout
|
|
)
|
|
|
|
# 使用 bash -l -c 强制加载登录环境 (Login Shell)
|
|
# 这样可以确保加载 /etc/profile 和 ~/.bash_profile，从而获取正确的 PATH
|
|
raw_cmd = (
|
|
"nvidia-smi --query-gpu=index,name,temperature.gpu,"
|
|
"utilization.gpu,utilization.memory,memory.total,"
|
|
"memory.used,memory.free --format=csv,noheader,nounits"
|
|
)
|
|
full_cmd = f'bash -l -c "{raw_cmd}"'
|
|
|
|
stdin, stdout, stderr = ssh.exec_command(full_cmd)
|
|
output = stdout.read().decode('utf-8').strip()
|
|
error = stderr.read().decode('utf-8').strip()
|
|
|
|
if error and not output:
|
|
logger.error(f"[{self.alias}] SSH Command Error: {error}")
|
|
return None
|
|
|
|
if not output:
|
|
logger.warning(f"[{self.alias}] No output received from nvidia-smi")
|
|
return None
|
|
|
|
# 解析 CSV 数据
|
|
lines = output.split('\n')
|
|
gpu_list = []
|
|
for line in lines:
|
|
if not line: continue
|
|
parts = [p.strip() for p in line.split(',')]
|
|
if len(parts) == 8:
|
|
gpu_list.append({
|
|
"index": int(parts[0]),
|
|
"name": parts[1],
|
|
"temp": int(parts[2]),
|
|
"util": int(parts[3]),
|
|
"mem_util": int(parts[4]),
|
|
"mem_total": int(parts[5]),
|
|
"mem_used": int(parts[6]),
|
|
"mem_free": int(parts[7])
|
|
})
|
|
|
|
return gpu_list
|
|
|
|
except paramiko.AuthenticationException:
|
|
logger.error(f"[{self.alias}] SSH Authentication failed for {self.ip}")
|
|
except paramiko.SSHException as e:
|
|
logger.error(f"[{self.alias}] SSH Exception: {e}")
|
|
except Exception as e:
|
|
logger.error(f"[{self.alias}] Unexpected error: {e}")
|
|
finally:
|
|
ssh.close()
|
|
|
|
return None
|