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