chore: only attach eula for model channel

Author: evakravi

src/sagemaker/estimator.py

@@ -2671,7 +2671,7 @@ def _get_train_args(cls, estimator, inputs, experiment_config, accept_eula=None)
             train_args["session_chaining_config"] = estimator.get_session_chaining_config()
 
         if accept_eula is not None:
-            cls._set_accept_eula_for_input_data_config(train_args, accept_eula)
+            cls._set_accept_eula_for_model_channel_input_data_config(train_args, accept_eula)
 
         return train_args
 
@@ -2696,11 +2696,11 @@ def _add_spot_checkpoint_args(cls, local_mode, estimator, train_args):
             train_args["checkpoint_local_path"] = estimator.checkpoint_local_path
 
     @classmethod
-    def _set_accept_eula_for_input_data_config(cls, train_args, accept_eula):
-        """Set the AcceptEula flag for all input data configurations.
+    def _set_accept_eula_for_model_channel_input_data_config(cls, train_args, accept_eula):
+        """Set the AcceptEula flag for model channel in input data configurations.
 
-        This method sets the AcceptEula flag in the ModelAccessConfig for all S3DataSources
-        in the input_config array. It handles cases where keys might not exist in the
+        This method sets the AcceptEula flag in the ModelAccessConfig for the model channel
+        S3DataSource in the input_config array. It handles cases where keys might not exist in the
         nested dictionary structure.
 
         Args:
@@ -2713,26 +2713,17 @@ def _set_accept_eula_for_input_data_config(cls, train_args, accept_eula):
         if accept_eula is None:
             return
 
-        eula_count = 0
-        s3_uris = []
-
         for idx in range(len(train_args["input_config"])):
-            if "DataSource" in train_args["input_config"][idx]:
+            if (
+                "DataSource" in train_args["input_config"][idx]
+                and train_args["input_config"][idx]["ChannelName"].lower().strip() == "model"
+            ):
                 data_source = train_args["input_config"][idx]["DataSource"]
                 if "S3DataSource" in data_source:
                     s3_data_source = data_source["S3DataSource"]
                     if "ModelAccessConfig" not in s3_data_source:
                         s3_data_source["ModelAccessConfig"] = {}
                     s3_data_source["ModelAccessConfig"]["AcceptEula"] = accept_eula
-                    eula_count += 1
-
-                    # Collect S3 URI if available
-                    if "S3Uri" in s3_data_source:
-                        s3_uris.append(s3_data_source["S3Uri"])
-
-        # Log info if more than one EULA needs to be accepted
-        if eula_count > 1:
-            logger.info("Accepting EULA for %d S3 data sources: %s", eula_count, ", ".join(s3_uris))
 
     @classmethod
     def _is_local_channel(cls, input_uri):

tests/unit/test_estimator.py

@@ -336,93 +336,119 @@ def training_job_description(sagemaker_session):
     return returned_job_description
 
 
-def test_set_accept_eula_for_input_data_config_no_input_data_config():
+def test_set_accept_eula_for_model_channel_input_data_config_no_input_data_config():
     """Test when input_config is not in train_args."""
     train_args = {}
     accept_eula = True
 
-    _TrainingJob._set_accept_eula_for_input_data_config(train_args, accept_eula)
+    _TrainingJob._set_accept_eula_for_model_channel_input_data_config(train_args, accept_eula)
 
     # Verify train_args remains unchanged
     assert train_args == {}
 
 
-def test_set_accept_eula_for_input_data_config_none_accept_eula():
+def test_set_accept_eula_for_model_channel_input_data_config_none_accept_eula():
     """Test when accept_eula is None."""
     train_args = {"input_config": [{"DataSource": {"S3DataSource": {}}}]}
     accept_eula = None
 
-    _TrainingJob._set_accept_eula_for_input_data_config(train_args, accept_eula)
+    _TrainingJob._set_accept_eula_for_model_channel_input_data_config(train_args, accept_eula)
 
     # Verify train_args remains unchanged
     assert train_args == {"input_config": [{"DataSource": {"S3DataSource": {}}}]}
 
 
-def test_set_accept_eula_for_input_data_config_single_data_source():
+def test_set_accept_eula_for_model_channel_input_data_config_single_data_source():
     """Test with a single S3DataSource."""
     with patch("sagemaker.estimator.logger") as logger:
         train_args = {
-            "input_config": [{"DataSource": {"S3DataSource": {"S3Uri": "s3://bucket/model"}}}]
+            "input_config": [
+                {
+                    "ChannelName": "model",
+                    "DataSource": {"S3DataSource": {"S3Uri": "s3://bucket/model"}},
+                }
+            ]
         }
         accept_eula = True
 
