构建安全的机器学习推理API：基于FastAPI的用户认证与管理实战

作为一名机器学习工程师，你已经成功训练了模型并将其部署到云端。然而，你创建的REST API端点并不安全——任何拥有该URL的人都可以访问。这带来了重大的安全风险。

那么，该如何解决这一问题呢？只添加一个静态API密钥就够了吗？答案是否定的。这远远不够。你需要实现一个完善的用户管理系统。

用户管理系统可以让你创建用户，并为其授予模型推理服务和其他功能的访问权限。这样，如果某个用户行为异常或其凭证被泄露，你可以轻松撤销其访问权限，而不会影响其他用户。这种方法能够为你的应用提供更好的控制与安全性。

在本教程中，我们将学习如何为机器学习应用设置认证机制，并搭建一个用户管理系统，使管理员可以根据需要添加或移除用户。最后，我们还会通过多种用例测试应用，确保各项功能都正确实现。

---

FastAPI应用规划

我们需要创建一个机器学习应用，管理员拥有高级权限，包括：

• 访问模型推理端点。

• 在系统中创建和删除用户。

系统设计如下：

1. 管理员API密钥

   ：管理员通过凭证生成API密钥。

2. 创建新用户

   ：管理员使用API密钥创建新用户。

3. 新用户登录

   ：新用户使用用户名和密码登录，生成自己的API密钥。

4. 访问模型

   ：用户使用API密钥访问模型信息并进行推理。

5. 移除用户

   ：如有需要，管理员可用API密钥移除某用户。

通过这样的设计，应用实现了安全且规范的基于角色的访问控制（RBAC）。管理员掌控用户管理，普通用户仅能访问推理相关功能。

---

通过认证与用户管理保护机器学习应用

FastAPI应用中添加认证与用户管理

在此前的“用FastAPI进行图像分类推理”项目中，我们已经训练并本地保存了简单的图像分类模型。现在，我们将创建一个机器学习推理应用：加载模型文件，允许用户上传图像文件并生成分类标签。

本次重点在于加入认证机制和用户管理系统，以确保安全访问和合理的权限分配。

主要功能包括：

• 环境设置

  ：加载管理员凭证和API配置的环境变量。

• 数据库

  ：使用SQLite存储用户凭证和API密钥，采用SQLAlchemy进行ORM管理。

• 用户管理

  ：管理员可安全地添加/移除用户，支持密码哈希和API密钥生成。

• 认证机制

  ：基于API密钥认证，保障端点安全。

• 机器学习模型

  ：加载针对CIFAR10微调的ResNet18模型，进行图像分类。

• 图像处理

  ：支持图像上传、预处理和推理，并有异常处理机制。

• API端点

  ：

  • /health

    ：健康检查

  • /model-info

    ：模型详情

  • /login

    ：用户认证并返回API密钥

  • /predict

    ：图像分类

  • /admin/add-user

     和 /admin/remove-user：仅管理员可用的用户管理

• 部署

  ：支持多进程（worker）的Uvicorn运行，具备扩展性。

---

主要代码结构如下：

（为保证可读性，代码部分分段翻译与整理，简述主流程并保留关键实现要点。）

1. 环境变量与密码哈希设置

# 加载环境变量
if not load_dotenv():raise ValueError("无法加载.env文件")
ADMIN_USERNAME = os.getenv("ADMIN_USERNAME")
ADMIN_PASSWORD = os.getenv("ADMIN_PASSWORD")
if not ADMIN_USERNAME or not ADMIN_PASSWORD:raise RuntimeError("必须在.env文件中设置ADMIN_USERNAME和ADMIN_PASSWORD。")# 密码哈希
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")

2. 数据库和用户模型

DATABASE_URL = "sqlite:///./data/database/app.db"
engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()class User(Base):__tablename__ = "users"id = Column(Integer, primary_key=True, index=True)username = Column(String, unique=True, index=True, nullable=False)password = Column(String, nullable=False)  # 存储哈希后的密码api_key = Column(String, unique=True, index=True, nullable=False)Base.metadata.create_all(bind=engine)

3. 密码处理与API密钥认证

def hash_password(password: str) -> str:return pwd_context.hash(password)def verify_password(plain_password: str, hashed_password: str) -> bool:return pwd_context.verify(plain_password, hashed_password)def authenticate_user(db: Session, username: str, password: str):# 支持老密码加密方式向新方式迁移...

4. FastAPI端点定义

• /health

  ：健康检查，返回API状态

• /model-info

  ：带认证，返回模型结构及类别

• /login

  ：用户登录，返回API密钥

• /predict

  ：带认证，上传图片返回预测类别或Top-K概率

• /admin/add-user

  、/admin/remove-user：管理员用户管理端点

示例：预测端点

@app.post("/predict", summary="预测图像类别", tags=["推理"])
async def predict(...):# 验证上传文件# 图像预处理# 推理与结果返回

5. 多进程部署

if __name__ == "__main__":uvicorn.run("cv_app:app", host="127.0.0.1", port=6565, workers=4, log_level="info")

---

创建.env文件示例

ADMIN_USERNAME=admin
ADMIN_PASSWORD=你的管理员密码

---

机器学习FastAPI应用测试

应用启动后，可通过 http://127.0.0.1:6565/docs 访问Swagger UI，交互测试各端点。

也可用脚本自动化测试各功能，包括：

1. 健康检查

   ：测试 /health，无需认证。

2. 管理员登录

   ：使用管理员账号获取API密钥。

3. 创建用户

   ：用管理员API密钥创建新用户。

4. 用户登录

   ：新用户登录，获取其API密钥。

5. 查询模型信息

   ：用用户API密钥访问 /model-info。

6. 模型推理

   ：上传图片，测试 /predict 端点。

7. 删除用户

   ：管理员删除新建用户。

自动化测试脚本片段：

def health_check():response = requests.get(f"{BASE_URL}/health")print(json.dumps(response.json(), indent=4))def admin_login():...
def create_test_user(admin_api_key):...
# 其他功能同理

运行结果示例：

=== Testing /health endpoint ===
{"status": "ok","message": "API is running","device": "cuda:0"
}
...
All tests completed successfully.

---

结论

为你的机器学习应用添加用户管理和认证机制变得前所未有的简单。你无需依赖第三方工具或外部服务器，仅需借助FastAPI框架和一些巧妙的实现，就能打造出强大而安全的生产级应用。

本项目实现了如下目标：

• 构建了包含用户管理的机器学习推理API。

• 实现了基于角色的访问控制（RBAC）：管理员可管理用户，普通用户仅能推理与查询模型。

• API密钥安全生成与管理，极大提升安全性。

• 代码简洁，易于扩展，适合真实生产环境部署。

这种架构不仅安全、灵活，还能便捷地适配各类实际业务需求，是部署和维护机器学习服务的理想选择。