mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-05 05:18:52 +08:00
59 lines
1.7 KiB
Python
59 lines
1.7 KiB
Python
import pdb
|
|
import logging
|
|
from pymongo import MongoClient
|
|
|
|
from .models.env import Env
|
|
from .models.hardware import Hardware
|
|
from .models.metric import Metric
|
|
from .models.server import Server
|
|
from .config import DB, UNIQUE_ID_COLLECTION, DOC_COLLECTION
|
|
from milvus_benchmark import config
|
|
|
|
|
|
# Initialize the mongoDB client
|
|
_client = MongoClient(config.MONGO_SERVER)
|
|
logger = logging.getLogger("milvus_benchmark.metric.api")
|
|
|
|
|
|
def insert_or_get(md5):
|
|
collection = _client[DB][UNIQUE_ID_COLLECTION]
|
|
found = collection.find_one({'md5': md5})
|
|
if not found:
|
|
return collection.insert_one({'md5': md5}).inserted_id
|
|
return found['_id']
|
|
|
|
|
|
def save(obj):
|
|
if not isinstance(obj, Metric):
|
|
logger.error("obj is not instance of Metric")
|
|
return False
|
|
|
|
if not isinstance(obj.server, Server):
|
|
logger.error("obj.server is not instance of Server")
|
|
return False
|
|
|
|
if not isinstance(obj.hardware, Hardware):
|
|
logger.error("obj.hardware is not instance of Hardware")
|
|
return False
|
|
|
|
if not isinstance(obj.env, Env):
|
|
logger.error("obj.env is not instance of Env")
|
|
return False
|
|
|
|
md5 = obj.server.json_md5()
|
|
server_doc_id = insert_or_get(md5)
|
|
obj.server = {"id": server_doc_id, "value": vars(obj.server)}
|
|
|
|
md5 = obj.hardware.json_md5()
|
|
hardware_doc_id = insert_or_get(md5)
|
|
obj.hardware = {"id": hardware_doc_id, "value": vars(obj.hardware)}
|
|
|
|
md5 = obj.env.json_md5()
|
|
env_doc_id = insert_or_get(md5)
|
|
obj.env = {"id": env_doc_id, "value": vars(obj.env)}
|
|
|
|
# insert data to mongoDB
|
|
logger.debug(vars(obj))
|
|
collection = _client[DB][DOC_COLLECTION]
|
|
collection.insert_one(vars(obj))
|