_metrics.py

# 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.session 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.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.
            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")