13649: Fixes default preemptible parameter when submitting committed child CRs

[arvados.git] / tools / crunchstat-summary / crunchstat_summary / reader.py

diff --git a/tools/crunchstat-summary/crunchstat_summary/reader.py b/tools/crunchstat-summary/crunchstat_summary/reader.py
index a0a838d6c9538c6231e4d284b6a2d2e5b7df77d6..98dda673d5a3ab70d65ab1d3989b49f539959b69 100644 (file)
--- a/tools/crunchstat-summary/crunchstat_summary/reader.py
+++ b/tools/crunchstat-summary/crunchstat_summary/reader.py
@@ -14,29 +14,41 @@ from crunchstat_summary import logger
 
 class CollectionReader(object):
     def __init__(self, collection_id):
-        logger.debug('load collection %s', collection_id)
-        collection = arvados.collection.CollectionReader(collection_id)
+        self._collection_id = collection_id
+        self._label = collection_id
+        self._reader = None
+
+    def __str__(self):
+        return self._label
+
+    def __iter__(self):
+        logger.debug('load collection %s', self._collection_id)
+        collection = arvados.collection.CollectionReader(self._collection_id)
         filenames = [filename for filename in collection]
         if len(filenames) == 1:
             filename = filenames[0]
         else:
             filename = 'crunchstat.txt'
+        self._label = "{}/{}".format(self._collection_id, filename)
         self._reader = collection.open(filename)
-        self._label = "{}/{}".format(collection_id, filename)
+        return iter(self._reader)
 
-    def __str__(self):
-        return self._label
+    def __enter__(self):
+        return self
 
-    def __iter__(self):
-        return iter(self._reader)
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self._reader:
+            self._reader.close()
+            self._reader = None
 
 
 class LiveLogReader(object):
     EOF = None
 
     def __init__(self, job_uuid):
-        logger.debug('load stderr events for job %s', job_uuid)
         self.job_uuid = job_uuid
+        self.event_types = (['stderr'] if '-8i9sb-' in job_uuid else ['crunchstat', 'arv-mount'])
+        logger.debug('load %s events for job %s', self.event_types, self.job_uuid)
 
     def __str__(self):
         return self.job_uuid
@@ -46,7 +58,7 @@ class LiveLogReader(object):
         last_id = 0
         filters = [
             ['object_uuid', '=', self.job_uuid],
-            ['event_type', '=', 'stderr']]
+            ['event_type', 'in', self.event_types]]
         try:
             while True:
                 page = arvados.api().logs().index(
@@ -83,3 +95,9 @@ class LiveLogReader(object):
             self._thread.join()
             raise StopIteration
         return line
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass