Merge remote-tracking branch 'upstream/master' into hmac-key-fix-v3

Author: aviruthen

sagemaker-train/src/sagemaker/ai_registry/evaluator.py

@@ -381,7 +381,7 @@ def _create_lambda_function(cls, name: str, source_file: str, role: Optional[str
 
         # Create Lambda function
         lambda_client = boto3.client("lambda")
-        function_name = f"SageMaker-evaluator-{name}"
+        function_name = f"SageMaker-evaluator-{name}-{datetime.now().strftime('%Y%m%d_%H%M%S')}"
         handler_name = f"{os.path.splitext(os.path.basename(source_file))[0]}.lambda_handler"
 
         try:

sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py

@@ -343,36 +343,27 @@ def _get_fine_tuning_options_and_model_arn(model_name: str, customization_techni
         recipes_with_template = [r for r in matching_recipes if r.get("SmtjRecipeTemplateS3Uri")]
 
         if not recipes_with_template:
-            raise ValueError(f"No recipes found with SmtjRecipeTemplateS3Uri for technique: {customization_technique}")
-
-        # If multiple recipes, filter by training_type (peft key)
-        if len(recipes_with_template) > 1:
-
-            if isinstance(training_type, TrainingType) and training_type == TrainingType.LORA:
-                # Filter recipes that have peft key for LORA
-                lora_recipes = [r for r in recipes_with_template if r.get("Peft")]
-                if lora_recipes:
-                    recipes_with_template = lora_recipes
-                elif len(recipes_with_template) > 1:
-                    raise ValueError(f"Multiple recipes found for LORA training but none have peft key")
-            elif isinstance(training_type, TrainingType) and training_type == TrainingType.FULL:
-                # For FULL training, if multiple recipes exist, throw error
-                if len(recipes_with_template) > 1:
-                    raise ValueError(f"Multiple recipes found for FULL training - cannot determine which to use")
-        
-        # If still multiple recipes after filtering, throw error
-        if len(recipes_with_template) > 1:
-            raise ValueError(f"Multiple recipes found after filtering - cannot determine which to use")
-        
-        recipe = recipes_with_template[0]
-        
-        if recipe and recipe.get("SmtjOverrideParamsS3Uri"):
+            raise ValueError(f"No recipes found with Smtj for technique: {customization_technique}")
+
+        # Select recipe based on training type
+        recipe = None
+        if (isinstance(training_type, TrainingType) and training_type == TrainingType.LORA) or training_type == "LORA":
+            recipe = next((r for r in recipes_with_template if r.get("Peft")), None)
+        elif (isinstance(training_type, TrainingType) and training_type == TrainingType.FULL) or training_type == "FULL":
+            recipe = next((r for r in recipes_with_template if not r.get("Peft")), None)
+
+        if not recipe:
+            raise ValueError(f"No recipes found with Smtj for technique: {customization_technique},training_type:{training_type}")
+
+        elif recipe and recipe.get("SmtjOverrideParamsS3Uri"):
             s3_uri = recipe["SmtjOverrideParamsS3Uri"]
             s3 = boto3.client("s3")
             bucket, key = s3_uri.replace("s3://", "").split("/", 1)
             obj = s3.get_object(Bucket=bucket, Key=key)
             options_dict = json.loads(obj["Body"].read())
             return FineTuningOptions(options_dict), model_arn, is_gated_model
+        else:
+            return FineTuningOptions({}), model_arn, is_gated_model
 
     except Exception as e:
         logger.error("Exception getting fine-tuning options: %s", e)
@@ -612,6 +603,9 @@ def _create_output_config(sagemaker_session,s3_output_path=None, kms_key_id=None
     # Use default S3 output path if none provided
     if s3_output_path is None:
         s3_output_path = _get_default_s3_output_path(sagemaker_session)
+    
+    # Validate S3 path exists
+    _validate_s3_path_exists(s3_output_path, sagemaker_session)
 
     return OutputDataConfig(
         s3_output_path=s3_output_path,
@@ -696,3 +690,43 @@ def _validate_eula_for_gated_model(model, accept_eula, is_gated_model):
         )
 
     return accept_eula
+
+
+def _validate_s3_path_exists(s3_path: str, sagemaker_session):
+    """Validate if S3 path exists and is accessible."""
+    if not s3_path.startswith("s3://"):
+        raise ValueError(f"Invalid S3 path format: {s3_path}")
+    
+    # Parse S3 URI
+    s3_parts = s3_path.replace("s3://", "").split("/", 1)
+    bucket_name = s3_parts[0]
+    prefix = s3_parts[1] if len(s3_parts) > 1 else ""
+    
+    s3_client = sagemaker_session.boto_session.client('s3')
+    
+    try:
+        # Check if bucket exists and is accessible
+        s3_client.head_bucket(Bucket=bucket_name)
+        
+        # If prefix is provided, check if it exists
+        if prefix:
+            response = s3_client.list_objects_v2(Bucket=bucket_name, Prefix=prefix, MaxKeys=1)
+            if 'Contents' not in response:
+                raise ValueError(f"S3 prefix '{prefix}' does not exist in bucket '{bucket_name}'")
+                
+    except Exception as e:
+        if "NoSuchBucket" in str(e):
+            raise ValueError(f"S3 bucket '{bucket_name}' does not exist or is not accessible")
+        raise ValueError(f"Failed to validate S3 path '{s3_path}': {str(e)}")
+
+
+def _validate_hyperparameter_values(hyperparameters: dict):
+    """Validate hyperparameter values for allowed characters."""
+    import re
+    allowed_chars = r"^[a-zA-Z0-9/_.:,\-\s'\"\[\]]*$"
+    for key, value in hyperparameters.items():
+        if isinstance(value, str) and not re.match(allowed_chars, value):
+            raise ValueError(
+                f"Hyperparameter '{key}' value '{value}' contains invalid characters. "
+                f"Only a-z, A-Z, 0-9, /, _, ., :, \\, -, space, ', \", [, ] and , are allowed."
+            )

sagemaker-train/src/sagemaker/train/common_utils/model_resolution.py

@@ -58,10 +58,10 @@ def __init__(self, sagemaker_session=None):
         
         Args:
             sagemaker_session: SageMaker session to use for API calls.
-                             If None, will be created with beta endpoint if configured.
+                             If None, will be created with endpoint if configured.
         """
         self.sagemaker_session = sagemaker_session
-        self._beta_endpoint = os.environ.get('SAGEMAKER_ENDPOINT')
+        self._endpoint = os.environ.get('SAGEMAKER_ENDPOINT')
 
     def resolve_model_info(
         self, 
@@ -188,8 +188,21 @@ def _resolve_model_package_object(self, model_package: 'ModelPackage') -> _Model
                 base_model_name = hub_content_name
             if hasattr(container.base_model, 'hub_content_arn'):
                 base_model_arn = container.base_model.hub_content_arn
+            
+            # If hub_content_arn is not present, construct it from hub_content_name and version
+            if not base_model_arn and hasattr(container.base_model, 'hub_content_version'):
+                hub_content_version = container.base_model.hub_content_version
+                model_pkg_arn = getattr(model_package, 'model_package_arn', None)
+                
+                if hub_content_name and hub_content_version and model_pkg_arn:
+                    # Extract region from model package ARN
+                    arn_parts = model_pkg_arn.split(':')
+                    if len(arn_parts) >= 4:
+                        region = arn_parts[3]
+                        # Construct hub content ARN for SageMaker public hub
+                        base_model_arn = f"arn:aws:sagemaker:{region}:aws:hub-content/SageMakerPublicHub/Model/{hub_content_name}/{hub_content_version}"
 
-        # If we couldn't extract base model ARN, this is not a supported model package
+        # If we couldn't extract or construct base model ARN, this is not a supported model package
         if not base_model_arn:
             raise ValueError(
                 f"NotSupported: Evaluation is only supported for model packages customized by SageMaker's fine-tuning flows. "
@@ -234,82 +247,23 @@ def _resolve_model_package_arn(self, model_package_arn: str) -> _ModelInfo:
             # Validate ARN format
             self._validate_model_package_arn(model_package_arn)
 
-            # TODO: Switch to sagemaker_core ModelPackage.get() once the bug is fixed
-            # Currently, ModelPackage.get() has a Pydantic validation issue where 
-            # the transform() function doesn't include model_package_name in the response,
-            # causing: "1 validation error for ModelPackage - model_package_name: Field required"
-            # Using boto3 directly as a workaround.
-            
-            # Use the sagemaker client from the session (which has the correct endpoint configured)
-            sm_client = session.sagemaker_client if hasattr(session, 'sagemaker_client') else session.boto_session.client('sagemaker')
-            response = sm_client.describe_model_package(ModelPackageName=model_package_arn)
-            
-            # Extract base model info from response
-            base_model_name = None
-            base_model_arn = None
-            hub_content_name = None
+            # Use sagemaker.core ModelPackage.get() to retrieve model package information
+            from sagemaker.core.resources import ModelPackage
 
-            # Check inference specification
-            if 'InferenceSpecification' not in response:
-                raise ValueError(
-                    f"NotSupported: Evaluation is only supported for model packages customized by SageMaker's fine-tuning flows. "
-                    f"The provided model package (ARN: {model_package_arn}) "
-                    f"does not have an inference_specification."
-                )
+            import logging
+            logger = logging.getLogger(__name__)
 
-            inf_spec = response['InferenceSpecification']
-            if 'Containers' not in inf_spec or len(inf_spec['Containers']) == 0:
-                raise ValueError(
-                    f"NotSupported: Evaluation is only supported for model packages customized by SageMaker's fine-tuning flows. "
-                    f"The provided model package (ARN: {model_package_arn}) "
-                    f"does not have any containers in its inference_specification."
-                )
-            
-            container = inf_spec['Containers'][0]
-            
-            # Extract base model info
-            if 'BaseModel' not in container:
-                raise ValueError(
-                    f"NotSupported: Evaluation is only supported for model packages customized by SageMaker's fine-tuning flows. "
-                    f"The provided model package (ARN: {model_package_arn}) "
-                    f"does not have base_model metadata in its inference_specification.containers[0]. "
-                    f"Please ensure the model was created using SageMaker's fine-tuning capabilities."
-                )
-            
-            base_model_info = container['BaseModel']
-            hub_content_name = base_model_info.get('HubContentName')
-            hub_content_version = base_model_info.get('HubContentVersion')
-            base_model_arn = base_model_info.get('HubContentArn')
-            
-            # If HubContentArn is None, construct it from HubContentName and version
-            # This handles cases where the API doesn't return the full ARN
-            if not base_model_arn and hub_content_name and hub_content_version:
-                # Extract region from model_package_arn
-                arn_parts = model_package_arn.split(':')
-                if len(arn_parts) >= 4:
-                    region = arn_parts[3]
-                    # Construct hub content ARN for SageMaker public hub
-                    base_model_arn = f"arn:aws:sagemaker:{region}:aws:hub-content/SageMakerPublicHub/Model/{hub_content_name}/{hub_content_version}"
-            
-            if not base_model_arn:
-                raise ValueError(
-                    f"NotSupported: Evaluation is only supported for model packages customized by SageMaker's fine-tuning flows. "
-                    f"The provided model package (ARN: {model_package_arn}) "
-                    f"does not have base_model metadata with HubContentArn or sufficient information to construct it. "
-                    f"Please ensure the model was created using SageMaker's fine-tuning capabilities."
-                )
+            # Get the model package using sagemaker.core
+            model_package = ModelPackage.get(
+                model_package_name=model_package_arn,
+                session=session.boto_session,
+                region=session.boto_session.region_name
+            )
 
-            # Use hub_content_name as base_model_name
-            base_model_name = hub_content_name if hub_content_name else response.get('ModelPackageGroupName', 'unknown')
+            logger.info(f"Retrieved ModelPackage in region: {session.boto_session.region_name}")
 
-            return _ModelInfo(
-                base_model_name=base_model_name,
-                base_model_arn=base_model_arn,
-                source_model_package_arn=model_package_arn,
-                model_type=_ModelType.FINE_TUNED,
-                hub_content_name=hub_content_name,
-                additional_metadata={}
-            )
+            # Now use the existing _resolve_model_package_object method to extract base model info
+            return self._resolve_model_package_object(model_package)
 
         except ValueError:
             # Re-raise ValueError as-is (our custom error messages)
@@ -342,7 +296,7 @@ def _validate_model_package_arn(self, arn: str) -> bool:
 
     def _get_session(self):
         """
-        Get or create SageMaker session with beta endpoint support.
+        Get or create SageMaker session with endpoint support.
         
         Returns:
             SageMaker session
@@ -352,12 +306,11 @@ def _get_session(self):
 
         from sagemaker.core.helper.session_helper import Session
 
-        # Check for beta endpoint in environment variable
-        if self._beta_endpoint:
+        # Check for endpoint in environment variable
+        if self._endpoint:
             sm_client = boto3.client(
                 'sagemaker',
-                endpoint_url=self._beta_endpoint,
-                region_name=os.environ.get('AWS_REGION', 'us-west-2')
+                endpoint_url=self._endpoint
             )
             return Session(sagemaker_client=sm_client)