Avoid deactivate/activate when feature flag for Aqua is enabled. (#797)

Author: VipulMascarenhas

ads/aqua/__init__.py

@@ -53,17 +53,10 @@ def set_log_level(log_level: str):
 if OCI_RESOURCE_PRINCIPAL_VERSION:
     set_auth("resource_principal")
 
-ODSC_MODEL_COMPARTMENT_OCID = os.environ.get("ODSC_MODEL_COMPARTMENT_OCID")
+ODSC_MODEL_COMPARTMENT_OCID = (
+    os.environ.get("ODSC_MODEL_COMPARTMENT_OCID") or fetch_service_compartment()
+)
 if not ODSC_MODEL_COMPARTMENT_OCID:
-    try:
-        ODSC_MODEL_COMPARTMENT_OCID = fetch_service_compartment()
-    except:
-        pass
-
-if not ODSC_MODEL_COMPARTMENT_OCID:
-    logger.error(
-        f"ODSC_MODEL_COMPARTMENT_OCID environment variable is not set for Aqua."
-    )
     if NB_SESSION_OCID:
         logger.error(
             f"Aqua is not available for this notebook session {NB_SESSION_OCID}."

ads/aqua/extension/common_handler.py

@@ -10,7 +10,7 @@
 from ads.aqua.decorator import handle_exceptions
 from ads.aqua.exception import AquaResourceAccessError
 from ads.aqua.extension.base_handler import AquaAPIhandler
-from ads.aqua.utils import known_realm
+from ads.aqua.utils import known_realm, fetch_service_compartment
 
 
 class ADSVersionHandler(AquaAPIhandler):
@@ -39,7 +39,7 @@ def get(self):
             AquaResourceAccessError: raised when aqua is not accessible in the given session/region.
 
         """
-        if ODSC_MODEL_COMPARTMENT_OCID:
+        if ODSC_MODEL_COMPARTMENT_OCID or fetch_service_compartment():
             return self.finish(dict(status="ok"))
         elif known_realm():
             return self.finish(dict(status="compatible"))

ads/aqua/utils.py

@@ -567,20 +567,27 @@ def get_container_image(
     return container_image
 
 
-def fetch_service_compartment():
-    """Loads the compartment mapping json from service bucket"""
+def fetch_service_compartment() -> Union[str, None]:
+    """Loads the compartment mapping json from service bucket. This json file has a service-model-compartment key which
+    contains a dictionary of namespaces and the compartment OCID of the service models in that namespace.
+    """
     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,
-    )
+    try:
+        config = load_config(
+            file_path=config_file_name,
+            config_file_name=CONTAINER_INDEX,
+        )
+    except AquaFileNotFoundError:
+        logger.error(
+            f"ODSC_MODEL_COMPARTMENT_OCID environment variable is not set for Aqua."
+        )
+        return
     compartment_mapping = config.get(COMPARTMENT_MAPPING_KEY)
     if compartment_mapping:
         return compartment_mapping.get(CONDA_BUCKET_NS)
-    return None
 
 
 def get_max_version(versions):