allow customizing tei deployment

ads/aqua/extension/deployment_handler.py

@@ -103,6 +103,8 @@ def post(self, *args, **kwargs):
         memory_in_gbs = input_data.get("memory_in_gbs")
         model_file = input_data.get("model_file")
         private_endpoint_id = input_data.get("private_endpoint_id")
+        container_image_uri = input_data.get("container_image_uri")
+        cmd_var = input_data.get("cmd_var")
 
         self.finish(
             AquaDeploymentApp().create(
@@ -126,6 +128,8 @@ def post(self, *args, **kwargs):
                 memory_in_gbs=memory_in_gbs,
                 model_file=model_file,
                 private_endpoint_id=private_endpoint_id,
+                container_image_uri=container_image_uri,
+                cmd_var=cmd_var,
             )
         )
 

ads/aqua/extension/model_handler.py

@@ -123,6 +123,7 @@ def post(self, *args, **kwargs):
         download_from_hf = (
             str(input_data.get("download_from_hf", "false")).lower() == "true"
         )
+        inference_container_uri = input_data.get("inference_container_uri")
 
         return self.finish(
             AquaModelApp().register(
@@ -134,6 +135,7 @@ def post(self, *args, **kwargs):
                 compartment_id=compartment_id,
                 project_id=project_id,
                 model_file=model_file,
+                inference_container_uri=inference_container_uri,
             )
         )
 

tests/unitary/with_extras/aqua/test_deployment_handler.py

@@ -130,6 +130,8 @@ def test_post(self, mock_create):
             ocpus=None,
             model_file=None,
             private_endpoint_id=None,
+            container_image_uri=None,
+            cmd_var=None,
         )
 
 

tests/unitary/with_extras/aqua/test_model_handler.py

@@ -10,6 +10,7 @@
 from huggingface_hub.hf_api import HfApi, ModelInfo
 from huggingface_hub.utils import GatedRepoError
 from notebook.base.handlers import IPythonHandler
+from parameterized import parameterized
 
 from ads.aqua.common.errors import AquaRuntimeError
 from ads.aqua.common.utils import get_hf_model_info
@@ -90,9 +91,25 @@ def test_list(self, mock_list):
                 compartment_id=None, project_id=None, model_type=None
             )
 
+    @parameterized.expand(
+        [
+            (None, None, False, None),
+            ("odsc-llm-fine-tuning", None, False, None),
+            (None, "test.gguf", True, None),
+            (None, None, True, "iad.ocir.io/<namespace>/<image>:<tag>"),
+        ],
+    )
     @patch("notebook.base.handlers.APIHandler.finish")
     @patch("ads.aqua.model.AquaModelApp.register")
-    def test_register(self, mock_register, mock_finish):
+    def test_register(
+        self,
+        finetuning_container,
+        model_file,
+        download_from_hf,
+        inference_container_uri,
+        mock_register,
+        mock_finish,
+    ):
         mock_register.return_value = AquaModel(
             id="test_id",
             inference_container="odsc-tgi-serving",
@@ -105,18 +122,23 @@ def test_register(self, mock_register, mock_finish):
                 model="test_model_name",
                 os_path="test_os_path",
                 inference_container="odsc-tgi-serving",
+                finetuning_container=finetuning_container,
+                model_file=model_file,
+                download_from_hf=download_from_hf,
+                inference_container_uri=inference_container_uri,
             )
         )
         result = self.model_handler.post()
         mock_register.assert_called_with(
             model="test_model_name",
             os_path="test_os_path",
             inference_container="odsc-tgi-serving",
-            finetuning_container=None,
+            finetuning_container=finetuning_container,
             compartment_id=None,
             project_id=None,
-            model_file=None,
-            download_from_hf=False,
+            model_file=model_file,
+            download_from_hf=download_from_hf,
+            inference_container_uri=inference_container_uri,
         )
         assert result["id"] == "test_id"
         assert result["inference_container"] == "odsc-tgi-serving"