AI 模型占了 10G 显存，服务重启却没释放？

聊聊 AI 后端那些事儿 · 第 1 篇 | 预估阅读：10 分钟

---

发现显存泄漏 - 「我有个朋友」系列

我有个朋友叫小禾，前些天，小禾踩了不少前端的坑。这几天，小禾开始对付后端。

后端用的是 FastAPI，跑在一台有 24G 显存的 GPU 服务器上。本地开发一切正常，部署上去也没问题。

直到有一天，小禾在服务器上按了 Ctrl+C 停掉服务，准备更新代码。

# 停止服务
^C
INFO: Shutting down...

# 重新启动
$ python run.py

然后，服务就崩了：

torch.cuda.OutOfMemoryError: CUDA out of memory.
Tried to allocate 2.00 GiB (GPU 0; 24.00 GiB total capacity;
22.10 GiB already allocated; 1.43 GiB free)

小禾懵了：我明明停掉服务了，为什么显存还被占着？

他打开 nvidia-smi 一看：

$ nvidia-smi
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|=============================================================================|
|    0   N/A  N/A     12345      C   python                          10240MiB |
+-----------------------------------------------------------------------------+

进程还在！

原来 Ctrl+C 虽然停了 uvicorn，但 Python 进程没有完全退出，模型还在显存里躺着。

小禾只好 kill -9 强制杀掉进程，然后重新启动。

"这也太蠢了吧，每次重启都要手动杀进程？"

---

FastAPI 的生命周期

小禾查了查文档，发现 FastAPI 有一套生命周期管理机制。

graph TD A[应用启动] B[startup 事件] C[加载 AI 模型
初始化数据库连接
预热缓存] D[处理请求...] E[shutdown 事件] F[释放 AI 模型
关闭数据库连接
清理临时文件] G[应用关闭] A --> B B --> C C --> D D --> E E --> F F --> G

原来 FastAPI 提供了 startup 和 shutdown 两个事件钩子，可以在应用启动和关闭时执行特定逻辑。

小禾之前只写了加载模型的代码，但从来没写过释放模型的代码。

难怪显存不释放------根本没人告诉程序要释放。

---

两种写法

FastAPI 支持两种生命周期管理方式。

写法一：装饰器方式（传统写法）

from fastapi import FastAPI

app = FastAPI()

@app.on_event("startup")
async def startup_event():
    print("应用启动")
    # 加载模型

@app.on_event("shutdown")
async def shutdown_event():
    print("应用关闭")
    # 释放模型

写法二：lifespan 上下文管理器（推荐，FastAPI 0.93+）

from contextlib import asynccontextmanager
from fastapi import FastAPI

@asynccontextmanager
async def lifespan(app: FastAPI):
    # 启动时执行
    print("应用启动")
    yield
    # 关闭时执行
    print("应用关闭")

app = FastAPI(lifespan=lifespan)

小禾选了第二种写法，因为它更清晰------启动和关闭的逻辑写在一起，一眼就能看出配对关系。

---

完整的生命周期管理

小禾开始重构代码：

# app/main.py
from contextlib import asynccontextmanager
from fastapi import FastAPI
from app.services.model_manager import ModelManager
from app.core.config import settings
from app.core.logger import logger

# 全局模型管理器
model_manager = ModelManager()

