milvus/tests/benchmark/milvus_benchmark/metrics/api.py

59 lines
1.7 KiB
Python
Raw Normal View History

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))