71eae519969bf1f34ec10b491c3885c1425596df
[arvados.git] / 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()