Compatibility check for Aqua extension and update compartment mapping config (#746)

Author: VipulMascarenhas

ads/aqua/__init__.py

@@ -6,7 +6,20 @@
 
 import logging
 import sys
+import os
+from ads.aqua.utils import fetch_service_compartment
+
 
 logger = logging.getLogger(__name__)
 handler = logging.StreamHandler(sys.stdout)
 logger.setLevel(logging.INFO)
+
+
+ODSC_MODEL_COMPARTMENT_OCID = os.environ.get("ODSC_MODEL_COMPARTMENT_OCID")
+if not ODSC_MODEL_COMPARTMENT_OCID:
+    try:
+        ODSC_MODEL_COMPARTMENT_OCID = fetch_service_compartment()
+    except:
+        logger.error(
+            "ODSC_MODEL_COMPARTMENT_OCID environment variable is not set for Aqua."
+        )

ads/aqua/extension/common_handler.py

@@ -7,6 +7,8 @@
 from importlib import metadata
 
 from ads.aqua.extension.base_handler import AquaAPIhandler
+from ads.aqua import ODSC_MODEL_COMPARTMENT_OCID
+from ads.aqua.exception import AquaResourceAccessError
 
 
 class ADSVersionHandler(AquaAPIhandler):
@@ -16,6 +18,19 @@ def get(self):
         self.finish({"data": metadata.version("oracle_ads")})
 
 
+class CompatibilityCheckHandler(AquaAPIhandler):
+    """The handler to check if the extension is compatible."""
+
+    def get(self):
+        if ODSC_MODEL_COMPARTMENT_OCID:
+            return self.finish(dict(status="ok"))
+        else:
+            raise AquaResourceAccessError(
+                f"The AI Quick actions extension is not compatible in the given region."
+            )
+
+
 __handlers__ = [
     ("ads_version", ADSVersionHandler),
+    ("hello", CompatibilityCheckHandler),
 ]

ads/aqua/model.py

@@ -47,10 +47,10 @@
     AQUA_SERVICE_MODELS_BUCKET,
     COMPARTMENT_OCID,
     CONDA_BUCKET_NS,
-    ODSC_MODEL_COMPARTMENT_OCID,
     PROJECT_OCID,
     TENANCY_OCID,
 )
+from ads.aqua import ODSC_MODEL_COMPARTMENT_OCID
 from ads.model import DataScienceModel
 from ads.model.model_metadata import MetadataTaxonomyKeys, ModelCustomMetadata
 from ads.telemetry import telemetry

ads/aqua/utils.py

@@ -29,7 +29,6 @@
 from ads.common.oci_resource import SEARCH_TYPE, OCIResource
 from ads.common.utils import get_console_link, upload_to_os
 from ads.config import (
-    AQUA_CONFIG_FOLDER,
     AQUA_SERVICE_MODELS_BUCKET,
     CONDA_BUCKET_NS,
     TENANCY_OCID,
@@ -45,6 +44,7 @@
 README = "README.md"
 LICENSE_TXT = "config/LICENSE.txt"
 DEPLOYMENT_CONFIG = "deployment_config.json"
+COMPARTMENT_MAPPING_KEY = "service-model-compartment"
 CONTAINER_INDEX = "container_index.json"
 EVALUATION_REPORT_JSON = "report.json"
 EVALUATION_REPORT_MD = "report.md"
@@ -574,6 +574,22 @@ def get_container_image(
     return container_image
 
 
+def fetch_service_compartment():
+    """Loads the compartment mapping json from service bucket"""
+    config_file_name = (
+        f"oci://{AQUA_SERVICE_MODELS_BUCKET}@{CONDA_BUCKET_NS}/service_models/config"
+    )
+
+    config = load_config(
+        file_path=config_file_name,
+        config_file_name=CONTAINER_INDEX,
+    )
+    compartment_mapping = config.get(COMPARTMENT_MAPPING_KEY)
+    if compartment_mapping:
+        return compartment_mapping.get(CONDA_BUCKET_NS)
+    return None
+
+
 def get_max_version(versions):
     """Takes in a list of versions and returns the higher version."""
     if not versions:

ads/config.py

@@ -83,17 +83,6 @@
 AQUA_SERVICE_NAME = "aqua"
 DATA_SCIENCE_SERVICE_NAME = "data-science"
 
-compartment_mapping = {
-    "idoaxnz5ar4s": "ocid1.compartment.oc1..aaaaaaaabxqyb6w4kj5y4fwwkc47auxdothlqc7jpm3zbah5aamflvp7th3q",
-    "id19sfcrra6z": "ocid1.compartment.oc1..aaaaaaaasdur4tm5apdm6qyapaba7apdiyv22n2mwt6zquyubqevk3uo7nha",
-}
-default_compartment = compartment_mapping.get(CONDA_BUCKET_NS)
-ODSC_MODEL_COMPARTMENT_OCID = (
-    default_compartment
-    if default_compartment
-    else os.environ.get("ODSC_MODEL_COMPARTMENT_OCID")
-)
-
 
 def export(
     uri: Optional[str] = DEFAULT_CONFIG_PATH,

tests/unitary/with_extras/aqua/test_model.py

@@ -78,13 +78,15 @@ def setUpClass(cls):
         os.environ["CONDA_BUCKET_NS"] = "test-namespace"
         os.environ["ODSC_MODEL_COMPARTMENT_OCID"] = TestDataset.SERVICE_COMPARTMENT_ID
         reload(ads.config)
+        reload(ads.aqua)
         reload(ads.aqua.model)
 
     @classmethod
     def tearDownClass(cls):
         os.environ.pop("CONDA_BUCKET_NS", None)
         os.environ.pop("ODSC_MODEL_COMPARTMENT_OCID", None)
         reload(ads.config)
+        reload(ads.aqua)
         reload(ads.aqua.model)
 
     def test_list_service_models(self):