Merge pull request #288 from univerone/add-model-retrieving

[MRG]Fix model retrieving from modelhub

modelci/types/models/mlmodel.py

@@ -7,45 +7,61 @@
 """
 import getpass
 import os
+import gridfs
 from datetime import datetime
 from pathlib import Path
 from typing import Union, Optional, Dict, List, Any
 
 from bson import ObjectId
-from gridfs import GridOut
 from pydantic import BaseModel, FilePath, DirectoryPath, PositiveInt, Field, root_validator
 
 from .common import Metric, IOShape, Framework, Engine, Task, ModelStatus, Status, PydanticObjectId, \
     named_enum_json_encoder
 from .pattern import as_form
 from ...hub.utils import parse_path_plain, generate_path_plain
+from modelci.config import db_settings
+from modelci.experimental.mongo_client import MongoClient
+
+_db = MongoClient()[db_settings.mongo_db]
+_fs = gridfs.GridFS(_db)
 
 
 class Weight(BaseModel):
     """TODO: Only works for MLModelIn"""
 
-    __slots__ = ('file',)
+    __slots__ = ('file', '_gridfs_out')
 
     __root__: Optional[PydanticObjectId]
 
     def __init__(self, __root__):
         if isinstance(__root__, Path):
             object.__setattr__(self, 'file', FilePath.validate(__root__))
+            object.__setattr__(self, '_gridfs_out', None)
             __root__ = None
 
-        self._grid_out: Optional[GridOut]
+        if isinstance(__root__, ObjectId):
+            if _fs.exists(__root__):
+                object.__setattr__(self, 'file', None)
+                object.__setattr__(self, '_gridfs_out', _fs.get(__root__))
 
         super().__init__(__root__=__root__)
 
     @property
     def filename(self):
         if self.file:
             return self.file.name
-        return ''
+        elif self._gridfs_out:
+            return self._gridfs_out.filename
+        else:
+            return ''
 
     def __bytes__(self):
         if self.file:
             return self.file.read_bytes()
+        elif self._gridfs_out:
+            return self._gridfs_out.read()
+        else:
+            return b''
 
 
 @as_form
@@ -85,7 +101,7 @@ def dict(self, use_enum_values: bool = False, **kwargs):
         # fix metric key as a Enum
         metric: dict = data.get('metric', None)
         if metric:
-            data['metric'] = {Metric(k).name: v for k, v in metric.items() }
+            data['metric'] = {Metric(k).name: v for k, v in metric.items()}
 
         return data
 
@@ -182,4 +198,4 @@ class ModelUpdateSchema(BaseModel):
     outputs: Optional[List[IOShape]] = Field(
         default_factory=list,
         example='[{"name": "output", "shape": [-1, 1000], "dtype": "TYPE_FP32"}]'
-    )
+    )