Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 3 additions & 10 deletions ads/aqua/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}."
Expand Down
4 changes: 2 additions & 2 deletions ads/aqua/extension/common_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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"))
Expand Down
21 changes: 14 additions & 7 deletions ads/aqua/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you add typing for this method and add mode details in the description?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

updated.

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):
Expand Down