<commit_before>#!/usr/bin/env python3
# -*- mode: python; -*-
# -*- coding: utf-8 -*-
# vim: set fileencoding=utf-8

"""
ibtopotool, a tool to do things with Infiniband topology.
Copyright (C) 2013-2020  <NAME>

  This Source Code Form is subject to the terms of the Mozilla Public
  License, v. 2.0. If a copy of the MPL was not distributed with this
  file, You can obtain one at http://mozilla.org/MPL/2.0/.
"""

import networkx as nx


def speed2weight(speed):
    """Convert an IB speed to an edge weight

    See e.g. https://en.wikipedia.org/wiki/InfiniBand

    """
    ss = speed.split('x')
    nlinks = int(ss[0])
    s1 = ss[1]
    if s1 == 'SDR':
        s = 2
    elif s1 == 'DDR':
        s = 4
    elif s1 == 'QDR':
        s = 8
    elif s1 == 'FDR10':
        s = 10 # ???
    elif s1 == 'FDR':
        s = 13.64
    elif s1 == 'EDR':
        s = 24.24
    elif s1 == 'HDR':
        s = 50 # roughly??
    elif s1 == 'NDR':
        s = 100 # roughly?
    elif s1 == 'XDR':
        s = 250 # the future?
    else:
        raise NotImplementedError('Support for Infiniband speed %s not implemented' % s1)
    return nlinks * s


def parse_ibtopo(topofile, shortlabel):
    """
    Parse an Infiniband topology file as generated by ibnetdiscover.

    Returns: A networkx graph representing the IB network
    """
    g = nx.DiGraph()
    switchidx = 0 # Index switches sequentially
    with open(topofile, 'r') as f:
        inblock = False # Switch or Host (Channel Adapter) block
        for line in f:
            if line.startswith('Switch'):
                inblock = True
                guid = line.split()[2][1:-1]
                i = line.index('#')
                s = line[i:].split('"')
                nodedesc = s[1]
                sid = "s%d" % switchidx
                if shortlabel:
                    label = sid
                else:
                    label = "%s\n%s" % (guid, nodedesc)
                g.add_node(guid, desc=nodedesc, type='Switch', label=label)
                switchidx += 1
            elif line.startswith('Ca'):
                inblock = True
                guid = line.split()[2][1:-1]
                i = line.index('#')
                s = line[i:].split('"')
                nodedesc = s[1].split()[0]
                g.add_node(guid, label=nodedesc, type='Host')
            elif len(line) == 0 or line.isspace():
                inblock = False
            elif inblock:
                ls = line.split()
                destguid = ls[1].split('"')[1]
                w = speed2weight(ls[-1])
                # If the edge already exists, add the weigth to it
                try:
                    g[guid][destguid]['weight'] += w
                    g[guid][destguid]['penwidth'] += 1
                except KeyError:
                    g.add_edge(guid, destguid, weight=w)
                    g[guid][destguid]['penwidth'] = 1
    return g

def gen_dot(graph, out):
    from networkx.drawing.nx_pydot import write_dot
    write_dot(graph, out)

def gen_slurm(g, out):
    """
    g: A networkx graph representing the IB network
    out: Output file-like object
    """
    try:
        import hostlist
    except ImportError:
        print("""To generate a slurm topology.conf, you need to install
python-hostlist, https://pypi.python.org/pypi/python-hostlist""")
        raise
    out.write('# topology.conf generated by ibtopo2dot.py\n')
    for n, nbrs in g.adjacency():
        if g.nodes[n]['type'] == 'Switch':
            switches = []
            nodes = []
            for nbr in nbrs:
                if g.nodes[nbr]['type'] == 'Switch':
                    switches.append(g.nodes[nbr]['label'])
                else:
                    nodename = g.nodes[nbr]['label']
                    nodes.append(nodename)
            switchstring = ""
            if len(switches) > 0:
                switches.sort()
                switchstring = " Switches=" + hostlist.collect_hostlist(switches)
            nodestr = ''
            if len(nodes) > 0:
                nodes.sort()
                nodestr = " Nodes=" + hostlist.collect_hostlist(nodes)
            out.write('SwitchName=%s%s%s\n' % (g.nodes[n]['label'],
                                               switchstring, nodestr))


