inference container uri validation check

Author: kumar-shivam-ranjan

ads/aqua/model/model.py

@@ -411,20 +411,33 @@ def edit_registered_model(
                 custom_metadata_list = ds_model.custom_metadata_list
                 freeform_tags = ds_model.freeform_tags
                 if inference_container:
-                    custom_metadata_list.add(
-                        key=ModelCustomMetadataFields.DEPLOYMENT_CONTAINER,
-                        value=inference_container,
-                        category=MetadataCustomCategory.OTHER,
-                        description="Deployment container mapping for SMC",
-                        replace=True,
-                    )
+                    if (inference_container==InferenceContainerTypeFamily.AQUA_TEI_CONTAINER_FAMILY
+                            and inference_container_uri is None):
+                        raise AquaRuntimeError(
+                            f"Failed to edit model:{id}. Inference container URI must be provided."
+                        )
+                    else:
+                        custom_metadata_list.add(
+                            key=ModelCustomMetadataFields.DEPLOYMENT_CONTAINER,
+                            value=inference_container,
+                            category=MetadataCustomCategory.OTHER,
+                            description="Deployment container mapping for SMC",
+                            replace=True,
+                        )
                 if inference_container_uri:
-                    custom_metadata_list.add(
-                        key=ModelCustomMetadataFields.DEPLOYMENT_CONTAINER_URI,
-                        value=inference_container_uri,
-                        category=MetadataCustomCategory.OTHER,
-                        description=f"Inference container URI for {ds_model.display_name}",
-                    )
+                    if inference_container==InferenceContainerTypeFamily.AQUA_TEI_CONTAINER_FAMILY or inference_container==None:
+                        custom_metadata_list.add(
+                            key=ModelCustomMetadataFields.DEPLOYMENT_CONTAINER_URI,
+                            value=inference_container_uri,
+                            category=MetadataCustomCategory.OTHER,
+                            description=f"Inference container URI for {ds_model.display_name}",
+                            replace=True
+                        )
+                    else:
+                        raise AquaRuntimeError(
+                            f"Failed to edit model:{id}. Inference container URI can be edited only with TEI container."
+                        )
+
                 if enable_finetuning is not None:
                     if enable_finetuning.lower() == "true":
                         custom_metadata_list.add(