text: #!/usr/bin/env python3
  
  # Copyright (c) Microsoft Corporation
  #
  # All rights reserved.
  #
  # MIT License
  #
  # Permission is hereby granted, free of charge, to any person obtaining a
  # copy of this software and associated documentation files (the "Software"),
  # to deal in the Software without restriction, including without limitation
  # the rights to use, copy, modify, merge, publish, distribute, sublicense,
  # and/or sell copies of the Software, and to permit persons to whom the
  # Software is furnished to do so, subject to the following conditions:
  #
  # The above copyright notice and this permission notice shall be included in
  # all copies or substantial portions of the Software.
  #
  # THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
  # FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
  # DEALINGS IN THE SOFTWARE.
  
  # stdlib imports
  import argparse
  import copy
  import datetime
  import json
  import pathlib
  import subprocess
  import sys
  # non-stdlib imports
  import azure.cosmosdb.table as azuretable
  
  # global defines
  _PARTITION_KEY = None
  _TABLE_NAME = None
  
  
  def _create_credentials(config: dict) -> azuretable.TableService:
      """Create authenticated clients
      :param dict config: configuration dict
      :rtype: azure.cosmosdb.table.TableService
      :return: table client
      """
      global _PARTITION_KEY, _TABLE_NAME
      _PARTITION_KEY = '{}${}'.format(
          config['credentials']['batch']['account'],
          config['pool_specification']['id'])
      try:
          sep = config['batch_shipyard']['storage_entity_prefix']
          if sep is None or len(sep) == 0:
              raise KeyError()
      except KeyError:
          sep = 'shipyard'
      _TABLE_NAME = sep + 'perf'
      ssel = config['batch_shipyard']['storage_account_settings']
      table_client = azuretable.TableService(
          account_name=config['credentials']['storage'][ssel]['account'],
          account_key=config['credentials']['storage'][ssel]['account_key'],
          endpoint_suffix=config['credentials']['storage'][ssel]['endpoint'])
      return table_client
  
  
  def _compute_delta_t(
          data: dict, nodeid: str, event1: str, event1_pos: int, event2: str,
          event2_pos: int) -> float:
      """Compute time delta between two events
      :param dict data: data
      :param str nodeid: node id
      :param str event1: event1
      :param int event1_pos: event1 position in stream
      :param str event2: event2
      :param int event2_pos: event2 position in stream
      :rtype: float
      :return: delta t of events
      """
      # attempt to get directly recorded diff
      try:
          return data[nodeid][event2][event2_pos]['message']['diff']
      except (TypeError, KeyError):
          return (data[nodeid][event2][event2_pos]['timestamp'] -
                  data[nodeid][event1][event1_pos]['timestamp']).total_seconds()
  
  
  def _parse_message(event: str, msg: str) -> dict:
      """Parse message
      :param str event: event
      :param str msg: message
      :rtype: dict
      :return: dict of message entries
      """
      parts = msg.split(',')
      m = {}
      for part in parts:
          tmp = part.split('=')
          if tmp[0] == 'size':
              if event == 'cascade:pull-end':
                  sz = tmp[1].split()
                  sz[0] = float(sz[0])
                  if sz[1] == 'kB':
                      sz[0] *= 1024
                  elif sz[1] == 'MB':
                      sz[0] *= 1024 * 1024
                  elif sz[1] == 'GB':
                      sz[0] *= 1024 * 1024 * 1024
                  elif sz[1] == 'TB':
                      sz[0] *= 1024 * 1024 * 1024 * 1024
                  tmp[1] = sz[0]
              m[tmp[0]] = int(tmp[1])
          elif tmp[0] == 'nglobalresources':
              m[tmp[0]] = int(tmp[1])
          elif tmp[0] == 'diff':
              m[tmp[0]] = float(tmp[1])
          else:
              m[tmp[0]] = tmp[1]
      return m
  
  
  def _diff_events(
          data: dict, nodeid: str, event: str, end_event: str, timing: dict,
          prefix: str, sizes: dict=None) -> None:
      """Diff start and end event
      :param dict data: data
      :param str nodeid: node id
      :param str event: start event
      :param str end_event: end event
      :param dict timing: timing dict
      :param str prefix: prefix
      :param dict sizes: sizes dict
      """
      for i in range(0, len(data[nodeid][event])):
          # torrent start -> load start may not always exist due to pull
          if (event == 'cascade:torrent-start' and
                  end_event == 'cascade:load-start' and
                  end_event not in data[nodeid]):
              return
          # find end event for this img
          subevent = data[nodeid][event][i]
          img = subevent['message']['img']
          found = False
          for j in range(0, len(data[nodeid][end_event])):
              pei = data[
                  nodeid][end_event][j]['message']['img']
              if pei == img:
                  timing[prefix + img] = _compute_delta_t(
                      data, nodeid, event, i, end_event, j)
                  if sizes is not None and img not in sizes:
                      try:
                          if event == 'cascade:load-start':
                              sizes[img] = data[
                                  nodeid][event][j]['message']['size']
                          else:
                              sizes[img] = data[
                                  nodeid][end_event][j]['message']['size']
                      except KeyError:
                          pass
                  found = True
                  break
          if not found and event != 'cascade:torrent-start':
              raise RuntimeError(
                  'could not find corresponding event for {}:{}'.format(
                      event, img))
  
  
  def coalesce_data(table_client: azuretable.TableService) -> tuple:
      """Coalesce perf data from table
      :param azure.cosmosdb.table.TableService table_client: table client
      :rtype: tuple
      :return: (timing, sizes, offer, sku)
      """
      print('graphing data from {} with pk={}'.format(
          _TABLE_NAME, _PARTITION_KEY))
      entities = table_client.query_entities(
          _TABLE_NAME, filter='PartitionKey eq \'{}\''.format(_PARTITION_KEY))
      data = {}
      # process events
      for ent in entities:
          nodeid = ent['NodeId']
          event = ent['Event']
          if nodeid not in data:
              data[nodeid] = {}
          if event not in data[nodeid]:
              data[nodeid][event] = []
          ev = {
              'timestamp': datetime.datetime.fromtimestamp(
                  float(ent['RowKey'])),
          }
          try:
              ev['message'] = _parse_message(event, ent['Message'])
          except KeyError:
              ev['message'] = None
          data[nodeid][event].append(ev)
      del entities
      sizes = {}
      offer = None
      sku = None
      for nodeid in data:
          if offer is None:
              offer = data[nodeid]['nodeprep:start'][0]['message']['offer']
              sku = data[nodeid]['nodeprep:start'][0]['message']['sku']
          # calculate dt timings
          timing = {
              'nodeprep': _compute_delta_t(
                  data, nodeid, 'nodeprep:start', 0, 'nodeprep:end', 0),
              'global_resources_loaded': _compute_delta_t(
                  data, nodeid, 'cascade:start', 0, 'cascade:gr-done', 0),
          }
          try:
              timing['docker_install'] = _compute_delta_t(
                  data, nodeid, 'nodeprep:start', 0, 'privateregistry:start', 0)
          except KeyError:
              # when no private registry setup exists, install time is
              # equivalent to nodeprep time
              timing['docker_install'] = timing['nodeprep']
          try:
              timing['private_registry_setup'] = _compute_delta_t(
                  data, nodeid, 'privateregistry:start', 0,
                  'privateregistry:end', 0)
          except KeyError:
              timing['private_registry_setup'] = 0
          try:
              timing['docker_shipyard_container_pull'] = _compute_delta_t(
                  data, nodeid, 'shipyard:pull-start', 0,
                  'shipyard:pull-end', 0)
          except KeyError:
              timing['docker_shipyard_container_pull'] = 0
          data[nodeid]['start'] = data[
              nodeid]['nodeprep:start'][0]['timestamp'].timestamp()
          data[nodeid].pop('nodeprep:start')
          data[nodeid].pop('nodeprep:end')
          data[nodeid].pop('privateregistry:start', None)
          data[nodeid].pop('privateregistry:end', None)
          data[nodeid].pop('shipyard:pull-start', None)
          data[nodeid].pop('shipyard:pull-end', None)
          data[nodeid].pop('cascade:start')
          data[nodeid].pop('cascade:gr-done')
          for event in data[nodeid]:
              # print(event, data[nodeid][event])
              if event == 'cascade:pull-start':
                  _diff_events(
                      data, nodeid, event, 'cascade:pull-end', timing, 'pull:',
                      sizes)
              elif event == 'cascade:save-start':
                  _diff_events(
                      data, nodeid, event, 'cascade:save-end', timing, 'save:',
                      sizes)
              elif event == 'cascade:torrent-start':
                  _diff_events(
                      data, nodeid, event, 'cascade:load-start', timing,
                      'torrent:')
              elif event == 'cascade:load-start':
                  _diff_events(
                      data, nodeid, event, 'cascade:load-end', timing,
                      'load:', sizes)
          data[nodeid].pop('cascade:pull-start', None)
          data[nodeid].pop('cascade:pull-end', None)
          data[nodeid].pop('cascade:save-start', None)
          data[nodeid].pop('cascade:save-end', None)
          data[nodeid].pop('cascade:torrent-start', None)
          data[nodeid].pop('cascade:load-start', None)
          data[nodeid].pop('cascade:load-end', None)
          data[nodeid]['timing'] = timing
      return data, sizes, offer, sku
  
  
  def graph_data(data: dict, sizes: dict, offer: str, sku: str):
      """Graph data via gnuplot
      :param dict data: timing data
      :param dict sizes: size data
      :param str offer: offer
      :param str sku: sku
      """
      print(sizes)
      # create data file
      dat_fname = _PARTITION_KEY.replace('$', '-') + '.dat'
      mintime = float(sys.maxsize)
      maxtime = 0.0
      rdata = {}
      for nodeid in data:
          start = data[nodeid]['start']
          if start in rdata:
              raise RuntimeError('cannot create reverse mapping')
          rdata[start] = nodeid
          if start < mintime:
              mintime = start
          if start > maxtime:
              maxtime = start
      print('nodeready variance:', maxtime - mintime)
      total_gr = 0
      total_ac = 0
      with open(dat_fname, 'w') as f:
          f.write(
              'NodePrepStartTime NodeId NodePrep+DockerInstall '
              'PrivateRegistrySetup ShipyardContainerPull GlobalResourcesLoad '
              'TotalPull TotalSave TotalLoad TotalTorrent\n')
          for start in sorted(rdata):
              nodeid = rdata[start]
              pull = 0
              save = 0
              load = 0
              torrent = 0
              for event in data[nodeid]['timing']:
                  if event.startswith('pull:'):
                      pull += data[nodeid]['timing'][event]
                  elif event.startswith('save:'):
                      save += data[nodeid]['timing'][event]
                  elif event.startswith('load:'):
                      load += data[nodeid]['timing'][event]
                  elif event.startswith('torrent:'):
                      torrent += data[nodeid]['timing'][event]
              acquisition = pull + torrent + load
              total_ac += acquisition
              print(nodeid, data[nodeid]['timing'])
              f.write(
                  ('{0} {1} {2} {3} {4} {5} {6:.5f} {7:.5f} {8:.5f} '
                   '{9:.5f}\n').format(
                       datetime.datetime.fromtimestamp(start).strftime(
                           '%Y-%m-%d-%H:%M:%S.%f'),
                       nodeid,
                       data[nodeid]['timing']['docker_install'],
                       data[nodeid]['timing']['private_registry_setup'],
                       data[nodeid]['timing']['docker_shipyard_container_pull'],
                       data[nodeid]['timing']['global_resources_loaded'],
                       pull,
                       save,
                       load,
                       torrent)
              )
              total_gr += data[nodeid]['timing']['global_resources_loaded']
      print('total gr: {} avg: {}'.format(total_gr, total_gr / len(data)))
      print('total acq: {} avg: {}'.format(total_ac, total_ac / len(data)))
      # create plot file
      plot_fname = _PARTITION_KEY.replace('$', '-') + '.plot'
      with open(plot_fname, 'w') as f:
          f.write('set terminal pngcairo enhanced transparent crop\n')
          f.write(
              ('set title "Shipyard Performance for {} ({} {})" '
               'font ", 10" \n').format(
                   _PARTITION_KEY.split('$')[-1], offer, sku))
          f.write(
              'set key top right horizontal autotitle columnhead '
              'font ", 7"\n')
          f.write('set xtics rotate by 45 right font ", 7"\n')
          f.write('set ytics font ", 8"\n')
          f.write('set xlabel "Node Prep Start Time" font ", 8"\n')
          f.write('set ylabel "Seconds" font ", 8"\n')
          f.write('set format x "%H:%M:%.3S"\n')
          f.write('set xdata time\n')
          f.write('set timefmt "%Y-%m-%d-%H:%M:%S"\n')
          f.write('set style fill solid\n')
          f.write('set boxwidth {0:.5f} absolute\n'.format(
              (maxtime - mintime) / 100.0))
          f.write('plot "{}" using 1:($3+$4+$5+$6) with boxes, \\\n'.format(
              dat_fname))
          f.write('\t"" using 1:($3+$4+$5) with boxes, \\\n')
          f.write('\t"" using 1:($3+$4) with boxes, \\\n')
          f.write('\t"" using 1:3 with boxes\n')
      png_fname = _PARTITION_KEY.replace('$', '-') + '.png'
      subprocess.check_call(
          'gnuplot {} > {}'.format(plot_fname, png_fname), shell=True)
  
  
  def merge_dict(dict1: dict, dict2: dict) -> dict:
      """Recursively merge dictionaries: dict2 on to dict1. This differs
      from dict.update() in that values that are dicts are recursively merged.
      Note that only dict value types are merged, not lists, etc.
  
      :param dict dict1: dictionary to merge to
      :param dict dict2: dictionary to merge with
      :rtype: dict
      :return: merged dictionary
      """
      if not isinstance(dict1, dict) or not isinstance(dict2, dict):
          raise ValueError('dict1 or dict2 is not a dictionary')
      result = copy.deepcopy(dict1)
      for k, v in dict2.items():
          if k in result and isinstance(result[k], dict):
              result[k] = merge_dict(result[k], v)
          else:
              result[k] = copy.deepcopy(v)
      return result
  
  
  def main():
      """Main function"""
      # get command-line args
      args = parseargs()
  
      if args.configdir is not None:
          if args.credentials is None:
              args.credentials = str(
                  pathlib.Path(args.configdir, 'credentials.json'))
          if args.config is None:
              args.config = str(pathlib.Path(args.configdir, 'config.json'))
          if args.pool is None:
              args.pool = str(pathlib.Path(args.configdir, 'pool.json'))
  
      if args.credentials is None:
          raise ValueError('credentials json not specified')
      if args.config is None:
          raise ValueError('config json not specified')
      if args.pool is None:
          raise ValueError('pool json not specified')
  
      with open(args.credentials, 'r') as f:
          config = json.load(f)
      with open(args.config, 'r') as f:
          config = merge_dict(config, json.load(f))
      with open(args.pool, 'r') as f:
          config = merge_dict(config, json.load(f))
  
      # create storage credentials
      table_client = _create_credentials(config)
      # graph data
      data, sizes, offer, sku = coalesce_data(table_client)
      graph_data(data, sizes, offer, sku)
  
  
  def parseargs():
      """Parse program arguments
      :rtype: argparse.Namespace
      :return: parsed arguments
      """
      parser = argparse.ArgumentParser(
          description='Batch Shipyard perf graph generator')
      parser.add_argument(
          '--configdir', help='json config dir')
      parser.add_argument(
          '--credentials', help='credentials json config')
      parser.add_argument(
          '--config', help='general json config for option')
      parser.add_argument(
          '--pool', help='pool json config')
      return parser.parse_args()
  
  
  if __name__ == '__main__':
      main()
  