-        _TrainingJob._set_accept_eula_for_input_data_config(train_args, accept_eula)
+        _TrainingJob._set_accept_eula_for_model_channel_input_data_config(train_args, accept_eula)
 
         # Verify ModelAccessConfig and AcceptEula are set correctly
         assert train_args["input_config"][0]["DataSource"]["S3DataSource"]["ModelAccessConfig"] == {
             "AcceptEula": True
         }
 
-        # Verify no logging occurred since there's only one data source
+        # Verify no logging occurred
         logger.info.assert_not_called()
 
 
-def test_set_accept_eula_for_input_data_config_multiple_data_sources():
-    """Test with multiple S3DataSources."""
+def test_set_accept_eula_for_model_channel_input_data_config_single_data_source():
+    """Test with a single S3DataSource."""
     with patch("sagemaker.estimator.logger") as logger:
         train_args = {
             "input_config": [
-                {"DataSource": {"S3DataSource": {"S3Uri": "s3://bucket/model1"}}},
-                {"DataSource": {"S3DataSource": {"S3Uri": "s3://bucket/model2"}}},
+                {
+                    "ChannelName": "NotModel",
+                    "DataSource": {"S3DataSource": {"S3Uri": "s3://bucket/not-model"}},
+                }
             ]
         }
         accept_eula = True
 
-        _TrainingJob._set_accept_eula_for_input_data_config(train_args, accept_eula)
+        _TrainingJob._set_accept_eula_for_model_channel_input_data_config(train_args, accept_eula)
+
+        # Verify ModelAccessConfig and AcceptEula are set correctly
+        assert (
+            train_args["input_config"][0]["DataSource"]["S3DataSource"].get("ModelAccessConfig")
+            == None
+        )
 
-        # Verify ModelAccessConfig and AcceptEula are set correctly for both data sources
-        assert train_args["input_config"][0]["DataSource"]["S3DataSource"]["ModelAccessConfig"] == {
-            "AcceptEula": True
-        }
-        assert train_args["input_config"][1]["DataSource"]["S3DataSource"]["ModelAccessConfig"] == {
-            "AcceptEula": True
-        }
 
-        # Verify logging occurred with correct information
-        logger.info.assert_called_once()
-        args = logger.info.call_args[0]
-        assert args[0] == "Accepting EULA for %d S3 data sources: %s"
-        assert args[1] == 2
-        assert args[2] == "s3://bucket/model1, s3://bucket/model2"
+def test_set_accept_eula_for_model_channel_input_data_config_multiple_model_channels():
+    """Test with multiple model channels."""
+    train_args = {
+        "input_config": [
+            {
+                "ChannelName": "model",
+                "DataSource": {"S3DataSource": {"S3Uri": "s3://bucket/model1"}},
+            },
+            {
+                "ChannelName": "model",
+                "DataSource": {"S3DataSource": {"S3Uri": "s3://bucket/model2"}},
+            },
+        ]
+    }
+    accept_eula = True
+
+    _TrainingJob._set_accept_eula_for_model_channel_input_data_config(train_args, accept_eula)
+
+    # Verify ModelAccessConfig and AcceptEula are set correctly for both model channels
+    assert train_args["input_config"][0]["DataSource"]["S3DataSource"]["ModelAccessConfig"] == {
+        "AcceptEula": True
+    }
+    assert train_args["input_config"][1]["DataSource"]["S3DataSource"]["ModelAccessConfig"] == {
+        "AcceptEula": True
+    }
 
 
-def test_set_accept_eula_for_input_data_config_existing_model_access_config():
+def test_set_accept_eula_for_model_channel_input_data_config_existing_model_access_config():
     """Test when ModelAccessConfig already exists."""
     train_args = {
         "input_config": [
             {
+                "ChannelName": "model",
                 "DataSource": {
                     "S3DataSource": {
                         "S3Uri": "s3://bucket/model",
                         "ModelAccessConfig": {"OtherSetting": "value"},
                     }
-                }
+                },
             }
         ]
     }
     accept_eula = True
 
-    _TrainingJob._set_accept_eula_for_input_data_config(train_args, accept_eula)
+    _TrainingJob._set_accept_eula_for_model_channel_input_data_config(train_args, accept_eula)
 
     # Verify AcceptEula is added to existing ModelAccessConfig
     assert train_args["input_config"][0]["DataSource"]["S3DataSource"]["ModelAccessConfig"] == {
@@ -431,51 +457,52 @@ def test_set_accept_eula_for_input_data_config_existing_model_access_config():
     }
 
 
