more realtime log fetching

Author: savingoyal

metaflow/metaflow_environment.py

@@ -6,7 +6,7 @@
 from . import metaflow_version
 from metaflow.exception import MetaflowException
 from metaflow.extension_support import dump_module_info
-from metaflow.mflog import BASH_MFLOG
+from metaflow.mflog import BASH_MFLOG, BASH_FLUSH_LOGS
 from . import R
 
 
@@ -159,6 +159,7 @@ def _get_install_dependencies_cmd(self, datastore_type):
     def get_package_commands(self, code_package_url, datastore_type):
         cmds = [
             BASH_MFLOG,
+            BASH_FLUSH_LOGS,
             "mflog 'Setting up task environment.'",
             self._get_install_dependencies_cmd(datastore_type),
             "mkdir metaflow",
@@ -176,6 +177,7 @@ def get_package_commands(self, code_package_url, datastore_type):
             "fi" % code_package_url,
             "TAR_OPTIONS='--warning=no-timestamp' tar xf job.tar",
             "mflog 'Task is starting.'",
+            "flush_mflogs",
         ]
         return cmds
 

metaflow/mflog/__init__.py

@@ -44,6 +44,8 @@
 BASH_SAVE_LOGS_ARGS = ["python", "-m", "metaflow.mflog.save_logs"]
 BASH_SAVE_LOGS = " ".join(BASH_SAVE_LOGS_ARGS)
 
+BASH_FLUSH_LOGS = "flush_mflogs(){ " f"{BASH_SAVE_LOGS}; " "}"
+
 
 # this function returns a bash expression that redirects stdout
 # and stderr of the given bash expression to mflog.tee
@@ -63,7 +65,7 @@ def bash_capture_logs(bash_expr, var_transform=None):
 # update_delay determines how often logs should be uploaded to S3
 # as a function of the task execution time
 
-MIN_UPDATE_DELAY = 1.0  # the most frequent update interval
+MIN_UPDATE_DELAY = 0.25  # the most frequent update interval
 MAX_UPDATE_DELAY = 30.0  # the least frequent update interval
 
 
@@ -110,7 +112,6 @@ def export_mflog_env_vars(
 
 def tail_logs(prefix, stdout_tail, stderr_tail, echo, has_log_updates):
     def _available_logs(tail, stream, echo, should_persist=False):
-        # print the latest batch of lines
         try:
             for line in tail:
                 if should_persist:
@@ -128,7 +129,7 @@ def _available_logs(tail, stream, echo, should_persist=False):
 
     start_time = time.time()
     next_log_update = start_time
-    log_update_delay = 1
+    log_update_delay = update_delay(0)
     while has_log_updates():
         if time.time() > next_log_update:
             _available_logs(stdout_tail, "stdout", echo)

metaflow/plugins/pypi/conda_environment.py

@@ -421,13 +421,15 @@ def bootstrap_commands(self, step_name, datastore_type):
         if id_:
             return [
                 "echo 'Bootstrapping virtual environment...'",
+                "flush_mflogs",
                 # We have to prevent the tracing module from loading,
                 # as the bootstrapping process uses the internal S3 client which would fail to import tracing
                 # due to the required dependencies being bundled into the conda environment,
                 # which is yet to be initialized at this point.
                 'DISABLE_TRACING=True python -m metaflow.plugins.pypi.bootstrap "%s" %s "%s"'
                 % (self.flow.name, id_, self.datastore_type),
                 "echo 'Environment bootstrapped.'",
+                "flush_mflogs",
                 # To avoid having to install micromamba in the PATH in micromamba.py, we add it to the PATH here.
                 "export PATH=$PATH:$(pwd)/micromamba/bin",
                 "export MF_ARCH=$(case $(uname)/$(uname -m) in Darwin/arm64)echo osx-arm64;;Darwin/*)echo osx-64;;Linux/aarch64)echo linux-aarch64;;*)echo linux-64;;esac)",