添加--模型管理方法添加删除模型功能

This commit is contained in:
haotian 2025-02-19 18:13:29 +08:00
parent 6cc0910112
commit 7b99da7251
6 changed files with 80 additions and 3 deletions

View File

@ -456,8 +456,26 @@ Response:
} }
``` ```
### 2.7 删除指定的训练好的模型. ### 2.7 删除指定的训练好的模型
```http
DELETE /api/models/{run_id}
Response:
{
"status": "success",
"message": "模型删除成功",
"details": {
"run_id": "7970364d490f4e0aa0375c2db26215f3",
"experiment_id": "656341556838275234",
"model_name": "XGBClassifier",
"deleted_artifacts": [
"models/7970364d490f4e0aa0375c2db26215f3/model.pkl",
"models/7970364d490f4e0aa0375c2db26215f3/requirements.txt",
"models/7970364d490f4e0aa0375c2db26215f3/conda.yaml"
]
}
}
```
### 2.8 获取训练状态 ### 2.8 获取训练状态
```http ```http

View File

@ -10,6 +10,7 @@
- 20250219~20250220 - 20250219~20250220
- [x] 实现模型训练和注册 完成 - [x] 实现模型训练和注册 完成
- [x] 实现获取已有模型 完成 - [x] 实现获取已有模型 完成
- [ ] 根据删除模型
- [ ] 根据已有模型加载模型预测 - [ ] 根据已有模型加载模型预测
- 2天 - 2天

View File

@ -210,6 +210,62 @@ class ModelManager:
except Exception as e: except Exception as e:
error_msg = f"Error getting experiments: {str(e)}" error_msg = f"Error getting experiments: {str(e)}"
self.logger.error(error_msg) self.logger.error(error_msg)
return {
'status': 'error',
'message': error_msg
}
def delete_model(self, run_id: str) -> Dict:
"""
删除指定的训练好的模型
Args:
run_id: MLflow运行ID
Returns:
删除操作的结果信息
"""
try:
# 获取运行信息
run = self.client.get_run(run_id)
if not run:
return {
'status': 'error',
'message': f'未找到运行ID为 {run_id} 的模型'
}
# 获取实验ID
experiment_id = run.info.experiment_id
# 获取模型信息
run_data = mlflow.get_run(run_id)
model_name = run_data.data.params.get('algorithm', 'Unknown')
# 获取工件列表
artifacts = []
for artifact in self.client.list_artifacts(run_id):
artifacts.append(artifact.path)
# 删除运行记录
self.client.delete_run(run_id)
# 记录日志
self.logger.info(f"已删除模型 - Run ID: {run_id}, 实验ID: {experiment_id}, 模型: {model_name}")
return {
'status': 'success',
'message': '模型删除成功',
'details': {
'run_id': run_id,
'experiment_id': experiment_id,
'model_name': model_name,
'deleted_artifacts': artifacts
}
}
except Exception as e:
error_msg = f"删除模型时发生错误: {str(e)}"
self.logger.error(error_msg)
return { return {
'status': 'error', 'status': 'error',
'message': error_msg 'message': error_msg

View File

@ -1,8 +1,9 @@
artifact_uri: mlflow-artifacts:/656341556838275234/7970364d490f4e0aa0375c2db26215f3/artifacts artifact_uri: mlflow-artifacts:/656341556838275234/7970364d490f4e0aa0375c2db26215f3/artifacts
deleted_time: 1739959968072
end_time: 1739954585256 end_time: 1739954585256
entry_point_name: '' entry_point_name: ''
experiment_id: '656341556838275234' experiment_id: '656341556838275234'
lifecycle_stage: active lifecycle_stage: deleted
run_id: 7970364d490f4e0aa0375c2db26215f3 run_id: 7970364d490f4e0aa0375c2db26215f3
run_name: colorful-stoat-263 run_name: colorful-stoat-263
run_uuid: 7970364d490f4e0aa0375c2db26215f3 run_uuid: 7970364d490f4e0aa0375c2db26215f3

View File

@ -1,8 +1,9 @@
artifact_uri: mlflow-artifacts:/656341556838275234/c88c5a1db98f49db9d8793ad19a08809/artifacts artifact_uri: mlflow-artifacts:/656341556838275234/c88c5a1db98f49db9d8793ad19a08809/artifacts
deleted_time: 1739959938908
end_time: 1739954359528 end_time: 1739954359528
entry_point_name: '' entry_point_name: ''
experiment_id: '656341556838275234' experiment_id: '656341556838275234'
lifecycle_stage: active lifecycle_stage: deleted
run_id: c88c5a1db98f49db9d8793ad19a08809 run_id: c88c5a1db98f49db9d8793ad19a08809
run_name: thoughtful-hawk-675 run_name: thoughtful-hawk-675
run_uuid: c88c5a1db98f49db9d8793ad19a08809 run_uuid: c88c5a1db98f49db9d8793ad19a08809