-def test_set_accept_eula_for_input_data_config_missing_s3_data_source():
+def test_set_accept_eula_for_model_channel_input_data_config_missing_s3_data_source():
     """Test when S3DataSource is missing."""
-    train_args = {"input_config": [{"DataSource": {"OtherDataSource": {}}}]}
+    train_args = {"input_config": [{"ChannelName": "model", "DataSource": {"OtherDataSource": {}}}]}
     accept_eula = True
 
-    _TrainingJob._set_accept_eula_for_input_data_config(train_args, accept_eula)
+    _TrainingJob._set_accept_eula_for_model_channel_input_data_config(train_args, accept_eula)
 
     # Verify train_args remains unchanged
-    assert train_args == {"input_config": [{"DataSource": {"OtherDataSource": {}}}]}
+    assert train_args == {
+        "input_config": [{"ChannelName": "model", "DataSource": {"OtherDataSource": {}}}]
+    }
 
 
-def test_set_accept_eula_for_input_data_config_missing_data_source():
+def test_set_accept_eula_for_model_channel_input_data_config_missing_data_source():
     """Test when DataSource is missing."""
     train_args = {"input_config": [{"OtherKey": {}}]}
     accept_eula = True
 
-    _TrainingJob._set_accept_eula_for_input_data_config(train_args, accept_eula)
+    _TrainingJob._set_accept_eula_for_model_channel_input_data_config(train_args, accept_eula)
 
     # Verify train_args remains unchanged
     assert train_args == {"input_config": [{"OtherKey": {}}]}
 
 
-def test_set_accept_eula_for_input_data_config_mixed_data_sources():
+def test_set_accept_eula_for_model_channel_input_data_config_mixed_data_sources():
     """Test with a mix of S3DataSource and other data sources."""
-    with patch("sagemaker.estimator.logger") as logger:
-        train_args = {
-            "input_config": [
-                {"DataSource": {"S3DataSource": {"S3Uri": "s3://bucket/model"}}},
-                {"DataSource": {"OtherDataSource": {}}},
-            ]
-        }
-        accept_eula = True
-
-        _TrainingJob._set_accept_eula_for_input_data_config(train_args, accept_eula)
+    train_args = {
+        "input_config": [
+            {
+                "ChannelName": "model",
+                "DataSource": {"S3DataSource": {"S3Uri": "s3://bucket/model"}},
+            },
+            {"ChannelName": "model", "DataSource": {"OtherDataSource": {}}},
+        ]
+    }
+    accept_eula = True
 
-        # Verify ModelAccessConfig and AcceptEula are set correctly for S3DataSource only
-        assert train_args["input_config"][0]["DataSource"]["S3DataSource"]["ModelAccessConfig"] == {
-            "AcceptEula": True
-        }
-        assert "ModelAccessConfig" not in train_args["input_config"][1]["DataSource"].get(
-            "OtherDataSource", {}
-        )
+    _TrainingJob._set_accept_eula_for_model_channel_input_data_config(train_args, accept_eula)
 
-        # Verify no logging occurred since there's only one S3 data source
-        logger.info.assert_not_called()
+    # Verify ModelAccessConfig and AcceptEula are set correctly for S3DataSource only
+    assert train_args["input_config"][0]["DataSource"]["S3DataSource"]["ModelAccessConfig"] == {
+        "AcceptEula": True
+    }
+    assert "ModelAccessConfig" not in train_args["input_config"][1]["DataSource"].get(
+        "OtherDataSource", {}
+    )
 
 
 def test_validate_smdistributed_unsupported_image_raises(sagemaker_session):
@@ -2709,6 +2736,8 @@ def test_fit_verify_job_name(strftime, sagemaker_session):
 def test_fit_verify_accept_eula(strftime, sagemaker_session):
     fw = DummyFramework(
         entry_point=SCRIPT_PATH,
+        model_uri="s3://mybucket/model",
+        image_uri=IMAGE_URI,
         role="DummyRole",
         sagemaker_session=sagemaker_session,
         instance_count=INSTANCE_COUNT,
@@ -2721,7 +2750,7 @@ def test_fit_verify_accept_eula(strftime, sagemaker_session):
     _, _, train_kwargs = sagemaker_session.train.mock_calls[0]
 
     assert (
-        train_kwargs["input_config"][0]
+        train_kwargs["input_config"][1]
         .get("DataSource", {})
         .get("S3DataSource", {})
         .get("ModelAccessConfig", {})