text: #!/usr/bin/env python3
  # coding: utf-8
  #
  # $Id: ibmap2slurm.py 3525 2021-11-02 11:06:57Z ltaulell $
  # SPDX-License-Identifier: BSD-2-Clause
  #
  
  #
  # some code from ibtoviz, Copyright (C) 2020 Jan de Cuveland (No license specified)
  #
  
  
  """
  create snippet for slurm topology.conf file from infiniband topology, hostname
  mapping and spines switchs.
  
  TODO: explain all data structures
  
  """
  
  import argparse
  import re
  import sys
  from pprint import pprint
  from ClusterShell.NodeSet import NodeSet
  
  debug = False
  
  
  def get_args():
      """ get arguments from CLI """
      parser = argparse.ArgumentParser(description="""Create snippet for slurm
                                       topology.conf from ibnetdiscover with
                                       files containing hostnames ('H-GUID
                                       hostname' format) and GUID's for spine
                                       switches (only 'S-GUID' format).
                                       """)
      parser.add_argument('-d', '--debug', action='store_true', help='toggle debug on')
      parser.add_argument('topofile', type=str, help='topology file')
      parser.add_argument('-m', '--mapfile', type=str, help='map file', default=None)
      parser.add_argument('-s', '--spinefile', type=str, help='spine file', default=None)
      parser.add_argument('-o', '--output', help='output file, if omitted stdout')
      return parser.parse_args()
  
  
  def read_map(map_file, debug=None):
      """ return a dict {'GUID': 'hostname'} from the map file"""
      d = {}
      with open(map_file, 'r') as f:
          for line in f:
              l = line.strip()
              if l.startswith('#') or len(l) == 0:
                  continue
              (key, val) = line.split()
              d[key] = val
      if debug:
          pprint(d)
      return d
  
  
  def read_spine(spine_file, debug=None):
      """ return a list ['GUID'] of spine switches from the spine file"""
      sl = []
      with open(spine_file, 'r') as f:
          for line in f:
              l = line.strip()
              if l.startswith('#') or len(l) == 0:
                  continue
              sl.append(l)
      if debug:
          pprint(sl)
      return sl
  
  
  def read_topology(topology_file, hostmap=None, spinemap=None, debug=None):
      """ Read topology file as generated by ibnetdiscover
  
      'map' step
  
      return a big dict
      """
      topology = {}
      current_node = ''
      with open(topology_file, mode='r', buffering=1) as f:
          for line in f:
              line = line.strip()
              if line:
                  # Read the name of nodes and the number of ports (Switches or HCAs)
                  m = re.search('^(\w+)\s+(\d+)\s+\"(.+?)\"\s+#\s+\"(.+?)\"', line)
                  if m:
                      current_node = m.groups()[2]
                      topology[current_node] = {}
                      topology[current_node]['number_of_ports'] = int(m.groups()[1])
                      topology[current_node]['label'] = m.groups()[3]
  
                      if m.groups()[0] == 'Switch':
                          topology[current_node]['node_type'] = 'switch'
                          topology[current_node]['label'] = current_node
  
                          if spinemap:
                              # mark if switch is a spine one
                              if current_node in spinemap:
                                  topology[current_node]['switch_type'] = 'spine'
                              else:
                                  topology[current_node]['switch_type'] = 'leaf'
                          # if no spinemap, we still need a switch_type, default to 'leaf'
                          else:
                              topology[current_node]['switch_type'] = 'leaf'
  
                      else:
                          topology[current_node]['node_type'] = 'hca'
  
                          if hostmap:
                              # if hca in mapfile, use hostname
                              if current_node in hostmap.keys():
                                  topology[current_node]['label'] = hostmap[current_node]
                              else:
                                  topology[current_node]['label'] = current_node
                          else:
                              topology[current_node]['label'] = current_node
  
                      topology[current_node]['ports'] = []
  
                  # Read the port lines
                  m = re.search('^\[(\d+)\].*?\"(.+?)\"\[(\d+)\]', line)
                  if m:
                      local_port = int(m.groups()[0])
                      remote = m.groups()[1]
                      remote_port = int(m.groups()[2])
  
                      topology[current_node]['ports'].append(
                          (local_port, remote, remote_port))
      return topology
  
  
  def prepare_slurm_structure(topology, debug=None):
      """ reduce the topology to a table
      SwitchName=SwitchID Nodes=NodeSet(nodes)  # for leaf switches
      SwitchName=SwitchID Switches=NodeSet(switches)  # for spine switches
  
      'reduce' step
  
      return a list
      """
      slurm_structure = []
      for node in topology.keys():
          node_info = topology[node]
          node_type = node_info['node_type']
          label = node_info['label']
          if node_type == "switch":
              if node_info['switch_type'] == 'spine':
                  # if debug:
                  #    pprint(node_info['ports'])
                  switches = NodeSet()
                  for (port1, remote, port1) in node_info['ports']:
                      switches.add(topology[remote]['label'])
  
                  slurm_structure.append('SwitchName={} Switches={}\n'.format(label, switches))
                  if debug:
                      print('SwitchName={} Switches={}'.format(label, switches))
  
              elif node_info['switch_type'] == 'leaf':
                  # if debug:
                  #    pprint(node_info['ports'])
                  nodes = NodeSet()
                  for (port1, host, port2) in node_info['ports']:
                      if topology[host]['node_type'] == 'hca':
                          nodes.add(topology[host]['label'])
  
                  if not nodes:  # if nodes is empty, we probably have a spine switch
                      slurm_structure.append('SwitchName={} Nodes="no nodes, probably a spine, please check"\n'.format(label))
                  else:
                      slurm_structure.append('SwitchName={} Nodes={}\n'.format(label, nodes))
  
                  if debug:
                      print('SwitchName={} Nodes={}'.format(label, nodes))
  
      return slurm_structure
  
  
  def write_slurm_file(slurm_structure, output):
      """ write the topology.conf (in slurm format) to output """
      output.write('# topology.conf snippet\n')
      output.write(''.join([str(elem) for elem in slurm_structure]))
      output.write('#\n')
  
  
  if __name__ == '__main__':
      """ """
      args = get_args()
  
      if args.debug:
          debug = True
          print(args)
  
      if args.mapfile:
          hostmap = read_map(args.mapfile, debug=debug)
      else:
          hostmap = None
  
      if args.spinefile:
          spinemap = read_spine(args.spinefile, debug=debug)
      else:
          spinemap = None
  
      topology = read_topology(args.topofile, hostmap=hostmap, spinemap=spinemap, debug=debug)
      if debug:
          pprint(topology)
  
      slurm_structure = prepare_slurm_structure(topology, debug=debug)
      if debug:
          pprint(slurm_structure)
  
      if args.output:
          with open(args.output, mode='w') as out:
              write_slurm_file(slurm_structure, out)
      else:
          out = sys.stdout
          write_slurm_file(slurm_structure, out)
  