From 6c7a72f5a3b268e1a25c5133371991032341a1fe Mon Sep 17 00:00:00 2001 From: Vipul Date: Mon, 9 Dec 2024 15:16:56 -0800 Subject: [PATCH] support input tags for evaluation --- ads/aqua/evaluation/entities.py | 6 +++ ads/aqua/evaluation/evaluation.py | 41 +++++++++++++++---- .../with_extras/aqua/test_evaluation.py | 17 ++++++-- 3 files changed, 53 insertions(+), 11 deletions(-) diff --git a/ads/aqua/evaluation/entities.py b/ads/aqua/evaluation/entities.py index ad49c88ae..bb165edd8 100644 --- a/ads/aqua/evaluation/entities.py +++ b/ads/aqua/evaluation/entities.py @@ -64,6 +64,10 @@ class CreateAquaEvaluationDetails(Serializable): The metrics for the evaluation. force_overwrite: (bool, optional). Defaults to `False`. Whether to force overwrite the existing file in object storage. + freeform_tags: (dict, optional) + Freeform tags for the evaluation model + defined_tags: (dict, optional) + Defined tags for the evaluation model """ evaluation_source_id: str @@ -85,6 +89,8 @@ class CreateAquaEvaluationDetails(Serializable): log_id: Optional[str] = None metrics: Optional[List[Dict[str, Any]]] = None force_overwrite: Optional[bool] = False + freeform_tags: Optional[dict] = None + defined_tags: Optional[dict] = None class Config: extra = "ignore" diff --git a/ads/aqua/evaluation/evaluation.py b/ads/aqua/evaluation/evaluation.py index cfd0029e6..c4585819d 100644 --- a/ads/aqua/evaluation/evaluation.py +++ b/ads/aqua/evaluation/evaluation.py @@ -297,6 +297,10 @@ def create( evaluation_mvs_freeform_tags = { Tags.AQUA_EVALUATION: Tags.AQUA_EVALUATION, } + evaluation_mvs_freeform_tags = { + **evaluation_mvs_freeform_tags, + **(create_aqua_evaluation_details.freeform_tags or {}), + } model_version_set = ( ModelVersionSet() @@ -307,6 +311,9 @@ def create( create_aqua_evaluation_details.experiment_description ) .with_freeform_tags(**evaluation_mvs_freeform_tags) + .with_defined_tags( + **(create_aqua_evaluation_details.defined_tags or {}) + ) # TODO: decide what parameters will be needed .create(**kwargs) ) @@ -369,6 +376,10 @@ def create( Tags.AQUA_EVALUATION: Tags.AQUA_EVALUATION, Tags.AQUA_EVALUATION_MODEL_ID: evaluation_model.id, } + evaluation_job_freeform_tags = { + **evaluation_job_freeform_tags, + **(create_aqua_evaluation_details.freeform_tags or {}), + } evaluation_job = Job(name=evaluation_model.display_name).with_infrastructure( DataScienceJob() @@ -379,6 +390,7 @@ def create( .with_shape_name(create_aqua_evaluation_details.shape_name) .with_block_storage_size(create_aqua_evaluation_details.block_storage_size) .with_freeform_tag(**evaluation_job_freeform_tags) + .with_defined_tag(**(create_aqua_evaluation_details.defined_tags or {})) ) if ( create_aqua_evaluation_details.memory_in_gbs @@ -425,6 +437,7 @@ def create( evaluation_job_run = evaluation_job.run( name=evaluation_model.display_name, freeform_tags=evaluation_job_freeform_tags, + defined_tags=(create_aqua_evaluation_details.defined_tags or {}), wait=False, ) logger.debug( @@ -444,13 +457,23 @@ def create( for metadata in evaluation_model_custom_metadata.to_dict()["data"] ] + evaluation_model_freeform_tags = { + Tags.AQUA_EVALUATION: Tags.AQUA_EVALUATION, + } + evaluation_model_freeform_tags = { + **evaluation_model_freeform_tags, + **(create_aqua_evaluation_details.freeform_tags or {}), + } + evaluation_model_defined_tags = ( + create_aqua_evaluation_details.defined_tags or {} + ) + self.ds_client.update_model( model_id=evaluation_model.id, update_model_details=UpdateModelDetails( custom_metadata_list=updated_custom_metadata_list, - freeform_tags={ - Tags.AQUA_EVALUATION: Tags.AQUA_EVALUATION, - }, + freeform_tags=evaluation_model_freeform_tags, + defined_tags=evaluation_model_defined_tags, ), ) @@ -520,10 +543,14 @@ def create( ), ), tags={ - "aqua_evaluation": Tags.AQUA_EVALUATION, - "evaluation_job_id": evaluation_job.id, - "evaluation_source": create_aqua_evaluation_details.evaluation_source_id, - "evaluation_experiment_id": experiment_model_version_set_id, + **{ + "aqua_evaluation": Tags.AQUA_EVALUATION, + "evaluation_job_id": evaluation_job.id, + "evaluation_source": create_aqua_evaluation_details.evaluation_source_id, + "evaluation_experiment_id": experiment_model_version_set_id, + }, + **evaluation_model_freeform_tags, + **evaluation_model_defined_tags, }, parameters=AquaEvalParams(), ) diff --git a/tests/unitary/with_extras/aqua/test_evaluation.py b/tests/unitary/with_extras/aqua/test_evaluation.py index 06f716061..ef3475184 100644 --- a/tests/unitary/with_extras/aqua/test_evaluation.py +++ b/tests/unitary/with_extras/aqua/test_evaluation.py @@ -475,6 +475,9 @@ def test_create_evaluation( self.app.ds_client.update_model = MagicMock() self.app.ds_client.update_model_provenance = MagicMock() + eval_model_freeform_tags = {"ftag1": "fvalue1", "ftag2": "fvalue2"} + eval_model_defined_tags = {"dtag1": "dvalue1", "dtag2": "dvalue2"} + create_aqua_evaluation_details = dict( evaluation_source_id="ocid1.datasciencemodel.oc1.iad.", evaluation_name="test_evaluation_name", @@ -486,6 +489,8 @@ def test_create_evaluation( experiment_name="test_experiment_name", memory_in_gbs=1, ocpus=1, + freeform_tags=eval_model_freeform_tags, + defined_tags=eval_model_defined_tags, ) aqua_evaluation_summary = self.app.create(**create_aqua_evaluation_details) @@ -516,10 +521,14 @@ def test_create_evaluation( "url": f"https://cloud.oracle.com/data-science/models/ocid1.datasciencemodel.oc1.iad.?region={self.app.region}", }, "tags": { - "aqua_evaluation": "aqua_evaluation", - "evaluation_experiment_id": f"{experiment.id}", - "evaluation_job_id": f"{mock_job_id.return_value}", - "evaluation_source": "ocid1.datasciencemodel.oc1.iad.", + **{ + "aqua_evaluation": "aqua_evaluation", + "evaluation_experiment_id": f"{experiment.id}", + "evaluation_job_id": f"{mock_job_id.return_value}", + "evaluation_source": "ocid1.datasciencemodel.oc1.iad.", + }, + **eval_model_freeform_tags, + **eval_model_defined_tags, }, "time_created": f"{oci_dsc_model.time_created}", }