Pytorch CPU OOM 杀死 linux 上的 ssh 服务器

Question

我遇到了一个问题，即当发生 CPU OOM 时，pytorch（使用 2.0.1+cu117 测试）不会正常失败。具体来说，我失去了所有 ssh 连接和 Xserver 对虚拟机或裸机的访问。

我没有测试过这种情况是否发生在任何其他操作系统上。

我找到的唯一解决方案是直接重新启动计算机（通过 vsphere 或仅通过电源按钮）。

我考虑过使用别名 python（考虑 conda env 切换）来将每个进程添加到 cgroup 中，这直接限制了内存使用，但有人建议我弄乱 cgroup 基本上是一个坏主意。

很难判断一个模型将占用多少内存，我需要一种优雅的方式来杀死它而不杀死我的 ssh 服务器。

一些注意事项：我已经在两台设备上尝试过这一点，一台虚拟机一台裸机，两台设备都是 Ubuntu 22.04，并且我相信两台设备都启用了操作系统级 OOM 杀手。这在 GPU OOM 中是无法重现的，因为这实际上会终止进程并返回典型的 GPU OOM 错误，但只能在 RAM 中重现。

我尝试按照此处所述设置 RLIMIT：https://www.geeksforgeeks.org/python-how-to-put-limits-on-memory-and-cpu-usage/。虽然这并没有解决我的问题。

我还考虑过在训练中添加一个条件来检查可用内存和中断，但这似乎是一个不干净的解决方案。

我没有尝试过 cgroup，因为如上所述。我对操作系统的事情也只是有点熟悉，我犹豫是否用它做任何事情，以免破坏一些我无法修复的东西。

Answer 1

您可以尝试调整

/proc/<pid>/oom_score_adj

以确保您的Python进程将被第一个杀死。

引用

man proc

：

坏度启发式为每个候选任务分配一个从 0（从不杀死）到 1000（总是杀死）的值，以确定哪个进程是目标。