我目前在外校的网络安全实验室线上实习，实验室有一台共享的 Linux 服务器，我至今为止的所有工作几乎都是在这台服务器上完成的。

尽管有一定的 Linux 使用经验，但是我还是在整个实习过程中遇到了比较多的问题。究其原因，使用一台属于自己的 Linux 本地机器，和作为一个普通用户访问共享的 Linux 服务器，这之间还是有较大的区别的。

这篇文章里，我想分享下我使用服务器中用到的一些工具，遇到的问题以及解决方案。

SSH 连接工具#

我找不出任何不推荐 VS Code Remote-SSH 插件的理由。

只需要简单的配置，你就可以使用本地的 VS Code 终端编辑远程服务器上的文件，同时，所有的插件和配置都能方便地同步到远程服务器上，体验上和本地写代码并没有区别。

最重要的是，Github Copilot 也能在远程服务器上方便地访问。使用过 AI Copilot 后，我很难回到没有 AI 插件的环境了，由俭入奢易，由奢入俭难啊。

并行化工具#

在服务器上，我经常需要同时运行多个任务。特别常见的场景是，同一段代码需要在不同的数据集或者不同的参数设置下进行评估。

这种情况如果在本地的机器，我可能会用 GNU Parallel；但是服务器上并没有安装这个工具。作为一个替代方案，我常用的方式是手动写一个 shell 脚本来管理任务。

例如，main.py 这个程序需要在 10 个数据集上运行，我写的 shell 脚本可能会像这样：

#!/bin/bash

for i in {1..10}
do
    nohup python main.py --dataset dataset_$i 2>&1 > output_$i.log &
done

这里的 nohup 是用来创建后台进程，如果不使用 nohup 直接运行 python main.py，在 SSH 连接中断后程序会停止运行；& 是用来将任务放到后台运行，和当前的终端分离。

当然，这是只是最粗糙的并行脚本，它的问题有很多，最容易想到的有两点：

  • 一些任务可能会占用太多的 CPU/GPU 资源，导致服务器上的其他任务无法正常运行。
  • 所有任务的输出都被 nohup 默认重定向到 nohup.out 文件，这样会导致输出混乱，不方便查看。

所以，我实际上会在脚本里手动控制并行任务的数量，并将输出重定向到不同的文件中。

#!/bin/bash

max_jobs=5
current_jobs=0

for i in {1..10}
do
    if [ $current_jobs -ge $max_jobs ]; then
        wait -n
        current_jobs=$((current_jobs-1))
    fi
    nohup python main.py --dataset $i 2>&1 > output_$i.log &
    current_jobs=$((current_jobs+1))
done

CPU 资源控制#

虽然上面展示的仍然是一个比较粗糙和简单的脚本，但是已经能满足我的需求了，我在几个月内都在使用这样的脚本来完成任务。

直到最近，共用服务器的其他老师反馈说我的任务占用了太多的 CPU 资源，导致他们的实验无法正常进行了，并提醒我要限制 numpy 的线程数。我打开 htop 一看，才发现所有的 CPU 都在满载状态，load average 高达 700~800（服务器有 112 个逻辑 CPU），这时我才意识到问题的严重性。

紧急 pkill 了所有任务来释放所有资源之后，我按照老师提示的方法设置了 numpy 的线程数，问题才得以解决。

事后复盘发现，我当时在实验的代码是一段 python 代码，其中的 numpy 模块运行时默认会尝试用所有 CPU 核心进行多线程优化，这是因为 numpy 调用的底层是一些高性能计算库，这些计算库会尽可能利用所有的 CPU. 为了限制这些高性能库榨干服务器的计算资源，需要在调用 numpy 之前先设置环境变量，使这些库认为只有少量的 CPU 可以使用。

import os
os.environ["OMP_NUM_THREADS"] = "4"  # OpenMP 线程数
os.environ["OPENBLAS_NUM_THREADS"] = "4"  # OpenBLAS 线程数
os.environ["MKL_NUM_THREADS"] = "4"  # MKL 线程数
os.environ["NUMEXPR_NUM_THREADS"] = "4"  # NumExpr 线程数
import numpy as np

在设置这些变量后重新运行一段时间后，服务器上出现了其他问题：内存占用过高，可用内存仅剩 10GB/400GB。经过 htop 查看和分析后，导致这些问题的仍然是我的并行任务…

对于内存的限制还是比较容易的，我在 shell 脚本的最上面加了一行 ulimit -v，限制整个脚本的内存使用量。

硬盘空间管理#

服务器上的硬盘空间是有限的，但是 home 文件夹的膨胀是没有止境的。

好在实验室有额外加装的硬盘，挂载在 /data 目录下。 最开始，我只是把不太常用的数据集放在了 /data 目录下。但后来随着使用时间的增加，home 文件夹下各类文件也越来越多。按照实验室建议，我开始把一些比较常用，但是占用了较大空间的文件夹都挪到了 /data 目录下，然后用 ln -s 重新软链接回 home 文件夹。

mv ~/folder /data/my_user_name/folder
ln -s /data/my_user_name/folder ~/folder

GPU 资源管理#

深度学习的任务十分依赖 GPU 加速。服务器上有若干 RTX 4090，但是由于服务器上有其他用户的实验，我们不能随意使用所有的 GPU 资源。

这一部分实际上分为两个话题，一个是我如何知道自己占用了多少显存，一个是如何合理地分配任务到不同的 GPU 上。

监控 GPU 资源#

查看 GPU 的状态可以使用 nividia 提供的 nvidia-smi 工具，这个工具可以查看 GPU 的使用情况，包括 GPU 的使用率、温度、显存使用情况等等。

这个工具当然很好用，不过，我需要看到这些显存具体是被哪个用户占用的。更具体地来说，我想要看到我自己占用了多少显存，同时，其他用户中是否有人也正在重度使用 GPU，我是否需要管理自己的任务，降低对服务器的影响。

因此我在 AI 工具辅助下编写了一个简单的脚本，这个脚本可以查看当前所有用户的显存使用情况，各个 GPU 剩余的显存，以及当前用户的内存使用情况。

为了方便高亮查看，我把整个脚本的核心代码贴在下面，去掉了 watch -n 1 的前缀。实际使用的脚本中，我使用 watch -n 1 来每秒刷新一次脚本显示，实现实时监控的效果。

#!/bin/sh

nvidia-smi --query-compute-apps=pid,used_memory --format=csv,noheader,nounits | \
while IFS=',' read -r pid memory; do 
  user=$(ps -o user= -p $pid); 
  echo "$user,$memory"; 
done | awk -F"," '"'"'{mem[$1]+=$2} END {for (u in mem) print u, mem[u] " MiB"}'"'"' | \
        awk -v current_user="'$(whoami)'" '"'"'{if ($1 == current_user) printf "%-10s %-10s %-10.2f GB %-10.2f Cards <==\n", $1, $2, $2/1024, $2/24576; else printf "%-10s %-10s %-10.2f GB %-10.2f Cards\n", $1, $2, $2/1024, $2/24576}'"'"' | sort -k3 -nr | column -t

echo ""
echo "=============================================="

echo "Memory usage by me:"
CURRENT_USER=$(whoami)

TOTAL_MEMORY=$(ps -u "$CURRENT_USER" -o rss= | awk '"'"'{sum+=$1} END {print sum}'"'"')

TOTAL_MEMORY_MB=$((TOTAL_MEMORY / 1024))
TOTAL_MEMORY_GB=$((TOTAL_MEMORY_MB / 1024))

echo "Total memory usage: $TOTAL_MEMORY_MB MB ($TOTAL_MEMORY_GB GB)"

echo ""
echo "=============================================="

echo "Free memory per GPU:"
nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits | awk '"'"'{print $1/1024 " GB"}'"'"' | column -t

当我运行脚本时，可以看到如下的输出：

Every 1.0s:                                 cxhpc: Thu Jan 23 19:22:23 2025

User_A      60806  59.38  GB  2.47  Cards
Me          49092  47.94  GB  2.00  Cards  <==

==============================================
Memory usage by me:
Total memory usage: 39969 MB (39 GB)

==============================================
Free memory per GPU:
7.61328  GB
21.5918  GB
3.69727  GB
3.69727  GB
9.93164  GB
21.7715  GB
10.1504  GB
2.45312  GB

为了直观显示显存占用的大小，我直接将显存换算为 RTX 4090 显存 24GB 的倍数，即大约占用了多少张显卡。

GPU 资源分配#

使用 pytorch 框架时，可以直接指定使用哪个 GPU，例如：

import torch
torch.cuda.set_device(1)

有一些场景下，需要保证所有的张量都在同一个 GPU 上，这时可以使用 Tensor.to() 方法，将某个张量转移到指定的 GPU 上。

import torch
device = torch.device("cuda:1") if torch.cuda.is_available() else "cpu"
x = torch.randn(5, 3).to(device)

然而，我的实验场景下，最常见的情况是同一个模型要同时在多个数据集上训练。如果在代码中指定某个特定的 GPU，那么无论并行的任务数量多少，所有的任务都会被分配到同一个 GPU 上，大大限制了任务的并行度。

我为此编写了一个简单的 python 函数，这个函数可以根据当前 GPU 的使用情况，自动选择一个较为空闲的 GPU。

import torch
import subprocess

def get_gpu_memory():
    result = subprocess.run(
        ['nvidia-smi', '--query-gpu=memory.total,memory.used', '--format=csv,nounits,noheader'],
        stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
    )
    if result.returncode != 0:
        raise RuntimeError(f"nvidia-smi error: {result.stderr}")
    return [tuple(map(int, x.split(', '))) for x in result.stdout.strip().split('\n')]

def get_lowest_memory_gpu():
    if not torch.cuda.is_available():
        return None
    gpu_memory = get_gpu_memory()
    free_memory = [(total - used, i) for i, (total, used) in enumerate(gpu_memory)]
    _, best_gpu = max(free_memory)
    if _ < 6000:
        raise RuntimeError("No free GPU available")
    print(f'Best GPU: {best_gpu}')
    return best_gpu

opti_device = torch.device(f'cuda:{get_lowest_memory_gpu()}')

这个函数会返回一个较为空闲的 GPU，并直接生成一个名为 opti_device 的 torch.device 对象，可以直接用于 pytorch 的代码中。

指定 GPU 为 opti_device 后，pytorch 会自动选择一个较为空闲的 GPU 进行计算。

使用这个 GPU 代码片段，shell 脚本并行任务需要做出一些修改。在提交每个任务需要调用 sleep 命令等待一段时间，使任务有足够的时间把需要的 GPU 显存分配完成。

#!/bin/bash

max_jobs=5
current_jobs=0

for i in {1..10}
do
    if [ $current_jobs -ge $max_jobs ]; then
        wait -n
        current_jobs=$((current_jobs-1))
    fi
    nohup python main.py --dataset $i 2>&1 > output_$i.log &
    current_jobs=$((current_jobs+1))
    sleep 30
done

如果不使用 sleep 命令，那么最开始的 max_jobs 个任务会在短时间内同时被提交，这时它们获取的最优 GPU 极有可能是同一个，导致单张卡的显存不足。