Revert "feature: allow setting the default bucket in Session (#1168)" (#1175)

ajaykarpur · web-flow · commit e98029119a07 · 2019-12-16T11:43:35.000-08:00
This reverts commit 98fa76d.
diff --git a/src/sagemaker/local/local_session.py b/src/sagemaker/local/local_session.py
@@ -379,7 +379,7 @@ def __init__(self, boto_session=None):
         if platform.system() == "Windows":
             logger.warning("Windows Support for Local Mode is Experimental")
 
-    def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client, default_bucket):
+    def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client):
         """Initialize this Local SageMaker Session.
 
         Args:
@@ -413,9 +413,6 @@ def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client,
 
             self.config = yaml.load(open(sagemaker_config_file, "r"))
 
-        self._default_bucket = None
-        self._desired_default_bucket_name = default_bucket
-
     def logs_for_job(self, job_name, wait=False, poll=5, log_type="All"):
         """
 
diff --git a/src/sagemaker/session.py b/src/sagemaker/session.py
@@ -76,13 +76,7 @@ class Session(object):  # pylint: disable=too-many-public-methods
     bucket based on a naming convention which includes the current AWS account ID.
     """
 
-    def __init__(
-        self,
-        boto_session=None,
-        sagemaker_client=None,
-        sagemaker_runtime_client=None,
-        default_bucket=None,
-    ):
+    def __init__(self, boto_session=None, sagemaker_client=None, sagemaker_runtime_client=None):
         """Initialize a SageMaker ``Session``.
 
         Args:
@@ -97,23 +91,15 @@ def __init__(
                 ``InvokeEndpoint`` calls to Amazon SageMaker (default: None). Predictors created
                 using this ``Session`` use this client. If not provided, one will be created using
                 this instance's ``boto_session``.
-            default_bucket (str): The default s3 bucket to be used by this session.
-                Ex: "sagemaker-us-west-2"
-
         """
         self._default_bucket = None
 
         # currently is used for local_code in local mode
         self.config = None
 
-        self._initialize(
-            boto_session=boto_session,
-            sagemaker_client=sagemaker_client,
-            sagemaker_runtime_client=sagemaker_runtime_client,
-            default_bucket=default_bucket,
-        )
+        self._initialize(boto_session, sagemaker_client, sagemaker_runtime_client)
 
-    def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client, default_bucket):
+    def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client):
         """Initialize this SageMaker Session.
 
         Creates or uses a boto_session, sagemaker_client and sagemaker_runtime_client.
@@ -140,12 +126,6 @@ def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client,
 
         prepend_user_agent(self.sagemaker_runtime_client)
 
-        self._default_bucket = None
-        self._desired_default_bucket_name = default_bucket
-
-        # Create default bucket on session init to verify that desired name, if specified, is valid
-        self.default_bucket()
-
         self.local_mode = False
 
     @property
@@ -334,14 +314,11 @@ def default_bucket(self):
         if self._default_bucket:
             return self._default_bucket
 
-        default_bucket = self._desired_default_bucket_name
         region = self.boto_session.region_name
-
-        if not default_bucket:
-            account = self.boto_session.client(
-                "sts", region_name=region, endpoint_url=sts_regional_endpoint(region)
-            ).get_caller_identity()["Account"]
-            default_bucket = "sagemaker-{}-{}".format(region, account)
+        account = self.boto_session.client(
+            "sts", region_name=region, endpoint_url=sts_regional_endpoint(region)
+        ).get_caller_identity()["Account"]
+        default_bucket = "sagemaker-{}-{}".format(region, account)
 
         s3 = self.boto_session.resource("s3")
         try:
diff --git a/tests/integ/test_local_mode.py b/tests/integ/test_local_mode.py
@@ -43,7 +43,7 @@ class LocalNoS3Session(LocalSession):
     def __init__(self):
         super(LocalSession, self).__init__()
 
-    def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client, default_bucket):
+    def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client):
         self.boto_session = boto3.Session(region_name=DEFAULT_REGION)
         if self.config is None:
             self.config = {"local": {"local_code": True, "region_name": DEFAULT_REGION}}
@@ -53,9 +53,6 @@ def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client,
         self.sagemaker_runtime_client = LocalSagemakerRuntimeClient(self.config)
         self.local_mode = True
 
-        self._default_bucket = None
-        self._desired_default_bucket_name = default_bucket
-
 
 @pytest.fixture(scope="module")
 def mxnet_model(sagemaker_local_session, mxnet_full_version):
