model/traffic/heron/prophet.py

# Copyright 2018 Twitter, Inc.
# Licensed under the Apache License, Version 2.0
# http://www.apache.org/licenses/LICENSE-2.0

""" This module contains methods and classes for forecasting the traffic from
spout instances in a Heron topology using Facebook's Prophet times series
prediction package: https://facebook.github.io/prophet/"""

import logging

import datetime as dt

from typing import Any, Dict, DefaultDict, Union, cast, List, Optional
from collections import defaultdict

import pandas as pd

from fbprophet import Prophet

from caladrius.common.heron import tracker
from caladrius.metrics.heron.client import HeronMetricsClient
from caladrius.model.traffic.heron.base import HeronTrafficModel
from caladrius.graph.gremlin.client import GremlinClient

LOG: logging.Logger = logging.getLogger(__name__)

INSTANCE_MODELS = DefaultDict[str, DefaultDict[int, Dict[str, Prophet]]]
COMPONENT_MODELS = DefaultDict[str, Dict[str, Prophet]]


def get_spout_emissions(metric_client: HeronMetricsClient, tracker_url: str,
                        topology_id: str, cluster: str, environ: str,
                        start: dt.datetime, end: dt.datetime) -> pd.DataFrame:

    emit_counts: pd.DataFrame = metric_client.get_emit_counts(
            topology_id, cluster, environ, start, end)

    lplan: Dict[str, Any] = tracker.get_logical_plan(tracker_url, cluster,
                                                     environ, topology_id)

    spout_emits: pd.DataFrame = \
        emit_counts[emit_counts.component.isin(lplan["spouts"].keys())]

    return spout_emits


def build_component_models(
        metric_client: HeronMetricsClient, tracker_url: str, topology_id: str,
        cluster: str, environ: str, start: dt.datetime = None, end: dt.datetime = None,
        spout_emits: Optional[pd.DataFrame]=None) -> DefaultDict[str, Dict[str, Prophet]]:

    LOG.info("Creating traffic models for spout components of topology %s",
             topology_id)

    if start and end and spout_emits is None:
        spout_emits = get_spout_emissions(metric_client, tracker_url,
                                          topology_id, cluster, environ, start,
                                          end)
    elif spout_emits is None and ((not start and end) or (start and not end)):
        err: str = ("Either start and end datetime instances or the spout "
                    "emits should be provided")
        LOG.error(err)
        raise RuntimeError(err)

    spout_comp_emits: pd.DataFrame = \
        (spout_emits.groupby(["component", "stream", "timestamp"])
         .mean()["emit_count"].reset_index())

    output: DefaultDict[str, Dict[str, Prophet]] = defaultdict(dict)

    for (spout_comp, stream), data in spout_comp_emits.groupby(["component",
                                                                "stream"]):

        LOG.info("Creating traffic model for spout %s stream %s", spout_comp,
                 stream)

        df: pd.DataFrame = (data[["timestamp", "emit_count"]]
                            .rename(index=str, columns={"timestamp": "ds",
                                                        "emit_count": "y"}))
        model: Prophet = Prophet()
        model.fit(df)

        output[spout_comp][stream] = model

    return output


def predict_per_component(metric_client: HeronMetricsClient, tracker_url: str,
                          topology_id: str, cluster: str, environ: str,
                          start: dt.datetime, end: dt.datetime,
                          future_mins: int) -> pd.DataFrame:

    models: DefaultDict[str, Dict[str, Prophet]] = \
        build_component_models(metric_client, tracker_url, topology_id,
                               cluster, environ, start, end)

    return run_per_component(models, future_mins)


def run_per_component(models: COMPONENT_MODELS, future_mins: int) -> pd.DataFrame:
    output: pd.DataFrame = None

    for spout_comp, stream_models in models.items():
        for stream, model in stream_models.items():

            future: pd.DataFrame = model.make_future_dataframe(
                periods=future_mins, freq='T', include_history=False)

            forecast: pd.DataFrame = model.predict(future)

            forecast["stream"] = stream
            forecast["component"] = spout_comp

            if output is None:
                output = forecast
            else:
                output = output.append(forecast)

    return output


