Model Cache Management Features

requirements.txt

@@ -23,3 +23,4 @@ onnxruntime==1.13.1
 pandas==1.5.1
 optimum==1.4.1
 opencv-python-headless==4.6.0.66
+psutil==5.9.4

setup.py

@@ -30,7 +30,8 @@
         "fastapi",
         "uvicorn[standard]",
         "fastapi_utils",
-        "opencv-python-headless"
+        "opencv-python-headless",
+        "psutil"
     ],
     name="marqo-engine",
     version="0.1.10",

src/marqo/errors.py

@@ -176,6 +176,11 @@ class IndexMaxFieldsError(__InvalidRequestError):
     code = "index_max_fields_error"
     status_code = HTTPStatus.BAD_REQUEST
 
+
+class ModelNotInCacheError(__InvalidRequestError):
+    code = "model_not_in_cache"
+    status_code = HTTPStatus.NOT_FOUND
+
 # ---MARQO INTERNAL ERROR---
 
 

src/marqo/s2_inference/errors.py

@@ -42,3 +42,7 @@ class RerankerImageError(S2InferenceError):
 
 class RerankerNameError(S2InferenceError):
     pass
+
+
+class ModelNotInCacheError(S2InferenceError):
+    pass

src/marqo/s2_inference/s2_inference.py

@@ -2,13 +2,13 @@
 The functions defined here would have endpoints, later on.
 """
 import numpy as np
-from marqo.s2_inference.errors import VectoriseError, InvalidModelPropertiesError, ModelLoadError, UnknownModelError
+from marqo.s2_inference.errors import VectoriseError, InvalidModelPropertiesError, ModelLoadError, UnknownModelError, ModelNotInCacheError
 from PIL import UnidentifiedImageError
 from marqo.s2_inference.model_registry import load_model_properties
 from marqo.s2_inference.configs import get_default_device, get_default_normalization, get_default_seq_length
 from marqo.s2_inference.types import *
 from marqo.s2_inference.logger import get_logger
-from timeit import default_timer as timer
+import torch
 
 logger = get_logger(__name__)
 
@@ -36,18 +36,17 @@ def vectorise(model_name: str, content: Union[str, List[str]], model_properties:
     Raises:
         VectoriseError: if the content can't be vectorised, for some reason.
     """
+
     validated_model_properties = _validate_model_properties(model_name, model_properties)
     model_cache_key = _create_model_cache_key(model_name, device, validated_model_properties)
 
     _update_available_models(model_cache_key, model_name, validated_model_properties, device, normalize_embeddings)
 
-
     try:
         vectorised = available_models[model_cache_key].encode(content, normalize=normalize_embeddings, **kwargs)
     except UnidentifiedImageError as e:
         raise VectoriseError from e
 
-
     return _convert_vectorized_output(vectorised)
 
 
