cocluto/PowerDiagram.py

'''
    The goal of this application is to generate a power diagram that will help system administrators to:
    - document the power supply architecture
    - easily spot potential power overloads (for example if the power consumption exceeds the capacity of a cable)
    
    This application takes its input from a database, currently in the form of an sql dump, but it could easily be adapted to read directly from a mysql database
'''

import re
import pygraphviz  # port install py-pygraphviz
from inventory import Inventory, MachineSpecIdNotFound
from SimpaDbUtil import SqlFile, SqlDatabaseReader
from Lib import SimpaDbUtil


def add_capacity_constraints(capacity1, capacity2):
    """
    combines 2 capacity constraints (max amperes) together
    
    :param float capacity1: max amperes for the first capacity, None if there are no constraints
    :param float capacity2: max amperes for the second capacity, None if there are no constraints
    :return float: max amperes for the combined capacity, None if there are no constraints
    """
    if capacity1 is None:
        return capacity2
    else:
        if capacity2 is None:
            return capacity1
        else:
            return min(capacity1, capacity2)


class Machine(object):
    """
    represents a device with input and output power plugs. It could represent a power consumer such as a server (in which case, it has no output plug), or a power distrubtion unit (in which case, it has one input plug and more than one output plugs), or even something else...
    """
    def __init__(self, name, power_config):
        self.name = name
        self.input_plugs = {}
        self.output_plugs = {}
        self.current_capacity_constraint = None  # the maximum amperes in this connection
        self.power_config = power_config
        self.power_consumption = 0.0
        self.rack_id = None

    def get_plug(self, plug_name):
        plugs = {'i': self.input_plugs, 'o': self.output_plugs}[plug_name[0]]
        if plug_name not in plugs:
            plugs[plug_name] = Plug(plug_name, self, self.power_config)
        return plugs[plug_name]

    def get_max_amperes(self):
        capacity = None
        if len(self.input_plugs) > 0:
            capacity = self.input_plugs.values()[0].get_max_amperes()
            capacity = add_capacity_constraints(capacity, self.current_capacity_constraint)
        return capacity

    def get_outgoing_connections(self):
        outgoing_connections = []
        for conn in self.power_config.connections:
            if conn.from_plug.machine == self:
                outgoing_connections.append(conn)
        return outgoing_connections

    def get_power_consumption(self):
        power_consumption = self.power_consumption
        for conn in self.get_outgoing_connections():
            power_consumption += conn.get_power_consumption()
            # print("machine %s : power_consumption += %f" % (self.name, conn.get_power_consumption()))
        return power_consumption

    def is_power_provider(self):
        return self.name == 'edf' or re.match('^ups[0-9]*$', self.name)



class Plug(object):
    """
    represents a power plug (input or output) of a device
    """
    def __init__(self, name, machine, power_config):
        self.name = name
        self.machine = machine
        self.current_capacity_constraint = None  # the maximum amperes in this connection
        self.power_config = power_config

    def __str__(self):
        return self.machine.name + '.' + self.name

    def get_incoming_connection(self):
        return self.power_config.get_connection_to(self)

#     def get_outgoing_connections(self):
#         return self.power_config.get_connections_from(self)

    def is_input_plug(self):
        return self.name[0] == 'i'

    def set_current_capacity_constraint(self, max_amps):
        self.current_capacity_constraint = max_amps

    def get_max_amperes(self):
        capacity = None

        debug = False
        if self.is_input_plug():
            in_con = self.get_incoming_connection()
            if in_con:
                # apply incoming connection amperes limitation
                capacity = add_capacity_constraints(capacity, in_con.get_max_amperes())
                if debug:
                    print(str(self)+ 'after incoming connection amperes limitation, capacity = ' + str(capacity))
        else:
            # apply the machine containing this plug's amperes limitation
            capacity = add_capacity_constraints(capacity, self.machine.get_max_amperes())
            if debug:
                print(str(self)+'apply the machine containing this plug s amperes limitation, capacity = ' + str(capacity))
                
        # apply this plug's amperes limitation
        capacity = add_capacity_constraints(capacity, self.current_capacity_constraint)
        if debug:
            print(str(self)+'after apply this plug s amperes limitation, capacity = ' + str(capacity), self.current_capacity_constraint)

        return capacity

#     def get_power_consumption(self):
#         power_consumption = 0.0
#         for conn in self.get_outgoing_connections():
#             power_consumption += conn.get_power_consumption()
#         power_consumption += self.get_power_consumption()
#         return self.from_plug.get_power_consumption()


class Connection(object):
    """
    a power cable connecting an input power plug to an output power plug
    """

    def __init__(self, from_plug, to_plug):
        self.from_plug = from_plug
        self.to_plug = to_plug
        self.current_capacity_constraint = None  # the maximum amperes in this connection

    def __str__(self):
        return str(self.from_plug) + ' -> ' + str(self.to_plug) + ' (' + str(self.from_plug.get_max_amperes()) + ' A, ' + str(self.get_max_amperes()) + ' A)'
    
    def get_max_amperes(self):
        # gLogger.debug('%s (%s A) -> %s (%s A): ' % (str(self.from_plug), str(self.from_plug.get_max_amperes()), str(self.to_plug), str(self.to_plug.current_capacity_constraint)))
        capacity = self.from_plug.get_max_amperes()
        capacity = add_capacity_constraints(capacity, self.to_plug.current_capacity_constraint)
        if self.current_capacity_constraint is not None:
            capacity = min(capacity, self.current_capacity_constraint)
        return capacity

    def is_redundancy_cable(self):
        to_machine = self.to_plug.machine
        my_power_provider = self.get_power_provider()
        # find the first sibling cable that has the same provider as self
        first_cable_with_same_provider = None
        for input_plug in to_machine.input_plugs.itervalues():
            sibling_cable = input_plug.get_incoming_connection()
            if sibling_cable.get_power_provider() == my_power_provider:
                first_cable_with_same_provider = sibling_cable
        if first_cable_with_same_provider is None:
            # no other connection with the same provider
            return False
        if first_cable_with_same_provider == self:
            # for each provider, the 1st cable amongst the connectors using this provider is considered to be the original (not redundant)
            return False
        else:
            # for each provider, all cable but the 1st one are considered as redundant
            return True

    def get_power_provider(self):
        from_machine = self.from_plug.machine
        if from_machine.is_power_provider():
            return from_machine
        input_plug_names = from_machine.input_plugs.keys()
        assert len(input_plug_names) == 1, "from_machine is supposed to be a power strip (which is expected to only have one input)"
        input_plug = from_machine.input_plugs[input_plug_names[0]]
        return input_plug.get_incoming_connection().get_power_provider()

    def get_power_consumption(self):
        # at the moment, this program doesn't handle redundant power supplies properly:
        # the energy dragged by a power supply depends whether both power supplies are connnected to the same provider or not
        if self.is_redundancy_cable():
            return 0.0  # consider secondary power supplies to drag 0 power
        else:
            return self.to_plug.machine.get_power_consumption()
            