def build_instance_models(
        metric_client: HeronMetricsClient, tracker_url: str, topology_id: str,
        cluster: str, environ: str, start: Optional[dt.datetime] = None,
        end: Optional[dt.datetime] = None,
        spout_emits: Optional[pd.DataFrame] = None) -> INSTANCE_MODELS:

    if start and end and spout_emits is None:
        spout_emits = get_spout_emissions(metric_client, tracker_url,
                                          topology_id, cluster, environ, start,
                                          end)
    elif spout_emits is None and ((not start and end) or (start and not end)):
        err: str = ("Both start and end datetimes should be provided if spout "
                    "emits DataFrame is not")
        LOG.error(err)
        raise RuntimeError(err)
    elif spout_emits is None and not start and not end:
        err = ("Either start and end datetimes or the spout emits should be "
               "provided")
        LOG.error(err)
        raise RuntimeError(err)
    else:
        spout_emits = cast(pd.DataFrame, spout_emits)

    spout_groups: pd.core.groupby.DataFrameGroupBy = \
        (spout_emits[["component", "task", "stream", "timestamp",
                      "emit_count"]]
         .groupby(["component", "task", "stream"]))

    output: INSTANCE_MODELS = defaultdict(lambda: defaultdict(dict))

    for (spout_comp, task, stream), data in spout_groups:
        df: pd.DataFrame = (data[["timestamp", "emit_count"]]
                            .rename(index=str, columns={"timestamp": "ds",
                                                        "emit_count": "y"}))

        model: Prophet = Prophet()
        model.fit(df)

        output[spout_comp][task][stream] = model

    return output


def run_per_instance_models(models: INSTANCE_MODELS,
                            future_mins: int) -> pd.DataFrame:

    output: pd.DataFrame = None

    for spout_comp, task_dict in models.items():
        for task, stream_models in task_dict.items():
            for stream, model in stream_models.items():

                future: pd.DataFrame = model.make_future_dataframe(
                    periods=future_mins, freq='T', include_history=False)

                forecast: pd.DataFrame = model.predict(future)

                forecast["stream"] = stream
                forecast["task"] = task
                forecast["component"] = spout_comp

                if output is None:
                    output = forecast
                else:
                    output = output.append(forecast)

    return output


def predict_per_instance(metric_client: HeronMetricsClient, tracker_url: str,
                         topology_id: str, cluster: str, environ: str,
                         start: dt.datetime, end: dt.datetime,
                         future_mins: int) -> pd.DataFrame:

    models = build_instance_models(metric_client, tracker_url, topology_id,
                                   cluster, environ, start, end)

    return run_per_instance_models(models, future_mins)



