add evaluator tagging for jumpstart models

Author: Narrohag

sagemaker-train/src/sagemaker/train/evaluate/base_evaluator.py

@@ -9,10 +9,11 @@
 
 import logging
 import re
-from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 from pydantic import BaseModel, validator
 
+from sagemaker.core.common_utils import TagsDict
 from sagemaker.core.resources import ModelPackageGroup
 from sagemaker.core.shapes import VpcConfig
 
@@ -411,6 +412,13 @@ def _source_model_package_arn(self) -> Optional[str]:
         """Get the resolved source model package ARN (None for JumpStart models)."""
         info = self._get_resolved_model_info()
         return info.source_model_package_arn if info else None
+
+    @property
+    def _is_jumpstart_model(self) -> bool:
+        """Determine if model is a JumpStart model"""
+        from sagemaker.train.common_utils.model_resolution import _ModelType
+        info = self._get_resolved_model_info()
+        return info.model_type == _ModelType.JUMPSTART
 
     def _infer_model_package_group_arn(self) -> Optional[str]:
         """Infer model package group ARN from source model package ARN.
@@ -795,6 +803,12 @@ def _start_execution(
             EvaluationPipelineExecution: Started execution object
         """
         from .execution import EvaluationPipelineExecution
+
+        tags: List[TagsDict] = []
+        
+        if self._is_jumpstart_model:
+            from sagemaker.core.jumpstart.utils import add_jumpstart_model_info_tags
+            tags = add_jumpstart_model_info_tags(tags, self.model, "*")
 
         execution = EvaluationPipelineExecution.start(
             eval_type=eval_type,
@@ -803,7 +817,8 @@ def _start_execution(
             role_arn=role_arn,
             s3_output_path=self.s3_output_path,
             session=self.sagemaker_session.boto_session if hasattr(self.sagemaker_session, 'boto_session') else None,
-            region=region
+            region=region,
+            tags=tags
         )
 
         return execution

sagemaker-train/src/sagemaker/train/evaluate/execution.py

@@ -16,6 +16,7 @@
 # Third-party imports
 from botocore.exceptions import ClientError
 from pydantic import BaseModel, Field
+from sagemaker.core.common_utils import TagsDict
 from sagemaker.core.helper.session_helper import Session
 from sagemaker.core.resources import Pipeline, PipelineExecution, Tag
 from sagemaker.core.telemetry.telemetry_logging import _telemetry_emitter
@@ -38,6 +39,7 @@ def _create_evaluation_pipeline(
     pipeline_definition: str,
     session: Optional[Any] = None,
     region: Optional[str] = None,
+    tags: Optional[List[TagsDict]] = [],
 ) -> Any:
     """Helper method to create a SageMaker pipeline for evaluation.
     
@@ -49,6 +51,7 @@ def _create_evaluation_pipeline(
         pipeline_definition (str): JSON pipeline definition (Jinja2 template).
         session (Optional[Any]): SageMaker session object.
         region (Optional[str]): AWS region.
+        tags (Optional[List[TagsDict]]): List of tags to include in pipeline
         
     Returns:
         Any: Created Pipeline instance (ready for execution).
@@ -65,9 +68,9 @@ def _create_evaluation_pipeline(
     resolved_pipeline_definition = template.render(pipeline_name=pipeline_name)
 
     # Create tags for the pipeline
-    tags = [
+    tags = tags.extend([
         {"key": _TAG_SAGEMAKER_MODEL_EVALUATION, "value": "true"}
-    ]
+    ])
 
     pipeline = Pipeline.create(
         pipeline_name=pipeline_name,
@@ -163,7 +166,8 @@ def _get_or_create_pipeline(
     pipeline_definition: str,
     role_arn: str,
     session: Optional[Session] = None,
-    region: Optional[str] = None
+    region: Optional[str] = None,
+    tags: Optional[List[TagsDict]] = [],
 ) -> Pipeline:
     """Get existing pipeline or create/update it.
     
@@ -177,6 +181,7 @@ def _get_or_create_pipeline(
         role_arn: IAM role ARN for pipeline execution
         session: Boto3 session (optional)
         region: AWS region (optional)
+        tags (Optional[List[TagsDict]]): List of tags to include in pipeline
         
     Returns:
         Pipeline instance (existing updated or newly created)
@@ -202,7 +207,7 @@ def _get_or_create_pipeline(
 
             # Get tags using Tag.get_all
             tags_list = Tag.get_all(resource_arn=pipeline_arn, session=session, region=region)
-            tags = {tag.key: tag.value for tag in tags_list}
+            tags = tags.extend({tag.key: tag.value for tag in tags_list})
 
             # Validate tag
             if tags.get(_TAG_SAGEMAKER_MODEL_EVALUATION) == "true":
@@ -505,7 +510,8 @@ def start(
         role_arn: str,
         s3_output_path: Optional[str] = None,
         session: Optional[Session] = None,
-        region: Optional[str] = None
+        region: Optional[str] = None,
+        tags: Optional[List[TagsDict]] = [],
     ) -> 'EvaluationPipelineExecution':
         """Create sagemaker pipeline execution. Optionally creates pipeline.
         
@@ -517,6 +523,7 @@ def start(
             s3_output_path (Optional[str]): S3 location where evaluation results are stored.
             session (Optional[Session]): Boto3 session for API calls.
             region (Optional[str]): AWS region for the pipeline.
+            tags (Optional[List[TagsDict]]): List of tags to include in pipeline
             
         Returns:
             EvaluationPipelineExecution: Started pipeline execution instance.
@@ -547,7 +554,8 @@ def start(
                 pipeline_definition=pipeline_definition,
                 role_arn=role_arn,
                 session=session,
-                region=region
+                region=region,
+                tags=tags,
             )
 
             # Start pipeline execution via boto3