diff --git a/tests/unit/sagemaker/huggingface/test_estimator.py b/tests/unit/sagemaker/huggingface/test_estimator.py index 3ad641a321..d6cf88d811 100644 --- a/tests/unit/sagemaker/huggingface/test_estimator.py +++ b/tests/unit/sagemaker/huggingface/test_estimator.py @@ -241,7 +241,7 @@ def test_huggingface( sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job( huggingface_training_version, f"pytorch{huggingface_pytorch_training_version}" diff --git a/tests/unit/sagemaker/training_compiler/test_huggingface_pytorch_compiler.py b/tests/unit/sagemaker/training_compiler/test_huggingface_pytorch_compiler.py index e1c21cf662..fee26668d2 100644 --- a/tests/unit/sagemaker/training_compiler/test_huggingface_pytorch_compiler.py +++ b/tests/unit/sagemaker/training_compiler/test_huggingface_pytorch_compiler.py @@ -402,7 +402,7 @@ def test_pytorchxla_distribution( sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job( huggingface_training_compiler_version, @@ -463,7 +463,7 @@ def test_default_compiler_config( sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job( huggingface_training_compiler_version, @@ -519,7 +519,7 @@ def test_debug_compiler_config( sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job( huggingface_training_compiler_version, @@ -575,7 +575,7 @@ def test_disable_compiler_config( sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job( huggingface_training_compiler_version, diff --git a/tests/unit/sagemaker/training_compiler/test_huggingface_tensorflow_compiler.py b/tests/unit/sagemaker/training_compiler/test_huggingface_tensorflow_compiler.py index e0d172f6e0..8bfb694040 100644 --- a/tests/unit/sagemaker/training_compiler/test_huggingface_tensorflow_compiler.py +++ b/tests/unit/sagemaker/training_compiler/test_huggingface_tensorflow_compiler.py @@ -349,7 +349,7 @@ def test_default_compiler_config( sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job( huggingface_training_compiler_version, @@ -407,7 +407,7 @@ def test_debug_compiler_config( sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job( huggingface_training_compiler_version, @@ -465,7 +465,7 @@ def test_disable_compiler_config( sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job( huggingface_training_compiler_version, diff --git a/tests/unit/sagemaker/training_compiler/test_pytorch_compiler.py b/tests/unit/sagemaker/training_compiler/test_pytorch_compiler.py index 34a1236a7f..cac31313cd 100644 --- a/tests/unit/sagemaker/training_compiler/test_pytorch_compiler.py +++ b/tests/unit/sagemaker/training_compiler/test_pytorch_compiler.py @@ -344,7 +344,7 @@ def test_pytorchxla_distribution( sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job( pytorch_training_compiler_version, @@ -403,7 +403,7 @@ def test_default_compiler_config( sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job( pytorch_training_compiler_version, @@ -458,7 +458,7 @@ def test_debug_compiler_config( sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job( pytorch_training_compiler_version, @@ -513,7 +513,7 @@ def test_disable_compiler_config( sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job( pytorch_training_compiler_version, diff --git a/tests/unit/sagemaker/training_compiler/test_tensorflow_compiler.py b/tests/unit/sagemaker/training_compiler/test_tensorflow_compiler.py index ac42bb53ab..6c76558568 100644 --- a/tests/unit/sagemaker/training_compiler/test_tensorflow_compiler.py +++ b/tests/unit/sagemaker/training_compiler/test_tensorflow_compiler.py @@ -289,7 +289,7 @@ def test_default( sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job( tensorflow_training_version, @@ -348,7 +348,7 @@ def test_byoc( sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job( tensorflow_training_version, @@ -399,7 +399,7 @@ def test_debug_compiler_config( sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job( tensorflow_training_version, @@ -450,7 +450,7 @@ def test_disable_compiler_config( sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job( tensorflow_training_version, diff --git a/tests/unit/test_chainer.py b/tests/unit/test_chainer.py index 8ad2ae0bab..dfaf13812e 100644 --- a/tests/unit/test_chainer.py +++ b/tests/unit/test_chainer.py @@ -354,7 +354,7 @@ def test_chainer(strftime, time, sagemaker_session, chainer_version, chainer_py_ sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job(chainer_version, chainer_py_version) expected_train_args["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] = inputs diff --git a/tests/unit/test_mxnet.py b/tests/unit/test_mxnet.py index 4a584dfae4..dbb11cd63e 100644 --- a/tests/unit/test_mxnet.py +++ b/tests/unit/test_mxnet.py @@ -360,7 +360,7 @@ def test_mxnet( sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] actual_train_args = sagemaker_session.method_calls[0][2] job_name = actual_train_args["job_name"] diff --git a/tests/unit/test_processing.py b/tests/unit/test_processing.py index 090a44f1dc..caac42b306 100644 --- a/tests/unit/test_processing.py +++ b/tests/unit/test_processing.py @@ -133,6 +133,14 @@ def uploaded_code( return UploadedCode(s3_prefix=s3_prefix, script_name=script_name) +@pytest.fixture(auto_use=True) +def mock_process_studio_metadata_file(): + with patch( + "sagemaker.user_agent.process_studio_metadata_file", MagicMock() + ) as process_studio_metadata_file: + yield process_studio_metadata_file + + @patch("sagemaker.utils._botocore_resolver") @patch("os.path.exists", return_value=True) @patch("os.path.isfile", return_value=True) @@ -277,7 +285,6 @@ def test_sklearn_with_all_parameters_via_run_args( sagemaker_session.process.assert_called_with(**expected_args) -@patch("sagemaker.user_agent.process_studio_metadata_file", return_value="TestApp") @patch("sagemaker.utils._botocore_resolver") @patch("os.path.exists", return_value=True) @patch("os.path.isfile", return_value=True) @@ -285,7 +292,6 @@ def test_sklearn_with_all_parameters_via_run_args_called_twice( exists_mock, isfile_mock, botocore_resolver, - process_studio_metadata_file_mock, sklearn_version, sagemaker_session, ): @@ -338,7 +344,6 @@ def test_sklearn_with_all_parameters_via_run_args_called_twice( sagemaker_session.process.assert_called_with(**expected_args) -@patch("sagemaker.user_agent.process_studio_metadata_file", return_value="TestApp") @patch("sagemaker.utils._botocore_resolver") @patch("os.path.exists", return_value=True) @patch("os.path.isfile", return_value=True) @@ -346,7 +351,6 @@ def test_pytorch_processor_with_required_parameters( exists_mock, isfile_mock, botocore_resolver, - process_studio_metadata_file_mock, sagemaker_session, pytorch_training_version, pytorch_training_py_version, @@ -418,14 +422,12 @@ def test_xgboost_processor_with_required_parameters( sagemaker_session.process.assert_called_with(**expected_args) -@patch("sagemaker.user_agent.process_studio_metadata_file", return_value="TestApp") @patch("sagemaker.utils._botocore_resolver") @patch("os.path.exists", return_value=True) @patch("os.path.isfile", return_value=True) def test_mxnet_processor_with_required_parameters( exists_mock, isfile_mock, - process_studio_metadata_file_mock, botocore_resolver, sagemaker_session, mxnet_training_version, @@ -468,7 +470,6 @@ def test_mxnet_processor_with_required_parameters( sagemaker_session.process.assert_called_with(**expected_args) -@patch("sagemaker.user_agent.process_studio_metadata_file", return_value="TestApp") @patch("sagemaker.utils._botocore_resolver") @patch("os.path.exists", return_value=True) @patch("os.path.isfile", return_value=True) @@ -476,7 +477,6 @@ def test_tensorflow_processor_with_required_parameters( exists_mock, isfile_mock, botocore_resolver, - process_studio_metadata_file_mock, sagemaker_session, tensorflow_training_version, tensorflow_training_py_version, diff --git a/tests/unit/test_pytorch.py b/tests/unit/test_pytorch.py index 618d0d7ea8..e35f17331c 100644 --- a/tests/unit/test_pytorch.py +++ b/tests/unit/test_pytorch.py @@ -337,7 +337,7 @@ def test_pytorch( sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job(pytorch_inference_version, pytorch_inference_py_version) expected_train_args["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] = inputs diff --git a/tests/unit/test_rl.py b/tests/unit/test_rl.py index d9c4129cf6..fb337f0063 100644 --- a/tests/unit/test_rl.py +++ b/tests/unit/test_rl.py @@ -335,7 +335,7 @@ def test_rl(time, strftime, sagemaker_session, coach_mxnet_version): sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job( RLToolkit.COACH.value, coach_mxnet_version, RLFramework.MXNET.value diff --git a/tests/unit/test_sklearn.py b/tests/unit/test_sklearn.py index b0df31fee1..68302ee572 100644 --- a/tests/unit/test_sklearn.py +++ b/tests/unit/test_sklearn.py @@ -332,7 +332,7 @@ def test_sklearn(time, strftime, sagemaker_session, sklearn_version): sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job(sklearn_version) expected_train_args["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] = inputs diff --git a/tests/unit/test_xgboost.py b/tests/unit/test_xgboost.py index 18eab98149..12d2354dad 100644 --- a/tests/unit/test_xgboost.py +++ b/tests/unit/test_xgboost.py @@ -330,7 +330,7 @@ def test_xgboost_cpu(time, strftime, sagemaker_session, xgboost_framework_versio sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job(xgboost_framework_version) expected_train_args["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] = inputs @@ -377,7 +377,7 @@ def test_xgboost_gpu(time, strftime, sagemaker_session, xgboost_gpu_framework_ve sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job( xgboost_gpu_framework_version, instance_type=GPU_INSTANCE_TYPE @@ -427,7 +427,7 @@ def test_distributed_training(time, strftime, sagemaker_session, xgboost_framewo sagemaker_call_names = [c[0] for c in sagemaker_session.method_calls] assert sagemaker_call_names == ["train", "logs_for_job"] boto_call_names = [c[0] for c in sagemaker_session.boto_session.method_calls] - assert boto_call_names == ["resource"] + assert boto_call_names == ["resource", "client"] expected_train_args = _create_train_job(xgboost_framework_version, DIST_INSTANCE_COUNT) expected_train_args["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] = inputs