change: allow ModelMonitor and Processor to take IAM role names (in addition to ARNs) (#1150)

Author: laurenyu

src/sagemaker/automl/automl.py

@@ -10,7 +10,7 @@
 # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
-"""A class for SageMaker AutoML Job."""
+"""A class for SageMaker AutoML Jobs."""
 from __future__ import absolute_import
 
 from six import string_types

src/sagemaker/model_monitor/model_monitoring.py

@@ -22,24 +22,17 @@
 import logging
 import uuid
 
-from six.moves.urllib.parse import urlparse
 from six import string_types
-
+from six.moves.urllib.parse import urlparse
 from botocore.exceptions import ClientError
 
+from sagemaker.exceptions import UnexpectedStatusException
+from sagemaker.model_monitor.monitoring_files import Constraints, ConstraintViolations, Statistics
 from sagemaker.network import NetworkConfig
+from sagemaker.processing import Processor, ProcessingInput, ProcessingJob, ProcessingOutput
 from sagemaker.s3 import S3Uploader
-
-from sagemaker.utils import name_from_base
 from sagemaker.session import Session
-from sagemaker.processing import Processor
-from sagemaker.processing import ProcessingJob
-from sagemaker.processing import ProcessingInput
-from sagemaker.processing import ProcessingOutput
-from sagemaker.model_monitor.monitoring_files import Constraints, ConstraintViolations
-from sagemaker.model_monitor.monitoring_files import Statistics
-from sagemaker.exceptions import UnexpectedStatusException
-from sagemaker.utils import retries
+from sagemaker.utils import name_from_base, retries
 
 _DEFAULT_MONITOR_IMAGE_URI_WITH_PLACEHOLDERS = (
     "{}.dkr.ecr.{}.amazonaws.com/sagemaker-model-monitor-analyzer"
@@ -331,7 +324,7 @@ def create_monitoring_schedule(
             max_runtime_in_seconds=self.max_runtime_in_seconds,
             environment=self.env,
             network_config=network_config_dict,
-            role_arn=self.role,
+            role_arn=self.sagemaker_session.expand_role(self.role),
             tags=self.tags,
         )
 
@@ -390,7 +383,7 @@ def update_monitoring_schedule(
             network_config (sagemaker.network.NetworkConfig): A NetworkConfig
                 object that configures network isolation, encryption of
                 inter-container traffic, security group IDs, and subnets.
-            role (str): An AWS IAM role. The Amazon SageMaker jobs use this role.
+            role (str): An AWS IAM role name or ARN. The Amazon SageMaker jobs use this role.
             image_uri (str): The uri of the image to use for the jobs started by
                 the Monitor.
 
@@ -478,7 +471,7 @@ def update_monitoring_schedule(
             max_runtime_in_seconds=max_runtime_in_seconds,
             environment=env,
             network_config=network_config_dict,
-            role_arn=role,
+            role_arn=self.sagemaker_session.expand_role(self.role),
         )
 
         self._wait_for_schedule_changes_to_apply()
@@ -988,7 +981,7 @@ def __init__(
         creating Amazon SageMaker Monitoring Schedules to monitor SageMaker endpoints.
 
         Args:
-            role (str): An AWS IAM role. The Amazon SageMaker jobs use this role.
+            role (str): An AWS IAM role name or ARN. The Amazon SageMaker jobs use this role.
             instance_count (int): The number of instances to run the jobs with.
             instance_type (str): Type of EC2 instance to use for the job, for example,
                 'ml.m5.xlarge'.
@@ -1292,7 +1285,7 @@ def create_monitoring_schedule(
             max_runtime_in_seconds=self.max_runtime_in_seconds,
             environment=normalized_env,
             network_config=network_config_dict,
-            role_arn=self.role,
+            role_arn=self.sagemaker_session.expand_role(self.role),
             tags=self.tags,
         )
 
@@ -1355,7 +1348,7 @@ def update_monitoring_schedule(
                 inter-container traffic, security group IDs, and subnets.
             enable_cloudwatch_metrics (bool): Whether to publish cloudwatch metrics as part of
                 the baselining or monitoring jobs.
-            role (str): An AWS IAM role. The Amazon SageMaker jobs use this role.
+            role (str): An AWS IAM role name or ARN. The Amazon SageMaker jobs use this role.
 
         """
         monitoring_inputs = None
@@ -1449,7 +1442,7 @@ def update_monitoring_schedule(
             max_runtime_in_seconds=max_runtime_in_seconds,
             environment=normalized_env,
             network_config=network_config_dict,
-            role_arn=role,
+            role_arn=self.sagemaker_session.expand_role(self.role),
         )
 
         self._wait_for_schedule_changes_to_apply()

src/sagemaker/processing.py

@@ -52,7 +52,7 @@ def __init__(
         SageMaker processing tasks.
 
         Args:
-            role (str): An AWS IAM role. The Amazon SageMaker training jobs
+            role (str): An AWS IAM role name or ARN. The Amazon SageMaker training jobs
                 and APIs that create Amazon SageMaker endpoints use this role
                 to access training data and model artifacts. After the endpoint
                 is created, the inference code might use the IAM role, if it
@@ -281,7 +281,7 @@ def __init__(
         handles Amazon SageMaker processing tasks for jobs using script mode.
 
         Args:
-            role (str): An AWS IAM role. The Amazon SageMaker training jobs
+            role (str): An AWS IAM role name or ARN. The Amazon SageMaker training jobs
                 and APIs that create Amazon SageMaker endpoints use this role
                 to access training data and model artifacts. After the endpoint
                 is created, the inference code might use the IAM role, if it
@@ -538,7 +538,7 @@ def start_new(cls, processor, inputs, outputs, experiment_config):
         else:
             process_request_args["network_config"] = None
 
-        process_request_args["role_arn"] = processor.role
+        process_request_args["role_arn"] = processor.sagemaker_session.expand_role(processor.role)
 
         process_request_args["tags"] = processor.tags
 

src/sagemaker/sklearn/processing.py

@@ -48,7 +48,7 @@ def __init__(
 
         Args:
             framework_version (str): The version of scikit-learn.
-            role (str): An AWS IAM role. The Amazon SageMaker training jobs
+            role (str): An AWS IAM role name or ARN. The Amazon SageMaker training jobs
                 and APIs that create Amazon SageMaker endpoints use this role
                 to access training data and model artifacts. After the endpoint
                 is created, the inference code might use the IAM role, if it

tests/integ/test_auto_ml.py

@@ -24,7 +24,6 @@
 from tests.integ import DATA_DIR, AUTO_ML_DEFAULT_TIMEMOUT_MINUTES, auto_ml_utils
 from tests.integ.timeout import timeout
 
-DEV_ACCOUNT = 142577830533
 ROLE = "SageMakerRole"
 PREFIX = "sagemaker/beta-automl-xgboost"
 HOSTING_INSTANCE_TYPE = "ml.c4.xlarge"
@@ -41,26 +40,10 @@
 # use a succeeded AutoML job to test describe and list candidates method, otherwise tests will run too long
 AUTO_ML_JOB_NAME = "python-sdk-integ-test-base-job"
 
-EXPECTED_DEFAULT_INPUT_CONFIG = [
-    {
-        "DataSource": {
-            "S3DataSource": {
-                "S3DataType": "S3Prefix",
-                "S3Uri": "s3://sagemaker-us-east-2-{}/{}/input/iris_training.csv".format(
-                    DEV_ACCOUNT, PREFIX
-                ),
-            }
-        },
-        "TargetAttributeName": TARGET_ATTRIBUTE_NAME,
-    }
-]
 EXPECTED_DEFAULT_JOB_CONFIG = {
     "CompletionCriteria": {"MaxCandidates": 3},
     "SecurityConfig": {"EnableInterContainerTrafficEncryption": False},
 }
-EXPECTED_DEFAULT_OUTPUT_CONFIG = {
-    "S3OutputPath": "s3://sagemaker-us-east-2-{}/".format(DEV_ACCOUNT)
-}
 
 
 @pytest.mark.skipif(
@@ -119,7 +102,7 @@ def test_auto_ml_input_object_fit(sagemaker_session):
     reason="AutoML is not supported in the region yet.",
 )
 def test_auto_ml_fit_optional_args(sagemaker_session):
-    output_path = "s3://sagemaker-us-east-2-{}/{}".format(DEV_ACCOUNT, "specified_ouput_path")
+    output_path = "s3://{}/{}".format(sagemaker_session.default_bucket(), "specified_ouput_path")
     problem_type = "MulticlassClassification"
     job_objective = {"MetricName": "Accuracy"}
     auto_ml = AutoML(
@@ -163,8 +146,24 @@ def test_auto_ml_invalid_target_attribute(sagemaker_session):
     reason="AutoML is not supported in the region yet.",
 )
 def test_auto_ml_describe_auto_ml_job(sagemaker_session):
-    auto_ml_utils.create_auto_ml_job_if_not_exist(sagemaker_session)
+    expected_default_input_config = [
+        {
+            "DataSource": {
+                "S3DataSource": {
+                    "S3DataType": "S3Prefix",
+                    "S3Uri": "s3://{}/{}/input/iris_training.csv".format(
+                        sagemaker_session.default_bucket(), PREFIX
+                    ),
+                }
+            },
+            "TargetAttributeName": TARGET_ATTRIBUTE_NAME,
+        }
+    ]
+    expected_default_output_config = {
+        "S3OutputPath": "s3://{}/".format(sagemaker_session.default_bucket())
+    }
 
+    auto_ml_utils.create_auto_ml_job_if_not_exist(sagemaker_session)
     auto_ml = AutoML(
         role=ROLE, target_attribute_name=TARGET_ATTRIBUTE_NAME, sagemaker_session=sagemaker_session
     )
@@ -173,9 +172,9 @@ def test_auto_ml_describe_auto_ml_job(sagemaker_session):
     assert desc["AutoMLJobName"] == AUTO_ML_JOB_NAME
     assert desc["AutoMLJobStatus"] == "Completed"
     assert isinstance(desc["BestCandidate"], dict)
-    assert desc["InputDataConfig"] == EXPECTED_DEFAULT_INPUT_CONFIG
+    assert desc["InputDataConfig"] == expected_default_input_config
     assert desc["AutoMLJobConfig"] == EXPECTED_DEFAULT_JOB_CONFIG
-    assert desc["OutputDataConfig"] == EXPECTED_DEFAULT_OUTPUT_CONFIG
+    assert desc["OutputDataConfig"] == expected_default_output_config
 
 
 @pytest.mark.skipif(

tests/integ/test_data_capture_config.py

@@ -17,23 +17,17 @@
 import sagemaker
 import tests.integ
 import tests.integ.timeout
-
-from tests.integ.retry import retries
-from sagemaker.model_monitor import NetworkConfig
-
-from sagemaker.model_monitor import DataCaptureConfig
+from sagemaker.model_monitor import DataCaptureConfig, NetworkConfig
 from sagemaker.tensorflow.serving import Model
 from sagemaker.utils import unique_name_from_base
+from tests.integ.retry import retries
 
 ROLE = "SageMakerRole"
-REGION = "us-west-2"
 SKLEARN_FRAMEWORK = "scikit-learn"
 
 INSTANCE_COUNT = 1
 INSTANCE_TYPE = "ml.m5.xlarge"
 VOLUME_SIZE_IN_GB = 20
-VOLUME_KMS_KEY = "arn:aws:kms:us-west-2:012345678901:key/volume-kms-key"
-OUTPUT_KMS_KEY = "arn:aws:kms:us-west-2:012345678901:key/output-kms-key"
 MAX_RUNTIME_IN_SECONDS = 2 * 60 * 60
 ENVIRONMENT = {"env_key_1": "env_value_1"}
 TAGS = [{"Key": "tag_key_1", "Value": "tag_value_1"}]

tests/integ/test_model_monitor.py

@@ -42,7 +42,7 @@
 from tests.integ.kms_utils import get_or_create_kms_key
 from tests.integ.retry import retries
 
-ROLE = "arn:aws:iam::142577830533:role/SageMakerRole"
+ROLE = "SageMakerRole"
 INSTANCE_COUNT = 1
 INSTANCE_TYPE = "ml.m5.xlarge"
 VOLUME_SIZE_IN_GB = 40
@@ -63,7 +63,7 @@
 DEFAULT_EXECUTION_MAX_RUNTIME_IN_SECONDS = 3600
 DEFAULT_IMAGE_SUFFIX = ".com/sagemaker-model-monitor-analyzer"
 
-UPDATED_ROLE = "arn:aws:iam::142577830533:role/SageMakerRole"
+UPDATED_ROLE = "SageMakerRole"
 UPDATED_INSTANCE_COUNT = 2
 UPDATED_INSTANCE_TYPE = "ml.m5.2xlarge"
 UPDATED_VOLUME_SIZE_IN_GB = 50
@@ -99,7 +99,7 @@ def predictor(sagemaker_session, tf_full_version):
     ):
         model = Model(
             model_data=model_data,
-            role="SageMakerRole",
+            role=ROLE,
             framework_version=tf_full_version,
             sagemaker_session=sagemaker_session,
         )
@@ -220,9 +220,10 @@ def byoc_monitoring_schedule_name(sagemaker_session, output_kms_key, volume_kms_
 
 @pytest.fixture(scope="module")
 def volume_kms_key(sagemaker_session):
+    role_arn = sagemaker_session.expand_role(ROLE)
     return get_or_create_kms_key(
         sagemaker_session=sagemaker_session,
-        role_arn=ROLE,
+        role_arn=role_arn,
         alias="integ-test-processing-volume-kms-key-{}".format(
             sagemaker_session.boto_session.region_name
         ),
@@ -231,9 +232,10 @@ def volume_kms_key(sagemaker_session):
 
 @pytest.fixture(scope="module")
 def output_kms_key(sagemaker_session):
+    role_arn = sagemaker_session.expand_role(ROLE)
     return get_or_create_kms_key(
         sagemaker_session=sagemaker_session,
-        role_arn=ROLE,
+        role_arn=role_arn,
         alias="integ-test-processing-output-kms-key-{}".format(
             sagemaker_session.boto_session.region_name
         ),
@@ -242,9 +244,10 @@ def output_kms_key(sagemaker_session):
 
 @pytest.fixture(scope="module")
 def updated_volume_kms_key(sagemaker_session):
+    role_arn = sagemaker_session.expand_role(ROLE)
     return get_or_create_kms_key(
         sagemaker_session=sagemaker_session,
-        role_arn=ROLE,
+        role_arn=role_arn,
         alias="integ-test-processing-volume-kms-key-updated-{}".format(
             sagemaker_session.boto_session.region_name
         ),
@@ -253,9 +256,10 @@ def updated_volume_kms_key(sagemaker_session):
 
 @pytest.fixture(scope="module")
 def updated_output_kms_key(sagemaker_session):
+    role_arn = sagemaker_session.expand_role(ROLE)
     return get_or_create_kms_key(
         sagemaker_session=sagemaker_session,
-        role_arn=ROLE,
+        role_arn=role_arn,
         alias="integ-test-processing-output-kms-key-updated-{}".format(
             sagemaker_session.boto_session.region_name
         ),
@@ -315,7 +319,7 @@ def test_default_monitor_suggest_baseline_and_create_monitoring_schedule_with_cu
         == volume_kms_key
     )
     assert DEFAULT_IMAGE_SUFFIX in baselining_job_description["AppSpecification"]["ImageUri"]
-    assert baselining_job_description["RoleArn"] == ROLE
+    assert ROLE in baselining_job_description["RoleArn"]
     assert (
         baselining_job_description["ProcessingInputs"][0]["InputName"] == "baseline_dataset_input"
     )
@@ -492,7 +496,7 @@ def test_default_monitor_suggest_baseline_and_create_monitoring_schedule_without
         is None
     )
     assert DEFAULT_IMAGE_SUFFIX in baselining_job_description["AppSpecification"]["ImageUri"]
-    assert baselining_job_description["RoleArn"] == ROLE
+    assert ROLE in baselining_job_description["RoleArn"]
     assert (
         baselining_job_description["ProcessingInputs"][0]["InputName"] == "baseline_dataset_input"
     )
@@ -1524,7 +1528,7 @@ def test_byoc_monitor_suggest_baseline_and_create_monitoring_schedule_with_custo
         == volume_kms_key
     )
     assert DEFAULT_IMAGE_SUFFIX in baselining_job_description["AppSpecification"]["ImageUri"]
-    assert baselining_job_description["RoleArn"] == ROLE
+    assert ROLE in baselining_job_description["RoleArn"]
     assert baselining_job_description["ProcessingInputs"][0]["InputName"] == "input-1"
     assert (
         baselining_job_description["ProcessingOutputConfig"]["Outputs"][0]["OutputName"]
@@ -1723,7 +1727,7 @@ def test_byoc_monitor_suggest_baseline_and_create_monitoring_schedule_without_cu
         is None
     )
     assert DEFAULT_IMAGE_SUFFIX in baselining_job_description["AppSpecification"]["ImageUri"]
-    assert baselining_job_description["RoleArn"] == ROLE
+    assert ROLE in baselining_job_description["RoleArn"]
     assert baselining_job_description["ProcessingInputs"][0]["InputName"] == "input-1"
     assert (
         baselining_job_description["ProcessingOutputConfig"]["Outputs"][0]["OutputName"]
@@ -2181,7 +2185,7 @@ def test_byoc_monitor_attach_followed_by_baseline_and_update_monitoring_schedule
         == volume_kms_key
     )
     assert DEFAULT_IMAGE_SUFFIX in baselining_job_description["AppSpecification"]["ImageUri"]
-    assert baselining_job_description["RoleArn"] == ROLE
+    assert ROLE in baselining_job_description["RoleArn"]
     assert baselining_job_description["ProcessingInputs"][0]["InputName"] == "input-1"
     assert (
         baselining_job_description["ProcessingOutputConfig"]["Outputs"][0]["OutputName"]

tests/integ/test_monitoring_files.py

@@ -19,13 +19,10 @@
 
 import tests.integ
 import tests.integ.timeout
-
 from sagemaker.model_monitor import Statistics, Constraints, ConstraintViolations
 from sagemaker.s3 import S3Uploader
 from tests.integ.kms_utils import get_or_create_kms_key
 
-ROLE = "arn:aws:iam::142577830533:role/SageMakerRole"
-
 
 @pytest.fixture(scope="module")
 def monitoring_files_kms_key(sagemaker_session):
@@ -40,7 +37,7 @@ def test_statistics_object_creation_from_file_path_with_customizations(
         kms_key=monitoring_files_kms_key,
         sagemaker_session=sagemaker_session,
     )
-    #
+
     assert statistics.file_s3_uri.startswith("s3://")
     assert statistics.file_s3_uri.endswith("statistics.json")