graph/builder/heron/builder.py

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

""" This module contains classes and methods for constructing graph
representations of Heron logical and physical plans within the Caladrius Graph
Database."""

import logging

import datetime as dt

from typing import List, Dict, Union, Any

from gremlin_python.process.traversal import P
from gremlin_python.process.graph_traversal import \
    GraphTraversalSource, out, outV, addE, inV
from gremlin_python.structure.graph import Vertex, Edge

from caladrius.common.heron import tracker
from caladrius.graph.gremlin.client import GremlinClient
from caladrius.graph.analysis.heron.routing_probabilities \
    import set_shuffle_routing_probs, set_fields_routing_probs
from caladrius.metrics.heron.client import HeronMetricsClient

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

# pylint: disable = too-many-arguments


def _create_stream_managers(graph_client: GremlinClient, topology_id: str,
                            topology_ref: str, physical_plan: Dict[str, Any]
                            ) -> None:

    LOG.info("Creating stream managers and container vertices")

    counter: int = 0

    for stream_manager in physical_plan["stmgrs"].values():

        # Create the stream manager vertex
        LOG.debug("Creating vertex for stream manager: %s",
                  stream_manager["id"])

        strmg: Vertex = (graph_client.graph_traversal
                         .addV("stream_manager")
                         .property("id", stream_manager["id"])
                         .property("host", stream_manager["host"])
                         .property("port", stream_manager["port"])
                         .property("topology_id", topology_id)
                         .property("topology_ref", topology_ref)
                         .next())

        # Create the stream manager vertex
        container: int = int(stream_manager["id"].split("-")[1])

        LOG.debug("Creating vertex for container: %d", container)

        cont: Vertex = (graph_client.graph_traversal
                        .addV("container")
                        .property("id", container)
                        .property("topology_id", topology_id)
                        .property("topology_ref", topology_ref)
                        .next())

        # Connect the stream manager to the container
        LOG.debug("Connecting stream manager %s to be within container %d",
                  stream_manager["id"], container)
        (graph_client.graph_traversal.V(strmg).addE("is_within").to(cont)
         .next())

        counter += 1

    LOG.info("Created %d container and stream manager vertices", counter)


def _create_spouts(graph_client: GremlinClient, topology_id: str,
                   topology_ref: str,
                   physical_plan: Dict[str, Any],
                   logical_plan: Dict[str, Any]) -> None:

    LOG.info("Creating spout instance vertices")

    # Create the spouts
    physical_spouts: Dict[str, List[str]] = physical_plan["spouts"]

    counter: int = 0

    for spout_name, spout_data in logical_plan["spouts"].items():
        LOG.debug("Creating vertices for instances of spout component: %s",
                  spout_name)
        for instance_name in physical_spouts[spout_name]:

            instance: Dict[str, Union[str, int]] = \
                    tracker.parse_instance_name(instance_name)

            LOG.debug("Creating vertex for spout instance: %s", instance_name)

            stream_manager_id: str = \
                physical_plan["instances"][instance_name]["stmgrId"]

            spout: Vertex = (graph_client.graph_traversal
                             .addV("spout")
                             .property("container", instance["container"])
                             .property("task_id", instance["task_id"])
                             .property("component", spout_name)
                             .property("stream_manager", stream_manager_id)
                             .property("spout_type",
                                       spout_data["spout_type"])
                             .property("spout_source",
                                       spout_data["spout_source"])
                             .property("topology_id", topology_id)
                             .property("topology_ref", topology_ref)
                             .next())

            # Connect the spout to its container vertex
            (graph_client.graph_traversal.V(spout).addE("is_within")
             .to(graph_client.graph_traversal.V()
                 .hasLabel("container")
                 .has("topology_id", topology_id)
                 .has("topology_ref", topology_ref)
                 .has("id", instance["container"])
                 )
             .next())

            counter += 1

    LOG.info("Created %d spout instances", counter)


