diff --git a/apps/setting/serializers/provider_serializers.py b/apps/setting/serializers/provider_serializers.py index ad4c7d265..2ac28343a 100644 --- a/apps/setting/serializers/provider_serializers.py +++ b/apps/setting/serializers/provider_serializers.py @@ -28,25 +28,30 @@ class ModelPullManage: @staticmethod def pull(model: Model, credential: Dict): - response = ModelProvideConstants[model.provider].value.down_model(model.model_type, model.model_name, - credential) - down_model_chunk = {} - timestamp = time.time() - for chunk in response: - down_model_chunk[chunk.digest] = chunk.to_dict() - if time.time() - timestamp > 5: - QuerySet(Model).filter(id=model.id).update(meta={"down_model_chunk": list(down_model_chunk.values())}) - timestamp = time.time() - status = Status.ERROR - message = "" - down_model_chunk_list = list(down_model_chunk.values()) - for chunk in down_model_chunk_list: - if chunk.get('status') == DownModelChunkStatus.success.value: - status = Status.SUCCESS - if chunk.get('status') == DownModelChunkStatus.error.value: - message = chunk.get("digest") - QuerySet(Model).filter(id=model.id).update(meta={"down_model_chunk": [], "message": message}, - status=status) + try: + response = ModelProvideConstants[model.provider].value.down_model(model.model_type, model.model_name, + credential) + down_model_chunk = {} + timestamp = time.time() + for chunk in response: + down_model_chunk[chunk.digest] = chunk.to_dict() + if time.time() - timestamp > 5: + QuerySet(Model).filter(id=model.id).update( + meta={"down_model_chunk": list(down_model_chunk.values())}) + timestamp = time.time() + status = Status.ERROR + message = "" + down_model_chunk_list = list(down_model_chunk.values()) + for chunk in down_model_chunk_list: + if chunk.get('status') == DownModelChunkStatus.success.value: + status = Status.SUCCESS + if chunk.get('status') == DownModelChunkStatus.error.value: + message = chunk.get("digest") + QuerySet(Model).filter(id=model.id).update(meta={"down_model_chunk": [], "message": message}, + 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):