class PowerConfig(object):
    """
    the description of how machines are connected together (in terms of electric power)
    """

    def __init__(self, simpa_db_sql_file_path):
        self.machines = {}
        self.connections = []
        
        sql_source = SqlFile(simpa_db_sql_file_path)
        sql_reader = SqlDatabaseReader(sql_source)
        inventory = Inventory(sql_reader)
        self._parse_from_inventory(inventory)
        
    def _parse_from_inventory(self, inventory):
        """
        :param Inventory inventory:
        """
        
        rows = inventory.query("SELECT * FROM machine_to_power")

        for row in rows:
            (to_plug_as_str, from_plug_as_str, powercordid) = row
            if to_plug_as_str != '':
                conn = self._add_connection(from_plug_as_str, to_plug_as_str)
                for plug in (conn.from_plug, conn.to_plug):
                    plug_capacity = inventory.read_plug_capacity(plug)
                    if plug_capacity is not None:
                        plug.set_current_capacity_constraint(plug_capacity)

        rows = inventory.query("SELECT * FROM power_output_specs")

        for row in rows:
            # print row
            # ('ups1_o1', 16.0),
            (to_plug_as_str, max_amps_as_str) = row
            to_plug = self._get_plug(to_plug_as_str)
            to_plug.set_current_capacity_constraint(float(max_amps_as_str))
            
        for machine in self.machines.values():

            machine_name = machine.name

            # find its rack location
            try:
                rack_id, rack_slot_index = inventory.get_machine_rack_location(machine_name)
                machine.rack_id = rack_id
                machine.rack_slot_index = rack_slot_index
            except SimpaDbUtil.TableAttrNotFound:
                pass

            if re.match('[a-z]+.._..', machine_name):
                # machine_name is a group of machines such as physix80_83
                # in this case, we use a hack : the type and power consumption is based on the first machine of this group (in this example physix80)
                machine_name = '_'.join(machine_name.split('_')[0:-1])

            # print(machine_name)

            machine_spec_id = None
            try:
                # assert machine_spec_id != '', 'non-empty value expected for machine_spec_id for machine %s' % machine_name
                machine_spec_id = inventory.machine_name_to_machine_spec_id(machine_name)
                if machine_spec_id == '':
                    machine_spec_id = None  # some simple 'machines' such as powerext003 have undefined machine_spec_id
            except MachineSpecIdNotFound:
                pass
                        
            if machine_spec_id is not None:
                power_consumption = inventory.machine_spec_id_to_power_consumption(machine_spec_id)
                if power_consumption is not None:
                    machine.power_consumption = power_consumption


    def get_connection_to(self, to_plug):
        for connection in self.connections:
            if connection.to_plug == to_plug:
                return connection
        return None

    def _get_machine(self, machine_name):
        if machine_name not in self.machines:
            self.machines[machine_name] = Machine(machine_name, self)
        return self.machines[machine_name]
        
    def _get_plug(self, plug_as_str):
        elements = plug_as_str.split('_')
        plug_name = elements[-1]
        machine_name = plug_as_str[0:-(len(plug_name) + 1)]
        machine = self._get_machine(machine_name)
        return machine.get_plug(plug_name)
        
    def _add_connection(self, from_plug_as_str, to_plug_as_str):
        from_plug = self._get_plug(from_plug_as_str)
        to_plug = self._get_plug(to_plug_as_str)
        conn = Connection(from_plug, to_plug)
        self.connections.append(conn)
        return conn
        
    def __str__(self):
        s = ''
        for c in self.connections:
            s += str(c) + '\n'
        return s

class CableColorer(object):

    def get_cable_color(self, cable):
        """
        :param Connection cable:
        """
        raise NotImplementedError

class SimpleColorer(CableColorer):

    def get_cable_color(self, cable):
        """
        :param Connection cable:
        """
        power_consumption = cable.get_power_consumption()
        amperes = power_consumption / 220.0
        capacity = cable.get_max_amperes()
        saturation = amperes / capacity
        if saturation > 1.0:
            color = '/svg/red'
        elif saturation > 0.75:
            color = '/svg/orange'
        elif saturation < 0.001:
            color = '/svg/black'  # probably an error
        else:
            color = '/svg/green'
        return color

class RampColorer(CableColorer):

    @staticmethod
    def hotness_to_hsv_color(hotness):
        """
        :param float hotness: temperature of the wire ratio (0.0 : cold -> 1.0 : hot)
        """
        clamped_hotness = max(min(hotness, 1.0), 0.0)
        return "%f, 1.0, 0.8" % ((clamped_hotness) * 0.1 + 0.23)

    def get_cable_color(self, cable):
        """
        :param Connection cable:
        """
        power_consumption = cable.get_power_consumption()
        amperes = power_consumption / 220.0
        capacity = cable.get_max_amperes()
        saturation = amperes / capacity
        color = None
        # print(cable.from_plug.machine.name, cable.to_plug.machine.name, cable.get_power_provider().name)
        power_is_backed_up = cable.get_power_provider().name == 'ups3'
        if power_is_backed_up:
            # greenish colors
            color = RampColorer.hotness_to_hsv_color(pow(saturation, 4.0))
        else:
            # blueish colors using the pretty ylgnbu9 color palette
            clamped_saturation = max(min(saturation, 1.0), 0.0)
            color = '/ylgnbu9/%d' % (int(clamped_saturation * 5) + 4)

        if saturation < 0.001:
            color = '/svg/black'  # probably an error
        elif saturation > 1.0:
            color = '/svg/red'  # draw overloaded cables in red
        return color