class ProphetTrafficModel(HeronTrafficModel):
    """ This model class provides traffic forecasting for Heron topologies
    using the Facebook Prophet time series forecasting package:
    https://facebook.github.io/prophet/
    """

    name: str = "prophet"

    description: str = ("This model uses the Facebook Prophet time series "
                        "prediction package to forecast the spout traffic for "
                        "a given amount of time into the future.")

    def __init__(self, config: dict, metrics_client: HeronMetricsClient,
                 graph_client: GremlinClient) -> None:

        super().__init__(config, metrics_client, graph_client)

        self.metrics_client: HeronMetricsClient

        if "heron.tracker.url" in config:

            self.tracker_url: str = \
                config["heron.tracker.url"]

        else:
            tracker_msg: str = (
                "The Heron Tracker URL was not supplied in the configuration "
                "dictionary given to the Prophet traffic model. Please check "
                "the keys under the 'heron.traffic.model.config' in the main "
                "configuration file")
            LOG.error(tracker_msg)
            raise RuntimeError(tracker_msg)

        if "prophet.model.default.source.hours" in config:

            self.default_source_hours: float = \
                config["prophet.model.default.source.hours"]

        else:
            self.default_source_hours = 24
            LOG.warning("Default source hours were not supplied via "
                        "configuration file. Setting to %d hours.",
                        self.default_source_hours)

        if "prophet.model.default.future.mins" in config:
            self.default_future_minutes: int = \
                config["prophet.model.default.future.mins"]

        else:
            self.default_future_minutes = 30
            LOG.warning("Default future minutes parameter was not supplied "
                        "via configuration file. Setting to %d minutes.",
                        self.default_future_minutes)

        if "prophet.model.default.metrics_sample_period" in config:
            self.default_metrics_sample_period: int = \
                config["prophet.model.default.metrics_sample_period"]

        else:
            self.default_metrics_sample_period = 60
            LOG.warning("Default metrics sample period parameter was not supplied "
                        "via configuration file. Setting to %d minutes.",
                        self.default_metrics_sample_period)


        if "prophet.model.quantiles" in config:
            self.quantiles: List[int] = config["prophet.model.quantiles"]
        else:
            LOG.warning("Quantile values were not set via configuration file "
                        " using: 10, 90, 95, 99 as defaults.")
            self.quantiles = [10, 90, 95, 99]

    def predict_traffic(self, topology_id: str, cluster: str, environ: str,
                        **kwargs: Union[str, int, float]) -> Dict[str, Any]:

        if "source_hours" not in kwargs:
            LOG.warning("source_hours parameter (indicating how many hours of "
                        "historical data to summarise) was not provided, "
                        "using default value of %d hours",
                        self.default_source_hours)
            source_hours: float = self.default_source_hours
        else:
            source_hours = cast(float, float(kwargs["source_hours"]))

        source_end: dt.datetime = \
            dt.datetime.utcnow().replace(tzinfo=dt.timezone.utc)
        source_start: dt.datetime = (source_end -
                                     dt.timedelta(hours=source_hours))

        if "future_mins" not in kwargs:
            LOG.warning("future_mins parameter (indicating how many minutes "
                        "into the future traffic should be predicted was not "
                        "provided, using default value of %d minutes",
                        self.default_future_minutes)
            future_mins: int = self.default_future_minutes
        else:
            future_mins = cast(int, int(kwargs["future_mins"]))

        LOG.info("Predicting traffic over the next %d minutes for topology %s "
                 "using a Prophet model trained on metrics from a %f hour "
                 "period from %s to %s", future_mins, topology_id,
                 (source_end-source_start).total_seconds() / 3600,
                 source_start.isoformat(), source_end.isoformat())

        if "metrics_sample_period" in kwargs:
            time_period_sec: float = \
                cast(float, float(kwargs["metrics_sample_period"]))
        else:
            LOG.warning("metrics_sample_period (indicating the period of time the metrics client retrieves metrics for)"
                        " was not supplied. Using default value of %d seconds.", self.default_metrics_sample_period)
            time_period_sec: int = self.default_metrics_sample_period

        output: Dict[str, Any] = {}

        # Per component predictions

        component_traffic: pd.DataFrame = predict_per_component(
            self.metrics_client, self.tracker_url, topology_id,
            cluster, environ, source_start, source_end, future_mins)

        traffic_by_component: pd.core.groupby.DataFrameGroupBy = \
            component_traffic.groupby(["component", "stream"])

        components: DefaultDict[str, DefaultDict[str, Dict[str, float]]] = \
            defaultdict(lambda: defaultdict(dict))

        for (spout_component, stream), data in traffic_by_component:

            components["mean"][spout_component][stream] = \
                (float(data.yhat.mean()) / time_period_sec)

            components["median"][spout_component][stream] = \
                (float(data.yhat.median()) / time_period_sec)

            components["max"][spout_component][stream] = \
                (float(data.yhat.max()) / time_period_sec)

            components["min"][spout_component][stream] = \
                (float(data.yhat.min()) / time_period_sec)

            for quantile in self.quantiles:
                components[f"{quantile}-quantile"][spout_component][stream] = \
                    (float(data.yhat.quantile(quantile/100)) /
                     time_period_sec)

        output["components"] = components

        # Per instance predictions

        instance_traffic: pd.DataFrame = predict_per_instance(
            self.metrics_client, self.tracker_url, topology_id, cluster,
            environ, source_start, source_end, future_mins)

        traffic_by_task: pd.core.groupby.DataFrameGroupBy = \
            instance_traffic.groupby(["task", "stream"])

        instances: DefaultDict[str, DefaultDict[str, Dict[str, float]]] = \
            defaultdict(lambda: defaultdict(dict))

        for (task_id, stream), data in traffic_by_task:

            instances["mean"][str(task_id)][stream] = \
                (float(data.yhat.mean()) / time_period_sec)

            instances["median"][str(task_id)][stream] = \
                (float(data.yhat.median()) / time_period_sec)

            instances["max"][str(task_id)][stream] = \
                (float(data.yhat.max()) / time_period_sec)

            instances["min"][str(task_id)][stream] = \
                (float(data.yhat.min()) / time_period_sec)

            for quantile in self.quantiles:
                instances[f"{quantile}-quantile"][str(task_id)][stream] = \
                    (float(data.yhat.quantile(quantile/100)) /
                     time_period_sec)

        output["instances"] = instances
        return output