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