diff --git a/sagemaker-core/src/sagemaker/core/local/local_session.py b/sagemaker-core/src/sagemaker/core/local/local_session.py index 9d6ab7b199..363585c1d1 100644 --- a/sagemaker-core/src/sagemaker/core/local/local_session.py +++ b/sagemaker-core/src/sagemaker/core/local/local_session.py @@ -670,8 +670,8 @@ def _initialize( validate_sagemaker_config(sagemaker_config) if self.s3_endpoint_url is not None: - self.s3_resource = boto_session.resource("s3", endpoint_url=self.s3_endpoint_url) - self.s3_client = boto_session.client("s3", endpoint_url=self.s3_endpoint_url) + self.s3_resource = self.boto_session.resource("s3", endpoint_url=self.s3_endpoint_url) + self.s3_client = self.boto_session.client("s3", endpoint_url=self.s3_endpoint_url) self.sagemaker_config = ( sagemaker_config if sagemaker_config diff --git a/sagemaker-core/tests/unit/local/test_local_session.py b/sagemaker-core/tests/unit/local/test_local_session.py index 91f4e4feaa..d20706d1b9 100644 --- a/sagemaker-core/tests/unit/local/test_local_session.py +++ b/sagemaker-core/tests/unit/local/test_local_session.py @@ -509,6 +509,26 @@ def test_local_session_with_s3_endpoint(self): assert session.s3_endpoint_url == "http://localhost:9000" + @patch("boto3.Session") + def test_local_session_with_s3_endpoint_and_default_boto_session(self, mock_boto_session_class): + """Test custom S3 endpoint when no boto session is passed in""" + mock_boto_session = Mock() + mock_boto_session.region_name = "us-west-2" + mock_boto_session_class.return_value = mock_boto_session + + with patch("sagemaker.core.local.local_session.load_sagemaker_config", return_value={}): + with patch( + "sagemaker.core.local.local_session.load_local_mode_config", + return_value={"local": {}}, + ): + session = LocalSession(s3_endpoint_url="http://localhost:9000") + + mock_boto_session.resource.assert_called_once_with( + "s3", endpoint_url="http://localhost:9000" + ) + mock_boto_session.client.assert_any_call("s3", endpoint_url="http://localhost:9000") + assert session.s3_resource is mock_boto_session.resource.return_value + @patch("boto3.Session") def test_local_session_no_region_raises(self, mock_boto_session_class): """Test that missing region raises error"""