Onboard BYOR feature

Author: kumar-shivam-ranjan

ads/aqua/extension/ui_handler.py

@@ -15,7 +15,7 @@
 from ads.aqua.extension.utils import validate_function_parameters
 from ads.aqua.model.entities import ImportModelDetails
 from ads.aqua.ui import AquaUIApp
-from ads.config import COMPARTMENT_OCID
+from ads.config import COMPARTMENT_OCID, IS_BYOR_ENABLED
 
 
 @dataclass
@@ -82,6 +82,10 @@ def get(self, id=""):
             return self.is_bucket_versioned()
         elif paths.startswith("aqua/containers"):
             return self.list_containers()
+        elif paths.startswith("aqua/capacityreservations/enabled"):
+            return self.is_capacity_reservations_enabled()
+        elif paths.startswith("aqua/capacityreservations"):
+            return self.list_capacity_reservations()
         else:
             raise HTTPError(400, f"The request {self.request.path} is invalid.")
 
@@ -103,6 +107,17 @@ def list_log_groups(self, **kwargs):
             AquaUIApp().list_log_groups(compartment_id=compartment_id, **kwargs)
         )
 
+    def is_capacity_reservations_enabled(self, **kwargs):
+        return self.finish({"status": str(IS_BYOR_ENABLED).strip().lower() == "true"})
+
+    def list_capacity_reservations(self, **kwargs):
+        compartment_id = self.get_argument("compartment_id", default=COMPARTMENT_OCID)
+        return self.finish(
+            AquaUIApp().list_capacity_reservations(
+                compartment_id=compartment_id, **kwargs
+            )
+        )
+
     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))
@@ -279,4 +294,5 @@ def post(self, *args, **kwargs):
     ("bucket/versioning/?([^/]*)", AquaUIHandler),
     ("containers/?([^/]*)", AquaUIHandler),
     ("cli/?([^/]*)", AquaCLIHandler),
+    ("capacityreservations/?([^/]*)", AquaUIHandler),
 ]

ads/aqua/ui.py

@@ -90,6 +90,16 @@ def list_logs(self, **kwargs) -> str:
         res = self.logging_client.list_logs(log_group_id=log_group_id, **kwargs).data
         return sanitize_response(oci_client=self.logging_client, response=res)
 
+    @telemetry(entry_point="plugin=ui&action=list_capacity_reservations", name="aqua")
+    def list_capacity_reservations(self, **kwargs):
+        compartment_id = kwargs.pop("compartment_id", COMPARTMENT_OCID)
+        logger.info(f"Loading Capacity reservations from compartment: {compartment_id}")
+        compute_client = oc.OCIClientFactory(**default_signer()).compute
+        reservations = compute_client.list_compute_capacity_reservations(
+            compartment_id=compartment_id
+        )
+        return reservations.data
+
     @telemetry(entry_point="plugin=ui&action=list_compartments", name="aqua")
     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

ads/common/oci_client.py

@@ -1,5 +1,4 @@
 #!/usr/bin/env python
-# -*- coding: utf-8; -*-
 
 # Copyright (c) 2021, 2024 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
@@ -9,20 +8,20 @@
 import oci.artifacts
 from oci.ai_language import AIServiceLanguageClient
 from oci.artifacts import ArtifactsClient
+from oci.core import ComputeClient, VirtualNetworkClient
 from oci.data_catalog import DataCatalogClient
 from oci.data_flow import DataFlowClient
 from oci.data_labeling_service import DataLabelingManagementClient
 from oci.data_labeling_service_dataplane import DataLabelingClient
 from oci.data_science import DataScienceClient
 from oci.identity import IdentityClient
+from oci.limits import LimitsClient
+from oci.logging import LoggingManagementClient
 from oci.marketplace import MarketplaceClient
 from oci.object_storage import ObjectStorageClient
 from oci.resource_search import ResourceSearchClient
 from oci.secrets import SecretsClient
 from oci.vault import VaultsClient
-from oci.logging import LoggingManagementClient
-from oci.core import VirtualNetworkClient
-from oci.limits import LimitsClient
 
 logger = logging.getLogger(__name__)
 
@@ -69,6 +68,7 @@ def _client_impl(client):
             "secret": SecretsClient,
             "vault": VaultsClient,
             "identity": IdentityClient,
+            "compute": ComputeClient,
             "ai_language": AIServiceLanguageClient,
             "data_labeling_dp": DataLabelingClient,
             "data_labeling_cp": DataLabelingManagementClient,
@@ -114,6 +114,10 @@ def create_client(self, client_name):
     def object_storage(self):
         return self.create_client("object_storage")
 
+    @property
+    def compute(self):
+        return self.create_client("compute")
+
     @property
     def identity(self):
         return self.create_client("identity")

ads/config.py

@@ -14,6 +14,7 @@
 OCI_IDENTITY_SERVICE_ENDPOINT = os.environ.get("OCI_IDENTITY_SERVICE_ENDPOINT")
 NB_SESSION_COMPARTMENT_OCID = os.environ.get("NB_SESSION_COMPARTMENT_OCID")
 PROJECT_OCID = os.environ.get("PROJECT_OCID") or os.environ.get("PIPELINE_PROJECT_OCID")
+IS_BYOR_ENABLED = os.environ.get("ALLOWLISTED_FOR_BYOR")
 NB_SESSION_OCID = os.environ.get("NB_SESSION_OCID")
 USER_OCID = os.environ.get("USER_OCID")
 OCI_RESOURCE_PRINCIPAL_VERSION = os.environ.get("OCI_RESOURCE_PRINCIPAL_VERSION")