Support backward compatibility for TEI register command

Author: VipulMascarenhas

ads/aqua/model/constants.py

@@ -17,6 +17,7 @@ class ModelCustomMetadataFields(str, metaclass=ExtendedEnumMeta):
     DEPLOYMENT_CONTAINER = "deployment-container"
     EVALUATION_CONTAINER = "evaluation-container"
     FINETUNE_CONTAINER = "finetune-container"
+    DEPLOYMENT_CONTAINER_URI = "deployment-container-uri"
 
 
 class ModelTask(str, metaclass=ExtendedEnumMeta):

ads/aqua/model/entities.py

@@ -98,6 +98,7 @@ class AquaModel(AquaModelSummary, DataClassSerializable):
 
     model_card: str = None
     inference_container: str = None
+    inference_container_uri: str = None
     finetuning_container: str = None
     evaluation_container: str = None
     artifact_location: str = None

ads/aqua/model/model.py

@@ -232,6 +232,12 @@ def get(self, model_id: str, load_model_card: Optional[bool] = True) -> "AquaMod
             ModelCustomMetadataFields.DEPLOYMENT_CONTAINER,
             ModelCustomMetadataItem(key=ModelCustomMetadataFields.DEPLOYMENT_CONTAINER),
         ).value
+        inference_container_uri = ds_model.custom_metadata_list.get(
+            ModelCustomMetadataFields.DEPLOYMENT_CONTAINER_URI,
+            ModelCustomMetadataItem(
+                key=ModelCustomMetadataFields.DEPLOYMENT_CONTAINER_URI
+            ),
+        ).value
         evaluation_container = ds_model.custom_metadata_list.get(
             ModelCustomMetadataFields.EVALUATION_CONTAINER,
             ModelCustomMetadataItem(key=ModelCustomMetadataFields.EVALUATION_CONTAINER),
@@ -250,6 +256,7 @@ def get(self, model_id: str, load_model_card: Optional[bool] = True) -> "AquaMod
             project_id=ds_model.project_id,
             model_card=model_card,
             inference_container=inference_container,
+            inference_container_uri=inference_container_uri,
             finetuning_container=finetuning_container,
             evaluation_container=evaluation_container,
             artifact_location=artifact_location,
@@ -696,23 +703,26 @@ def _create_model_catalog_entry(
                 description=f"Inference container mapping for {model_name}",
                 category="Other",
             )
+            if inference_container_uri:
+                metadata.add(
+                    key=AQUA_DEPLOYMENT_CONTAINER_URI_METADATA_NAME,
+                    value=inference_container_uri,
+                    description=f"Inference container URI for {model_name}",
+                    category="Other",
+                )
+
+            inference_containers = (
+                AquaContainerConfig.from_container_index_json().inference
+            )
+            smc_container_set = {
+                container.family for container in inference_containers.values()
+            }
+            # only add cmd vars if inference container is not an SMC
             if (
-                inference_container
+                inference_container not in smc_container_set
+                and inference_container
                 == InferenceContainerTypeFamily.AQUA_TEI_CONTAINER_FAMILY
             ):
-                if not inference_container_uri:
-                    logger.warn(
-                        f"Proceeding with model registration without the inference container URI for "
-                        f"{inference_container}. You can still add this configuration during model deployment."
-                    )
-                else:
-                    metadata.add(
-                        key=AQUA_DEPLOYMENT_CONTAINER_URI_METADATA_NAME,
-                        value=inference_container_uri,
-                        description=f"Inference container URI for {model_name}",
-                        category="Other",
-                    )
-
                 cmd_vars = generate_tei_cmd_var(os_path)
                 metadata.add(
                     key=AQUA_DEPLOYMENT_CONTAINER_CMD_VAR_METADATA_NAME,
@@ -1231,20 +1241,23 @@ def register(
         # registered model will always have inference and evaluation container, but
         # fine-tuning container may be not set
         inference_container = ds_model.custom_metadata_list.get(
-            ModelCustomMetadataFields.DEPLOYMENT_CONTAINER
+            ModelCustomMetadataFields.DEPLOYMENT_CONTAINER,
+            ModelCustomMetadataItem(key=ModelCustomMetadataFields.DEPLOYMENT_CONTAINER),
+        ).value
+        inference_container_uri = ds_model.custom_metadata_list.get(
+            ModelCustomMetadataFields.DEPLOYMENT_CONTAINER_URI,
+            ModelCustomMetadataItem(
+                key=ModelCustomMetadataFields.DEPLOYMENT_CONTAINER_URI
+            ),
+        ).value
+        evaluation_container = ds_model.custom_metadata_list.get(
+            ModelCustomMetadataFields.EVALUATION_CONTAINER,
+            ModelCustomMetadataItem(key=ModelCustomMetadataFields.EVALUATION_CONTAINER),
+        ).value
+        finetuning_container: str = ds_model.custom_metadata_list.get(
+            ModelCustomMetadataFields.FINETUNE_CONTAINER,
+            ModelCustomMetadataItem(key=ModelCustomMetadataFields.FINETUNE_CONTAINER),
         ).value
-        try:
-            evaluation_container = ds_model.custom_metadata_list.get(
-                ModelCustomMetadataFields.EVALUATION_CONTAINER,
-            ).value
-        except Exception:
-            evaluation_container = None
-        try:
-            finetuning_container = ds_model.custom_metadata_list.get(
-                ModelCustomMetadataFields.FINETUNE_CONTAINER,
-            ).value
-        except Exception:
-            finetuning_container = None
 
         aqua_model_attributes = dict(
             **self._process_model(ds_model, self.region),
@@ -1256,6 +1269,7 @@ def register(
                 )
             ),
             inference_container=inference_container,
+            inference_container_uri=inference_container_uri,
             finetuning_container=finetuning_container,
             evaluation_container=evaluation_container,
             artifact_location=artifact_path,