Make datastore.__init__(mode='w') a cheap operation (#448)

* Make datastore.__init__(mode='w') cheap

* Make datastore.__init__(mode='w') cheap

Author: savingoyal

metaflow/datastore/datastore.py

@@ -323,15 +323,18 @@ def __init__(self,
                  event_logger=None,
                  monitor=None,
                  data_obj=None,
-                 artifact_cache=None):
+                 artifact_cache=None,
+                 allow_unsuccessful=False):
         if run_id == 'data':
             raise DataException("Run ID 'data' is reserved. "
                                 "Try with a different --run-id.")
         if self.datastore_root is None:
             raise DataException("Datastore root not found. "
                                 "Specify with METAFLOW_DATASTORE_SYSROOT_%s "
                                 "environment variable." % self.TYPE.upper())
-        
+        # NOTE: calling __init__(mode='w') should be a cheap operation:
+        # no file system accesses are allowed. It is called frequently
+        # e.g. to resolve log file location.
         self.event_logger = event_logger if event_logger else NullEventLogger()
         self.monitor = monitor if monitor else NullMonitor()
         self.metadata = metadata
@@ -356,14 +359,7 @@ def __init__(self,
                                    task_id)
 
         self.attempt = attempt
-        if mode == 'w':
-            if run_id is not None:
-                # run_id may be None when datastore is used to save
-                # things not related to runs, e.g. the job package
-                self.save_metadata('attempt', {'time': time.time()})
-                self.objects = {}
-                self.info = {}
-        elif mode == 'r':
+        if mode == 'r':
             if data_obj is None:
                 # what is the latest attempt ID of this data store?
 
@@ -389,25 +385,37 @@ def __init__(self,
                         self.attempt = i
 
                 # was the latest attempt completed successfully?
-                if not self.is_done():
+                if self.is_done():
+                    # load the data from the latest attempt
+                    data_obj = self.load_metadata('data')
+                elif allow_unsuccessful and self.attempt is not None:
+                    # this mode can be used to load_logs, for instance
+                    data_obj = None
+                else:
                     raise DataException("Data was not found or not finished at %s"\
                                         % self.root)
 
-                # load the data from the latest attempt
-                data_obj = self.load_metadata('data')
-
-            self.origin = data_obj.get('origin')
-            self.objects = data_obj['objects']
-            self.info = data_obj.get('info', {})
+            if data_obj:
+                self.origin = data_obj.get('origin')
+                self.objects = data_obj['objects']
+                self.info = data_obj.get('info', {})
         elif mode == 'd':
             # Direct access mode used by the client. We effectively don't load any
             # objects and can only access things using the load_* functions
             self.origin = None
             self.objects = None
             self.info = None
-        else:
+        elif mode != 'w':
             raise DataException('Unknown datastore mode: %s' % mode)
 
+    def init_task(self):
+        # this method should be called once after datastore has been opened
+        # for task-related write operations
+        self.save_metadata('attempt', {'time': time.time()})
+        self.objects = {}
+        self.info = {}
+
+
     @property
     def pathspec(self):
         return '%s/%s/%s' % (self.run_id, self.step_name, self.task_id)

metaflow/metaflow_config.py

@@ -170,6 +170,9 @@ def from_conf(name, default=None):
 # increasing this limit has real performance implications for all tasks.
 # Decreasing this limit is very unsafe, as it can lead to wrong results
 # being read from old tasks.
+#
+# Note also that DataStoreSet resolves the latest attempt_id using
+# lexicographic ordering of attempts. This won't work if MAX_ATTEMPTS > 99.
 MAX_ATTEMPTS = 6
 
 

metaflow/runtime.py

@@ -529,6 +529,7 @@ def new_attempt(self):
                                    attempt=self.retries,
                                    event_logger=self.event_logger,
                                    monitor=self.monitor)
+        self._ds.init_task()
 
 
     def log(self, msg, system_msg=False, pid=None):

metaflow/task.py

@@ -215,6 +215,7 @@ def clone_only(self, step_name, run_id, task_id, clone_origin_task):
                                 attempt=0,
                                 event_logger=self.event_logger,
                                 monitor=self.monitor)
+        output.init_task()
         origin_run_id, origin_step_name, origin_task_id =\
             clone_origin_task.split('/')
         # 2. initialize origin datastore
@@ -280,6 +281,7 @@ def run_step(self,
                                 attempt=retry_count,
                                 event_logger=self.event_logger,
                                 monitor=self.monitor)
+        output.init_task()
 
         if input_paths:
             # 2. initialize input datastores

metaflow/util.py

@@ -29,11 +29,21 @@
     def unquote_bytes(x):
         return to_unicode(unquote(to_bytes(x)))
 
+    # this is used e.g. by datastore.save_logs to identify paths
+    class Path(object):
+
+        def __init__(path):
+            self.path = path
+
+        def __str__(self):
+            return self.path
+
 except:
     # python3
     unicode_type = str
     bytes_type = bytes
     from urllib.parse import quote, unquote
+    from pathlib import Path
 
     def unquote_bytes(x):
         return unquote(to_unicode(x))