support input tags for deployments

Author: VipulMascarenhas

ads/aqua/modeldeployment/deployment.py

@@ -110,6 +110,8 @@ def create(
         private_endpoint_id: Optional[str] = None,
         container_image_uri: Optional[None] = None,
         cmd_var: List[str] = None,
+        freeform_tags: Optional[dict] = None,
+        defined_tags: Optional[dict] = None,
     ) -> "AquaDeployment":
         """
         Creates a new Aqua deployment
@@ -163,6 +165,10 @@ def create(
             Required parameter for BYOC based deployments if this parameter was not set during model registration.
         cmd_var: List[str]
             The cmd of model deployment container runtime.
+        freeform_tags: dict
+            Freeform tags for the model deployment
+        defined_tags: dict
+            Defined tags for the model deployment
         Returns
         -------
         AquaDeployment
@@ -172,7 +178,11 @@ def create(
         # TODO validate if the service model has no artifact and if it requires import step before deployment.
         # Create a model catalog entry in the user compartment
         aqua_model = AquaModelApp().create(
-            model_id=model_id, compartment_id=compartment_id, project_id=project_id
+            model_id=model_id,
+            compartment_id=compartment_id,
+            project_id=project_id,
+            freeform_tags=freeform_tags,
+            defined_tags=defined_tags,
         )
 
         tags = {}
@@ -418,12 +428,14 @@ def create(
         if cmd_var:
             container_runtime.with_cmd(cmd_var)
 
+        tags = {**tags, **(freeform_tags or {})}
         # configure model deployment and deploy model on container runtime
         deployment = (
             ModelDeployment()
             .with_display_name(display_name)
             .with_description(description)
             .with_freeform_tags(**tags)
+            .with_defined_tags(**(defined_tags or {}))
             .with_infrastructure(infrastructure)
             .with_runtime(container_runtime)
         ).deploy(wait_for_completion=False)

ads/aqua/modeldeployment/entities.py

@@ -98,9 +98,12 @@ def from_oci_model_deployment(
             ),
         )
 
-        freeform_tags = oci_model_deployment.freeform_tags or UNKNOWN_DICT
-        aqua_service_model_tag = freeform_tags.get(Tags.AQUA_SERVICE_MODEL_TAG, None)
-        aqua_model_name = freeform_tags.get(Tags.AQUA_MODEL_NAME_TAG, UNKNOWN)
+        tags = {}
+        tags.update(oci_model_deployment.freeform_tags or UNKNOWN_DICT)
+        tags.update(oci_model_deployment.defined_tags or UNKNOWN_DICT)
+
+        aqua_service_model_tag = tags.get(Tags.AQUA_SERVICE_MODEL_TAG, None)
+        aqua_model_name = tags.get(Tags.AQUA_MODEL_NAME_TAG, UNKNOWN)
         private_endpoint_id = getattr(
             instance_configuration, "private_endpoint_id", UNKNOWN
         )
@@ -125,7 +128,7 @@ def from_oci_model_deployment(
                 ocid=oci_model_deployment.id,
                 region=region,
             ),
-            tags=freeform_tags,
+            tags=tags,
             environment_variables=environment_variables,
             cmd=cmd,
         )

tests/unitary/with_extras/aqua/test_deployment.py

@@ -469,6 +469,9 @@ def test_create_deployment_for_foundation_model(
 
         self.app.get_deployment_config = MagicMock(return_value=config)
 
+        freeform_tags = {"ftag1": "fvalue1", "ftag2": "fvalue2"}
+        defined_tags = {"dtag1": "dvalue1", "dtag2": "dvalue2"}
+
         container_index_json = os.path.join(
             self.curr_dir, "test_data/ui/container_index.json"
         )
@@ -483,6 +486,8 @@ def test_create_deployment_for_foundation_model(
         model_deployment_obj = ModelDeployment.from_yaml(uri=aqua_deployment)
         model_deployment_dsc_obj = copy.deepcopy(TestDataset.model_deployment_object[0])
         model_deployment_dsc_obj["lifecycle_state"] = "CREATING"
+        model_deployment_dsc_obj["defined_tags"] = defined_tags
+        model_deployment_dsc_obj["freeform_tags"].update(freeform_tags)
         model_deployment_obj.dsc_model_deployment = (
             oci.data_science.models.ModelDeploymentSummary(**model_deployment_dsc_obj)
         )
@@ -495,10 +500,16 @@ def test_create_deployment_for_foundation_model(
             log_group_id="ocid1.loggroup.oc1.<region>.<OCID>",
             access_log_id="ocid1.log.oc1.<region>.<OCID>",
             predict_log_id="ocid1.log.oc1.<region>.<OCID>",
+            freeform_tags=freeform_tags,
+            defined_tags=defined_tags,
         )
 
         mock_create.assert_called_with(
-            model_id=TestDataset.MODEL_ID, compartment_id=None, project_id=None
+            model_id=TestDataset.MODEL_ID,
+            compartment_id=None,
+            project_id=None,
+            freeform_tags=freeform_tags,
+            defined_tags=defined_tags,
         )
         mock_get_container_image.assert_called()
         mock_deploy.assert_called()
@@ -508,6 +519,8 @@ def test_create_deployment_for_foundation_model(
         assert set(actual_attributes) == set(expected_attributes), "Attributes mismatch"
         expected_result = copy.deepcopy(TestDataset.aqua_deployment_object)
         expected_result["state"] = "CREATING"
+        expected_result["tags"].update(freeform_tags)
+        expected_result["tags"].update(defined_tags)
         assert actual_attributes == expected_result
 
     @patch("ads.aqua.modeldeployment.deployment.get_container_config")
@@ -566,7 +579,11 @@ def test_create_deployment_for_fine_tuned_model(
         )
 
         mock_create.assert_called_with(
-            model_id=TestDataset.MODEL_ID, compartment_id=None, project_id=None
+            model_id=TestDataset.MODEL_ID,
+            compartment_id=None,
+            project_id=None,
+            freeform_tags=None,
+            defined_tags=None,
         )
         mock_get_container_image.assert_called()
         mock_deploy.assert_called()
@@ -638,7 +655,11 @@ def test_create_deployment_for_gguf_model(
         )
 
         mock_create.assert_called_with(
-            model_id=TestDataset.MODEL_ID, compartment_id=None, project_id=None
+            model_id=TestDataset.MODEL_ID,
+            compartment_id=None,
+            project_id=None,
+            freeform_tags=None,
+            defined_tags=None,
         )
         mock_get_container_image.assert_called()
         mock_deploy.assert_called()
@@ -654,83 +675,87 @@ def test_create_deployment_for_gguf_model(
         )
         assert actual_attributes == expected_result
 
-    # @patch("ads.aqua.modeldeployment.deployment.get_container_config")
-    # @patch("ads.aqua.model.AquaModelApp.create")
-    # @patch("ads.aqua.modeldeployment.deployment.get_container_image")
-    # @patch("ads.model.deployment.model_deployment.ModelDeployment.deploy")
-    # def test_create_deployment_for_tei_byoc_embedding_model(
-    #     self,
-    #     mock_deploy,
-    #     mock_get_container_image,
-    #     mock_create,
-    #     mock_get_container_config,
-    # ):
-    #     """Test to create a deployment for fine-tuned model"""
-    #     aqua_model = os.path.join(
-    #         self.curr_dir, "test_data/deployment/aqua_tei_byoc_embedding_model.yaml"
-    #     )
-    #     datascience_model = DataScienceModel.from_yaml(uri=aqua_model)
-    #     mock_create.return_value = datascience_model
-    #
-    #     config_json = os.path.join(
-    #         self.curr_dir, "test_data/deployment/deployment_config.json"
-    #     )
-    #     with open(config_json, "r") as _file:
-    #         config = json.load(_file)
-    #
-    #     self.app.get_deployment_config = MagicMock(return_value=config)
-    #
-    #     container_index_json = os.path.join(
-    #         self.curr_dir, "test_data/ui/container_index.json"
-    #     )
-    #     with open(container_index_json, "r") as _file:
-    #         container_index_config = json.load(_file)
-    #     mock_get_container_config.return_value = container_index_config
-    #
-    #     mock_get_container_image.return_value = TestDataset.DEPLOYMENT_IMAGE_NAME
-    #     aqua_deployment = os.path.join(
-    #         self.curr_dir, "test_data/deployment/aqua_create_embedding_deployment.yaml"
-    #     )
-    #     model_deployment_obj = ModelDeployment.from_yaml(uri=aqua_deployment)
-    #     model_deployment_dsc_obj = copy.deepcopy(
-    #         TestDataset.model_deployment_object_tei_byoc[0]
-    #     )
-    #     model_deployment_dsc_obj["lifecycle_state"] = "CREATING"
-    #     model_deployment_obj.dsc_model_deployment = (
-    #         oci.data_science.models.ModelDeploymentSummary(**model_deployment_dsc_obj)
-    #     )
-    #     mock_deploy.return_value = model_deployment_obj
-    #
-    #     result = self.app.create(
-    #         model_id=TestDataset.MODEL_ID,
-    #         instance_shape=TestDataset.DEPLOYMENT_SHAPE_NAME,
-    #         display_name="model-deployment-name",
-    #         log_group_id="ocid1.loggroup.oc1.<region>.<OCID>",
-    #         access_log_id="ocid1.log.oc1.<region>.<OCID>",
-    #         predict_log_id="ocid1.log.oc1.<region>.<OCID>",
-    #         container_family="odsc-tei-serving",
-    #         cmd_var=[],
-    #     )
-    #
-    #     mock_create.assert_called_with(
-    #         model_id=TestDataset.MODEL_ID, compartment_id=None, project_id=None
-    #     )
-    #     mock_get_container_image.assert_called()
-    #     mock_deploy.assert_called()
-    #
-    #     expected_attributes = set(AquaDeployment.__annotations__.keys())
-    #     actual_attributes = asdict(result)
-    #     assert set(actual_attributes) == set(expected_attributes), "Attributes mismatch"
-    #     expected_result = copy.deepcopy(TestDataset.aqua_deployment_object)
-    #     expected_result["state"] = "CREATING"
-    #     expected_result["shape_info"] = (
-    #         TestDataset.aqua_deployment_tei_byoc_embeddings_shape_info
-    #     )
-    #     expected_result["cmd"] = TestDataset.aqua_deployment_tei_byoc_embeddings_cmd
-    #     expected_result["environment_variables"] = (
-    #         TestDataset.aqua_deployment_tei_byoc_embeddings_env_vars
-    #     )
-    #     assert actual_attributes == expected_result
+    @patch("ads.aqua.modeldeployment.deployment.get_container_config")
+    @patch("ads.aqua.model.AquaModelApp.create")
+    @patch("ads.aqua.modeldeployment.deployment.get_container_image")
+    @patch("ads.model.deployment.model_deployment.ModelDeployment.deploy")
+    def test_create_deployment_for_tei_byoc_embedding_model(
+        self,
+        mock_deploy,
+        mock_get_container_image,
+        mock_create,
+        mock_get_container_config,
+    ):
+        """Test to create a deployment for fine-tuned model"""
+        aqua_model = os.path.join(
+            self.curr_dir, "test_data/deployment/aqua_tei_byoc_embedding_model.yaml"
+        )
+        datascience_model = DataScienceModel.from_yaml(uri=aqua_model)
+        mock_create.return_value = datascience_model
+
+        config_json = os.path.join(
+            self.curr_dir, "test_data/deployment/deployment_config.json"
+        )
+        with open(config_json, "r") as _file:
+            config = json.load(_file)
+
+        self.app.get_deployment_config = MagicMock(return_value=config)
+
+        container_index_json = os.path.join(
+            self.curr_dir, "test_data/ui/container_index.json"
+        )
+        with open(container_index_json, "r") as _file:
+            container_index_config = json.load(_file)
+        mock_get_container_config.return_value = container_index_config
+
+        mock_get_container_image.return_value = TestDataset.DEPLOYMENT_IMAGE_NAME
+        aqua_deployment = os.path.join(
+            self.curr_dir, "test_data/deployment/aqua_create_embedding_deployment.yaml"
+        )
+        model_deployment_obj = ModelDeployment.from_yaml(uri=aqua_deployment)
+        model_deployment_dsc_obj = copy.deepcopy(
+            TestDataset.model_deployment_object_tei_byoc[0]
+        )
+        model_deployment_dsc_obj["lifecycle_state"] = "CREATING"
+        model_deployment_obj.dsc_model_deployment = (
+            oci.data_science.models.ModelDeploymentSummary(**model_deployment_dsc_obj)
+        )
+        mock_deploy.return_value = model_deployment_obj
+
+        result = self.app.create(
+            model_id=TestDataset.MODEL_ID,
+            instance_shape=TestDataset.DEPLOYMENT_SHAPE_NAME,
+            display_name="model-deployment-name",
+            log_group_id="ocid1.loggroup.oc1.<region>.<OCID>",
+            access_log_id="ocid1.log.oc1.<region>.<OCID>",
+            predict_log_id="ocid1.log.oc1.<region>.<OCID>",
+            container_family="odsc-tei-serving",
+            cmd_var=[],
+        )
+
+        mock_create.assert_called_with(
+            model_id=TestDataset.MODEL_ID,
+            compartment_id=None,
+            project_id=None,
+            freeform_tags=None,
+            defined_tags=None,
+        )
+        mock_get_container_image.assert_called()
+        mock_deploy.assert_called()
+
+        expected_attributes = set(AquaDeployment.__annotations__.keys())
+        actual_attributes = asdict(result)
+        assert set(actual_attributes) == set(expected_attributes), "Attributes mismatch"
+        expected_result = copy.deepcopy(TestDataset.aqua_deployment_object)
+        expected_result["state"] = "CREATING"
+        expected_result["shape_info"] = (
+            TestDataset.aqua_deployment_tei_byoc_embeddings_shape_info
+        )
+        expected_result["cmd"] = TestDataset.aqua_deployment_tei_byoc_embeddings_cmd
+        expected_result["environment_variables"] = (
+            TestDataset.aqua_deployment_tei_byoc_embeddings_env_vars
+        )
+        assert actual_attributes == expected_result
 
     @parameterized.expand(
         [