<commit_before>#!/usr/bin/env python3

# Generate a graphviz dot file from IB topology

def add_lidpair(graph, srclid, targetlid):
    if srclid in graph:
        if targetlid in graph[srclid]:
            graph[srclid][targetlid] += 1
        else:
            graph[srclid][targetlid] = 1
    else:
        graph[srclid] = {targetlid:1}

def parse_ibtopo(topofile):
    desc = {} # dict with {lid:description} mappings
    # Connectivity graph, nested dict
    # {lid:{target_lid:numlinks,...},...}
    graph = {} 
    with open(topofile, 'r') as f:
        sblock = False # Switch block
        hblock = False # Host (Channel Adapter) block
        for line in f:
            if line.startswith('Switch'):
                sblock = True
                i = line.index('lid')
                s = line[i:].split()
                lid = int(s[1])
                i = line.index('#')
                s = line[i:].split('"')
                nodedesc = s[1]
                desc[lid] = {'type':'Switch', 'desc':nodedesc} 
            elif line.startswith('Ca'):
                hblock = True
                i = line.index('#')
                s = line[i:].split('"')
                nodedesc = s[1]
            elif len(line) == 0 or line.isspace():
                sblock = False
                hblock = False
            elif hblock:
                i = line.index('#')
                s = line[i:].split()
                lid = int(s[2])
                desc[lid] = {'type':'Host', 'desc':nodedesc}
                i = line.rindex('lid')
                s = line[i:].split()
                tlid = int(s[1])
                add_lidpair(graph, lid, tlid)
            elif sblock:
                i = line.index('lid')
                s = line[i:].split()
                tlid = int(s[1])
                add_lidpair(graph, lid, tlid)
    return desc, graph

def gen_dot(graph, desc, switches, out):
    out.write('graph ibtopo {\n//For twopi\nranksep=14;\n')
    for k in desc:
        d = desc[k]['desc']
        if desc[k]['type'] == 'Switch':
            c = 'lightblue'
        else:
            c = 'lightgreen'
        if switches:
            if desc[k]['type'] != 'Switch':
                continue
            nh = 0
            for k2 in graph[k]:
                if desc[k2]['type'] == 'Host':
                    nh += 1
            d += '\\nhosts: %d' % nh
        out.write('l%d [label = "%s", style=filled, fillcolor=%s];\n' 
                  % (k, d, c))
    for k in graph:
        if switches and desc[k]['type'] != 'Switch':
            continue
        for k2 in graph[k]:
            if switches and desc[k2]['type'] != 'Switch':
                continue
            if k2 < k:
                continue # Avoid double counting
            numlinks = graph[k][k2]
            if numlinks == 1:
                out.write('l%d -- l%d [weight=10];\n' % (k, k2))
            else:
                out.write('l%d -- l%d [label="%dx",color=red];\n' % (k, k2, numlinks))
    out.write('}\n')

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', 
                      action='store_true',
                      help='Include only switch nodes')
    parser.add_option('-o', '--output', dest='output',
                      help='Output file, if omitted stdout')
    (options, args) = parser.parse_args()
    desc, graph = parse_ibtopo(args[0])
    if options.output:
        out = open(options.output, 'w')
    else:
        out = sys.stdout
    gen_dot(graph, desc, options.switches, out)<commit_msg>Initial work towards generating slurm topology.conf<commit_after>#!/usr/bin/env python3

# Generate a graphviz dot file, or slurm topology.conf from IB topology

def add_lidpair(graph, srclid, targetlid):
    if srclid in graph:
        if targetlid in graph[srclid]:
            graph[srclid][targetlid] += 1
        else:
            graph[srclid][targetlid] = 1
    else:
        graph[srclid] = {targetlid:1}

