[ODSC-52993] Set up caching for listing models and compartments (#592)

Author: VipulMascarenhas

THIRD_PARTY_LICENSES.txt

@@ -405,6 +405,12 @@ delta
 * Source code: https://github.com/delta-io/delta/
 * Project home: https://delta.io/
 
+cachetools
+* Copyright (c) 2014-2024 Thomas Kemmer
+* License: The MIT License (MIT)
+* Source code: https://github.com/tkem/cachetools/
+* Project home: https://cachetools.readthedocs.io/
+
 =============================== Licenses ===============================
 ------------------------------------------------------------------------
 

ads/aqua/extension/model_handler.py

@@ -3,6 +3,8 @@
 # Copyright (c) 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
+from urllib.parse import urlparse
+from tornado.web import HTTPError
 from ads.aqua.decorator import handle_exceptions
 from ads.aqua.extension.base_handler import AquaAPIhandler
 from ads.aqua.model import AquaModelApp
@@ -14,7 +16,6 @@ class AquaModelHandler(AquaAPIhandler):
  @handle_exceptions
  def get(self, model_id=""):
  """Handle GET request."""
-
  if not model_id:
  return self.list()
  return self.read(model_id)
@@ -24,6 +25,16 @@ def read(self, model_id):
  """Read the information of an Aqua model."""
  return self.finish(AquaModelApp().get(model_id))
 
+ @handle_exceptions
+ def delete(self, id=""):
+ """Handles DELETE request for clearing cache"""
+ url_parse = urlparse(self.request.path)
+ paths = url_parse.path.strip("/")
+ if paths.startswith("aqua/model/cache"):
+ return self.finish(AquaModelApp().clear_model_list_cache())
+ else:
+ raise HTTPError(400, f"The request {self.request.path} is invalid.")
+
  @handle_exceptions
  def list(self):
  """List Aqua models."""

ads/aqua/extension/ui_handler.py

@@ -4,9 +4,7 @@
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 from urllib.parse import urlparse
-
 from tornado.web import HTTPError
-
 from ads.aqua.decorator import handle_exceptions
 from ads.aqua.extension.base_handler import AquaAPIhandler
 from ads.aqua.ui import AquaUIApp
@@ -54,26 +52,34 @@ def get(self, id=""):
  else:
  raise HTTPError(400, f"The request {self.request.path} is invalid.")
 
+ @handle_exceptions
+ def delete(self, id=""):
+ """Handles DELETE request for clearing cache"""
+ # todo: added for dev work, to be deleted if there's no feature to refresh cache in Aqua
+ url_parse = urlparse(self.request.path)
+ paths = url_parse.path.strip("/")
+ if paths.startswith("aqua/compartments/cache"):
+ return self.finish(AquaUIApp().clear_compartments_list_cache())
+ else:
+ raise HTTPError(400, f"The request {self.request.path} is invalid.")
+
  @handle_exceptions
  def list_log_groups(self, **kwargs):
  """Lists all log groups for the specified compartment or tenancy."""
  compartment_id = self.get_argument("compartment_id", default=COMPARTMENT_OCID)
-
  return self.finish(
  AquaUIApp().list_log_groups(compartment_id=compartment_id, **kwargs)
  )
 
  @handle_exceptions
  def list_logs(self, log_group_id: str, **kwargs):
  """Lists the specified log group's log objects."""
-
  return self.finish(AquaUIApp().list_logs(log_group_id=log_group_id, **kwargs))
 
  @handle_exceptions
- def list_compartments(self, **kwargs):
+ def list_compartments(self):
  """Lists the compartments in a compartment specified by ODSC_MODEL_COMPARTMENT_OCID env variable."""
-
- return self.finish(AquaUIApp().list_compartments(**kwargs))
+ return self.finish(AquaUIApp().list_compartments())
 
  @handle_exceptions
  def get_default_compartment(self):

ads/aqua/model.py

@@ -6,9 +6,11 @@
 from dataclasses import dataclass
 from enum import Enum
 from typing import List, Union
+from datetime import datetime, timedelta
+from threading import Lock
+from cachetools import TTLCache
 
 import oci
-
 from ads.aqua import logger
 from ads.aqua.base import AquaApp
 from ads.aqua.exception import AquaRuntimeError
@@ -80,13 +82,20 @@ class AquaModelApp(AquaApp):
  Retrieves details of an Aqua model by its unique identifier.
  list(compartment_id: str = None, project_id: str = None, **kwargs) -> List[AquaModelSummary]:
  Lists all Aqua models within a specified compartment and/or project.
+ clear_model_list_cache()
+ Allows clear list model cache items from the service models compartment.
 
  Note:
  This class is designed to work within the Oracle Cloud Infrastructure
  and requires proper configuration and authentication set up to interact
  with OCI services.
  """
 
+ _service_models_cache = TTLCache(
+ maxsize=10, ttl=timedelta(hours=5), timer=datetime.now
+ )
+ _cache_lock = Lock()
+
  def create(
  self, model_id: str, project_id: str, compartment_id: str = None, **kwargs
  ) -> DataScienceModel:
@@ -134,10 +143,7 @@ def create(
  .with_defined_metadata_list(service_model.defined_metadata_list)
  .with_provenance_metadata(service_model.provenance_metadata)
  # TODO: decide what kwargs will be needed.
- .create(
- model_by_reference=True, 
- **kwargs
- )
+ .create(model_by_reference=True, **kwargs)
  )
  logger.debug(
  f"Aqua Model {custom_model.id} created with the service model {model_id}"
@@ -177,8 +183,9 @@ def list(
  ) -> List["AquaModelSummary"]:
  """Lists all Aqua models within a specified compartment and/or project.
  If `compartment_id` is not specified, the method defaults to returning
- the service models within the pre-configured default compartment.
-
+ the service models within the pre-configured default compartment. By default, the list
+ of models in the service compartment are cached. Use clear_model_list_cache() to invalidate
+ the cache.
 
  Parameters
  ----------
@@ -199,6 +206,11 @@ def list(
  logger.info(f"Fetching custom models from compartment_id={compartment_id}.")
  models = self._rqs(compartment_id)
  else:
+ if ODSC_MODEL_COMPARTMENT_OCID in self._service_models_cache.keys():
+ logger.info(
+ f"Returning service models list in {ODSC_MODEL_COMPARTMENT_OCID} from cache."
+ )
+ return self._service_models_cache.get(ODSC_MODEL_COMPARTMENT_OCID)
  logger.info(
  f"Fetching service models from compartment_id={ODSC_MODEL_COMPARTMENT_OCID}"
  )
@@ -213,7 +225,6 @@ def list(
  )
 
  aqua_models = []
- # TODO: build index.json for service model as caching if needed.
 
  for model in models:
  aqua_models.append(
@@ -223,8 +234,35 @@ def list(
  )
  )
 
+ if not compartment_id:
+ self._service_models_cache.__setitem__(
+ key=ODSC_MODEL_COMPARTMENT_OCID, value=aqua_models
+ )
+
  return aqua_models
 
+ def clear_model_list_cache(
+ self,
+ ):
+ """
+ Allows user to clear list model cache items from the service models compartment.
+ Returns
+ -------
+ dict with the key used, and True if cache has the key that needs to be deleted.
+ """
+ res = {}
+ logger.info(f"Clearing _service_models_cache")
+ with self._cache_lock:
+ if ODSC_MODEL_COMPARTMENT_OCID in self._service_models_cache.keys():
+ self._service_models_cache.pop(key=ODSC_MODEL_COMPARTMENT_OCID)
+ res = {
+ "key": {
+ "compartment_id": ODSC_MODEL_COMPARTMENT_OCID,
+ },
+ "cache_deleted": True,
+ }
+ return res
+
  def _process_model(
  self, model: Union["ModelSummary", "Model", "ResourceSummary"], region: str
  ) -> dict:

ads/aqua/ui.py

@@ -5,6 +5,9 @@
 
 from oci.exceptions import ServiceError
 from oci.identity.models import Compartment
+from datetime import datetime, timedelta
+from threading import Lock
+from cachetools import TTLCache
 
 from ads.aqua import logger
 from ads.aqua.base import AquaApp
@@ -30,6 +33,11 @@ class AquaUIApp(AquaApp):
 
  """
 
+ _compartments_cache = TTLCache(
+ maxsize=10, ttl=timedelta(hours=2), timer=datetime.now
+ )
+ _cache_lock = Lock()
+
  def list_log_groups(self, **kwargs) -> str:
  """Lists all log groups for the specified compartment or tenancy. This is a pass through the OCI list_log_groups
  API.
@@ -72,16 +80,10 @@ def list_logs(self, **kwargs) -> str:
  log_group_id=log_group_id, **kwargs
  ).data.__repr__()
 
- def list_compartments(self, **kwargs) -> str:
- """Lists the compartments in a compartment specified by TENANCY_OCID env variable. This is a pass through the OCI list_compartments
+ def list_compartments(self) -> str:
+ """Lists the compartments in a tenancy specified by TENANCY_OCID env variable. This is a pass through the OCI list_compartments
  API.
 
- Parameters
- ----------
- kwargs
- Keyword arguments, such as compartment_id,
- for `list_compartments <https://docs.oracle.com/en-us/iaas/tools/python/2.119.1/api/logging/client/oci.identity.IdentityClient.html#oci.identity.IdentityClient.list_compartments>`_
-
  Returns
  -------
  str:
@@ -93,6 +95,13 @@ def list_compartments(self, **kwargs) -> str:
  f"TENANCY_OCID must be available in environment"
  " variables to list the sub compartments."
  )
