Collect shape from configuration entry. (#1089)

Author: mrDzurb

ads/aqua/modeldeployment/utils.py

@@ -135,25 +135,28 @@ def _fetch_deployment_configs_concurrently(
     def _extract_model_shape_gpu(
         self, deployment_configs: Dict[str, AquaDeploymentConfig]
     ):
-        """Extracts shape and GPU count details from deployment configurations."""
+        """Extracts shape and GPU count details from deployment configurations.
+        Supported shapes for multi model deployment will be collected from `configuration` entry in deployment config.
+        """
         model_shape_gpu = {}
         deployment = {}
 
         for model_id, config in deployment_configs.items():
+            multi_deployment_shape = list(config.configuration.keys())
             model_shape_gpu[model_id] = {
                 shape: [
                     item.gpu_count
                     for item in config.configuration.get(
                         shape, ConfigurationItem()
                     ).multi_model_deployment
                 ]
-                for shape in config.shape
+                for shape in multi_deployment_shape
             }
             deployment[model_id] = {
-                "shape": config.shape,
+                "shape": multi_deployment_shape,
                 "configuration": {
                     shape: config.configuration.get(shape, ConfigurationItem())
-                    for shape in config.shape
+                    for shape in multi_deployment_shape
                 },
             }
 

tests/unitary/with_extras/aqua/test_deployment.py

@@ -499,10 +499,10 @@ class TestDataset:
         "deployment_config": {
             "model_a": {
                 "shape": [
-                    "VM.GPU.A10.2",
-                    "VM.GPU.A10.4",
                     "BM.GPU.A100-v2.8",
                     "BM.GPU.H100.8",
+                    "VM.GPU.A10.2",
+                    "VM.GPU.A10.4",
                 ],
                 "configuration": {
                     "VM.GPU.A10.2": {
@@ -815,13 +815,9 @@ def test_get_multimodel_deployment_config(
 
     @parameterized.expand(
         [
-            [
-                "shape",
-                "Unable to determine a valid GPU allocation for the selected models based on their current configurations. Please try selecting a different set of models.",
-            ],
             [
                 "configuration",
-                "Unable to determine a valid GPU allocation for the selected models based on their current configurations. Please select a different set of models.",
+                "Unable to determine a valid GPU allocation for the selected models based on their current configurations. Please try selecting a different set of models.",
             ],
         ]
     )