Rollback fix in src code

Author: mollyheamazon

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

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

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

@@ -112,7 +112,6 @@ 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"],
         )
 
@@ -216,14 +215,19 @@ def test_base_model_evaluation_uses_correct_weights(self):
             if execution.status.failure_reason:
                 logger.error(f"  Failure reason: {execution.status.failure_reason}")
 
-            # Log step failures
+            # Log step failures with detailed information
             if execution.status.step_details:
-                logger.error("\nFailed steps:")
+                logger.error("\n" + "=" * 80)
+                logger.error("DETAILED STEP FAILURE INFORMATION:")
+                logger.error("=" * 80)
                 for step in execution.status.step_details:
-                    if "failed" in step.status.lower():
-                        logger.error(f"  {step.name}: {step.status}")
-                        if step.failure_reason:
-                            logger.error(f"    Reason: {step.failure_reason}")
+                    logger.error(f"\nStep: {step.name}")
+                    logger.error(f"  Status: {step.status}")
+                    logger.error(f"  Start Time: {step.start_time}")
+                    logger.error(f"  End Time: {step.end_time}")
+                    if step.failure_reason:
+                        logger.error(f"  ❌ FAILURE REASON: {step.failure_reason}")
+                logger.error("=" * 80)
 
             # Re-raise to fail the test
             raise
@@ -248,7 +252,6 @@ 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
         )