diff --git a/ads/aqua/__init__.py b/ads/aqua/__init__.py index 45f1a8b03..e75a773f4 100644 --- a/ads/aqua/__init__.py +++ b/ads/aqua/__init__.py @@ -18,17 +18,10 @@ 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}." diff --git a/ads/aqua/extension/common_handler.py b/ads/aqua/extension/common_handler.py index 3529f12b2..dd56284fe 100644 --- a/ads/aqua/extension/common_handler.py +++ b/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")) diff --git a/ads/aqua/utils.py b/ads/aqua/utils.py index d3369ba33..578449b4e 100644 --- a/ads/aqua/utils.py +++ b/ads/aqua/utils.py @@ -578,20 +578,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):