Merge branch 'master' into custom

Author: ahsan-z-khan

CHANGELOG.md

@@ -1,10 +1,23 @@
 # Changelog
 
+## v2.29.1 (2021-03-09)
+
+### Bug Fixes and Other Changes
+
+ * return all failed row indices in feature_group.ingest
+ * move service-role path parsing for AmazonSageMaker-ExecutionRole for get_execution_role() into except block of IAM get_role() call and add warning message
+ * add description parameter for RegisterModelStep
+ * add type annotations for Lineage
+
+### Documentation Changes
+
+ * remove ellipsis from CHANGELOG.md
+
 ## v2.29.0 (2021-03-04)
 
 ### Features
 
- * add support for TensorFlow 2.4.1 for training, inference and…
+ * add support for TensorFlow 2.4.1 for training, inference and data parallel
  * Support profiler config in the pipeline training job step
  * support PyTorch 1.7.1 training, inference and data parallel
 

VERSION

@@ -1 +1 @@
-2.29.1.dev0
+2.29.2.dev0

setup.py

@@ -42,12 +42,11 @@ def read_version():
     "smdebug_rulesconfig==1.0.1",
     "importlib-metadata>=1.4.0",
     "packaging>=20.0",
+    "pandas",
 ]
 
 # Specific use case dependencies
 extras = {
-    "analytics": ["pandas"],
-    "feature-store": ["pandas"],
     "local": [
         "urllib3>=1.21.1,!=1.25,!=1.25.1",
         "docker-compose>=1.25.2",

src/sagemaker/feature_store/feature_group.py

@@ -160,6 +160,7 @@ class IngestionManagerPandas:
     data_frame: DataFrame = attr.ib()
     max_workers: int = attr.ib(default=1)
     _futures: Dict[Any, Any] = attr.ib(init=False, factory=dict)
+    _failed_indices: List[int] = attr.ib(factory=list)
 
     @staticmethod
     def _ingest_single_batch(
@@ -168,7 +169,7 @@ def _ingest_single_batch(
         sagemaker_session: Session,
         start_index: int,
         end_index: int,
-    ):
+    ) -> List[int]:
         """Ingest a single batch of DataFrame rows into FeatureStore.
 
         Args:
@@ -177,19 +178,38 @@ def _ingest_single_batch(
             sagemaker_session (Session): session instance to perform boto calls.
             start_index (int): starting position to ingest in this batch.
             end_index (int): ending position to ingest in this batch.
+
+        Returns:
+            List of row indices that failed to be ingested.
         """
         logger.info("Started ingesting index %d to %d", start_index, end_index)
-        for row in data_frame[start_index:end_index].itertuples(index=False):
+        failed_rows = list()
+        for row in data_frame[start_index:end_index].itertuples():
             record = [
                 FeatureValue(
-                    feature_name=data_frame.columns[index], value_as_string=str(row[index])
+                    feature_name=data_frame.columns[index - 1], value_as_string=str(row[index])
                 )
-                for index in range(len(row))
+                for index in range(1, len(row))
                 if pd.notna(row[index])
             ]
-            sagemaker_session.put_record(
-                feature_group_name=feature_group_name, record=[value.to_dict() for value in record]
-            )
+            try:
+                sagemaker_session.put_record(
+                    feature_group_name=feature_group_name,
+                    record=[value.to_dict() for value in record],
+                )
+            except Exception as e:  # pylint: disable=broad-except
+                logger.error("Failed to ingest row %d: %s", row[0], e)
+                failed_rows.append(row[0])
+        return failed_rows
+
+    @property
+    def failed_rows(self) -> List[int]:
+        """Get rows that failed to ingest
+
+        Returns:
+            List of row indices that failed to be ingested.
+        """
+        return self._failed_indices
 
     def wait(self, timeout=None):
         """Wait for the ingestion process to finish.
@@ -198,18 +218,17 @@ def wait(self, timeout=None):
             timeout (Union[int, float]): ``concurrent.futures.TimeoutError`` will be raised
                 if timeout is reached.
         """
-        failed = False
+        self._failed_indices = list()
         for future in as_completed(self._futures, timeout=timeout):
             start, end = self._futures[future]
-            try:
-                future.result()
-            except Exception as e:  # pylint: disable=broad-except
-                failed = True
-                logger.error("Failed to ingest row %d to %d: %s", start, end, e)
+            result = future.result()
+            if result:
+                logger.error("Failed to ingest row %d to %d", start, end)
             else:
                 logger.info("Successfully ingested row %d to %d", start, end)
+            self._failed_indices += result
 
-        if failed:
+        if len(self._failed_indices) > 0:
             raise RuntimeError(
                 f"Failed to ingest some data into FeatureGroup {self.feature_group_name}"
             )

src/sagemaker/session.py

@@ -3498,14 +3498,6 @@ def get_caller_identity_arn(self):
             endpoint_url=sts_regional_endpoint(self.boto_region_name),
         ).get_caller_identity()["Arn"]
 
-        if "AmazonSageMaker-ExecutionRole" in assumed_role:
-            role = re.sub(
-                r"^(.+)sts::(\d+):assumed-role/(.+?)/.*$",
-                r"\1iam::\2:role/service-role/\3",
-                assumed_role,
-            )
-            return role
-
         role = re.sub(r"^(.+)sts::(\d+):assumed-role/(.+?)/.*$", r"\1iam::\2:role/\3", assumed_role)
 
         # Call IAM to get the role's path
@@ -3518,6 +3510,24 @@ def get_caller_identity_arn(self):
                 role_name,
             )
 