@@ -62,15 +61,17 @@ def _create_model_cache_key(model_name: str, device: str, model_properties: dict
     Returns:
         str: _description_
     """
+    # Changing the format of model cache key will also need to change eject_model api
+
     if model_properties is None:
         model_properties = dict()
 
-    model_cache_key = (model_name
-                       + model_properties.get('name', '')
-                       + str(model_properties.get('dimensions', ''))
-                       + model_properties.get('type', '')
-                       + str(model_properties.get('tokens', ''))
-                       + device)
+    model_cache_key = (model_name + "||" +
+                       model_properties.get('name', '') + "||" +
+                       str(model_properties.get('dimensions', '')) + "||" +
+                       model_properties.get('type', '') + "||" +
+                       str(model_properties.get('tokens', '')) + "||" +
+                       device)
 
     return model_cache_key
 
@@ -293,6 +294,35 @@ def _load_model(model_name: str, model_properties: dict, device: str = get_defau
 
     return model
 
+
+def get_available_models():
+    return available_models
+
+
+def eject_model(model_name:str, device:str):
+
+    model_cache_keys = available_models.keys()
+
+    model_cache_key = None
+
+    # we can't handle the situation where there are two models with the same name and device
+    # but different properties.
+    for key in model_cache_keys:
+        if key.startswith(model_name) and key.endswith(device):
+            model_cache_key = key
+            break
+
+    if model_cache_key is None:
+        raise ModelNotInCacheError(f"The model_name \`{model_name}\` device \`{device}\` is not cached or found")
+
+    if model_cache_key in available_models:
+        del available_models[model_cache_key]
+        if device.startswith("cuda"):
+            torch.cuda.empty_cache()
+        return {"result": "success", "message": f"successfully eject model_name \`{model_name}\` from device \`{device}\`"}
+    else:
+        raise ModelNotInCacheError(f"The model_name \`{model_name}\` device \`{device}\` is not cached or found")
+
 # def normalize(inputs):
 
 #     is_valid = False

src/marqo/tensor_search/api.py

@@ -212,6 +212,28 @@ def check_health(marqo_config: config.Config = Depends(generate_config)):
 def get_indexes(marqo_config: config.Config = Depends(generate_config)):
     return tensor_search.get_indexes(config=marqo_config)
 
+
+@app.get("/models")
+def get_loaded_models():
+    return tensor_search.get_loaded_models()
+
+
+@app.delete("/models")
+def eject_model(model_name:str, model_device:str):
+    return tensor_search.eject_model(model_name = model_name, device = model_device)
+
+
+@app.get("/device/cpu")
+def get_cpu_info():
+    return tensor_search.get_cpu_info()
+
+
+@app.get("/device/cuda")
+def get_cuda_info():
+    return tensor_search.get_cuda_info()
+
+
+
 # try these curl commands:
 
 # ADD DOCS:
@@ -282,3 +304,25 @@ def get_indexes(marqo_config: config.Config = Depends(generate_config)):
 curl -XDELETE http://localhost:8882/indexes/my-irst-ix
 """
 
+# check cpu info
+"""
+curl -XGET http://localhost:8882/device/cpu
+"""
+
+# check cuda info
+"""
+curl -XGET http://localhost:8882/device/cuda
+"""
+
+# check the loaded models
+"""
+curl -XGET http://localhost:8882/models
+"""
+
+# eject a model
+"""
+curl -X DELETE 'http://localhost:8882/models?model_name=ViT-L/14&model_device=cuda'
+curl -X DELETE 'http://localhost:8882/models?model_name=ViT-L/14&model_device=cpu'
+curl -X DELETE 'http://localhost:8882/models?model_name=hf/all_datasets_v4_MiniLM-L6&model_device=cuda' 
+curl -X DELETE 'http://localhost:8882/models?model_name=hf/all_datasets_v4_MiniLM-L6&model_device=cpu' 
+"""

src/marqo/tensor_search/tensor_search.py

@@ -55,7 +55,8 @@
 from marqo.s2_inference.clip_utils import _is_image
 from marqo.s2_inference.reranking import rerank
 from marqo.s2_inference import s2_inference
-
+import torch.cuda
+import psutil
 # We depend on _httprequests.py for now, but this may be replaced in the future, as
 # _httprequests.py is designed for the client
 from marqo._httprequests import HttpRequests
@@ -1330,3 +1331,41 @@ def _get_model_properties(index_info):
                 f"Please provide model_properties if the model is a custom model and is not supported by default")
 
     return model_properties
+
+def get_loaded_models() -> dict:
+    available_models = s2_inference.get_available_models()
+    message = {
+        'models' : [
+            {"model_name": ix.split("||")[0], "model_device": ix.split("||")[-1]} for ix in available_models.keys()
+        ]
+    }
+    return message
+
+
+def eject_model(model_name: str, device: str) -> dict:
+    try:
+       result = s2_inference.eject_model(model_name, device)
+    except s2_inference_errors.ModelNotInCacheError as e:
+        raise errors.ModelNotInCacheError(message=str(e))
+    return result
+
+
+def get_cpu_info() -> dict:
+    return {
+        "cpu_usage_percent": f"{psutil.cpu_percent(1)} %", # The number 1 is a time interval for CPU usage calculation.
+        "memory_used_percent": f"{psutil.virtual_memory()[2]} %",  # The number 2 is just a index number to get the expected results
+        "memory_used_gb": f"{round(psutil.virtual_memory()[3]/1000000000,1)}", # The number 3 is just a index number to get the expected results
+    }
+
+
+def get_cuda_info() -> dict:
+    if torch.cuda.is_available():
+        return {"cuda_devices": [{"device_id" : _device_id, "device_name" : torch.cuda.get_device_name(_device_id),
+                "memory_used":f"{round(torch.cuda.memory_allocated(_device_id) / 1024**3, 1)} GiB",
+                "total_memory": f"{round(torch.cuda.get_device_properties(_device_id).total_memory/ 1024**3, 1)} GiB"}
+                for _device_id in range(torch.cuda.device_count())]}
+
+    else:
+        raise errors.HardwareCompatabilityError(message=str(
+            "ERROR: cuda is not supported in your machine!!"
+        ))

tests/s2_inference/test_utils.py

@@ -39,11 +39,11 @@ def test_create_model_cache_key(self):
                 assert (
                             _create_model_cache_key(name, device, model_properties)
                             == (
-                               name
-                               + model_properties.get('name', '')
-                               + str(model_properties.get('dimensions', ''))
-                               + model_properties.get('type', '')
-                               + str(model_properties.get('tokens', ''))
+                               name + "||"
+                               + model_properties.get('name', '') + "||"
+                               + str(model_properties.get('dimensions', '')) + "||"
+                               + model_properties.get('type', '') + "||"
+                               + str(model_properties.get('tokens', '')) + "||"
                                + device)
                 )