diff --git a/tests/integ/test_processing.py b/tests/integ/test_processing.py
@@ -14,10 +14,7 @@
 
 import os
 
-import boto3
 import pytest
-from botocore.config import Config
-from sagemaker import Session
 from sagemaker.fw_registry import default_framework_uri
 
 from sagemaker.processing import ProcessingInput, ProcessingOutput, ScriptProcessor, Processor
@@ -26,35 +23,6 @@
 from tests.integ.kms_utils import get_or_create_kms_key
 
 ROLE = "SageMakerRole"
-DEFAULT_REGION = "us-west-2"
-CUSTOM_BUCKET_PATH = "sagemaker-custom-bucket"
-
-
-@pytest.fixture(scope="module")
-def sagemaker_session_with_custom_bucket(
-    boto_config, sagemaker_client_config, sagemaker_runtime_config
-):
-    boto_session = (
-        boto3.Session(**boto_config) if boto_config else boto3.Session(region_name=DEFAULT_REGION)
-    )
-    sagemaker_client_config.setdefault("config", Config(retries=dict(max_attempts=10)))
-    sagemaker_client = (
-        boto_session.client("sagemaker", **sagemaker_client_config)
-        if sagemaker_client_config
-        else None
-    )
-    runtime_client = (
-        boto_session.client("sagemaker-runtime", **sagemaker_runtime_config)
-        if sagemaker_runtime_config
-        else None
-    )
-
-    return Session(
-        boto_session=boto_session,
-        sagemaker_client=sagemaker_client,
-        sagemaker_runtime_client=runtime_client,
-        default_bucket=CUSTOM_BUCKET_PATH,
-    )
 
 
 @pytest.fixture(scope="module")
