Merge branch 'master' into fix-pr-checks

Author: aviruthen

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

@@ -352,13 +352,18 @@ def _get_fine_tuning_options_and_model_arn(model_name: str, customization_techni
         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 recipe and recipe.get("SmtjOverrideParamsS3Uri"):
+        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)
@@ -598,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,
@@ -682,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)
 

sagemaker-train/src/sagemaker/train/dpo_trainer.py

@@ -17,7 +17,8 @@
     _create_serverless_config,
     _create_mlflow_config,
     _create_model_package_config,
-    _validate_eula_for_gated_model
+    _validate_eula_for_gated_model,
+    _validate_hyperparameter_values
 )
 from sagemaker.core.telemetry.telemetry_logging import _telemetry_emitter
 from sagemaker.core.telemetry.constants import Feature
@@ -137,8 +138,38 @@ def __init__(
 
                                                                                     ))
 
+        # Process hyperparameters
+        self._process_hyperparameters()
+        
         # Validate and set EULA acceptance
         self.accept_eula = _validate_eula_for_gated_model(model, accept_eula, is_gated_model)
+
+    def _process_hyperparameters(self):
+        """Remove hyperparameter keys that are handled by constructor inputs."""
+        if self.hyperparameters:
+            # Remove keys that are handled by constructor inputs
+            if hasattr(self.hyperparameters, 'data_path'):
+                delattr(self.hyperparameters, 'data_path')
+                self.hyperparameters._specs.pop('data_path', None)
+            if hasattr(self.hyperparameters, 'output_path'):
+                delattr(self.hyperparameters, 'output_path')
+                self.hyperparameters._specs.pop('output_path', None)
+            if hasattr(self.hyperparameters, 'data_s3_path'):
+                delattr(self.hyperparameters, 'data_s3_path')
+                self.hyperparameters._specs.pop('data_s3_path', None)
+            if hasattr(self.hyperparameters, 'output_s3_path'):
+                delattr(self.hyperparameters, 'output_s3_path')
+                self.hyperparameters._specs.pop('output_s3_path', None)
+            if hasattr(self.hyperparameters, 'training_data_name'):
+                delattr(self.hyperparameters, 'training_data_name')
+                self.hyperparameters._specs.pop('training_data_name', None)
+            if hasattr(self.hyperparameters, 'validation_data_name'):
+                delattr(self.hyperparameters, 'validation_data_name')
+                self.hyperparameters._specs.pop('validation_data_name', None)
+            if hasattr(self.hyperparameters, 'validation_data_path'):
+                delattr(self.hyperparameters, 'validation_data_path')
+                self.hyperparameters._specs.pop('validation_data_path', None)
+
     @_telemetry_emitter(feature=Feature.MODEL_CUSTOMIZATION, func_name="DPOTrainer.train")
     def train(self,
               training_dataset: Optional[Union[str, DataSet]] = None,
@@ -198,6 +229,7 @@ def train(self,
         )
 
         final_hyperparameters = self.hyperparameters.to_dict()
+        _validate_hyperparameter_values(final_hyperparameters)
 
         model_package_config = _create_model_package_config(
             model_package_group_name=self.model_package_group_name,

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

@@ -546,6 +546,8 @@ def _get_or_create_artifact_arn(self, source_uri: str, region: str) -> str:
                 properties['HubContentArn'] = source_uri
             else:
                 properties['SourceUri'] = source_uri
+
+            _logger.info(f"source_uri: {source_uri}, region: {region}, properties: {properties}")
 
             # Create artifact using Artifact.create()
             artifact = Artifact.create(

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

@@ -308,6 +308,10 @@ def _get_custom_scorer_template_additions(self, evaluator_config: dict) -> dict:
             'evaluator_arn': evaluator_config['evaluator_arn'],
         }
 
+        # Add lambda_type for Nova models
+        if is_nova:
+            custom_scorer_context['lambda_type'] = 'rft'
+        
         # Add preset_reward_function if present
         if evaluator_config['preset_reward_function']:
             custom_scorer_context['preset_reward_function'] = evaluator_config['preset_reward_function']

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

@@ -632,7 +632,8 @@
                     "task": "{{ task }}",
                     "strategy": "{{ strategy }}"{% if metric is defined %},
                     "metric": "{{ metric }}"{% elif evaluation_metric is defined %},
-                    "evaluation_metric": "{{ evaluation_metric }}"{% endif %}{% if max_new_tokens is defined %},
+                    "evaluation_metric": "{{ evaluation_metric }}"{% endif %}{% if lambda_type is defined %},
+                    "lambda_type": "{{ lambda_type }}"{% endif %}{% if max_new_tokens is defined %},
                     "max_new_tokens": "{{ max_new_tokens }}"{% endif %}{% if temperature is defined %},
                     "temperature": "{{ temperature }}"{% endif %}{% if top_k is defined %},
                     "top_k": "{{ top_k }}"{% endif %}{% if top_p is defined %},
@@ -694,7 +695,8 @@
                     "task": "{{ task }}",
                     "strategy": "{{ strategy }}"{% if metric is defined %},
                     "metric": "{{ metric }}"{% elif evaluation_metric is defined %},
-                    "evaluation_metric": "{{ evaluation_metric }}"{% endif %}{% if max_new_tokens is defined %},
+                    "evaluation_metric": "{{ evaluation_metric }}"{% endif %}{% if lambda_type is defined %},
+                    "lambda_type": "{{ lambda_type }}"{% endif %}{% if max_new_tokens is defined %},
                     "max_new_tokens": "{{ max_new_tokens }}"{% endif %}{% if temperature is defined %},
                     "temperature": "{{ temperature }}"{% endif %}{% if top_k is defined %},
                     "top_k": "{{ top_k }}"{% endif %}{% if top_p is defined %},
@@ -872,7 +874,8 @@
                     "task": "{{ task }}",
                     "strategy": "{{ strategy }}"{% if metric is defined %},
                     "metric": "{{ metric }}"{% elif evaluation_metric is defined %},
-                    "evaluation_metric": "{{ evaluation_metric }}"{% endif %}{% if max_new_tokens is defined %},
+                    "evaluation_metric": "{{ evaluation_metric }}"{% endif %}{% if lambda_type is defined %},
+                    "lambda_type": "{{ lambda_type }}"{% endif %}{% if max_new_tokens is defined %},
                     "max_new_tokens": "{{ max_new_tokens }}"{% endif %}{% if temperature is defined %},
                     "temperature": "{{ temperature }}"{% endif %}{% if top_k is defined %},
                     "top_k": "{{ top_k }}"{% endif %}{% if top_p is defined %},