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_powered_machines(self):
        powered_machines = {}
        powered_machines[self.name] = self
        for conn in self.get_outgoing_connections():
            conn_powered_machines = conn.to_plug.machine.get_powered_machines()
            for machine in conn_powered_machines.itervalues():
                if machine not in powered_machines.keys():
                    powered_machines[machine.name] = machine
        return powered_machines

    def get_power_consumption(self, worst_case_scenario=False):
        """
        Returns the number of watts going through this 'machine' (which could just be a powerstrip)

        :param bool worst_case_scenario: if True, computes the number of watts going through this cable in the worst case scenario, in which this cable has to provide the power supply assumung all other backup powers are dead
        """
        if worst_case_scenario:
            # this machine has to provide the full power to all machine it powers, assuming all the power backup providers are dead
            power_consumption = 0.0
            for machine in self.get_powered_machines().itervalues():
                power_consumption += machine.power_consumption
        else:
            power_consumption = self.power_consumption
            for conn in self.get_outgoing_connections():
                power_consumption += conn.get_power_consumption(worst_case_scenario)
                # 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

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, worst_case_scenario=False):
        """
        Returns the number of watts going through this cable

        :param bool worst_case_scenario: if True, computes the number of watts going through this cable in the worst case scenario, in which this cable has to provide the power supply assumung all other backup powers are dead
        """
        to_machine = self.to_plug.machine
        power_consumption = None
        if worst_case_scenario:
            if False: # self.is_redundancy_cable():
                power_consumption = 0.0
            else:
                power_consumption = to_machine.get_power_consumption(worst_case_scenario)
        else:
            num_input_cables = 0
            for input_plug in to_machine.input_plugs.itervalues():
                input_cable = input_plug.get_incoming_connection()
                assert input_cable
                num_input_cables += 1
            assert num_input_cables > 0
            power_consumption = to_machine.get_power_consumption(worst_case_scenario) / num_input_cables
        return 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, worst_case_scenario):
        """
        :param Connection cable:
        """
        raise NotImplementedError

class SimpleColorer(CableColorer):

    def get_cable_color(self, cable, worst_case_scenario):
        """
        :param Connection cable:
        """
        power_consumption = cable.get_power_consumption(worst_case_scenario)
        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, worst_case_scenario):
        """
        :param Connection cable:
        """
        power_consumption = cable.get_power_consumption(worst_case_scenario)
        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, worst_case_scenario=True):
    """
    creates a svg diagram representing the input power configuration
    
    :param PowerConfig power_config: the input power config
    """
    graph = pygraphviz.AGraph(strict=False)  # strict=False allows more than one connection between 2 nodes
    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)

    for machine in power_config.machines.itervalues():
        graph.add_node(machine.name)
        node = graph.get_node(machine.name)
        machine_total_power_consumption = int(machine.get_power_consumption(worst_case_scenario=worst_case_scenario))
        # node.attr['label'] = '%s (%d W)' % (machine.name, machine_total_power_consumption)

        node.attr['shape'] = 'plaintext'

        node.attr['label'] = '<\
            <table border="1" cellborder="0" cellspacing="0">\
                <tr>\
                    <td color="red">%s</td>\
                    <td bgcolor="black"><font color="white">%s W</font></td>\
                </tr>\
            </table>>' % (machine.name, machine_total_power_consumption)


    if False:
        x = 0.0
        for rack in racks.itervalues():
            y = 0.0
            for machine in rack['machines']:
                node = graph.get_node(machine.name)
                node.attr['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(worst_case_scenario=worst_case_scenario)
            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, worst_case_scenario=worst_case_scenario)
                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="%s:%s" % (color, wsc_color), label=label, penwidth="%s:%s" % (penwidth, penwidth))
            graph.add_edge(con.from_plug.machine.name, con.to_plug.machine.name, color=color, label=label, penwidth=penwidth)

    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, style='rounded')
        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)