sagemaker-core 1.0.62__py3-none-any.whl → 2.1.1__py3-none-any.whl

sagemaker/core/experiments/_helper.py

@@ -0,0 +1,294 @@
+# 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 helper classes for SageMaker Experiment."""
+from __future__ import absolute_import
+
+import json
+import logging
+import os
+
+import botocore
+
+from sagemaker.core import s3
+from sagemaker.core.experiments._utils import is_already_exist_error
+
+logger = logging.getLogger(__name__)
+
+
+_DEFAULT_ARTIFACT_PREFIX = "trial-component-artifacts"
+_DEFAULT_ARTIFACT_TYPE = "Tracker"
+
+
+class _ArtifactUploader(object):
+    """Artifact uploader"""
+
+    def __init__(
+        self,
+        trial_component_name,
+        sagemaker_session,
+        artifact_bucket=None,
+        artifact_prefix=_DEFAULT_ARTIFACT_PREFIX,
+    ):
+        """Initialize a `_ArtifactUploader` instance.
+
+        Args:
+            trial_component_name (str): The name of the trial component,
+                which is used to generate the S3 path to upload the artifact to.
+            sagemaker_session (sagemaker.core.helper.session_helper.Session): Session object which
+                manages interactions with Amazon SageMaker APIs and any other
+                AWS services needed.
+            artifact_bucket (str): The S3 bucket to upload the artifact to.
+                If not specified, the default bucket defined in `sagemaker_session`
+                will be used.
+            artifact_prefix (str): The S3 key prefix used to generate the S3 path
+                to upload the artifact to (default: "trial-component-artifacts").
+        """
+        self.sagemaker_session = sagemaker_session
+        self.trial_component_name = trial_component_name
+        self.artifact_bucket = artifact_bucket
+        self.artifact_prefix = artifact_prefix
+        self._s3_client = self.sagemaker_session.boto_session.client("s3")
+
+    def upload_artifact(self, file_path, extra_args=None):
+        """Upload an artifact file to S3.
+
+        Args:
+            file_path (str): the file path of the artifact
+            extra_args (dict): Optional extra arguments that may be passed to the upload operation.
+                Similar to ExtraArgs parameter in S3 upload_file function. Please refer to the
+                ExtraArgs parameter documentation here:
+                https://boto3.amazonaws.com/v1/documentation/api/latest/guide/s3-uploading-files.html#the-extraargs-parameter
+
+        Returns:
+            (str, str): The s3 URI of the uploaded file and the etag of the file.
+
+        Raises:
+            ValueError: If file does not exist.
+        """
+        file_path = os.path.expanduser(file_path)
+        if not os.path.isfile(file_path):
+            raise ValueError(
+                "{} does not exist or is not a file. Please supply a file path.".format(file_path)
+            )
+
+        # If self.artifact_bucket is falsy, it will be set to sagemaker_session.default_bucket.
+        # In that case, and if sagemaker_session.default_bucket_prefix exists, self.artifact_prefix
+        # needs to be updated too (because not updating self.artifact_prefix would result in
+        # different behavior the 1st time this method is called vs the 2nd).
+        self.artifact_bucket, self.artifact_prefix = s3.determine_bucket_and_prefix(
+            bucket=self.artifact_bucket,
+            key_prefix=self.artifact_prefix,
+            sagemaker_session=self.sagemaker_session,
+        )
+
+        artifact_name = os.path.basename(file_path)
+        artifact_s3_key = "{}/{}/{}".format(
+            self.artifact_prefix, self.trial_component_name, artifact_name
+        )
+        self._s3_client.upload_file(
+            file_path,
+            self.artifact_bucket,
+            artifact_s3_key,
+            ExtraArgs=extra_args,
+        )
+        etag = self._try_get_etag(artifact_s3_key)
+        return "s3://{}/{}".format(self.artifact_bucket, artifact_s3_key), etag
+
+    def upload_object_artifact(self, artifact_name, artifact_object, file_extension=None):
+        """Upload an artifact object to S3.
+
+        Args:
+            artifact_name (str): the name of the artifact.
+            artifact_object (obj): the object of the artifact
+            file_extension (str): Optional file extension.
+
+        Returns:
+            str: The s3 URI of the uploaded file and the version of the file.
+        """
+
+        # If self.artifact_bucket is falsy, it will be set to sagemaker_session.default_bucket.
+        # In that case, and if sagemaker_session.default_bucket_prefix exists, self.artifact_prefix
+        # needs to be updated too (because not updating self.artifact_prefix would result in
+        # different behavior the 1st time this method is called vs the 2nd).
+        self.artifact_bucket, self.artifact_prefix = s3.determine_bucket_and_prefix(
+            bucket=self.artifact_bucket,
+            key_prefix=self.artifact_prefix,
+            sagemaker_session=self.sagemaker_session,
+        )
+
+        if file_extension:
+            artifact_name = (
+                artifact_name + ("" if file_extension.startswith(".") else ".") + file_extension
+            )
+        artifact_s3_key = "{}/{}/{}".format(
+            self.artifact_prefix, self.trial_component_name, artifact_name
+        )
+        self._s3_client.put_object(
+            Body=json.dumps(artifact_object), Bucket=self.artifact_bucket, Key=artifact_s3_key
+        )
+        etag = self._try_get_etag(artifact_s3_key)
+        return "s3://{}/{}".format(self.artifact_bucket, artifact_s3_key), etag
+
+    def _try_get_etag(self, key):
+        """Get ETag of given key and return None if not allowed
+
+        Args:
+            key (str): The S3 object key.
+
+        Returns:
+            str: The S3 object ETag if it allows, otherwise return None.
+        """
+        try:
+            response = self._s3_client.head_object(Bucket=self.artifact_bucket, Key=key)
+            return response["ETag"]
+        except botocore.exceptions.ClientError as error:
+            # requires read permissions
+            logger.warning("Failed to get ETag of %s due to %s", key, error)
+        return None
+
+
+class _LineageArtifactManager(object):
+    """A helper class to manage Lineage Artifacts"""
+
+    def __init__(
+        self,
+        name,
+        source_uri,
+        etag,
+        source_arn=None,
+        dest_arn=None,
+        artifact_type=_DEFAULT_ARTIFACT_TYPE,
+    ):
+        """Initialize a `_LineageArtifactManager` instance.
+
+        Args:
+            name (str): The name of the Lineage artifact to be created.
+            source_uri (str): The source URI used to create the Lineage artifact.
+            etag (str): The S3 Etag used to create the Lineage artifact.
+            source_arn (str): The source ARN of a trail component to associate
+                this Lineage artifact with (default: None).
+            dest_arn (str): The destination ARN of a trial component to associate
+                this Lineage artifact with (default: None).
+            artifact_type (str): The type of the Lineage artifact (default: "Tracker").
+        """
+        self.name = name
+        self.source_uri = source_uri
+        self.etag = etag
+        self.source_arn = source_arn
+        self.dest_arn = dest_arn
+        self.artifact_arn = None
+        self.artifact_type = artifact_type
+
+    def create_artifact(self, sagemaker_session):
+        """Create the artifact by calling `CreateArtifact` API
+
+        Args:
+            sagemaker_session (sagemaker.core.helper.session_helper.Session): Session object which
+                manages interactions with Amazon SageMaker APIs and any other
+                AWS services needed.
+        """
+        source_ids = []
+        if self.etag:
+            source_ids.append({"SourceIdType": "S3ETag", "Value": self.etag})
+
+        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)
+
+        Args:
+            sagemaker_session (sagemaker.core.helper.session_helper.Session): Session object which
+                manages interactions with Amazon SageMaker APIs and any other
+                AWS services needed.
+        """
+        source_arn = self.source_arn if self.source_arn else self.artifact_arn
+        dest_arn = self.dest_arn if self.dest_arn else self.artifact_arn
+        # 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"
+        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):
+    """Lineage Artifact Tracker"""
+
+    def __init__(self, trial_component_arn, sagemaker_session):
+        """Initialize a `_LineageArtifactTracker` instance.
+
+        Args:
+            trial_component_arn (str): The ARN of the trial component to be
+                associated with the input/output artifacts.
+            sagemaker_session (sagemaker.core.helper.session_helper.Session): Session object which
+                manages interactions with Amazon SageMaker APIs and any other
+                AWS services needed.
+        """
+        self.trial_component_arn = trial_component_arn
+        self.sagemaker_session = sagemaker_session
+        self.artifacts = []
+
+    def add_input_artifact(self, name, source_uri, etag, artifact_type):
+        """Add a Lineage input artifact locally
+
+        Args:
+            name (str): The name of the Lineage input artifact to be added.
+            source_uri (str): The source URI used to create the Lineage input artifact.
+            etag (str): The S3 Etag used to create the Lineage input artifact.
+            artifact_type (str): The type of the Lineage input artifact.
+        """
+        artifact = _LineageArtifactManager(
+            name, source_uri, etag, dest_arn=self.trial_component_arn, artifact_type=artifact_type
+        )
+        self.artifacts.append(artifact)
+
+    def add_output_artifact(self, name, source_uri, etag, artifact_type):
+        """Add a Lineage output artifact locally
+
+        Args:
+            name (str): The name of the Lineage output artifact to be added.
+            source_uri (str): The source URI used to create the Lineage output artifact.
+            etag (str): The S3 Etag used to create the Lineage output artifact.
+            artifact_type (str): The type of the Lineage output artifact.
+        """
+        artifact = _LineageArtifactManager(
+            name, source_uri, etag, source_arn=self.trial_component_arn, artifact_type=artifact_type
+        )
+        self.artifacts.append(artifact)
+
+    def save(self):
+        """Persist any artifact data saved locally"""
+        for artifact in self.artifacts:
+            artifact.create_artifact(self.sagemaker_session)
+            artifact.add_association(self.sagemaker_session)

sagemaker/core/experiments/_metrics.py

@@ -0,0 +1,333 @@
+# 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 classes to manage metrics for Sagemaker Experiment"""
+from __future__ import absolute_import
+
+import datetime
+import logging
+import os
+import time
+import threading
+import queue
+
+import dateutil.tz
+
+from sagemaker.core.helper.session_helper import Session
+
+METRICS_DIR = os.environ.get("SAGEMAKER_METRICS_DIRECTORY", "")
+METRIC_TS_LOWER_BOUND_TO_NOW = 1209600  # on seconds
+METRIC_TS_UPPER_BOUND_FROM_NOW = 7200  # on seconds
+
+BATCH_SIZE = 10
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class _RawMetricData(object):
+    """A Raw Metric Data Object"""
+
+    MetricName = None
+    Value = None
+    Timestamp = None
+    Step = None
+
+    def __init__(self, metric_name, value, timestamp=None, step=None):
+        """Construct a `_RawMetricData` instance.
+
+        Args:
+            metric_name (str): The name of the metric.
+            value (float): The value of the metric.
+            timestamp (datetime.datetime or float or str): Timestamp of the metric.
+                If not specified, the current UTC time will be used.
+            step (int):  Iteration number of the metric (default: None).
+        """
+        if timestamp is None:
+            timestamp = time.time()
+        elif isinstance(timestamp, datetime.datetime):
+            # If the input is a datetime then convert it to UTC time.
+            # Assume a naive datetime is in local timezone
+            if not timestamp.tzinfo:
+                timestamp = timestamp.replace(tzinfo=dateutil.tz.tzlocal())
+            timestamp = (timestamp - timestamp.utcoffset()).replace(tzinfo=datetime.timezone.utc)
+            timestamp = timestamp.timestamp()
+        else:
+            timestamp = float(timestamp)
+
+        if timestamp < (time.time() - METRIC_TS_LOWER_BOUND_TO_NOW) or timestamp > (
+            time.time() + METRIC_TS_UPPER_BOUND_FROM_NOW
+        ):
+            raise ValueError(
+                "Supplied timestamp %f is invalid."
+                " Timestamps must be between two weeks before and two hours from now." % timestamp
+            )
+        value = float(value)
+
+        self.MetricName = metric_name
+        self.Value = float(value)
+        self.Timestamp = timestamp
+        if step is not None:
+            if not isinstance(step, int):
+                raise ValueError("step must be int.")
+            self.Step = step
+
+    def to_record(self):
+        """Convert the `_RawMetricData` object to dict"""
+        return self.__dict__
+
+    def to_raw_metric_data(self):
+        """Converts the metric data to a BatchPutMetrics RawMetricData item"""
+        # Convert timestamp from float to timestamp str.
+        # Otherwise will get ParamValidationError
+        raw_metric_data = {
+            "MetricName": self.MetricName,
+            "Value": self.Value,
+            "Timestamp": str(int(self.Timestamp)),
+        }
+        if self.Step is not None:
+            raw_metric_data["Step"] = int(self.Step)
+        return raw_metric_data
+
+    def __str__(self):
+        """String representation of the `_RawMetricData` object."""
+        return repr(self)
+
+    def __repr__(self):
+        """Return a string representation of this _RawMetricData` object."""
+        return "{}({})".format(
+            type(self).__name__,
+            ",".join(["{}={}".format(k, repr(v)) for k, v in vars(self).items()]),
+        )
+
+
+class _MetricsManager(object):
+    """Collects metrics and sends them directly to SageMaker Metrics data plane APIs."""
+
+    def __init__(self, trial_component_name: str, sagemaker_session: Session, sink=None) -> None:
+        """Initialize a `_MetricsManager` instance
+
+        Args:
+            trial_component_name (str): The Name of the Trial Component to log metrics to
+            sagemaker_session (sagemaker.core.helper.session_helper.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.
+            sink (object): The metrics sink to use.
+        """
+        if sink is None:
+            self.sink = _SyncMetricsSink(
+                trial_component_name, sagemaker_session.sagemaker_metrics_client
+            )
+        else:
+            self.sink = sink
+
+    def log_metric(self, metric_name, value, timestamp=None, step=None):
+        """Sends a metric to metrics service."""
+
+        metric_data = _RawMetricData(metric_name, value, timestamp, step)
+        self.sink.log_metric(metric_data)
+
+    def __enter__(self):
+        """Return self"""
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        """Execute self.close()"""
+        self.sink.close()
+
+    def close(self):
+        """Close the metrics object."""
+        self.sink.close()
+
+
+class _SyncMetricsSink(object):
+    """Collects metrics and sends them directly to metrics service."""
+
+    def __init__(self, trial_component_name, metrics_client) -> None:
+        """Initialize a `_SyncMetricsSink` instance
+
+        Args:
+            trial_component_name (str): The Name of the Trial Component to log metrics.
+            metrics_client (boto3.client): boto client for metrics service
+        """
+        self._trial_component_name = trial_component_name
+        self._metrics_client = metrics_client
+        self._buffer = []
+
+    def log_metric(self, metric_data):
+        """Sends a metric to metrics service."""
+
+        # this is a simplistic solution which calls BatchPutMetrics
+        # on the same thread as the client code
+        self._buffer.append(metric_data)
+        self._drain()
+
+    def _drain(self, close=False):
+        """Pops off all metrics in the buffer and starts sending them to metrics service."""
+
+        if not self._buffer:
+            return
+
+        if len(self._buffer) < BATCH_SIZE and not close:
+            return
+
+        # pop all the available metrics
+        available_metrics, self._buffer = self._buffer, []
+
+        self._send_metrics(available_metrics)
+
+    def _send_metrics(self, metrics):
+        """Calls BatchPutMetrics directly on the metrics service."""
+        while metrics:
+            batch, metrics = (
+                metrics[:BATCH_SIZE],
+                metrics[BATCH_SIZE:],
+            )
+            request = self._construct_batch_put_metrics_request(batch)
+            response = self._metrics_client.batch_put_metrics(**request)
+            errors = response["Errors"] if "Errors" in response else None
+            if errors:
+                error_code = errors[0]["Code"]
+                raise Exception(f'{len(errors)} errors with error code "{error_code}"')
+
+    def _construct_batch_put_metrics_request(self, batch):
+        """Creates dictionary object used as request to metrics service."""
+        return {
+            "TrialComponentName": self._trial_component_name.lower(),
+            "MetricData": list(map(lambda x: x.to_raw_metric_data(), batch)),
+        }
+
+    def close(self):
+        """Drains any remaining metrics."""
+        self._drain(close=True)
+
+
+class _MetricQueue(object):
+    """A thread safe queue for sending metrics to SageMaker.
+
+    Args:
+        trial_component_name (str): the ARN of the resource
+        metric_name (str): the name of the metric
+        metrics_client (boto_client): the boto client for SageMaker Metrics service
+    """
+
+    _CONSUMER_SLEEP_SECONDS = 5
+
+    def __init__(self, trial_component_name, metric_name, metrics_client):
+        # infinite queue size
+        self._queue = queue.Queue()
+        self._buffer = []
+        self._thread = threading.Thread(target=self._run)
+        self._started = False
+        self._finished = False
+        self._trial_component_name = trial_component_name
+        self._metrics_client = metrics_client
+        self._metric_name = metric_name
+        self._logged_metrics = 0
+
+    def log_metric(self, metric_data):
+        """Adds a metric data point to the queue"""
+        self._buffer.append(metric_data)
+
+        if len(self._buffer) < BATCH_SIZE:
+            return
+
+        self._enqueue_all()
+
+        if not self._started:
+            self._thread.start()
+            self._started = True
+
+    def _run(self):
+        """Starts the metric thread which sends metrics to SageMaker in batches"""
+
+        while not self._queue.empty() or not self._finished:
+            if self._queue.empty():
+                time.sleep(self._CONSUMER_SLEEP_SECONDS)
+            else:
+                batch = self._queue.get()
+                self._send_metrics(batch)
+
+    def _send_metrics(self, metrics_batch):
+        """Calls BatchPutMetrics directly on the metrics service."""
+        request = self._construct_batch_put_metrics_request(metrics_batch)
+        self._logged_metrics += len(metrics_batch)
+        self._metrics_client.batch_put_metrics(**request)
+
+    def _construct_batch_put_metrics_request(self, batch):
+        """Creates dictionary object used as request to metrics service."""
+
+        return {
+            "TrialComponentName": self._trial_component_name,
+            "MetricData": list(map(lambda x: x.to_raw_metric_data(), batch)),
+        }
+
+    def _enqueue_all(self):
+        """Enqueue all buffered metrics to be sent to SageMaker"""
+
+        available_metrics, self._buffer = self._buffer, []
+        if available_metrics:
+            self._queue.put(available_metrics)
+
+    def close(self):
+        """Flushes any buffered metrics"""
+
+        self._enqueue_all()
+        self._finished = True
+
+    def is_active(self):
+        """Is the thread active (still draining metrics to SageMaker)"""
+
+        return self._thread.is_alive()
+
+
+class _AsyncMetricsSink(object):
+    """Collects metrics and sends them directly to metrics service."""
+
+    _COMPLETE_SLEEP_SECONDS = 1.0
+
+    def __init__(self, trial_component_name, metrics_client) -> None:
+        """Initialize a `_AsyncMetricsSink` instance
+
+        Args:
+            trial_component_name (str): The Name of the Trial Component to log metrics to.
+            metrics_client (boto3.client): boto client for metrics service
+        """
+        self._trial_component_name = trial_component_name
+        self._metrics_client = metrics_client
+        self._buffer = []
+        self._is_draining = False
+        self._metric_queues = {}
+
+    def log_metric(self, metric_data):
+        """Sends a metric to metrics service."""
+
+        if metric_data.MetricName in self._metric_queues:
+            self._metric_queues[metric_data.MetricName].log_metric(metric_data)
+        else:
+            cur_metric_queue = _MetricQueue(
+                self._trial_component_name, metric_data.MetricName, self._metrics_client
+            )
+            self._metric_queues[metric_data.MetricName] = cur_metric_queue
+            cur_metric_queue.log_metric(metric_data)
+
+    def close(self):
+        """Closes the metric file."""
+        logging.debug("Closing")
+        for q in self._metric_queues.values():
+            q.close()
+
+        # TODO should probably use join
+        while any(map(lambda x: x.is_active(), self._metric_queues.values())):
+            time.sleep(self._COMPLETE_SLEEP_SECONDS)
+        logging.debug("Closed")

sagemaker/core/experiments/_run_context.py

@@ -0,0 +1,58 @@
+# 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 SageMaker Experiment _RunContext class."""
+from __future__ import absolute_import
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from sagemaker.core.experiments import Run
+
+
+class _RunContext:
+    """A static context variable to keep track of the current Run object"""
+
+    _context_run = None
+
+    @classmethod
+    def add_run_object(cls, run: "Run"):
+        """Keep track of the current executing Run object
+
+        by adding it to a class static variable.
+
+        Args:
+            run (Run): The current Run object to be tracked.
+        """
+        cls._context_run = run
+
+    @classmethod
+    def drop_current_run(cls) -> "Run":
+        """Drop the Run object tracked in the global static variable
+
+        as its execution finishes (its "with" block ends).
+
+        Return:
+            Run: the dropped Run object.
+        """
+        current_run = cls._context_run
+        cls._context_run = None
+        return current_run
+
+    @classmethod
+    def get_current_run(cls) -> "Run":
+        """Return the current Run object without dropping it.
+
+        Return:
+            Run: the current Run object to be returned.
+        """
+        return cls._context_run