+            # This conditional has been present since the inception of SageMaker
+            # Guessing this conditional's purpose was to handle lack of IAM permissions
+            # https://github.com/aws/sagemaker-python-sdk/issues/2089#issuecomment-791802713
+            if "AmazonSageMaker-ExecutionRole" in assumed_role:
+                LOGGER.warning(
+                    "Assuming role was created in SageMaker AWS console, "
+                    "as the name contains `AmazonSageMaker-ExecutionRole`. "
+                    "Defaulting to Role ARN with service-role in path. "
+                    "If this Role ARN is incorrect, please add "
+                    "IAM read permissions to your role or supply the "
+                    "Role Arn directly."
+                )
+                role = re.sub(
+                    r"^(.+)sts::(\d+):assumed-role/(.+?)/.*$",
+                    r"\1iam::\2:role/service-role/\3",
+                    assumed_role,
+                )
+
         return role
 
     def logs_for_job(  # noqa: C901 - suppress complexity warning for this method

tests/integ/test_feature_store.py

@@ -197,6 +197,7 @@ def test_create_feature_store(
             data_frame=pandas_data_frame, max_workers=3, wait=False
         )
         ingestion_manager.wait()
+        assert 0 == len(ingestion_manager.failed_rows)
 
         # Query the integrated Glue table.
         athena_query = feature_group.athena_query()

tests/unit/sagemaker/feature_store/test_feature_store.py

@@ -274,19 +274,20 @@ def test_ingestion_manager_run_success():
 
 @patch(
     "sagemaker.feature_store.feature_group.IngestionManagerPandas._ingest_single_batch",
-    MagicMock(side_effect=Exception("Failed!")),
+    MagicMock(return_value=[1]),
 )
 def test_ingestion_manager_run_failure():
     df = pd.DataFrame({"float": pd.Series([2.0], dtype="float64")})
     manager = IngestionManagerPandas(
         feature_group_name="MyGroup",
         sagemaker_session=sagemaker_session_mock,
         data_frame=df,
-        max_workers=10,
+        max_workers=1,
     )
     with pytest.raises(RuntimeError) as error:
         manager.run()
     assert "Failed to ingest some data into FeatureGroup MyGroup" in str(error)
+    assert manager.failed_rows == [1]
 
 
 @pytest.fixture

tests/unit/test_session.py

@@ -398,12 +398,30 @@ def test_get_caller_identity_arn_from_a_role(sts_regional_endpoint, boto_session
 @patch("os.path.exists", side_effect=mock_exists(NOTEBOOK_METADATA_FILE, False))
 @patch("sagemaker.session.sts_regional_endpoint", return_value=STS_ENDPOINT)
 def test_get_caller_identity_arn_from_an_execution_role(sts_regional_endpoint, boto_session):
+    sess = Session(boto_session)
+    sts_arn = "arn:aws:sts::369233609183:assumed-role/AmazonSageMaker-ExecutionRole-20171129T072388/SageMaker"
+    sess.boto_session.client("sts", endpoint_url=STS_ENDPOINT).get_caller_identity.return_value = {
+        "Arn": sts_arn
+    }
+    iam_arn = "arn:aws:iam::369233609183:role/AmazonSageMaker-ExecutionRole-20171129T072388"
+    sess.boto_session.client("iam").get_role.return_value = {"Role": {"Arn": iam_arn}}
+
+    actual = sess.get_caller_identity_arn()
+    assert actual == iam_arn
+
+
+@patch("os.path.exists", side_effect=mock_exists(NOTEBOOK_METADATA_FILE, False))
+@patch("sagemaker.session.sts_regional_endpoint", return_value=STS_ENDPOINT)
+def test_get_caller_identity_arn_from_a_sagemaker_execution_role_with_iam_client_error(
+    sts_regional_endpoint, boto_session
+):
     sess = Session(boto_session)
     arn = "arn:aws:sts::369233609183:assumed-role/AmazonSageMaker-ExecutionRole-20171129T072388/SageMaker"
     sess.boto_session.client("sts", endpoint_url=STS_ENDPOINT).get_caller_identity.return_value = {
         "Arn": arn
     }
-    sess.boto_session.client("iam").get_role.return_value = {"Role": {"Arn": arn}}
+
+    sess.boto_session.client("iam").get_role.side_effect = ClientError({}, {})
 
     actual = sess.get_caller_identity_arn()
     assert (