predict.py

import time
import os
import stomp
import threading
from src.model_predict import predict
from amq_message_python_library import *
from src.influxdb_predictions import InfluxdbPredictionsSender
from timeloop import Timeloop
from datetime import timedelta
import json
import sys
import pandas as pd
import logging
from src.dataset_maker import CSVData


METHOD = os.environ.get("METHOD", "tft")
START_TOPIC = f"start_forecasting.{METHOD}"
STOP_TOPIC = f"stop_forecasting.{METHOD}"
PRED_TOPIC_PREF = f"intermediate_prediction.{METHOD}"
PREDICTION_CYCLE = 1  # minutes
APP_NAME = os.environ.get("APP_NAME", "demo")
AMQ_USER = os.environ.get("AMQ_USER", "admin")
AMQ_PASSWORD = os.environ.get("AMQ_PASSWORD", "admin")


class CustomListener(stomp.ConnectionListener):
    """Custom listener, parameters:
    - conn (stomp connector)
    - topic_name, name of topic to subscribe,
    - start , if start is set to be true recived metrics
    are added to predicted metric, otherwise recived
    metrics are removed from predicted metric (start
    mode corresponds to start_prediction)"""

    def __init__(self, conn, topic_name, start=True):
        self.conn = conn
        self.topic_name = topic_name
        self.start = start

    def on_error(self, frame):
        print('received an error "%s"' % frame.body)

    def on_message(self, frame):
        global predicted_metrics, lock, number_of_forward_predictions, metrics_info
        lock.acquire()
        try:
            metrics = set(json.loads(frame.body)["metrics"])

            if self.start:
                predicted_metrics = predicted_metrics.union(metrics)
                for metric in metrics:
                    frame_body = json.loads(frame.body)
                    logging.debug(frame_body)
                    number_of_forward_predictions[metric] = frame_body[
                        "number_of_forward_predictions"
                    ]
                    self.prediction_cycle = frame_body["prediction_horizon"]
            else:
                predicted_metrics = predicted_metrics.difference(metrics)

        finally:
            lock.release()


def main():
    logging.debug(f"metrics to predict {predicted_metrics}")
    start_conn = morphemic.Connection(AMQ_USER, AMQ_PASSWORD)
    start_conn.connect()

    stop_conn = morphemic.Connection(AMQ_USER, AMQ_PASSWORD)
    stop_conn.connect()

    start_conn.conn.subscribe(f"/topic/{START_TOPIC}", "1", ack="auto")
    stop_conn.conn.subscribe(f"/topic/{STOP_TOPIC}", "2", ack="auto")

    start_conn.set_listener(
        "1", CustomListener(start_conn.conn, START_TOPIC, start=True)
    )
    stop_conn.set_listener("2", CustomListener(stop_conn.conn, STOP_TOPIC, start=False))

    dataset_preprocessor = CSVData(APP_NAME)
    dataset_preprocessor.prepare_csv()
    logging.debug("dataset downloaded")

    influxdb_conn = InfluxdbPredictionsSender(METHOD, APP_NAME)

    logging.debug(
        f"waiting {(msg['epoch_start'] - int(time.time()) * 1000 - prediction_cycle * 1000) // 1000} seconds"
    )

    time.sleep(
        max(0, (msg["epoch_start"] - int(time.time()) * 1000 - prediction_cycle * 1000))
        // 1000
    )

    tl = Timeloop()

    dataset_preprocessor.prepare_csv()
    for metric in predicted_metrics:
        predictions = None
        global time_0
        time_0 = time_0 + prediction_cycle * 1000
        for i in range(number_of_forward_predictions[metric]):
            prediction_msgs, prediction = predict(
                metric,
                (prediction_cycle * 1000) // msg["publish_rate"],
                extra_data=predictions,
                m=i + 1,
                prediction_hor=prediction_horizon,
                timestamp=time_0,
            )
            if predictions is not None:
                predictions = pd.concat([predictions, prediction], ignore_index=True)
            else:
                predictions = prediction

            if prediction_msgs:
                dest = f"{PRED_TOPIC_PREF}.{metric}"
                start_conn.send_to_topic(dest, prediction_msgs[metric])
                influxdb_conn.send_to_influxdb(metric, prediction_msgs)

    @tl.job(interval=timedelta(seconds=prediction_cycle))
    def metric_predict():
        logging.debug("prediction")
        dataset_preprocessor.prepare_csv()
        for metric in predicted_metrics:
            predictions = None
            global time_0
            time_0 = time_0 + prediction_cycle * 1000
            for i in range(number_of_forward_predictions[metric]):
                prediction_msgs, prediction = predict(
                    metric,
                    (prediction_cycle * 1000) // msg["publish_rate"],
                    extra_data=predictions,
                    m=i + 1,
                    prediction_hor=prediction_horizon,
                    timestamp=time_0,
                )
                if predictions is not None:
                    predictions = pd.concat(
                        [predictions, prediction], ignore_index=True
                    )
                else:
                    predictions = prediction

                if prediction_msgs:
                    dest = f"{PRED_TOPIC_PREF}.{metric}"
                    start_conn.send_to_topic(dest, prediction_msgs[metric])
                    influxdb_conn.send_to_influxdb(metric, prediction_msgs)

    tl.start(block=True)

    while True:
        pass


if __name__ == "__main__":
    msg = json.loads(sys.argv[1])
    metrics_info = {
        m["metric"]: {
            "level": m["level"],
            "publish_rate": m["publish_rate"],
        }
        for m in msg["all_metrics"]
    }
    time_0 = msg["epoch_start"]
    prediction_horizon = msg["prediction_horizon"] * 1000
    predicted_metrics = set(msg["metrics"])
    prediction_cycle = msg["prediction_horizon"]

    logging.debug(f"Predicted metrics: {predicted_metrics}")
    number_of_forward_predictions = {
        metric: msg["number_of_forward_predictions"] for metric in predicted_metrics
    }  # deafult number of forward predictions
    lock = threading.Lock()

    main()