Added API to get multi model deployment config (#1055)

Author: VipulMascarenhas

ads/aqua/extension/deployment_handler.py

@@ -1,7 +1,8 @@
 #!/usr/bin/env python
 # Copyright (c) 2024, 2025 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
-import logging
+
+from typing import List, Union
 from urllib.parse import urlparse
 
 from tornado.web import HTTPError
@@ -20,7 +21,7 @@ class AquaDeploymentHandler(AquaAPIhandler):
 
     Methods
     -------
-    get(self, id="")
+    get(self, id: Union[str, List[str]])
         Retrieves a list of AQUA deployments or model info or logs by ID.
     post(self, *args, **kwargs)
         Creates a new AQUA deployment.
@@ -37,7 +38,7 @@ class AquaDeploymentHandler(AquaAPIhandler):
     """
 
     @handle_exceptions
-    def get(self, id=""):
+    def get(self, id: Union[str, List[str]] = None):
         """Handle GET request."""
         url_parse = urlparse(self.request.path)
         paths = url_parse.path.strip("/")
@@ -47,6 +48,16 @@ def get(self, id=""):
                     400, f"The request {self.request.path} requires model id."
                 )
             return self.get_deployment_config(id)
+        elif paths.startswith("aqua/deployments/modelconfig"):
+            if isinstance(id, list):
+                return self.get_multimodel_compatible_shapes(id)
+            elif isinstance(id, str):
+                return self.get_deployment_config(id)
+            else:
+                raise HTTPError(
+                    400,
+                    f"The request {self.request.path} requires either a model id or a list of model ids.",
+                )
         elif paths.startswith("aqua/deployments"):
             if not id:
                 return self.list()
@@ -121,6 +132,15 @@ def get_deployment_config(self, model_id):
         """Gets the deployment config for Aqua model."""
         return self.finish(AquaDeploymentApp().get_deployment_config(model_id=model_id))
 
+    def get_multimodel_compatible_shapes(self, model_ids: List[str]):
+        """Gets the multi model deployment config and optimal GPU allocations for Aqua models."""
+        primary_model_id = self.get_argument("primary_model_id", default=None)
+        return self.finish(
+            AquaDeploymentApp().get_multimodel_compatible_shapes(
+                model_ids=model_ids, primary_model_id=primary_model_id
+            )
+        )
+
 
 class AquaDeploymentInferenceHandler(AquaAPIhandler):
     @staticmethod
@@ -237,6 +257,7 @@ def post(self, *args, **kwargs):  # noqa: ARG002
 __handlers__ = [
     ("deployments/?([^/]*)/params", AquaDeploymentParamsHandler),
     ("deployments/config/?([^/]*)", AquaDeploymentHandler),
+    ("deployments/modelconfig/?([^/]*)", AquaDeploymentHandler),
     ("deployments/?([^/]*)", AquaDeploymentHandler),
     ("deployments/?([^/]*)/activate", AquaDeploymentHandler),
     ("deployments/?([^/]*)/deactivate", AquaDeploymentHandler),

ads/aqua/modeldeployment/deployment.py

@@ -2,6 +2,7 @@
 # Copyright (c) 2024, 2025 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
+import copy
 import shlex
 from typing import Dict, List, Union
 
@@ -41,9 +42,14 @@
 from ads.aqua.model import AquaModelApp
 from ads.aqua.modeldeployment.entities import (
     AquaDeployment,
+    AquaDeploymentConfig,
     AquaDeploymentDetail,
     CreateModelDeploymentDetails,
+    GPUModelAllocation,
+    GPUShapeAllocation,
+    ModelDeploymentConfigSummary,
 )
+from ads.aqua.modeldeployment.utils import get_combinations
 from ads.aqua.ui import ModelFormat
 from ads.common.object_storage_details import ObjectStorageDetails
 from ads.common.utils import get_log_links
@@ -621,6 +627,183 @@ def get_deployment_config(self, model_id: str) -> Dict:
             )
         return config
 
+    @telemetry(
+        entry_point="plugin=deployment&action=get_multimodel_compatible_shapes",
+        name="aqua",
+    )
+    def get_multimodel_compatible_shapes(
+        self, model_ids: List[str], primary_model_id: str = None
+    ) -> ModelDeploymentConfigSummary:
+        """Gets the deployment config of multiple Aqua models and calculate the gpu allocations for all compatible shapes.
+        If no primary Aqua model id provided, gpu count for each compatible shape will be evenly allocated.
+        If provided, gpu count for each compatible shape will be prioritized for primary model.
+
+        For example, there is one compatible shape "BM.GPU.H100.8" for three models A, B, C, and each model has a gpu count as below:
+
+        A - BM.GPU.H100.8 - 1, 2, 4, 8
+        B - BM.GPU.H100.8 - 1, 2, 4, 8
+        C - BM.GPU.H100.8 - 1, 2, 4, 8
+
+        If no primary model is provided, the gpu allocation for A, B, C could be [2, 4, 2], [2, 2, 4] or [4, 2, 2]
+        If B is the primary model, the gpu allocation is [2, 4, 2] as B always gets the maximum gpu count.
+
+        Parameters
+        ----------
+        model_ids: List[str]
+            A list of OCID of the Aqua model.
+        primary_model_id: str
+            The OCID of the primary Aqua model
+
+        Returns
+        -------
+        ModelDeploymentConfigSummary:
+            An instance of ModelDeploymentConfigSummary.
+        """
+        deployment = {}
+        model_shape_gpu = {}
+        for model_id in model_ids:
+            deployment_config = AquaDeploymentConfig(
+                **self.get_deployment_config(model_id=model_id)
+            )
+            model_shape_gpu[model_id] = {
+                shape: [
+                    item.gpu_count
+                    for item in deployment_config.configuration[
+                        shape
+                    ].multi_model_deployment
+                ]
+                for shape in deployment_config.shape
+            }
+
+            deployment.update(
+                {
+                    model_id: {
+                        "shape": deployment_config.shape,
+                        "configuration": {
+                            shape: deployment_config.configuration[shape]
+                            for shape in deployment_config.shape
+                        },
+                    }
+                }
+            )
+
+        common_shapes = []
+        for shape_gpu in model_shape_gpu.values():
+            if not common_shapes:
+                common_shapes = list(shape_gpu.keys())
+            else:
+                common_shapes = [
+                    shape for shape in common_shapes if shape in list(shape_gpu.keys())
+                ]
+
+        if not common_shapes:
+            raise AquaValueError(
+                "There are no available shapes for models selected at this moment, please select different model to deploy."
+            )
+
+        gpu_allocation = {}
+        for common_shape in common_shapes:
+            model_gpu = {
+                model: shape_gpu[common_shape]
+                for model, shape_gpu in model_shape_gpu.items()
+            }
+            is_compatible, maximum_gpu_count, combination = self._verify_compatibility(
+                model_gpu, primary_model_id
+            )
+            if is_compatible:
+                gpu_allocation[common_shape] = GPUShapeAllocation(
+                    models=combination, total_gpus_available=maximum_gpu_count
+                )
+
+        if not gpu_allocation:
+            raise AquaValueError(
+                "There are no available gpu allocations for models selected at this moment, please select different model to deploy."
+            )
+
+        return ModelDeploymentConfigSummary(
+            deployment_config=deployment, gpu_allocation=gpu_allocation
+        )
+
+    @staticmethod
+    def _verify_compatibility(
+        model_gpu_dict: Dict, primary_model_id: str = None
+    ) -> tuple:
+        """Calculates the gpu allocations for all compatible shapes.
+        If no primary Aqua model id provided, gpu count for each compatible shape will be evenly allocated.
+        If provided, gpu count for each compatible shape will be prioritized for primary model.
+
+        For example, there is one compatible shape "BM.GPU.H100.8" for three models A, B, C, and each model has a gpu count as below:
+
+        A - BM.GPU.H100.8 - 1, 2, 4, 8
+        B - BM.GPU.H100.8 - 1, 2, 4, 8
+        C - BM.GPU.H100.8 - 1, 2, 4, 8
+
+        If no primary model is provided, the gpu allocation for A, B, C could be [2, 4, 2], [2, 2, 4] or [4, 2, 2]
+        If B is the primary model, the gpu allocation is [2, 4, 2] as B always gets the maximum gpu count.
+
+        Parameters
+        ----------
+        model_gpu_dict: Dict
+            A dict of Aqua model and its gpu counts.
+        primary_model_id: str
+            The OCID of the primary Aqua model
+
+        Returns
+        -------
+        tuple:
+            A tuple of gpu count allocation result.
+        """
+        maximum_gpu_count = max([sorted(gpus)[-1] for gpus in model_gpu_dict.values()])
+        model_gpu_dict_copy = copy.deepcopy(model_gpu_dict)
+        if primary_model_id:
+            primary_model_gpu_list = sorted(model_gpu_dict_copy.pop(primary_model_id))
+            for gpu_count in reversed(primary_model_gpu_list):
+                combinations = get_combinations(model_gpu_dict_copy)
+                for combination in combinations:
+                    if (
+                        len(combination) == len(model_gpu_dict_copy)
+                        and sum(combination.values()) == maximum_gpu_count - gpu_count
+                    ):
+                        combination[primary_model_id] = gpu_count
+                        return (
+                            True,
+                            maximum_gpu_count,
+                            [
+                                GPUModelAllocation(ocid=ocid, gpu_count=gpu_count)
+                                for ocid, gpu_count in combination.items()
+                            ],
+                        )
+
+        else:
+            combinations = get_combinations(model_gpu_dict_copy)
+            minimal_difference = float("inf")  # gets the positive infinity
+            optimal_combination = []
+            for combination in combinations:
+                if (
+                    len(combination) == len(model_gpu_dict_copy)
+                    and sum(combination.values()) == maximum_gpu_count
+                ):
+                    difference = max(combination.values()) - min(combination.values())
+                    if difference < minimal_difference:
+                        minimal_difference = difference
+                        optimal_combination = combination
+
+                        # find the optimal combination, no need to continue
+                        if minimal_difference == 0:
+                            break
+
+            if optimal_combination:
+                return (
+                    True,
+                    maximum_gpu_count,
+                    [
+                        GPUModelAllocation(ocid=ocid, gpu_count=gpu_count)
+                        for ocid, gpu_count in optimal_combination.items()
+                    ],
+                )
+
+        return (False, 0, [])
+
     def get_deployment_default_params(
         self,
         model_id: str,
@@ -671,9 +854,8 @@ def get_deployment_default_params(
             ).get(instance_shape, UNKNOWN_DICT)
 
             if "multi_model_deployment" in instance_shape_config and gpu_count:
-                gpu_params = (
-                    instance_shape_config
-                    .get("multi_model_deployment", UNKNOWN_DICT)
+                gpu_params = instance_shape_config.get(
+                    "multi_model_deployment", UNKNOWN_DICT
                 )
 
                 for gpu_config in gpu_params:
@@ -684,11 +866,9 @@ def get_deployment_default_params(
                         break
 
             else:
-                config_params = (
-                    instance_shape_config
-                    .get("parameters", UNKNOWN_DICT)
-                    .get(get_container_params_type(container_type_key), UNKNOWN)
-                )
+                config_params = instance_shape_config.get(
+                    "parameters", UNKNOWN_DICT
+                ).get(get_container_params_type(container_type_key), UNKNOWN)
 
             if config_params:
                 params_list = get_params_list(config_params)