Further fix for the same LLM as judge integ test failure

Author: mollyheamazon

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

@@ -701,10 +701,18 @@ def _get_base_template_context(
         Returns:
             dict: Base template context dictionary
         """
+        # Generate default mlflow_experiment_name if not provided
+        # This is required by AWS when ModelPackageGroupArn is not provided in training jobs
+        mlflow_experiment_name = self.mlflow_experiment_name
+        if not mlflow_experiment_name and self.mlflow_resource_arn:
+            # Use pipeline_name as default experiment name
+            mlflow_experiment_name = '{{ pipeline_name }}'
+            _logger.info("No mlflow_experiment_name provided, using pipeline_name as default")
+        
         return {
             'role_arn': role_arn,
             'mlflow_resource_arn': self.mlflow_resource_arn,
-            'mlflow_experiment_name': self.mlflow_experiment_name,
+            'mlflow_experiment_name': mlflow_experiment_name,
             'mlflow_run_name': self.mlflow_run_name,
             'model_package_group_arn': model_package_group_arn,
             'source_model_package_arn': self._source_model_package_arn,

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

@@ -337,7 +337,8 @@
             "Name": "EvaluateBaseInferenceModel",
             "Type": "Training",
             "Arguments": {
-                "TrainingJobName": "BaseInference",
+                "TrainingJobName": "BaseInference",{% if mlflow_experiment_name %}
+                "MlflowExperimentName": "{{ mlflow_experiment_name }}",{% endif %}
                 "RoleArn": "{{ role_arn }}",
                 "ServerlessJobConfig": {
                     "BaseModelArn": "{{ base_model_arn }}",
@@ -1007,7 +1008,8 @@
             "Name": "EvaluateBaseInferenceModel",
             "Type": "Training",
             "Arguments": {
-                "TrainingJobName": "BaseInference",
+                "TrainingJobName": "BaseInference",{% if mlflow_experiment_name %}
+                "MlflowExperimentName": "{{ mlflow_experiment_name }}",{% endif %}
                 "RoleArn": "{{ role_arn }}",
                 "ServerlessJobConfig": {
                     "BaseModelArn": "{{ base_model_arn }}",

sagemaker-train/tests/integ/train/test_llm_as_judge_base_model_fix.py

@@ -112,6 +112,7 @@ def test_base_model_evaluation_uses_correct_weights(self):
             builtin_metrics=TEST_CONFIG["builtin_metrics"],
             custom_metrics=TEST_CONFIG["custom_metrics_json"],
             s3_output_path=TEST_CONFIG["s3_output_path"],
+            mlflow_resource_arn=TEST_CONFIG["mlflow_tracking_server_arn"],
             evaluate_base_model=TEST_CONFIG["evaluate_base_model"],
         )
 
@@ -247,6 +248,7 @@ def test_base_model_false_still_works(self):
             dataset=TEST_CONFIG["dataset_s3_uri"],
             builtin_metrics=TEST_CONFIG["builtin_metrics"],
             s3_output_path=TEST_CONFIG["s3_output_path"],
+            mlflow_resource_arn=TEST_CONFIG["mlflow_tracking_server_arn"],
             evaluate_base_model=False,  # Only evaluate custom model
         )