def treeify(g, rootfile):
    """Generate a DAG with roots given in the file rootfile"""
    roots = []
    with open(rootfile, 'r') as f:
        for line in f:
            l = line.strip()
            if l.startswith('#') or len(l) == 0:
                continue
            ii = l.find('#')
            if ii >= 1:
                l = l[:ii].rstrip()
            roots.append(l)

    for root in roots:
        # Mark the roots with color for graphviz
        g.nodes[root]['fillcolor'] = 'red'
        g.nodes[root]['style'] = 'filled'
        # Mark the roots as roots for graphviz
        g.nodes[root]['root'] = 'true'
        g.nodes[root]['rank'] = 0

    # Calculate distance from roots for all nodes
    for n in g.nodes():
        if n in roots:
            continue
        l = []
        for root in roots:
            l.append(nx.shortest_path_length(g, n, root))
        g.nodes[n]['rank'] = min(l)

    # Drop all edges that go towards the roots, based on the ranks we
    # just computed
    todel = []
    for n, nbrs in g.adjacency():
        for nbr in nbrs:
            if g.nodes[n]['rank'] > g.nodes[nbr]['rank']:
                todel.append((n, nbr))
    g.remove_edges_from(todel)
    return g

def only_switches(g):
    """Filter out nodes that are not switches"""
    return g.subgraph([n for n in g.nodes_iter() if n['type']
                       == 'Switch'])

def relabel_switch_tree(g):
    """If shortlabels and treeify is in effect, relabel switches taking
into account the rank (distance from root(s)) in the tree.

    """
    srl = {} # rank:labelindex dict
    for n in g.nodes():
        if g.nodes[n]['type'] == 'Switch':
            r = g.nodes[n]['rank']
            if not r in srl:
                srl[r] = 0
            g.nodes[n]['label'] = 's%d-%d' % (r, srl[r])
            srl[r] += 1

