Merge pull request #4 from aws/master

update with aws:master

Author: chuyang-deng

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):
+    def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client, default_bucket):
         """Initialize this Local SageMaker Session.
 
         Args:
@@ -413,6 +413,9 @@ 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"):
         """
 

src/sagemaker/processing.py

@@ -369,8 +369,9 @@ def run(
         """
         self._current_job_name = self._generate_current_job_name(job_name=job_name)
 
-        user_script_name = self._get_user_script_name(code)
-        user_code_s3_uri = self._upload_code(code)
+        user_code_s3_uri = self._handle_user_code_url(code)
+        user_script_name = self._get_user_code_name(code)
+
         inputs_with_code = self._convert_code_and_add_to_inputs(inputs, user_code_s3_uri)
 
         self._set_entrypoint(self.command, user_script_name)
@@ -389,25 +390,59 @@ def run(
         if wait:
             self.latest_job.wait(logs=logs)
 
-    def _get_user_script_name(self, code):
-        """Finds the user script name using the provided code file,
-        directory, or script name.
+    def _get_user_code_name(self, code):
+        """Gets the basename of the user's code from the URL the customer provided.
 
         Args:
-            code (str): This can be an S3 uri or a local path to either
-                a directory or a file.
+            code (str): A URL to the user's code.
+
+        Returns:
+            str: The basename of the user's code.
+
+        """
+        code_url = urlparse(code)
+        return os.path.basename(code_url.path)
+
+    def _handle_user_code_url(self, code):
+        """Gets the S3 URL containing the user's code.
+
+           Inspects the scheme the customer passed in ("s3://" for code in S3, "file://" or nothing
+           for absolute or local file paths. Uploads the code to S3 if the code is a local file.
+
+        Args:
+            code (str): A URL to the customer's code.
 
         Returns:
-            str: The script name from the S3 uri or from the file found
-                on the user's local machine.
+            str: The S3 URL to the customer's code.
+
         """
-        if os.path.isdir(code) is None or not os.path.splitext(code)[1]:
+        code_url = urlparse(code)
+        if code_url.scheme == "s3":
+            user_code_s3_uri = code
+        elif code_url.scheme == "" or code_url.scheme == "file":
+            # Validate that the file exists locally and is not a directory.
+            if not os.path.exists(code):
+                raise ValueError(
+                    """code {} wasn't found. Please make sure that the file exists.
+                    """.format(
+                        code
+                    )
+                )
+            if not os.path.isfile(code):
+                raise ValueError(
+                    """code {} must be a file, not a directory. Please pass a path to a file.
+                    """.format(
+                        code
+                    )
+                )
+            user_code_s3_uri = self._upload_code(code)
+        else:
             raise ValueError(
-                """'code' must be a file, not a directory. Please pass a path to a file, not a
-                directory.
-                """
+                "code {} url scheme {} is not recognized. Please pass a file path or S3 url".format(
+                    code, code_url.scheme
+                )
             )
-        return os.path.basename(code)
+        return user_code_s3_uri
 
     def _upload_code(self, code):
         """Uploads a code file or directory specified as a string

src/sagemaker/session.py

@@ -76,7 +76,13 @@ 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):
+    def __init__(
+        self,
+        boto_session=None,
+        sagemaker_client=None,
+        sagemaker_runtime_client=None,
+        default_bucket=None,
+    ):
         """Initialize a SageMaker ``Session``.
 
         Args:
@@ -91,15 +97,23 @@ def __init__(self, boto_session=None, sagemaker_client=None, sagemaker_runtime_c
                 ``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, sagemaker_client, sagemaker_runtime_client)
+        self._initialize(
+            boto_session=boto_session,
+            sagemaker_client=sagemaker_client,
+            sagemaker_runtime_client=sagemaker_runtime_client,
+            default_bucket=default_bucket,
+        )
 
-    def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client):
+    def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client, default_bucket):
         """Initialize this SageMaker Session.
 
         Creates or uses a boto_session, sagemaker_client and sagemaker_runtime_client.
@@ -126,6 +140,12 @@ 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
@@ -314,11 +334,14 @@ 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
-        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)
+
+        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)
 
         s3 = self.boto_session.resource("s3")
         try:

src/sagemaker/sklearn/processing.py

@@ -26,6 +26,8 @@
 class SKLearnProcessor(ScriptProcessor):
     """Handles Amazon SageMaker processing tasks for jobs using scikit-learn."""
 
+    _valid_framework_versions = ["0.20.0"]
+
     def __init__(
         self,
         framework_version,
@@ -84,6 +86,13 @@ def __init__(
         session = sagemaker_session or Session()
         region = session.boto_region_name
 
+        if framework_version not in self._valid_framework_versions:
+            raise ValueError(
+                "scikit-learn version {} is not supported. Supported versions are {}".format(
+                    framework_version, self._valid_framework_versions
+                )
+            )
+
         if not command:
             command = ["python3"]
 

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):
+    def _initialize(self, boto_session, sagemaker_client, sagemaker_runtime_client, default_bucket):
         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,6 +53,9 @@ 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):

tests/integ/test_processing.py

@@ -14,7 +14,10 @@
 
 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
@@ -23,6 +26,35 @@
 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")
@@ -170,6 +202,90 @@ 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
 ):

tests/unit/test_create_deploy_entities.py

@@ -34,6 +34,13 @@
 @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