From afd6b5ae3bc4543ed78280d5680f73d72fbf45c5 Mon Sep 17 00:00:00 2001
From: Lucas Jia <lucasjia@amazon.com>
Date: Thu, 4 Jun 2026 16:25:19 -0700
Subject: [PATCH] fix: conditionally render MlflowConfig in pipeline templates
 when ARN is None

When mlflow_resource_arn is None, Jinja2 renders the literal string "None"
into the pipeline JSON, causing SageMaker to reject it with a validation
error. Wrap MlflowConfig blocks with {% if mlflow_resource_arn %} so the
section is omitted entirely when MLflow tracking is unavailable.
---
 .../train/evaluate/pipeline_templates.py      | 24 +++++++++----------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/sagemaker-train/src/sagemaker/train/evaluate/pipeline_templates.py b/sagemaker-train/src/sagemaker/train/evaluate/pipeline_templates.py
index aa457ac0a1..d6650f8e9d 100644
--- a/sagemaker-train/src/sagemaker/train/evaluate/pipeline_templates.py
+++ b/sagemaker-train/src/sagemaker/train/evaluate/pipeline_templates.py
@@ -8,11 +8,11 @@
 
 DETERMINISTIC_TEMPLATE = """{
     "Version": "2020-12-01",
-    "Metadata": {},
+    "Metadata": {},{% if mlflow_resource_arn %}
     "MlflowConfig": {
         "MlflowResourceArn": "{{ mlflow_resource_arn }}"{% if mlflow_experiment_name %},
         "MlflowExperimentName": "{{ mlflow_experiment_name }}"{% endif %}
-    },
+    },{% endif %}
     "Parameters": [],
     "Steps": [
         {
@@ -325,11 +325,11 @@
 # LLM-as-a-Judge Template for Base Model Only - 2-Phase Evaluation with optional ModelPackageConfig
 LLMAJ_TEMPLATE_BASE_MODEL_ONLY = """{
     "Version": "2020-12-01",
-    "Metadata": {},
+    "Metadata": {},{% if mlflow_resource_arn %}
     "MlflowConfig": {
         "MlflowResourceArn": "{{ mlflow_resource_arn }}"{% if mlflow_experiment_name %},
         "MlflowExperimentName": "{{ mlflow_experiment_name }}"{% endif %}
-    },
+    },{% endif %}
     "Parameters": [],
     "Steps": [
         {
@@ -460,11 +460,11 @@
 # Deterministic Template for Base Model Only - Single step evaluation with optional ModelPackageConfig
 DETERMINISTIC_TEMPLATE_BASE_MODEL_ONLY = """{
     "Version": "2020-12-01",
-    "Metadata": {},
+    "Metadata": {},{% if mlflow_resource_arn %}
     "MlflowConfig": {
         "MlflowResourceArn": "{{ mlflow_resource_arn }}"{% if mlflow_experiment_name %},
         "MlflowExperimentName": "{{ mlflow_experiment_name }}"{% endif %}
-    },
+    },{% endif %}
     "Parameters": [],
     "Steps": [
         {
@@ -531,11 +531,11 @@
 # Custom Scorer Template with Jinja2 Placeholders - Similar to Deterministic but uses CustomScorerEvaluation
 CUSTOM_SCORER_TEMPLATE = """{
     "Version": "2020-12-01",
-    "Metadata": {},
+    "Metadata": {},{% if mlflow_resource_arn %}
     "MlflowConfig": {
         "MlflowResourceArn": "{{ mlflow_resource_arn }}"{% if mlflow_experiment_name %},
         "MlflowExperimentName": "{{ mlflow_experiment_name }}"{% endif %}
-    },
+    },{% endif %}
     "Parameters": [],
     "Steps": [
         {
@@ -850,11 +850,11 @@
 # Custom Scorer Template for Base Model Only - Single step evaluation with optional ModelPackageConfig
 CUSTOM_SCORER_TEMPLATE_BASE_MODEL_ONLY = """{
     "Version": "2020-12-01",
-    "Metadata": {},
+    "Metadata": {},{% if mlflow_resource_arn %}
     "MlflowConfig": {
         "MlflowResourceArn": "{{ mlflow_resource_arn }}"{% if mlflow_experiment_name %},
         "MlflowExperimentName": "{{ mlflow_experiment_name }}"{% endif %}
-    },
+    },{% endif %}
     "Parameters": [],
     "Steps": [
         {
@@ -924,11 +924,11 @@
 # Phase 2: Use judge model to evaluate responses with built-in and custom metrics
 LLMAJ_TEMPLATE = """{
     "Version": "2020-12-01",
-    "Metadata": {},
+    "Metadata": {},{% if mlflow_resource_arn %}
     "MlflowConfig": {
         "MlflowResourceArn": "{{ mlflow_resource_arn }}"{% if mlflow_experiment_name %},
         "MlflowExperimentName": "{{ mlflow_experiment_name }}"{% endif %}
-    },
+    },{% endif %}
     "Parameters": [],
     "Steps": [
         {