@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理"""

    # ===== 启动阶段 =====
    logger.info(f"正在启动 {settings.PROJECT_NAME}")
    logger.info(f"运行环境: {settings.ENVIRONMENT}")
    logger.info(f"设备: {settings.DEVICE}")

    try:
        # 加载 AI 模型
        model_manager.load_model()
        logger.info(f"模型 ({settings.IMAGE_MODEL}) 加载完成")

    except Exception as e:
        logger.error(f"启动失败: {str(e)}")
        raise  # 启动失败要抛出异常，阻止服务启动

    # ===== 运行阶段 =====
    yield  # 服务运行中...

    # ===== 关闭阶段 =====
    logger.info("正在关闭服务...")

    try:
        # 释放模型资源
        model_manager.cleanup()
        logger.info("模型资源已释放")

    except Exception as e:
        logger.error(f"清理时出错: {str(e)}")

    logger.info("服务已关闭")

app = FastAPI(
    title=settings.PROJECT_NAME,
    lifespan=lifespan
)

关键在于 cleanup() 方法，它负责释放 GPU 显存。

---

模型清理的正确姿势

小禾最初写的清理方法是这样的：

def cleanup(self):
    del self.model
    self.model = None

结果发现：显存还是没释放。

原来 PyTorch 的显存管理比较特殊，单纯 del 是不够的。小禾查了资料，找到了正确的清理流程：

# app/services/model_manager.py
import gc
import torch

class ModelManager:
    """模型管理器"""

    def __init__(self):
        self.pipe = None
        self.device = settings.DEVICE

    def load_model(self):
        """加载模型到 GPU"""
        if self.pipe is not None:
            return  # 已加载，不重复加载

        logger.info("正在加载模型...")
        self.pipe = StableDiffusionXLPipeline.from_pretrained(
            settings.MODEL_PATH,
            torch_dtype=torch.float16,
        ).to(self.device)
        logger.info("模型加载完成")

    def cleanup(self):
        """释放模型资源"""
        if self.pipe is None:
            return

        logger.info("正在释放模型资源...")

        # 第一步：把模型从 GPU 移到 CPU
        self.pipe.to("cpu")

        # 第二步：删除引用
        del self.pipe
        self.pipe = None

        # 第三步：强制垃圾回收
        gc.collect()

        # 第四步：清理 CUDA 缓存
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()  # 等待 GPU 操作完成

        logger.info("模型资源已释放")

四步缺一不可：

1. 移到 CPU：先把张量从 GPU 搬走
2. 删除引用：解除 Python 对象的引用
3. 垃圾回收：触发 Python 的 GC
4. 清理缓存：告诉 CUDA 释放已分配但未使用的显存

小禾测试了一下：

# 启动前
$ nvidia-smi
GPU Memory: 2GB / 24GB

# 服务运行中
$ nvidia-smi
GPU Memory: 12GB / 24GB

# Ctrl+C 停止服务
$ nvidia-smi
GPU Memory: 2GB / 24GB  # 释放了！

终于正常了。

---

但是，Ctrl+C 有时候不管用

小禾高兴了没两天，问题又来了。

有时候用 Ctrl+C 停止服务，shutdown 事件根本没触发，显存还是被占着。

他研究了一下，发现问题出在信号处理上。

当你按 Ctrl+C 时，系统发送的是 SIGINT 信号。正常情况下，uvicorn 会捕获这个信号，优雅地关闭服务。

但如果程序在执行某些阻塞操作（比如 AI 推理），可能来不及处理信号就被强制终止了。

更糟糕的是 kill -9，它发送的是 SIGKILL 信号，根本不给程序任何清理的机会。

小禾的解决方案是：注册信号处理器。

import signal
import sys

def signal_handler(signum, frame):
    """处理终止信号"""
    logger.info(f"收到信号 {signum}，准备清理...")

    try:
        model_manager.cleanup()
    except Exception as e:
        logger.error(f"清理失败: {e}")

    sys.exit(0)

# 注册信号处理器
signal.signal(signal.SIGINT, signal_handler)   # Ctrl+C
signal.signal(signal.SIGTERM, signal_handler)  # kill 命令

这样即使 FastAPI 的 shutdown 事件没触发，信号处理器也能兜底。

---

Uvicorn 的优雅关闭配置

小禾还发现 uvicorn 有个配置项可以控制优雅关闭的超时时间：

# run.py
import uvicorn

if __name__ == "__main__":
    uvicorn.run(
        "app.main:app",
        host="0.0.0.0",
        port=8000,
        reload=False,  # 生产环境不要开 reload
        workers=1,     # AI 应用一般单进程
        timeout_graceful_shutdown=30,  # 给 30 秒清理时间
    )

timeout_graceful_shutdown 参数告诉 uvicorn：收到停止信号后，等待最多 30 秒让程序完成清理工作。如果 30 秒内没清理完，才强制终止。

---

调试技巧：监控显存使用

为了方便调试，小禾写了个显存监控函数：

def log_gpu_memory():
    """打印当前 GPU 显存使用"""
    if torch.cuda.is_available():
        allocated = torch.cuda.memory_allocated() / 1024**3
        reserved = torch.cuda.memory_reserved() / 1024**3
        logger.info(f"GPU 显存: 已分配 {allocated:.2f}GB, 已预留 {reserved:.2f}GB")

在关键位置调用：

@asynccontextmanager
async def lifespan(app: FastAPI):
    log_gpu_memory()  # 启动前
    model_manager.load_model()
    log_gpu_memory()  # 加载后

    yield

    log_gpu_memory()  # 清理前
    model_manager.cleanup()
    log_gpu_memory()  # 清理后

输出类似这样：

GPU 显存: 已分配 0.00GB, 已预留 0.00GB
正在加载模型...
模型加载完成
GPU 显存: 已分配 9.87GB, 已预留 10.12GB
...
正在释放模型资源...
模型资源已释放
GPU 显存: 已分配 0.00GB, 已预留 0.00GB

一目了然。

---

健康检查端点

小禾还加了个健康检查端点，方便运维监控：

@app.get("/health")
async def health_check():
    gpu_info = {}
    if torch.cuda.is_available():
        gpu_info = {
            "allocated_gb": round(torch.cuda.memory_allocated() / 1024**3, 2),
            "reserved_gb": round(torch.cuda.memory_reserved() / 1024**3, 2),
        }

    return {
        "status": "healthy",
        "gpu": gpu_info,
        "model_loaded": model_manager.pipe is not None
    }

请求 /health 就能看到当前状态：

{
    "status": "healthy",
    "gpu": {
        "allocated_gb": 9.87,
        "reserved_gb": 10.12
    },
    "model_loaded": true
}

---

生命周期清单

经过这番折腾，小禾总结了一份清单：

阶段	动作
startup	加载模型、初始化连接、预热缓存
运行中	处理请求、监控资源
shutdown	释放模型、关闭连接、清理缓存

关键点：

1. 使用 lifespan 上下文管理器，启动和关闭逻辑配对
2. cleanup() 要彻底：移到 CPU、删引用、gc、empty_cache
3. 注册信号处理器应对异常关闭
4. 设置 timeout_graceful_shutdown 给清理留时间
5. 添加健康检查端点方便监控

---

小禾的感悟

显存不会说谎，
不释放就是不释放。

启动时加载，
关闭时释放，
听起来很简单，
做起来全是坑。

del 不等于释放，
Ctrl+C 不等于关闭，
kill -9 更是不讲武德。

lifespan 是救星，
信号处理是兜底，
四步清理要牢记：
移到 CPU、删引用、gc、empty_cache。

GPU 很贵，
显存更贵，
别让它们白白浪费。

小禾看着 nvidia-smi 显示的 2GB 占用，心情舒畅。

终于可以放心地重启服务了。

---

下一篇预告：接口报 500 了，日志里却啥都没有

找了三小时的 bug，原来是异常被吞了。

敬请期待。

---

#FastAPI #Python #GPU #显存管理 #生命周期 #AI应用