tools/crunchstat-summary/crunchstat_summary/reader.py

   1 # Copyright (C) The Arvados Authors. All rights reserved.
   2 #
   3 # SPDX-License-Identifier: AGPL-3.0
   4
   5 import arvados
   6 import itertools
   7 import queue
   8 import threading
   9
  10 from crunchstat_summary import logger
  11
  12
  13 class CollectionReader(object):
  14     def __init__(self, collection_id):
  15         self._collection_id = collection_id
  16         self._label = collection_id
  17         self._readers = []
  18
  19     def __str__(self):
  20         return self._label
  21
  22     def __iter__(self):
  23         logger.debug('load collection %s', self._collection_id)
  24         collection = arvados.collection.CollectionReader(self._collection_id)
  25         filenames = [filename for filename in collection]
  26         # Crunch2 has multiple stats files
  27         if len(filenames) > 1:
  28             filenames = ['crunchstat.txt', 'arv-mount.txt']
  29         for filename in filenames:
  30             try:
  31                 self._readers.append(collection.open(filename))
  32             except IOError:
  33                 logger.warn('Unable to open %s', filename)
  34         self._label = "{}/{}".format(self._collection_id, filenames[0])
  35         return itertools.chain(*[iter(reader) for reader in self._readers])
  36
  37     def __enter__(self):
  38         return self
  39
  40     def __exit__(self, exc_type, exc_val, exc_tb):
  41         if self._readers:
  42             for reader in self._readers:
  43                 reader.close()
  44             self._readers = []
  45
  46
  47 class LiveLogReader(object):
  48     EOF = None
  49
  50     def __init__(self, job_uuid):
  51         self.job_uuid = job_uuid
  52         self.event_types = (['stderr'] if '-8i9sb-' in job_uuid else ['crunchstat', 'arv-mount'])
  53         logger.debug('load %s events for job %s', self.event_types, self.job_uuid)
  54
  55     def __str__(self):
  56         return self.job_uuid
  57
  58     def _get_all_pages(self):
  59         got = 0
  60         last_id = 0
  61         filters = [
  62             ['object_uuid', '=', self.job_uuid],
  63             ['event_type', 'in', self.event_types]]
  64         try:
  65             while True:
  66                 page = arvados.api().logs().index(
  67                     limit=1000,
  68                     order=['id asc'],
  69                     filters=filters + [['id','>',str(last_id)]],
  70                     select=['id', 'properties'],
  71                 ).execute(num_retries=2)
  72                 got += len(page['items'])
  73                 logger.debug(
  74                     '%s: received %d of %d log events',
  75                     self.job_uuid, got,
  76                     got + page['items_available'] - len(page['items']))
  77                 for i in page['items']:
  78                     for line in i['properties']['text'].split('\n'):
  79                         self._queue.put(line+'\n')
  80                     last_id = i['id']
  81                 if (len(page['items']) == 0 or
  82                     len(page['items']) >= page['items_available']):
  83                     break
  84         finally:
  85             self._queue.put(self.EOF)
  86
  87     def __iter__(self):
  88         self._queue = queue.Queue()
  89         self._thread = threading.Thread(target=self._get_all_pages)
  90         self._thread.daemon = True
  91         self._thread.start()
  92         return self
  93
  94     def __next__(self):
  95         line = self._queue.get()
  96         if line is self.EOF:
  97             self._thread.join()
  98             raise StopIteration
  99         return line
 100
 101     next = __next__ # for Python 2
 102
 103     def __enter__(self):
 104         return self
 105
 106     def __exit__(self, exc_type, exc_val, exc_tb):
 107         pass