def parse_ibtopo(topofile):
    desc = {} # dict with {lid:description} mappings
    # Connectivity graph, nested dict
    # {lid:{target_lid:numlinks,...},...}
    graph = {} 
    with open(topofile, 'r') as f:
        sblock = False # Switch block
        hblock = False # Host (Channel Adapter) block
        for line in f:
            if line.startswith('Switch'):
                sblock = True
                i = line.index('lid')
                s = line[i:].split()
                lid = int(s[1])
                i = line.index('#')
                s = line[i:].split('"')
                nodedesc = s[1]
                desc[lid] = {'type':'Switch', 'desc':nodedesc} 
            elif line.startswith('Ca'):
                hblock = True
                i = line.index('#')
                s = line[i:].split('"')
                nodedesc = s[1]
            elif len(line) == 0 or line.isspace():
                sblock = False
                hblock = False
            elif hblock:
                i = line.index('#')
                s = line[i:].split()
                lid = int(s[2])
                desc[lid] = {'type':'Host', 'desc':nodedesc}
                i = line.rindex('lid')
                s = line[i:].split()
                tlid = int(s[1])
                add_lidpair(graph, lid, tlid)
            elif sblock:
                i = line.index('lid')
                s = line[i:].split()
                tlid = int(s[1])
                add_lidpair(graph, lid, tlid)
    return desc, graph

def gen_dot(graph, desc, switches, out):
    out.write('graph ibtopo {\n//For twopi\nranksep=14;\n')
    for k in desc:
        d = desc[k]['desc']
        if desc[k]['type'] == 'Switch':
            c = 'lightblue'
        else:
            c = 'lightgreen'
        if switches:
            if desc[k]['type'] != 'Switch':
                continue
            nh = 0
            for k2 in graph[k]:
                if desc[k2]['type'] == 'Host':
                    nh += 1
            d += '\\nhosts: %d' % nh
        out.write('l%d [label = "%s", style=filled, fillcolor=%s];\n' 
                  % (k, d, c))
    for k in graph:
        if switches and desc[k]['type'] != 'Switch':
            continue
        for k2 in graph[k]:
            if switches and desc[k2]['type'] != 'Switch':
                continue
            if k2 < k:
                continue # Avoid double counting
            numlinks = graph[k][k2]
            if numlinks == 1:
                out.write('l%d -- l%d [weight=10];\n' % (k, k2))
            else:
                out.write('l%d -- l%d [label="%dx",color=red];\n' % (k, k2, numlinks))
    out.write('}\n')

def gen_slurm(graph, desc, out):
    """
    graph: Connectivity graph
    desc: {lid:Description} pairs
    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 k, v in desc.items():
        if v['type'] == 'Switch':
            switches = []
            nodes = []
            for k2, v2 in graph[k].items():
                if desc[k2]['type'] == 'Switch':
                    switches.append("s%d" % k2)
                else:
                    nodename = desc[k2]['desc'].split()[0]
                    nodes.append(nodename)
            switchstring = ""
            if len(switches) > 0:
                switches.sort()
                switchstring = " Switches=" + hostlist.collect_hostlist(switches)
            if len(nodes) > 0:
                nodes.sort()
                nodestr = " Nodes=" + hostlist.collect_hostlist(nodes)
            out.write('SwitchName=s%d%s%s\n' % (k, switchstring, nodestr))


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', 
                      action='store_true',
                      help='Include only switch nodes')
    parser.add_option('-o', '--output', dest='output',
                      help='Output file, if omitted stdout')
    parser.add_option('--slurm', dest='slurm', action='store_true',
                      help='Output in slurm topology.conf format')
    (options, args) = parser.parse_args()
    desc, graph = parse_ibtopo(args[0])
    if options.output:
        out = open(options.output, 'w')
    else:
        out = sys.stdout
    if options.slurm:
        gen_slurm(graph, desc, out)
    else:
        gen_dot(graph, desc, options.switches, out)
{'fraction_non_alphanumeric': 0.14301700984780663, 'fraction_numerical': 0.0070501342882721574, 'mean_word_length': 1.5732795853728765, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7268043', 'n_tokens_mistral': 2881, 'n_tokens_neox': 2732, 'n_words': 678}