+
+ if TENANCY_OCID in self._compartments_cache.keys():
+ logger.info(
+ f"Returning compartments list in {TENANCY_OCID} from cache."
+ )
+ return self._compartments_cache.get(TENANCY_OCID)
+
  compartments = []
  # User may not have permissions to list compartment.
  try:
@@ -132,7 +141,13 @@ def list_compartments(self, **kwargs) -> str:
  0,
  Compartment(id=TENANCY_OCID, name=" ** Root - Name N/A **"),
  )
- return compartments.__repr__()
+ # convert the string of the results flattened as a dict
+ res = compartments.__repr__()
+
+ # cache compartment results
+ self._compartments_cache.__setitem__(key=TENANCY_OCID, value=res)
+
+ return res
 
  # todo : update this once exception handling is set up
  except ServiceError as se:
@@ -150,6 +165,25 @@ def get_default_compartment(self) -> dict:
  logger.error("No compartment id found from environment variables.")
  return dict(compartment_id=COMPARTMENT_OCID)
 
+ def clear_compartments_list_cache(self) -> dict:
+ """Allows caller to clear compartments list cache
+ Returns
+ -------
+ dict with the key used, and True if cache has the key that needs to be deleted.
+ """
+ res = {}
+ logger.info(f"Clearing list_compartments cache")
+ with self._cache_lock:
+ if TENANCY_OCID in self._compartments_cache.keys():
+ self._compartments_cache.pop(key=TENANCY_OCID)
+ res = {
+ "key": {
+ "tenancy_ocid": TENANCY_OCID,
+ },
+ "cache_deleted": True,
+ }
+ return res
+
  def list_model_version_sets(self, **kwargs) -> str:
  """Lists all model version sets for the specified compartment or tenancy.
 

pyproject.toml

@@ -163,7 +163,7 @@ pii = [
  "spacy==3.6.1",
 ]
 llm = ["langchain>=0.0.295", "evaluate>=0.4.0"]
-aqua = ["fire"]
+aqua = ["fire", "cachetools"]
 
 [project.urls]
 "Github" = "https://github.com/oracle/accelerated-data-science"