change: Update integ tests for Exp Plus M1 changes (aws#741)

Co-authored-by: Dewen Qi <[email protected]>

Author: qidewenwhen

src/sagemaker/experiments/_api_types.py

@@ -13,6 +13,7 @@
 """Contains API objects for SageMaker experiments."""
 from __future__ import absolute_import
 
+import enum
 import numbers
 
 from sagemaker.apiutils import _base_types
@@ -102,6 +103,14 @@ def __init__(self, value=None, media_type=None, **kwargs):
         super(TrialComponentArtifact, self).__init__(value=value, media_type=media_type, **kwargs)
 
 
+class _TrialComponentStatusType(enum.Enum):
+    """The type of trial component status"""
+
+    InProgress = "InProgress"
+    Completed = "Completed"
+    Failed = "Failed"
+
+
 class TrialComponentStatus(_base_types.ApiObject):
     """Status of the trial component.
 

src/sagemaker/experiments/_environment.py

@@ -28,7 +28,7 @@
 logger = logging.getLogger(__name__)
 
 
-class EnvironmentType(enum.Enum):
+class _EnvironmentType(enum.Enum):
     """SageMaker jobs which data can be pulled from the environment."""
 
     SageMakerTrainingJob = 1
@@ -42,7 +42,7 @@ def __init__(self, environment_type, source_arn):
         """Init for _RunEnvironment.
 
         Args:
-            environment_type (EnvironmentType): The environment type.
+            environment_type (_EnvironmentType): The environment type.
             source_arn (str): The ARN of the current job.
         """
         self.environment_type = environment_type
@@ -65,12 +65,13 @@ def load(
         Returns:
             _RunEnvironment: Job data loaded from the environment. None if config does not exist.
         """
+        # TODO: enable to determine transform job env
         if training_job_arn_env in os.environ:
-            environment_type = EnvironmentType.SageMakerTrainingJob
+            environment_type = _EnvironmentType.SageMakerTrainingJob
             source_arn = os.environ.get(training_job_arn_env)
             return _RunEnvironment(environment_type, source_arn)
         if os.path.exists(processing_job_config_path):
-            environment_type = EnvironmentType.SageMakerProcessingJob
+            environment_type = _EnvironmentType.SageMakerProcessingJob
             source_arn = json.loads(open(processing_job_config_path).read())["ProcessingJobArn"]
             return _RunEnvironment(environment_type, source_arn)
         return None

src/sagemaker/experiments/_helper.py

@@ -19,6 +19,8 @@
 
 import botocore
 
+from sagemaker.experiments._utils import is_already_exist_error
+
 logger = logging.getLogger(__name__)
 
 
@@ -171,12 +173,20 @@ def create_artifact(self, sagemaker_session):
         if self.etag:
             source_ids.append({"SourceIdType": "S3ETag", "Value": self.etag})
 
-        response = sagemaker_session.sagemaker_client.create_artifact(
-            ArtifactName=self.name,
-            ArtifactType=self.artifact_type,
-            Source={"SourceUri": self.source_uri, "SourceTypes": source_ids},
-        )
-        self.artifact_arn = response["ArtifactArn"]
+        try:
+            response = sagemaker_session.sagemaker_client.create_artifact(
+                ArtifactName=self.name,
+                ArtifactType=self.artifact_type,
+                Source={"SourceUri": self.source_uri, "SourceTypes": source_ids},
+            )
+            self.artifact_arn = response["ArtifactArn"]
+        except botocore.exceptions.ClientError as err:
+            err_info = err.response["Error"]
+            if not is_already_exist_error(err_info):
+                raise
+            logger.warning(
+                "Skip creating the artifact since it already exists: %s", err_info["Message"]
+            )
 
     def add_association(self, sagemaker_session):
         """Associate the artifact with a source/destination ARN (e.g. trial component arn)
@@ -191,9 +201,17 @@ def add_association(self, sagemaker_session):
         # if the trial component (job) is the source then it produced the artifact,
         # otherwise the artifact contributed to the trial component (job)
         association_edge_type = "Produced" if self.source_arn else "ContributedTo"
-        sagemaker_session.sagemaker_client.add_association(
-            SourceArn=source_arn, DestinationArn=dest_arn, AssociationType=association_edge_type
-        )
+        try:
+            sagemaker_session.sagemaker_client.add_association(
+                SourceArn=source_arn, DestinationArn=dest_arn, AssociationType=association_edge_type
+            )
+        except botocore.exceptions.ClientError as err:
+            err_info = err.response["Error"]
+            if not is_already_exist_error(err_info):
+                raise
+            logger.warning(
+                "Skip associating since the association already exists: %s", err_info["Message"]
+            )
 
 
 class _LineageArtifactTracker(object):
@@ -246,87 +264,3 @@ def save(self):
         for artifact in self.artifacts:
             artifact.create_artifact(self.sagemaker_session)
             artifact.add_association(self.sagemaker_session)
-
-
-class _ArtifactConverter(object):
-    """Converts data to easily consumed by Studio."""
-
-    @classmethod
-    def convert_dict_to_fields(cls, values):
-        """Converts a dictionary to list of field types.
-
-        Args:
-            values (dict): The values of the dictionary.
-
-        Returns:
-            dict: Dictionary of fields.
-        """
-        fields = []
-        for key in values:
-            fields.append({"name": key, "type": "string"})
-        return fields
-
-    @classmethod
-    def convert_data_frame_to_values(cls, data_frame):
-        """Converts a pandas data frame to a dictionary in the table artifact format.
-
-        Args:
-            data_frame (DataFrame): The pandas data frame to convert.
-
-        Returns:
-            dict: dictionary of values in the format needed to log the artifact.
-        """
-        df_dict = data_frame.to_dict()
-        new_df = {}
-        for key in df_dict:
-            col_value = df_dict[key]
-            values = []
-
-            for row_key in col_value:
-                values.append(col_value[row_key])
-
-            new_df[key] = values
-
-        return new_df
-
-    @classmethod
-    def convert_data_frame_to_fields(cls, data_frame):
-        """Converts a dataframe to a dictionary describing the type of fields.
-
-        Args:
-            data_frame(DataFrame): The data frame to convert.
-
-        Returns:
-            dict: Dictionary of fields.
-        """
-        fields = []
-
-        for key in data_frame:
-            col_type = data_frame.dtypes[key]
-            fields.append(
-                {"name": key, "type": _ArtifactConverter.convert_df_type_to_simple_type(col_type)}
-            )
-        return fields
-
-    @classmethod
-    def convert_df_type_to_simple_type(cls, data_frame_type):
-        """Converts a dataframe type to a type for rendering a table in Studio.
-
-        Args:
-            data_frame_type (str): The pandas type.
-
-        Returns:
-            str: The type of the table field.
-        """
-
-        type_pairs = [
-            ("datetime", "datetime"),
-            ("float", "number"),
-            ("int", "number"),
-            ("uint", "number"),
-            ("boolean", "boolean"),
-        ]
-        for pair in type_pairs:
-            if str(data_frame_type).lower().startswith(pair[0]):
-                return pair[1]
-        return "string"

src/sagemaker/experiments/_utils.py

@@ -83,3 +83,13 @@ def wrapper(*args, **kwargs):
         return func(*args, **kwargs)
 
     return wrapper
+
+
+def is_already_exist_error(error):
+    """Check if the error indicates resource already exists
+
+    Args:
+        error (dict): The "Error" field in the response of the
+            `botocore.exceptions.ClientError`
+    """
+    return error["Code"] == "ValidationException" and "already exists" in error["Message"]