def power_config_to_svg(power_config, svg_file_path):

    """
    creates a svg diagram representing the input power configuration
    
    :param PowerConfig power_config: the input power config
    """
    graph = pygraphviz.AGraph()
    graph.graph_attr['overlap'] = 'false'
    graph.graph_attr['splines'] = 'true'
    graph.graph_attr['rankdir'] = 'LR'  # to get hrizontal tree rather than vertical
    
    graph.edge_attr['colorscheme'] = 'rdylgn9'  # 'brbg11'
    graph.node_attr['shape'] = 'box'
    graph.node_attr['height'] = 0.3  # default 0.5 inches
    graph.node_attr['fontname'] = 'Helvetica'  # default : Times-Roman
    graph.edge_attr['fontsize'] = 10  # default : 14 pt
    graph.edge_attr['len'] = 1.5  # default : 1.0
    

    # graph.add_node('a')
    # graph.add_node('b')
    # graph.add_edge(u'a',u'b',color='blue')
    # graph.add_edge(u'b',u'a',color='blue')
    racks = {}

    for con in power_config.connections:
        for machine in [con.from_plug.machine, con.to_plug.machine]:
            rack_id = machine.rack_id
            if rack_id is not None:
                if rack_id not in racks.keys():
                    rack = {}
                    rack['name'] = rack_id
                    rack['machines'] = []
                    racks[rack_id] = rack
                rack = racks[rack_id]
                if machine.name not in rack['machines']:
                    rack['machines'].append(machine)

    if False:
        x = 0.0
        for rack in racks.itervalues():
            y = 0.0
            for machine in rack['machines']:
                graph.add_node(machine.name, pos='%f,%f!' % (x, y))  # https://observablehq.com/@magjac/placing-graphviz-nodes-in-fixed-positions
                # print(machine.name, x, y)
                y += 1.0
            x += 1.0

    cable_colorer = RampColorer()

    for con in power_config.connections:
        # print(con.from_plug.machine.name, con.to_plug.machine.name)
        if not con.is_redundancy_cable():  # don't display redundancy cables, as they might overlap and hide the main one
            power_consumption = con.get_power_consumption()
            amperes = power_consumption / 220.0
            color = '/svg/green'
            capacity = con.get_max_amperes()
            penwidth_scaler = 0.25
            if capacity is None:
                max_amp = '? A'
                color = '/svg/red'
                penwidth = 100.0 * penwidth_scaler  # make the problem clearly visible
            else:
                max_amp = str(capacity) + 'A'
                color = cable_colorer.get_cable_color(con)
                penwidth = capacity * penwidth_scaler
            label = "%.1f/%s" % (amperes, max_amp)
            # color='//%d' % int(9.0-amperes/capacity*8)
            
            graph.add_edge(con.from_plug.machine.name, con.to_plug.machine.name, color=color, label=label, penwidth=penwidth)

    if True:
        for rack_id, rack in racks.iteritems():
            # sub = graph.add_subgraph(rack, name='cluster_%s' % rack_id, rank='same')
            machine_names = list(machine.name for machine in rack['machines'])
            sub = graph.add_subgraph(machine_names, name='cluster_%s' % rack_id)
            sub.graph_attr['label'] = rack_id
            # sub.graph_attr['rank']='same'
    # assert False
    #graph.layout(prog='twopi')
    with open('./toto.dot', 'w') as f:
        f.write(graph.string())


    graph.layout(prog='dot')
    graph.draw(svg_file_path)