作为一名机器学习工程师，你已经成功训练了模型并将其部署到云端。然而，你创建的REST API端点并不安全------任何拥有该URL的人都可以访问。这带来了重大的安全风险。
那么，该如何解决这一问题呢？只添加一个静态API密钥就够了吗？答案是否定的。这远远不够。你需要实现一个完善的用户管理系统。
用户管理系统可以让你创建用户，并为其授予模型推理服务和其他功能的访问权限。这样，如果某个用户行为异常或其凭证被泄露，你可以轻松撤销其访问权限，而不会影响其他用户。这种方法能够为你的应用提供更好的控制与安全性。
在本教程中，我们将学习如何为机器学习应用设置认证机制，并搭建一个用户管理系统，使管理员可以根据需要添加或移除用户。最后，我们还会通过多种用例测试应用，确保各项功能都正确实现。
FastAPI应用规划
我们需要创建一个机器学习应用，管理员拥有高级权限，包括：
-
访问模型推理端点。
-
在系统中创建和删除用户。
系统设计如下：
-
管理员API密钥
：管理员通过凭证生成API密钥。
-
创建新用户
：管理员使用API密钥创建新用户。
-
新用户登录
：新用户使用用户名和密码登录，生成自己的API密钥。
-
访问模型
：用户使用API密钥访问模型信息并进行推理。
-
移除用户
：如有需要，管理员可用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，交互测试各端点。
也可用脚本自动化测试各功能，包括：
-
健康检查
：测试
/health，无需认证。 -
管理员登录
：使用管理员账号获取API密钥。
-
创建用户
：用管理员API密钥创建新用户。
-
用户登录
：新用户登录，获取其API密钥。
-
查询模型信息
：用用户API密钥访问
/model-info。 -
模型推理
：上传图片，测试
/predict端点。 -
删除用户
：管理员删除新建用户。
自动化测试脚本片段：
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密钥安全生成与管理，极大提升安全性。
-
代码简洁，易于扩展，适合真实生产环境部署。
这种架构不仅安全、灵活，还能便捷地适配各类实际业务需求，是部署和维护机器学习服务的理想选择。