有这样的一个场景，本人是在一家 AI 公司，有一台机器是公用的，可能经常有被人使用的可能，但是我想在空闲的时候使用，那么就有了这样的一个自动化脚，本，通过命令来连接主机，并通过查询显存用量，如果有的话就在桌面上发送通知。
ssh 和 node-ssh
SSH（Secure Shell）是一个加密的网络协议，用于在不安全的网络上安全地执行远程命令和管理系统。SSH 最常见的用途是通过加密的连接远程登录计算机，但它也可以用于传输文件、隧道其他协议、以及执行其他网络服务。
SSH 在工作时，通常涉及三个步骤：
-
建立连接：客户端和服务器通过 SSH 协议建立一个加密连接。
-
身份验证：客户端向服务器提供凭据（如用户名、密码、或公钥），服务器验证这些凭据后，允许访问。
-
数据传输：所有的数据（命令、输出、文件等）都通过加密通道传输，保证数据的机密性和完整性。
它最常用的 SSH 命令主要有以下这些：
-
ssh username@hostname：远程登录到指定的主机。
-
scp file.txt username@hostname:/remote/directory/：将文件复制到远程服务器。
-
ssh -L 8080:localhost:80 username@hostname：将本地的 8080 端口转发到远程主机的 80 端口。
而 node-ssh 是一个用于 Node.js 的库，封装了 SSH 协议的功能，使得开发者可以通过 JavaScript 或 TypeScript 代码在 Node.js 应用程序中实现 SSH 功能。node-ssh 库的底层使用了流行的 ssh2 库，并提供了更易用的 API。
下面是一个基本的使用案例：
const { NodeSSH } = require("node-ssh");
const ssh = new NodeSSH();
ssh
.connect({
host: "127.0.0.1",
username: "moment",
password: "moment",
})
.then(() => {
// 执行远程命令
ssh.execCommand("uptime").then((result) => {
console.log("STDOUT: " + result.stdout);
console.log("STDERR: " + result.stderr);
});
// 上传文件
ssh
.putFile("local-path.txt", "/remote-path.txt")
.then(() => {
console.log("文件上传成功");
})
.catch((err) => {
console.error("文件上传失败:", err);
});
// 下载文件
ssh
.getFile("local-path.txt", "/remote-path.txt")
.then(() => {
console.log("文件下载成功");
})
.catch((err) => {
console.error("文件下载失败:", err);
});
})
.catch((err) => {
console.error("SSH连接失败:", err);
});
如何在电脑上发送通知
在 nodejs 中有一个 node-notifier 包，它能够在桌面上发送通知，是因为它借用了操作系统原生的通知系统或工具，并在不同平台上通过这些工具来实现通知功能。这样，node-notifier 可以在 macOS、Windows 和 Linux 等不同操作系统上都实现相同的效果，同时屏蔽了各平台之间的差异，使得开发者可以使用统一的 API 来发送桌面通知。
以下是一个简单的使用示例：
const notifier = require("node-notifier");
// 基本通知
notifier.notify({
title: "下班辣下班辣",
message: "还不下班，都什么时候了",
sound: true, // 可选：播放通知声音
wait: true, // 可选：等待用户响应
});
最终效果如图所示：
最终实现
首先我们先来上完整的代码，但是隐藏了关键信息：
const { NodeSSH } = require("node-ssh");
const cron = require("node-cron");
const notifier = require("node-notifier");
const ssh = new NodeSSH();
const remoteHost = ""; // 远程机器的IP地址
const remoteUser = ""; // SSH用户名
const remotePassword = ""; // SSH密码
const remotePort = 6667; // SSH端口
async function checkGpuMemory() {
try {
await ssh.connect({
host: remoteHost,
port: remotePort,
username: remoteUser,
password: remotePassword,
});
console.log("SSH连接成功");
// 执行nvidia-smi命令来检查显存总量和使用情况
const result = await ssh.execCommand(
"nvidia-smi --query-gpu=memory.total,memory.used,memory.free --format=csv,noheader,nounits"
);
if (result.stderr) {
console.error("Error executing command:", result.stderr);
return;
}
const lines = result.stdout
.split("\n")
.filter((line) => line.trim() !== "");
const gpuInfo = lines.map((line) => {
const [total, used, free] = line
.split(",")
.map((value) => parseInt(value.trim()));
return { total, used, free };
});
console.log("GPU显存信息:", gpuInfo);
gpuInfo.forEach((gpu, index) => {
console.log(
`GPU ${index}: 总显存 ${gpu.total} MiB, 已使用 ${gpu.used} MiB, 空闲 ${gpu.free} MiB`
);
if (gpu.used <= 10) {
showNotification(index, gpu.used, gpu.total, gpu.free);
}
});
} catch (err) {
console.error("SSH连接失败:", err);
} finally {
ssh.dispose();
}
}
function showNotification(gpuIndex, used, total, free) {
notifier.notify({
title: `GPU ${gpuIndex} 显存空闲`,
message: `GPU ${gpuIndex} 当前显存空闲情况: 总显存 ${total} MiB, 已使用 ${used} MiB, 空闲 ${free} MiB`,
sound: true, // 播放提示音
wait: true, // 用户点击通知时才消失
});
}
// 每隔1分钟检查一次显存使用情况
cron.schedule("* * * * *", checkGpuMemory);
console.log("显存监测脚本已启动，每分钟检查一次显存使用情况");
最终执行效果如下图所示：
这段代码的作用是通过 SSH 连接到远程服务器，每分钟检查一次服务器上每块 GPU 的显存使用情况，并在显存使用较低时通过桌面通知提醒用户。
首先，代码使用 node-ssh 模块建立与远程服务器的 SSH 连接，通过 nvidia-smi 命令获取 GPU 的显存信息，包括总显存、已使用显存和空闲显存。然后，它将命令的输出解析成每块 GPU 的显存数据，并检查每块 GPU 的已使用显存。如果某块 GPU 的已使用显存小于或等于 10 MiB，代码会通过 node-notifier 模块发送桌面通知，提醒用户该 GPU 几乎处于空闲状态。
代码使用 node-cron 模块设置了一个定时任务，每分钟执行一次 checkGpuMemory 函数，以定期监控 GPU 的显存使用情况。如果连接失败或者命令执行出错，代码会在控制台输出相应的错误信息。
总结
借助 nodejs，它能够与操作系统交互，我们可以编写一些自动化脚本来帮助我们完成日常需求。
最后分享两个我的两个开源项目,它们分别是:
这两个项目都会一直维护的，如果你想参与或者交流学习，可以加我微信 yunmz777 如果你也喜欢,欢迎 star 🚗🚗🚗
