[Ready to merge] Onboarding BYOR feature in AQUA (#1179)

Author: kumar-shivam-ranjan

ads/aqua/app.py

@@ -64,6 +64,7 @@ def __init__(self) -> None:
             set_auth("resource_principal")
         self._auth = default_signer({"service_endpoint": OCI_ODSC_SERVICE_ENDPOINT})
         self.ds_client = oc.OCIClientFactory(**self._auth).data_science
+        self.compute_client = oc.OCIClientFactory(**default_signer()).compute
         self.logging_client = oc.OCIClientFactory(**default_signer()).logging_management
         self.identity_client = oc.OCIClientFactory(**default_signer()).identity
         self.region = extract_region(self._auth)

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,19 @@ def list_log_groups(self, **kwargs):
             AquaUIApp().list_log_groups(compartment_id=compartment_id, **kwargs)
         )
 
+    def is_capacity_reservations_enabled(self):
+        """Checks if the tenant is whitelisted for BYOR (Bring your own reservation) feature."""
+        return self.finish({"status": str(IS_BYOR_ENABLED).strip().lower() == "true"})
+
+    def list_capacity_reservations(self, **kwargs):
+        """Lists users compute reservations in a specified compartment."""
+        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 +296,5 @@ def post(self, *args, **kwargs):
     ("bucket/versioning/?([^/]*)", AquaUIHandler),
     ("containers/?([^/]*)", AquaUIHandler),
     ("cli/?([^/]*)", AquaCLIHandler),
+    ("capacityreservations/?([^/]*)", AquaUIHandler),
 ]

ads/aqua/ui.py

@@ -90,6 +90,26 @@ 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) -> list:
+        """
+        Lists users compute reservations in a specified compartment
+
+        Returns
+        -------
+            json representation of `oci.core.models.ComputeCapacityReservationSummary`.
+
+        """
+        compartment_id = kwargs.pop("compartment_id", COMPARTMENT_OCID)
+        logger.info(f"Loading Capacity reservations from compartment: {compartment_id}")
+
+        reservations = self.compute_client.list_compute_capacity_reservations(
+            compartment_id=compartment_id, **kwargs
+        )
+        return sanitize_response(
+            oci_client=self.compute_client, response=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", False)
 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")

tests/unitary/with_extras/aqua/test_data/ui/capacity_reservations_list.json

@@ -0,0 +1,15 @@
+[
+  {
+    "availability_domain": "<AD>",
+    "compartment_id": "ocid1.compartment.oc1..<OCID>",
+    "defined_tags": {},
+    "display_name": "test-reservation",
+    "freeform_tags": {},
+    "id": "ocid1.capacityreservation.oc1.iad...<OCID>",
+    "is_default_reservation": false,
+    "lifecycle_state": "ACTIVE",
+    "reserved_instance_count": 0,
+    "time_created": "2025-05-14T07:17:53.795000Z",
+    "used_instance_count": 0
+  }
+]

tests/unitary/with_extras/aqua/test_ui.py

@@ -97,6 +97,46 @@ def tearDownClass(cls):
         reload(ads.aqua)
         reload(ads.aqua.ui)
 
+    def test_list_capacity_reservations(self):
+        capacity_reservations_list = os.path.join(
+            self.curr_dir, "test_data/ui/capacity_reservations_list.json"
+        )
+        with open(capacity_reservations_list, "r") as _file:
+            capacity_reservations = json.load(_file)
+
+        self.app.compute_client.list_compute_capacity_reservations = MagicMock(
+            return_value=oci.response.Response(
+                status=200,
+                request=MagicMock(),
+                headers=MagicMock(),
+                data=[
+                    oci.core.models.ComputeCapacityReservationSummary(
+                        **capacity_reservation
+                    )
+                    for capacity_reservation in capacity_reservations
+                ],
+            )
+        )
+        results = self.app.list_capacity_reservations()
+        expected_attributes = {
+            "id",
+            "compartmentId",
+            "displayName",
+            "definedTags",
+            "freeformTags",
+            "lifecycleState",
+            "availabilityDomain",
+            "reservedInstanceCount",
+            "usedInstanceCount",
+            "isDefaultReservation",
+            "timeCreated",
+        }
+        for result in results:
+            self.assertTrue(
+                expected_attributes.issuperset(set(result)), "Attributes mismatch"
+            )
+        assert len(results) == len(capacity_reservations)
+
     def test_list_log_groups(self):
         """Test to lists all log groups for the specified compartment or tenancy"""
         log_groups_list = os.path.join(

tests/unitary/with_extras/aqua/test_ui_handler.py

@@ -50,6 +50,14 @@ def tearDownClass(cls):
         reload(ads.aqua)
         reload(ads.aqua.extension.ui_handler)
 
+    @patch("ads.aqua.ui.AquaUIApp.list_capacity_reservations")
+    def test_list_capacity_reservations(self, mock_list_reservations):
+        self.ui_handler.request.path = "aqua/capacityreservations"
+        self.ui_handler.get(id="")
+        mock_list_reservations.assert_called_with(
+            compartment_id=TestDataset.USER_COMPARTMENT_ID
+        )
+
     @patch("ads.aqua.ui.AquaUIApp.list_log_groups")
     def test_list_log_groups(self, mock_list_log_groups):
         """Test get method to fetch log groups"""