zhuchaowe/maxkb
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: 下载模型出错处理

Browse Source
This commit is contained in:
shaohuzhang1 2024-03-26 16:55:41 +08:00
parent b891d0612f
commit 8d318d2a41
1 changed files with 24 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
apps/setting/serializers/provider_serializers.py
View File

@ -28,6 +28,7 @@ class ModelPullManage:
@staticmethod @staticmethod
def pull(model: Model, credential: Dict): def pull(model: Model, credential: Dict):
try:
response = ModelProvideConstants[model.provider].value.down_model(model.model_type, model.model_name, response = ModelProvideConstants[model.provider].value.down_model(model.model_type, model.model_name,
credential) credential)
down_model_chunk = {} down_model_chunk = {}
@ -35,7 +36,8 @@ class ModelPullManage:
for chunk in response: for chunk in response:
down_model_chunk[chunk.digest] = chunk.to_dict() down_model_chunk[chunk.digest] = chunk.to_dict()
if time.time() - timestamp > 5: if time.time() - timestamp > 5:
QuerySet(Model).filter(id=model.id).update(meta={"down_model_chunk": list(down_model_chunk.values())}) QuerySet(Model).filter(id=model.id).update(
meta={"down_model_chunk": list(down_model_chunk.values())})
timestamp = time.time() timestamp = time.time()
status = Status.ERROR status = Status.ERROR
message = "" message = ""
@ -47,6 +49,9 @@ class ModelPullManage:
message = chunk.get("digest") message = chunk.get("digest")
QuerySet(Model).filter(id=model.id).update(meta={"down_model_chunk": [], "message": message}, QuerySet(Model).filter(id=model.id).update(meta={"down_model_chunk": [], "message": message},
status=status) status=status)
except Exception as e:
QuerySet(Model).filter(id=model.id).update(meta={"down_model_chunk": [], "message": str(e)},
status=Status.ERROR)
class ModelSerializer(serializers.Serializer): class ModelSerializer(serializers.Serializer):