if __name__ == '__main__':
    from optparse import OptionParser
    import sys
    usage = """%prog [options] ibtopofile

ibtopofile is a file containing the output of 'ibnetdiscover'."""
    parser = OptionParser(usage)
    parser.add_option('-s', '--switches', dest='switches', <commit_msg>proper iterator<commit_after>#!/usr/bin/env python3
# -*- mode: python; -*-
# -*- coding: utf-8 -*-
# vim: set fileencoding=utf-8

"""
ibtopotool, a tool to do things with Infiniband topology.
Copyright (C) 2013-2020  <NAME>

  This Source Code Form is subject to the terms of the Mozilla Public
  License, v. 2.0. If a copy of the MPL was not distributed with this
  file, You can obtain one at http://mozilla.org/MPL/2.0/.
"""

import networkx as nx


def speed2weight(speed):
    """Convert an IB speed to an edge weight

    See e.g. https://en.wikipedia.org/wiki/InfiniBand

    """
    ss = speed.split('x')
    nlinks = int(ss[0])
    s1 = ss[1]
    if s1 == 'SDR':
        s = 2
    elif s1 == 'DDR':
        s = 4
    elif s1 == 'QDR':
        s = 8
    elif s1 == 'FDR10':
        s = 10 # ???
    elif s1 == 'FDR':
        s = 13.64
    elif s1 == 'EDR':
        s = 24.24
    elif s1 == 'HDR':
        s = 50 # roughly??
    elif s1 == 'NDR':
        s = 100 # roughly?
    elif s1 == 'XDR':
        s = 250 # the future?
    else:
        raise NotImplementedError('Support for Infiniband speed %s not implemented' % s1)
    return nlinks * s


def parse_ibtopo(topofile, shortlabel):
    """
    Parse an Infiniband topology file as generated by ibnetdiscover.

    Returns: A networkx graph representing the IB network
    """
    g = nx.DiGraph()
    switchidx = 0 # Index switches sequentially
    with open(topofile, 'r') as f:
        inblock = False # Switch or Host (Channel Adapter) block
        for line in f:
            if line.startswith('Switch'):
                inblock = True
                guid = line.split()[2][1:-1]
                i = line.index('#')
                s = line[i:].split('"')
                nodedesc = s[1]
                sid = "s%d" % switchidx
                if shortlabel:
                    label = sid
                else:
                    label = "%s\n%s" % (guid, nodedesc)
                g.add_node(guid, desc=nodedesc, type='Switch', label=label)
                switchidx += 1
            elif line.startswith('Ca'):
                inblock = True
                guid = line.split()[2][1:-1]
                i = line.index('#')
                s = line[i:].split('"')
                nodedesc = s[1].split()[0]
                g.add_node(guid, label=nodedesc, type='Host')
            elif len(line) == 0 or line.isspace():
                inblock = False
            elif inblock:
                ls = line.split()
                destguid = ls[1].split('"')[1]
                w = speed2weight(ls[-1])
                # If the edge already exists, add the weigth to it
                try:
                    g[guid][destguid]['weight'] += w
                    g[guid][destguid]['penwidth'] += 1
                except KeyError:
                    g.add_edge(guid, destguid, weight=w)
                    g[guid][destguid]['penwidth'] = 1
    return g

def gen_dot(graph, out):
    from networkx.drawing.nx_pydot import write_dot
    write_dot(graph, out)

def gen_slurm(g, out):
    """
    g: A networkx graph representing the IB network
    out: Output file-like object
    """
    try:
        import hostlist
    except ImportError:
        print("""To generate a slurm topology.conf, you need to install
python-hostlist, https://pypi.python.org/pypi/python-hostlist""")
        raise
    out.write('# topology.conf generated by ibtopo2dot.py\n')
    for n, nbrs in g.adjacency():
        if g.nodes[n]['type'] == 'Switch':
            switches = []
            nodes = []
            for nbr in nbrs:
                if g.nodes[nbr]['type'] == 'Switch':
                    switches.append(g.nodes[nbr]['label'])
                else:
                    nodename = g.nodes[nbr]['label']
                    nodes.append(nodename)
            switchstring = ""
            if len(switches) > 0:
                switches.sort()
                switchstring = " Switches=" + hostlist.collect_hostlist(switches)
            nodestr = ''
            if len(nodes) > 0:
                nodes.sort()
                nodestr = " Nodes=" + hostlist.collect_hostlist(nodes)
            out.write('SwitchName=%s%s%s\n' % (g.nodes[n]['label'],
                                               switchstring, nodestr))


def treeify(g, rootfile):
    """Generate a DAG with roots given in the file rootfile"""
    roots = []
    with open(rootfile, 'r') as f:
        for line in f:
            l = line.strip()
            if l.startswith('#') or len(l) == 0:
                continue
            ii = l.find('#')
            if ii >= 1:
                l = l[:ii].rstrip()
            roots.append(l)

    for root in roots:
        # Mark the roots with color for graphviz
        g.nodes[root]['fillcolor'] = 'red'
        g.nodes[root]['style'] = 'filled'
        # Mark the roots as roots for graphviz
        g.nodes[root]['root'] = 'true'
        g.nodes[root]['rank'] = 0

    # Calculate distance from roots for all nodes
    for n in g.nodes():
        if n in roots:
            continue
        l = []
        for root in roots:
            l.append(nx.shortest_path_length(g, n, root))
        g.nodes[n]['rank'] = min(l)

    # Drop all edges that go towards the roots, based on the ranks we
    # just computed
    todel = []
    for n, nbrs in g.adjacency():
        for nbr in nbrs:
            if g.nodes[n]['rank'] > g.nodes[nbr]['rank']:
                todel.append((n, nbr))
    g.remove_edges_from(todel)
    return g

def only_switches(g):
    """Filter out nodes that are not switches"""
    return g.subgraph([n for n in g if n['type']
                       == 'Switch'])

def relabel_switch_tree(g):
    """If shortlabels and treeify is in effect, relabel switches taking
into account the rank (distance from root(s)) in the tree.

    """
    srl = {} # rank:labelindex dict
    for n in g.nodes():
        if g.nodes[n]['type'] == 'Switch':
            r = g.nodes[n]['rank']
            if not r in srl:
                srl[r] = 0
            g.nodes[n]['label'] = 's%d-%d' % (r, srl[r])
            srl[r] += 1

if __name__ == '__main__':
    from optparse import OptionParser
    import sys
    usage = """%prog [options] ibtopofile

ibtopofile is a file containing the output of 'ibnetdiscover'."""
    parser = OptionParser(usage)
    parser.add_option('-s', '--switches', dest='switches', 
{'fraction_non_alphanumeric': 0.1316509102273615, 'fraction_numerical': 0.012188452222829908, 'mean_word_length': 1.6896406808152973, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 13, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5241307', 'n_tokens_mistral': 4145, 'n_tokens_neox': 3915, 'n_words': 1147}