def _create_bolts(graph_client: GremlinClient, topology_id: str,
                  topology_ref: str,
                  physical_plan: Dict[str, Any],
                  logical_plan: Dict[str, Any]) -> None:

    LOG.info("Creating bolt instance vertices")

    physical_bolts: Dict[str, List[str]] = physical_plan["bolts"]

    counter: int = 0

    for bolt_name in logical_plan["bolts"]:
        LOG.debug("Creating vertices for instances of bolt component: %s",
                  bolt_name)
        for instance_name in physical_bolts[bolt_name]:

            instance: Dict[str, Union[str, int]] = \
                tracker.parse_instance_name(instance_name)

            LOG.debug("Creating vertex for bolt instance: %s",
                      instance_name)

            stream_manager_id: str = \
                physical_plan["instances"][instance_name]["stmgrId"]

            bolt: Vertex = (graph_client.graph_traversal
                            .addV("bolt")
                            .property("container", instance["container"])
                            .property("task_id", instance["task_id"])
                            .property("component", bolt_name)
                            .property("stream_manager", stream_manager_id)
                            .property("topology_id", topology_id)
                            .property("topology_ref", topology_ref)
                            .next())

            # Connect the bolt to its container vertex
            (graph_client.graph_traversal.V(bolt).addE("is_within")
             .to(graph_client.graph_traversal.V()
                 .hasLabel("container")
                 .has("topology_id", topology_id)
                 .has("topology_ref", topology_ref)
                 .has("id", instance["container"])
                 )
             .next())

            counter += 1

    LOG.info("Created %d bolt instances", counter)


def _create_logical_connections(graph_client: GremlinClient, topology_id: str,
                                topology_ref: str,
                                logical_plan: Dict[str, Any]
                                ) -> None:

    # Add all the logical connections between the topology's instances
    LOG.info("Adding logical connections to topology %s instances",
             topology_id)

    topo_traversal: GraphTraversalSource = \
        graph_client.topology_subgraph(topology_id, topology_ref)

    counter: int = 0

    for bolt_name, bolt_data in logical_plan["bolts"].items():

        LOG.debug("Adding logical connections for instances of "
                  "destination bolt: %s", bolt_name)

        # Get a list of all instance vertices for this bolt
        destination_instances: List[Vertex] = (
            topo_traversal.V()
            .has("component", bolt_name)
            .toList())

        for incoming_stream in bolt_data["inputs"]:
            source_instances: List[Vertex] = (
                topo_traversal.V()
                .has("component", incoming_stream["component_name"])
                .toList())

            for destination in destination_instances:
                for source in source_instances:
                    (topo_traversal.V(source)
                     .addE("logically_connected")
                     .property("stream",
                               incoming_stream["stream_name"])
                     .property("grouping", incoming_stream["grouping"])
                     .to(destination).next())

                    counter += 1

    LOG.info("Created %d logical connections", counter)


def _create_physical_connections(graph_client: GremlinClient, topology_id: str,
                                 topology_ref: str) -> None:

    LOG.info("Creating physical connections for topology: %s, reference: "
             "%s", topology_id, topology_ref)

    topo_traversal: GraphTraversalSource = \
        graph_client.topology_subgraph(topology_id, topology_ref)

    # First get all logically connected pairs of vertex and their associated
    # containers and stream managers
    logical_edges: List[Dict[str, Union[Vertex, Edge]]] = (
        topo_traversal.V().hasLabel(P.within("bolt", "spout"))
        .outE("logically_connected")
        .project("source_instance", "source_container",
                 "source_stream_manager", "l_edge", "destination_instance",
                 "destination_container", "destination_stream_manager")
        .by(outV())
        .by(outV().out("is_within"))
        .by(outV().out("is_within").in_("is_within")
            .hasLabel("stream_manager"))
        .by()
        .by(inV())
        .by(inV().out("is_within"))
        .by(inV().out("is_within").in_("is_within")
            .hasLabel("stream_manager"))
        .toList())

    LOG.debug("Processing %d logical connected vertices", len(logical_edges))

    for logical_edge in logical_edges:
        source: Vertex = logical_edge["source_instance"]
        source_container: Vertex = logical_edge["source_container"]
        source_stream_manager: Vertex = logical_edge["source_stream_manager"]
        destination: Vertex = logical_edge["destination_instance"]
        destination_container: Vertex = logical_edge["destination_container"]
        destination_stream_manager: Vertex = \
            logical_edge["destination_stream_manager"]
        l_edge: Edge = logical_edge["l_edge"]

        # Connect the source instance to its stream manager, checking first
        # if the connection already exists
        (graph_client.graph_traversal.V(source)
         .coalesce(out("physically_connected").is_(source_stream_manager),
                   addE("physically_connected").to(source_stream_manager))
         .next())

        if source_container == destination_container:

            # If the source and destination instances are in the same
            # container then they share the same stream manager so just use
            # the source stream manager found above. Connect the source
            # stream manager to the destination instance

            (graph_client.graph_traversal.V(source_stream_manager)
             .coalesce(out("physically_connected").is_(destination),
                       addE("physically_connected").to(destination))
             .next())

            # Set the logical edge for this pair to "local"
            graph_client.graph_traversal.E(l_edge).property("type",
                                                            "local").next()

        else:
            # Connect the two stream managers (if they aren't already)
            (graph_client.graph_traversal.V(source_stream_manager)
             .coalesce(
                 out("physically_connected").is_(destination_stream_manager),
                 addE("physically_connected").to(destination_stream_manager))
             .next())

            (graph_client.graph_traversal.V(destination_stream_manager)
             .coalesce(out("physically_connected").is_(destination),
                       addE("physically_connected").to(destination))
             .next())

            # Set the logical edge for this pair to "remote"
            graph_client.graph_traversal.E(l_edge).property("type",
                                                            "remote").next()