@@ -202,90 +170,6 @@ def test_sklearn_with_customizations(
     assert job_description["StoppingCondition"] == {"MaxRuntimeInSeconds": 3600}
 
 
-def test_sklearn_with_custom_default_bucket(
-    sagemaker_session_with_custom_bucket,
-    image_uri,
-    sklearn_full_version,
-    cpu_instance_type,
-    output_kms_key,
-):
-
-    input_file_path = os.path.join(DATA_DIR, "dummy_input.txt")
-
-    sklearn_processor = SKLearnProcessor(
-        framework_version=sklearn_full_version,
-        role=ROLE,
-        command=["python3"],
-        instance_type=cpu_instance_type,
-        instance_count=1,
-        volume_size_in_gb=100,
-        volume_kms_key=None,
-        output_kms_key=output_kms_key,
-        max_runtime_in_seconds=3600,
-        base_job_name="test-sklearn-with-customizations",
-        env={"DUMMY_ENVIRONMENT_VARIABLE": "dummy-value"},
-        tags=[{"Key": "dummy-tag", "Value": "dummy-tag-value"}],
-        sagemaker_session=sagemaker_session_with_custom_bucket,
-    )
-
-    sklearn_processor.run(
-        code=os.path.join(DATA_DIR, "dummy_script.py"),
-        inputs=[
-            ProcessingInput(
-                source=input_file_path,
-                destination="/opt/ml/processing/input/container/path/",
-                input_name="dummy_input",
-                s3_data_type="S3Prefix",
-                s3_input_mode="File",
-                s3_data_distribution_type="FullyReplicated",
-                s3_compression_type="None",
-            )
-        ],
-        outputs=[
-            ProcessingOutput(
-                source="/opt/ml/processing/output/container/path/",
-                output_name="dummy_output",
-                s3_upload_mode="EndOfJob",
-            )
-        ],
-        arguments=["-v"],
-        wait=True,
-        logs=True,
-    )
-
-    job_description = sklearn_processor.latest_job.describe()
-
-    assert job_description["ProcessingInputs"][0]["InputName"] == "dummy_input"
-    assert CUSTOM_BUCKET_PATH in job_description["ProcessingInputs"][0]["S3Input"]["S3Uri"]
-
-    assert job_description["ProcessingInputs"][1]["InputName"] == "code"
-    assert CUSTOM_BUCKET_PATH in job_description["ProcessingInputs"][1]["S3Input"]["S3Uri"]
-
-    assert job_description["ProcessingJobName"].startswith("test-sklearn-with-customizations")
-
-    assert job_description["ProcessingJobStatus"] == "Completed"
-
-    assert job_description["ProcessingOutputConfig"]["KmsKeyId"] == output_kms_key
-    assert job_description["ProcessingOutputConfig"]["Outputs"][0]["OutputName"] == "dummy_output"
-
-    assert job_description["ProcessingResources"] == {
-        "ClusterConfig": {"InstanceCount": 1, "InstanceType": "ml.m4.xlarge", "VolumeSizeInGB": 100}
-    }
-
-    assert job_description["AppSpecification"]["ContainerArguments"] == ["-v"]
-    assert job_description["AppSpecification"]["ContainerEntrypoint"] == [
-        "python3",
-        "/opt/ml/processing/input/code/dummy_script.py",
-    ]
-    assert job_description["AppSpecification"]["ImageUri"] == image_uri
-
-    assert job_description["Environment"] == {"DUMMY_ENVIRONMENT_VARIABLE": "dummy-value"}
-
-    assert ROLE in job_description["RoleArn"]
-
-    assert job_description["StoppingCondition"] == {"MaxRuntimeInSeconds": 3600}
-
-
 def test_sklearn_with_no_inputs_or_outputs(
     sagemaker_session, image_uri, sklearn_full_version, cpu_instance_type
 ):
diff --git a/tests/unit/test_create_deploy_entities.py b/tests/unit/test_create_deploy_entities.py
@@ -34,13 +34,6 @@
 @pytest.fixture()
 def sagemaker_session():
     boto_mock = Mock(name="boto_session", region_name=REGION)
-    client_mock = Mock()
-    client_mock.get_caller_identity.return_value = {
-        "UserId": "mock_user_id",
-        "Account": "012345678910",
-        "Arn": "arn:aws:iam::012345678910:user/mock-user",
-    }
-    boto_mock.client.return_value = client_mock
     ims = sagemaker.Session(boto_session=boto_mock)
     ims.expand_role = Mock(return_value=EXPANDED_ROLE)
     return ims
diff --git a/tests/unit/test_default_bucket.py b/tests/unit/test_default_bucket.py
@@ -25,13 +25,6 @@
 @pytest.fixture()
 def sagemaker_session():
     boto_mock = Mock(name="boto_session", region_name=REGION)
-    client_mock = Mock()
-    client_mock.get_caller_identity.return_value = {
-        "UserId": "mock_user_id",
-        "Account": "012345678910",
-        "Arn": "arn:aws:iam::012345678910:user/mock-user",
-    }
-    boto_mock.client.return_value = client_mock
     boto_mock.client("sts").get_caller_identity.return_value = {"Account": ACCOUNT_ID}
     ims = sagemaker.Session(boto_session=boto_mock)
     return ims
@@ -55,13 +48,11 @@ def test_default_already_cached(sagemaker_session):
     existing_default = "mydefaultbucket"
     sagemaker_session._default_bucket = existing_default
 
-    before_create_calls = sagemaker_session.boto_session.resource().create_bucket.mock_calls
-
     bucket_name = sagemaker_session.default_bucket()
 
-    after_create_calls = sagemaker_session.boto_session.resource().create_bucket.mock_calls
+    create_calls = sagemaker_session.boto_session.resource().create_bucket.mock_calls
     assert bucket_name == existing_default
-    assert before_create_calls == after_create_calls
+    assert create_calls == []
 
 
 def test_default_bucket_exists(sagemaker_session):
@@ -87,42 +78,22 @@ def test_concurrent_bucket_modification(sagemaker_session):
     assert bucket_name == DEFAULT_BUCKET_NAME
 
 
-def test_bucket_creation_client_error():
+def test_bucket_creation_client_error(sagemaker_session):
     with pytest.raises(ClientError):
-        boto_mock = Mock(name="boto_session", region_name=REGION)
-        client_mock = Mock()
-        client_mock.get_caller_identity.return_value = {
-            "UserId": "mock_user_id",
-            "Account": "012345678910",
-            "Arn": "arn:aws:iam::012345678910:user/mock-user",
-        }
-        boto_mock.client.return_value = client_mock
-        boto_mock.client("sts").get_caller_identity.return_value = {"Account": ACCOUNT_ID}
-
         error = ClientError(
             error_response={"Error": {"Code": "SomethingWrong", "Message": "message"}},
             operation_name="foo",
         )
-        boto_mock.resource().create_bucket.side_effect = error
+        sagemaker_session.boto_session.resource().create_bucket.side_effect = error
 
-        session = sagemaker.Session(boto_session=boto_mock)
-        assert session._default_bucket is None
+        sagemaker_session.default_bucket()
+    assert sagemaker_session._default_bucket is None
 
 
-def test_bucket_creation_other_error():
+def test_bucket_creation_other_error(sagemaker_session):
     with pytest.raises(RuntimeError):
-        boto_mock = Mock(name="boto_session", region_name=REGION)
-        client_mock = Mock()
-        client_mock.get_caller_identity.return_value = {
-            "UserId": "mock_user_id",
-            "Account": "012345678910",
-            "Arn": "arn:aws:iam::012345678910:user/mock-user",
-        }
-        boto_mock.client.return_value = client_mock
-        boto_mock.client("sts").get_caller_identity.return_value = {"Account": ACCOUNT_ID}
-
         error = RuntimeError()
-        boto_mock.resource().create_bucket.side_effect = error
+        sagemaker_session.boto_session.resource().create_bucket.side_effect = error
 
-        session = sagemaker.Session(boto_session=boto_mock)
-        assert session._default_bucket is None
+        sagemaker_session.default_bucket()
+    assert sagemaker_session._default_bucket is None
diff --git a/tests/unit/test_endpoint_from_job.py b/tests/unit/test_endpoint_from_job.py
@@ -43,13 +43,6 @@
 @pytest.fixture()
 def sagemaker_session():
     boto_mock = Mock(name="boto_session", region_name=REGION)
-    client_mock = Mock()
-    client_mock.get_caller_identity.return_value = {
-        "UserId": "mock_user_id",
-        "Account": "012345678910",
-        "Arn": "arn:aws:iam::012345678910:user/mock-user",
-    }
-    boto_mock.client.return_value = client_mock
     ims = sagemaker.Session(sagemaker_client=Mock(name="sagemaker_client"), boto_session=boto_mock)
     ims.sagemaker_client.describe_training_job = Mock(
         name="describe_training_job", return_value=TRAINING_JOB_RESPONSE
diff --git a/tests/unit/test_endpoint_from_model_data.py b/tests/unit/test_endpoint_from_model_data.py
@@ -36,13 +36,6 @@
 @pytest.fixture()
 def sagemaker_session():
     boto_mock = Mock(name="boto_session", region_name=REGION)
-    client_mock = Mock()
-    client_mock.get_caller_identity.return_value = {
-        "UserId": "mock_user_id",
-        "Account": "012345678910",
-        "Arn": "arn:aws:iam::012345678910:user/mock-user",
-    }
-    boto_mock.client.return_value = client_mock
     ims = sagemaker.Session(sagemaker_client=Mock(name="sagemaker_client"), boto_session=boto_mock)
     ims.sagemaker_client.describe_model = Mock(
         name="describe_model", side_effect=_raise_does_not_exist_client_error
diff --git a/tests/unit/test_exception_on_bad_status.py b/tests/unit/test_exception_on_bad_status.py
@@ -29,13 +29,7 @@ def get_sagemaker_session(returns_status):
     client_mock.describe_model_package = MagicMock(
         return_value={"ModelPackageStatus": returns_status}
     )
-    client_mock.get_caller_identity.return_value = {
-        "UserId": "mock_user_id",
-        "Account": "012345678910",
-        "Arn": "arn:aws:iam::012345678910:user/mock-user",
-    }
     client_mock.describe_endpoint = MagicMock(return_value={"EndpointStatus": returns_status})
-    boto_mock.client.return_value = client_mock
     ims = sagemaker.Session(boto_session=boto_mock, sagemaker_client=client_mock)
     ims.expand_role = Mock(return_value=EXPANDED_ROLE)
     return ims
diff --git a/tests/unit/test_local_session.py b/tests/unit/test_local_session.py
@@ -473,12 +473,5 @@ def test_file_input_content_type():
 
 def test_local_session_is_set_to_local_mode():
     boto_session = Mock(region_name="us-west-2")
-    client_mock = Mock()
-    client_mock.get_caller_identity.return_value = {
-        "UserId": "mock_user_id",
-        "Account": "012345678910",
-        "Arn": "arn:aws:iam::012345678910:user/mock-user",
-    }
-    boto_session.client.return_value = client_mock
     local_session = sagemaker.local.local_session.LocalSession(boto_session=boto_session)
     assert local_session.local_mode
diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
diff --git a/tests/unit/test_upload_data.py b/tests/unit/test_upload_data.py
