tools/crunchstat-summary/crunchstat_summary/command.py

   1 # Copyright (C) The Arvados Authors. All rights reserved.
   2 #
   3 # SPDX-License-Identifier: AGPL-3.0
   4
   5 import argparse
   6 import gzip
   7 from io import open
   8 import logging
   9 import sys
  10 import arvados
  11
  12 from crunchstat_summary import logger, summarizer, reader
  13 from crunchstat_summary._version import __version__
  14
  15
  16 class ArgumentParser(argparse.ArgumentParser):
  17     def __init__(self):
  18         super(ArgumentParser, self).__init__(
  19             description='Summarize resource usage of an Arvados Crunch job')
  20         src = self.add_mutually_exclusive_group()
  21         src.add_argument(
  22             '--job', '--container-request',
  23             type=str, metavar='UUID',
  24             help='Look up the specified job or container request '
  25             'and read its log data from Keep (or from the Arvados event log, '
  26             'if the job is still running)')
  27         src.add_argument(
  28             '--container',
  29             type=str, metavar='UUID',
  30             help='[Deprecated] Look up the specified container find its container request '
  31             'and read its log data from Keep (or from the Arvados event log, '
  32             'if the job is still running)')
  33         src.add_argument(
  34             '--log-file', type=str,
  35             help='Read log data from a regular file')
  36         self.add_argument(
  37             '--skip-child-jobs', action='store_true',
  38             help='Do not include stats from child jobs/containers')
  39         self.add_argument(
  40             '--format', type=str, choices=('html', 'text'), default='text',
  41             help='Report format')
  42         self.add_argument(
  43             '--threads', type=int, default=8,
  44             help='Maximum worker threads to run')
  45         self.add_argument(
  46             '--verbose', '-v', action='count', default=0,
  47             help='Log more information (once for progress, twice for debug)')
  48         self.add_argument('--version', action='version',
  49                          version="%s %s" % (sys.argv[0], __version__),
  50                          help='Print version and exit.')
  51
  52
  53 class UTF8Decode(object):
  54     '''Wrap a file-like iterable to decode UTF-8 bytes into a strings
  55     '''
  56     def __init__(self, fh):
  57         self.fh = fh
  58
  59     def __enter__(self):
  60         return self
  61
  62     def __exit__(self, exc_type, exc_val, exc_tb):
  63         self.close()
  64
  65     def __iter__(self):
  66         return self
  67
  68     def __next__(self):
  69         return next(self.fh).decode('utf-8')
  70
  71     next = __next__
  72
  73     def close(self):
  74         # mimic Gzip behavior and don't close underlying object
  75         pass
  76
  77
  78 class Command(object):
  79     def __init__(self, args):
  80         self.args = args
  81         logger.setLevel(logging.WARNING - 10 * args.verbose)
  82
  83     def run(self):
  84         kwargs = {
  85             'skip_child_jobs': self.args.skip_child_jobs,
  86             'threads': self.args.threads,
  87             'arv': arvados.api('v1')
  88         }
  89         if self.args.job:
  90             self.summer = summarizer.NewSummarizer(self.args.job, **kwargs)
  91         elif self.args.container:
  92             self.summer = summarizer.NewSummarizer(self.args.container, **kwargs)
  93         elif self.args.log_file:
  94             if self.args.log_file.endswith('.gz'):
  95                 fh = UTF8Decode(gzip.open(self.args.log_file))
  96             else:
  97                 fh = open(self.args.log_file, mode = 'r', encoding = 'utf-8')
  98             self.summer = summarizer.Summarizer(reader.StubReader(fh), **kwargs)
  99         else:
 100             self.summer = summarizer.Summarizer(reader.StubReader(sys.stdin), **kwargs)
 101         return self.summer.run()
 102
 103     def report(self):
 104         if self.args.format == 'html':
 105             return self.summer.html_report()
 106         elif self.args.format == 'text':
 107             return self.summer.text_report()