feature: Add experiment plus Run class (aws#691)

* feature: Add Experiment helper classes (aws#646)

* feature: Add Experiment helper classes

feature: Add helper class _RunEnvironment

* change: Change sleep retry to backoff retry for get TC

* minor fixes in backoff retry

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

* feature: Add helper classes and methods for Run class (aws#660)

* feature: Add helper classes and methods for Run class

* Add Parent class to address comment

* fix docstyle check

* Add arg docstrings in _helper

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

* feature: Add Experiment Run class (aws#651)

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

* change: Add integ tests for Run (aws#673)

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

* Update run log metric to use MetricsManager (aws#678)

* Update run.log_metric to use _MetricsManager

* fix several metrics issues

* Add doc strings to metrics.py

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

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

Author: 3 people

.gitignore

@@ -31,4 +31,7 @@ env/
 **/tmp
 .python-version
 **/_repack_model.py
-**/_repack_script_launcher.sh
+**/_repack_script_launcher.sh
+tests/data/experiment/docker/boto
+tests/data/experiment/docker/sagemaker-dev.tar.gz
+tests/data/**/_repack_model.py

requirements/extras/test_requirements.txt

@@ -20,3 +20,4 @@ requests==2.27.1
 sagemaker-experiments==0.1.35
 Jinja2==3.0.3
 pandas>=1.3.5,<1.5
+scikit-learn==1.0.2

src/sagemaker/apiutils/_base_types.py

@@ -173,8 +173,10 @@ def _search(
                 search_items = search_method_response.get("Results", [])
                 next_token = search_method_response.get(boto_next_token_name)
                 for item in search_items:
-                    if cls.__name__ in item:
-                        yield search_item_factory(item[cls.__name__])
+                    # _TrialComponent class in experiments module is not public currently
+                    class_name = cls.__name__.lstrip("_")
+                    if class_name in item:
+                        yield search_item_factory(item[class_name])
                 if not next_token:
                     break
         except StopIteration:

src/sagemaker/apiutils/_boto_functions.py

@@ -68,7 +68,9 @@ def from_boto(boto_dict, boto_name_to_member_name, member_name_to_type):
             api_type, is_collection = member_name_to_type[member_name]
             if is_collection:
                 if isinstance(boto_value, dict):
-                    member_value = api_type.from_boto(boto_value)
+                    member_value = {
+                        key: api_type.from_boto(value) for key, value in boto_value.items()
+                    }
                 else:
                     member_value = [api_type.from_boto(item) for item in boto_value]
             else:

src/sagemaker/dataset_definition/inputs.py

@@ -124,8 +124,10 @@ class DatasetDefinition(ApiObject):
     """DatasetDefinition input."""
 
     _custom_boto_types = {
-        "redshift_dataset_definition": (RedshiftDatasetDefinition, True),
-        "athena_dataset_definition": (AthenaDatasetDefinition, True),
+        # RedshiftDatasetDefinition and AthenaDatasetDefinition are not collection
+        # Instead they are singleton objects. Thus, set the is_collection flag to False.
+        "redshift_dataset_definition": (RedshiftDatasetDefinition, False),
+        "athena_dataset_definition": (AthenaDatasetDefinition, False),
     }
 
     def __init__(

src/sagemaker/experiments/__init__.py

@@ -0,0 +1,226 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# 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.
+"""Contains API objects for SageMaker experiments."""
+from __future__ import absolute_import
+
+import numbers
+
+from sagemaker.apiutils import _base_types
+
+
+class TrialComponentMetricSummary(_base_types.ApiObject):
+    """Summary model of a trial component.
+
+    Attributes:
+        metric_name (str): The name of the metric.
+        source_arn (str):  The ARN of the source.
+        time_stamp (datetime): Metric last updated value.
+        max (float): The max value of the metric.
+        min (float):  The min value of the metric.
+        last (float):  The last value of the metric.
+        count (float):  The number of samples used to generate the metric.
+        avg (float):  The average value of the metric.
+        std_dev (float):  The standard deviation of the metric.
+    """
+
+    metric_name = None
+    source_arn = None
+    time_stamp = None
+    max = None
+    min = None
+    last = None
+    count = None
+    avg = None
+    std_dev = None
+
+    def __init__(self, metric_name=None, source_arn=None, **kwargs):
+        super(TrialComponentMetricSummary, self).__init__(
+            metric_name=metric_name, source_arn=source_arn, **kwargs
+        )
+
+
+class TrialComponentParameters(_base_types.ApiObject):
+    """A dictionary of TrialComponentParameterValues"""
+
+    @classmethod
+    def from_boto(cls, boto_dict, **kwargs):
+        """Converts a boto dict to a dictionary of TrialComponentParameterValues
+
+        Args:
+            boto_dict (dict): boto response dictionary.
+            **kwargs:  Arbitrary keyword arguments.
+
+        Returns:
+            dict: Dictionary of parameter values.
+        """
+        return_map = {}
+        for key, value in boto_dict.items():
+            return_map[key] = value.get("NumberValue", value.get("StringValue", None))
+        return return_map
+
+    @classmethod
+    def to_boto(cls, parameters):
+        """Converts TrialComponentParameters to dict.
+
+        Args:
+            parameters (TrialComponentParameters): Dictionary to convert.
+
+        Returns:
+            dict: Dictionary of trial component parameters in boto format.
+        """
+        boto_map = {}
+        for key, value in parameters.items():
+            if isinstance(value, numbers.Number):
+                boto_map[key] = {"NumberValue": value}
+            else:
+                boto_map[key] = {"StringValue": str(value)}
+        return boto_map
+
+
+class TrialComponentArtifact(_base_types.ApiObject):
+    """Trial component artifact.
+
+    Attributes:
+        value (str): The artifact value.
+        media_type (str): The media type.
+    """
+
+    value = None
+    media_type = None
+
+    def __init__(self, value=None, media_type=None, **kwargs):
+        super(TrialComponentArtifact, self).__init__(value=value, media_type=media_type, **kwargs)
+
+
+class TrialComponentStatus(_base_types.ApiObject):
+    """Status of the trial component.
+
+    Attributes:
+        primary_status (str): The status of a trial component.
+        message (str): Status message.
+    """
+
+    primary_status = None
+    message = None
+
+    def __init__(self, primary_status=None, message=None, **kwargs):
+        super(TrialComponentStatus, self).__init__(
+            primary_status=primary_status, message=message, **kwargs
+        )
+
+
+class TrialComponentSummary(_base_types.ApiObject):
+    """Summary model of a trial component.
+
+    Attributes:
+        trial_component_name (str): Name of trial component.
+        trial_component_arn (str): ARN of the trial component.
+        display_name (str): Friendly display name in UI.
+        source_arn (str): ARN of the trial component source.
+        status (str): Status.
+        start_time (datetime): Start time.
+        end_time (datetime): End time.
+        creation_time (datetime): Creation time.
+        created_by (str): Created by.
+        last_modified_time (datetime): Date last modified.
+        last_modified_by (datetime): User last modified.
+    """
+
+    _custom_boto_types = {
+        "status": (TrialComponentStatus, False),
+    }
+    trial_component_name = None
+    trial_component_arn = None
+    display_name = None
+    source_arn = None
+    status = None
+    start_time = None
+    end_time = None
+    creation_time = None
+    created_by = None
+    last_modified_time = None
+    last_modified_by = None
+
+
+class TrialComponentSource(_base_types.ApiObject):
+    """Trial Component Source
+
+    Attributes:
+        source_arn (str): The ARN of the source.
+    """
+
+    source_arn = None
+
+    def __init__(self, source_arn=None, **kwargs):
+        super(TrialComponentSource, self).__init__(source_arn=source_arn, **kwargs)
+
+
+class Parent(_base_types.ApiObject):
+    """The trial/experiment/run that a trial component is associated with.
+
+    Attributes:
+        trial_name (str): Name of the trial.
+        experiment_name (str): Name of the experiment.
+        run_name (str): Name of the run.
+    """
+
+    trial_name = None
+    experiment_name = None
+    run_name = None
+
+
+class TrialComponentSearchResult(_base_types.ApiObject):
+    """Summary model of an Trial Component search result.
+
+    Attributes:
+        trial_component_arn (str): ARN of the trial component.
+        trial_component_name (str): Name of the trial component.
+        display_name (str): Display name of the trial component for UI display.
+        source (dict): The source of the trial component.
+        status (dict): The status of the trial component.
+        start_time (datetime): Start time.
+        end_time (datetime): End time.
+        creation_time (datetime): Creation time.
+        created_by (str): Created by.
+        last_modified_time (datetime): Date last modified.
+        last_modified_by (datetime): User last modified.
+        parameters (dict): The hyperparameters of the component.
+        input_artifacts (dict): The input artifacts of the component.
+        output_artifacts (dict): The output artifacts of the component.
+        metrics (list): The metrics for the component.
+        source_detail (dict): The source of the trial component.
+        tags (list): The list of tags that are associated with the trial component.
+        parents (list[Parent]): The parent of trial component.
+    """
+
+    _custom_boto_types = {
+        "parents": (Parent, True),  # parents is a collection (list) of Parent objects
+    }
+    trial_component_arn = None
+    trial_component_name = None
+    display_name = None
+    source = None
+    status = None
+    start_time = None
+    end_time = None
+    creation_time = None
+    created_by = None
+    last_modified_time = None
+    last_modified_by = None
+    parameters = None
+    input_artifacts = None
+    output_artifacts = None
+    metrics = None
+    source_detail = None
+    tags = None
+    parents = None

src/sagemaker/experiments/_api_types.py

@@ -0,0 +1,112 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# 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.
+"""Contains the _RunEnvironment class."""
+from __future__ import absolute_import
+
+import enum
+import json
+import logging
+import os
+
+from sagemaker.experiments import trial_component
+from sagemaker.utils import retry_with_backoff
+
+TRAINING_JOB_ARN_ENV = "TRAINING_JOB_ARN"
+PROCESSING_JOB_CONFIG_PATH = "/opt/ml/config/processingjobconfig.json"
+MAX_RETRY_ATTEMPTS = 7
+
+logger = logging.getLogger(__name__)
+
+
+class EnvironmentType(enum.Enum):
+    """SageMaker jobs which data can be pulled from the environment."""
+
+    SageMakerTrainingJob = 1
+    SageMakerProcessingJob = 2
+
+
+class _RunEnvironment(object):
+    """Retrieves job specific data from the environment."""
+
+    def __init__(self, environment_type, source_arn):
+        """Init for _RunEnvironment.
+
+        Args:
+            environment_type (EnvironmentType): The environment type.
+            source_arn (str): The ARN of the current job.
+        """
+        self.environment_type = environment_type
+        self.source_arn = source_arn
+
+    @classmethod
+    def load(
+        cls,
+        training_job_arn_env=TRAINING_JOB_ARN_ENV,
+        processing_job_config_path=PROCESSING_JOB_CONFIG_PATH,
+    ):
+        """Loads source arn of current job from environment.
+
+        Args:
+            training_job_arn_env (str): The environment key for training job ARN
+                (default: `TRAINING_JOB_ARN`).
+            processing_job_config_path (str): The processing job config path
+                (default: `/opt/ml/config/processingjobconfig.json`).
+
+        Returns:
+            _RunEnvironment: Job data loaded from the environment. None if config does not exist.
+        """
+        if training_job_arn_env in os.environ:
+            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
+            source_arn = json.loads(open(processing_job_config_path).read())["ProcessingJobArn"]
+            return _RunEnvironment(environment_type, source_arn)
+        return None
+
+    def get_trial_component(self, sagemaker_session):
+        """Retrieves the trial component from the job in the environment.
+
+        Args:
+            sagemaker_session (sagemaker.session.Session): Session object which
+                manages interactions with Amazon SageMaker APIs and any other
+                AWS services needed. If not specified, one is created using the
+                default AWS configuration chain.
+
+        Returns:
+            _TrialComponent: The trial component created from the job. None if not found.
+        """
+
+        def _get_trial_component():
+            summaries = list(
+                trial_component._TrialComponent.list(
+                    source_arn=self.source_arn, sagemaker_session=sagemaker_session
+                )
+            )
+            if summaries:
+                summary = summaries[0]
+                return trial_component._TrialComponent.load(
+                    trial_component_name=summary.trial_component_name,
+                    sagemaker_session=sagemaker_session,
+                )
+            return None
+
+        job_tc = None
+        try:
+            job_tc = retry_with_backoff(_get_trial_component, MAX_RETRY_ATTEMPTS)
+        except Exception as ex:  # pylint: disable=broad-except
+            logger.error(
+                "Failed to get trail component in the current environment due to %s", str(ex)
+            )
+        return job_tc