def create_physical_graph(graph_client: GremlinClient,
                          topology_id: str, topology_ref: str,
                          logical_plan: Dict[str, Any],
                          physical_plan: Dict[str, Any]) -> None:
    """ This method will build the physical graph of the specified topology
    in the caladrius graph database. It will attach the supplied reference
    to all vertices of this topology physical graph.

    Arguments:
        graph_client (GremlinClient):   The client instance for the graph
                                        database.
        topology_id (str):  The topology identification string
        topology_ref (str): The unique reference string for this topology
                            physical graph.
        logical_plan (dict):    Dictionary describing the logical plan of the
                                topology. This should match the format of the
                                logical plan returned by the Heron tracker
                                API.
        physical_plan (dict):   Dictionary describing the physical plan of the
                                topology. This should match the format of the
                                physical plan returned by the Heron tracker
                                API.
    Raises:
        RuntimeError:   If the graph database already contains entries with
                        the supplied topology ID and reference.
    """

    if graph_client.topology_ref_exists(topology_id, topology_ref):
        msg: str = (f"A graph of topology {topology_id} with reference "
                    f"{topology_ref} is already present in the graph "
                    f"database.")
        LOG.error(msg)
        raise RuntimeError(msg)

    LOG.info("Building physical graph for topology %s with reference %s",
             topology_id, topology_ref)

    start: dt.datetime = dt.datetime.now()

    _create_stream_managers(graph_client, topology_id, topology_ref,
                            physical_plan)

    _create_spouts(graph_client, topology_id, topology_ref, physical_plan,
                   logical_plan)

    _create_bolts(graph_client, topology_id, topology_ref, physical_plan,
                  logical_plan)

    _create_logical_connections(graph_client, topology_id, topology_ref,
                                logical_plan)

    _create_physical_connections(graph_client, topology_id, topology_ref)

    LOG.info("Physical graph construction completed after %d seconds",
             (dt.datetime.now() - start).total_seconds())


def populate_physical_graph(graph_client: GremlinClient,
                            metrics_client: HeronMetricsClient,
                            topology_id: str, topology_ref: str,
                            start: dt.datetime, end: dt.datetime) -> None:
    """ Populates the specified graph with metrics gathered from the defined
    time period.

    Arguments:
        graph_client (GremlinClient):   The client instance for the graph
                                        database.
        metrics_client (HeronMetricsClient):    The client instance for the
                                                metrics database.
        start (dt.datetime):    UTC datetime instance for the start of the
                                metrics gathering period.
        end (dt.datetime):  UTC datetime instance for the end of the metrics
                            gathering period.
        topology_id (str):  The topology identification string
        topology_ref (str): The unique reference string for this topology
                            physical graph.

    Raises:
        RuntimeError:   If the graph database does not contain a graph with the
                        supplied ID and reference.
    """

    if not graph_client.topology_ref_exists(topology_id, topology_ref):
        msg: str = (f"A graph of topology {topology_id} with reference "
                    f"{topology_ref} is not present in the graph "
                    f"database and therefore cannot be populated")
        LOG.error(msg)
        raise RuntimeError(msg)

    LOG.info("Populating topology %s reference %s physical graph with "
             "with metrics data from a %d second window from %s and %s",
             topology_id, topology_ref, (end - start).total_seconds(),
             start.isoformat(), end.isoformat())

    set_shuffle_routing_probs(graph_client, topology_id, topology_ref)

    set_fields_routing_probs(graph_client, metrics_client, topology_id,
                             topology_ref, start, end)