fix: unit tests

Author: evakravi

src/sagemaker/estimator.py

@@ -2710,6 +2710,9 @@ def _set_accept_eula_for_input_data_config(cls, train_args, accept_eula):
         if "InputDataConfig" not in train_args:
             return
 
+        if accept_eula is None:
+            return
+
         eula_count = 0
         s3_uris = []
 

src/sagemaker/model.py

@@ -1805,7 +1805,7 @@ def deploy(
                     container_startup_health_check_timeout=container_startup_health_check_timeout,
                     explainer_config_dict=explainer_config_dict,
                     async_inference_config_dict=async_inference_config_dict,
-                    serverless_inference_config_dict=serverless_inference_config_dict,
+                    serverless_inference_config=serverless_inference_config_dict,
                     routing_config=routing_config,
                     inference_ami_version=inference_ami_version,
                 )

tests/unit/sagemaker/experiments/test_run_context.py

@@ -54,7 +54,7 @@ def test_auto_pass_in_exp_config_to_train_job(mock_start_job, run_obj, sagemaker
         assert _RunContext.get_current_run() == run_obj
 
     expected_exp_config = run_obj.experiment_config
-    mock_start_job.assert_called_once_with(estimator, _train_input_path, expected_exp_config)
+    mock_start_job.assert_called_once_with(estimator, _train_input_path, expected_exp_config, None)
 
     # _RunContext is cleaned up after exiting the with statement
     assert not _RunContext.get_current_run()
@@ -94,7 +94,7 @@ def test_auto_pass_in_exp_config_under_load_run(
         assert loaded_run.experiment_config == run_obj.experiment_config
 
     expected_exp_config = run_obj.experiment_config
-    mock_start_job.assert_called_once_with(estimator, _train_input_path, expected_exp_config)
+    mock_start_job.assert_called_once_with(estimator, _train_input_path, expected_exp_config, None)
 
     # _RunContext is cleaned up after exiting the with statement
     assert not _RunContext.get_current_run()
@@ -174,7 +174,7 @@ def test_user_supply_exp_config_to_train_job(mock_start_job, run_obj, sagemaker_
 
         assert _RunContext.get_current_run() == run_obj
 
-    mock_start_job.assert_called_once_with(estimator, _train_input_path, supplied_exp_cfg)
+    mock_start_job.assert_called_once_with(estimator, _train_input_path, supplied_exp_cfg, None)
 
     # _RunContext is cleaned up after exiting the with statement
     assert not _RunContext.get_current_run()

tests/unit/sagemaker/jumpstart/estimator/test_estimator.py

@@ -607,6 +607,7 @@ def test_gated_model_s3_uri_with_eula_in_fit(
             inputs=channels,
             wait=True,
             job_name="meta-textgeneration-llama-2-7b-f-8675309",
+            accept_eula=True,
         )
 
         assert hasattr(estimator, "model_access_config")
@@ -688,6 +689,7 @@ def test_gated_model_non_model_package_s3_uri(
             instance_count=1,
             image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pyt"
             "orch-training:2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04",
+            model_uri="s3://jumpstart-private-cache-prod-us-west-2/some/dummy/key",
             source_dir="s3://jumpstart-cache-prod-us-west-2/source-d"
             "irectory-tarballs/meta/transfer_learning/textgeneration/prepack/v1.0.1/sourcedir.tar.gz",
             entry_point="transfer_learning.py",
@@ -1346,7 +1348,7 @@ def test_jumpstart_estimator_kwargs_match_parent_class(self):
         and reach out to JumpStart team."""
 
         init_args_to_skip: Set[str] = set(["kwargs"])
-        fit_args_to_skip: Set[str] = set(["accept_eula"])
+        fit_args_to_skip: Set[str] = set([])
         deploy_args_to_skip: Set[str] = set(["kwargs"])
 
         parent_class_init = Estimator.__init__

tests/unit/sagemaker/jumpstart/test_types.py

@@ -333,7 +333,7 @@ def test_use_training_model_artifact():
     specs1 = JumpStartModelSpecs(BASE_SPEC)
     assert specs1.use_training_model_artifact()
     specs1.gated_bucket = True
-    assert not specs1.use_training_model_artifact()
+    assert specs1.use_training_model_artifact()
     specs1.gated_bucket = False
     specs1.training_model_package_artifact_uris = {"region1": "blah", "region2": "blah2"}
     assert not specs1.use_training_model_artifact()

tests/unit/test_estimator.py

@@ -336,6 +336,148 @@ def training_job_description(sagemaker_session):
     return returned_job_description
 
 
+def test_set_accept_eula_for_input_data_config_no_input_data_config():
+    """Test when InputDataConfig is not in train_args."""
+    train_args = {}
+    accept_eula = True
+
+    EstimatorBase._set_accept_eula_for_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():
+    """Test when accept_eula is None."""
+    train_args = {"InputDataConfig": [{"DataSource": {"S3DataSource": {}}}]}
+    accept_eula = None
+
+    EstimatorBase._set_accept_eula_for_input_data_config(train_args, accept_eula)
+
+    # Verify train_args remains unchanged
+    assert train_args == {"InputDataConfig": [{"DataSource": {"S3DataSource": {}}}]}
+
+
+def test_set_accept_eula_for_input_data_config_single_data_source():
+    """Test with a single S3DataSource."""
+    with patch("sagemaker.estimator.logger") as logger:
+        train_args = {
+            "InputDataConfig": [{"DataSource": {"S3DataSource": {"S3Uri": "s3://bucket/model"}}}]
+        }
+        accept_eula = True
+
+        EstimatorBase._set_accept_eula_for_input_data_config(train_args, accept_eula)
+
+        # Verify ModelAccessConfig and AcceptEula are set correctly
+        assert train_args["InputDataConfig"][0]["DataSource"]["S3DataSource"][
+            "ModelAccessConfig"
+        ] == {"AcceptEula": True}
+
+        # Verify no logging occurred since there's only one data source
+        logger.info.assert_not_called()
+
+
+def test_set_accept_eula_for_input_data_config_multiple_data_sources():
+    """Test with multiple S3DataSources."""
+    with patch("sagemaker.estimator.logger") as logger:
+        train_args = {
+            "InputDataConfig": [
+                {"DataSource": {"S3DataSource": {"S3Uri": "s3://bucket/model1"}}},
+                {"DataSource": {"S3DataSource": {"S3Uri": "s3://bucket/model2"}}},
+            ]
+        }
+        accept_eula = True
+
+        EstimatorBase._set_accept_eula_for_input_data_config(train_args, accept_eula)
+
+        # Verify ModelAccessConfig and AcceptEula are set correctly for both data sources
+        assert train_args["InputDataConfig"][0]["DataSource"]["S3DataSource"][
+            "ModelAccessConfig"
+        ] == {"AcceptEula": True}
+        assert train_args["InputDataConfig"][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_input_data_config_existing_model_access_config():
+    """Test when ModelAccessConfig already exists."""
+    train_args = {
+        "InputDataConfig": [
+            {
+                "DataSource": {
+                    "S3DataSource": {
+                        "S3Uri": "s3://bucket/model",
+                        "ModelAccessConfig": {"OtherSetting": "value"},
+                    }
+                }
+            }
+        ]
+    }
+    accept_eula = True
+
+    EstimatorBase._set_accept_eula_for_input_data_config(train_args, accept_eula)
+
+    # Verify AcceptEula is added to existing ModelAccessConfig
+    assert train_args["InputDataConfig"][0]["DataSource"]["S3DataSource"]["ModelAccessConfig"] == {
+        "OtherSetting": "value",
+        "AcceptEula": True,
+    }
+
+
+def test_set_accept_eula_for_input_data_config_missing_s3_data_source():
+    """Test when S3DataSource is missing."""
+    train_args = {"InputDataConfig": [{"DataSource": {"OtherDataSource": {}}}]}
+    accept_eula = True
+
+    EstimatorBase._set_accept_eula_for_input_data_config(train_args, accept_eula)
+
+    # Verify train_args remains unchanged
+    assert train_args == {"InputDataConfig": [{"DataSource": {"OtherDataSource": {}}}]}
+
+
+def test_set_accept_eula_for_input_data_config_missing_data_source():
+    """Test when DataSource is missing."""
+    train_args = {"InputDataConfig": [{"OtherKey": {}}]}
+    accept_eula = True
+
+    EstimatorBase._set_accept_eula_for_input_data_config(train_args, accept_eula)
+
+    # Verify train_args remains unchanged
+    assert train_args == {"InputDataConfig": [{"OtherKey": {}}]}
+
+
+def test_set_accept_eula_for_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 = {
+            "InputDataConfig": [
+                {"DataSource": {"S3DataSource": {"S3Uri": "s3://bucket/model"}}},
+                {"DataSource": {"OtherDataSource": {}}},
+            ]
+        }
+        accept_eula = True
+
+        EstimatorBase._set_accept_eula_for_input_data_config(train_args, accept_eula)
+
+        # Verify ModelAccessConfig and AcceptEula are set correctly for S3DataSource only
+        assert train_args["InputDataConfig"][0]["DataSource"]["S3DataSource"][
+            "ModelAccessConfig"
+        ] == {"AcceptEula": True}
+        assert "ModelAccessConfig" not in train_args["InputDataConfig"][1]["DataSource"].get(
+            "OtherDataSource", {}
+        )
+
+        # Verify no logging occurred since there's only one S3 data source
+        logger.info.assert_not_called()
+
+
 def test_validate_smdistributed_unsupported_image_raises(sagemaker_session):
     # Test unsupported image raises error.
     for unsupported_image in DummyFramework.UNSUPPORTED_DLC_IMAGE_FOR_SM_PARALLELISM: