Support user specified python executables for remote bootstrapping and save logs functionality

Author: talsperre

metaflow/metaflow_environment.py

@@ -42,7 +42,7 @@ def decospecs(self):
         """
         return ()
 
-    def bootstrap_commands(self, step_name, datastore_type):
+    def bootstrap_commands(self, step_name, datastore_type, default_executable_path=None):
         """
         A list of shell commands to bootstrap this environment in a remote runtime.
         """

metaflow/mflog/__init__.py

@@ -41,21 +41,23 @@
     " }" % TASK_LOG_SOURCE
 )
 
-BASH_SAVE_LOGS_ARGS = ["python", "-m", "metaflow.mflog.save_logs"]
-BASH_SAVE_LOGS = " ".join(BASH_SAVE_LOGS_ARGS)
+
+def _get_bash_capture_log(default_python_executable="python"):
+    bash_save_log_args = [default_python_executable, "-m", "metaflow.mflog.save_logs"]
+    return bash_save_log_args, " ".join(bash_save_log_args)
 
 
 # this function returns a bash expression that redirects stdout
 # and stderr of the given bash expression to mflog.tee
-def bash_capture_logs(bash_expr, var_transform=None):
+def bash_capture_logs(bash_expr, var_transform=None, default_python_executable_path="python"):
     if var_transform is None:
         var_transform = lambda s: "$%s" % s
 
-    cmd = "python -m metaflow.mflog.tee %s %s"
+    cmd = "%s -m metaflow.mflog.tee %s %s"
     parts = (
         bash_expr,
-        cmd % (TASK_LOG_SOURCE, var_transform("MFLOG_STDOUT")),
-        cmd % (TASK_LOG_SOURCE, var_transform("MFLOG_STDERR")),
+        cmd % (default_python_executable_path, TASK_LOG_SOURCE, var_transform("MFLOG_STDOUT")),
+        cmd % (default_python_executable_path, TASK_LOG_SOURCE, var_transform("MFLOG_STDERR")),
     )
     return "(%s) 1>> >(%s) 2>> >(%s >&2)" % parts
 

metaflow/mflog/save_logs_periodically.py

@@ -5,7 +5,7 @@
 from threading import Thread
 
 from metaflow.sidecar import MessageTypes
-from . import update_delay, BASH_SAVE_LOGS_ARGS
+from . import update_delay, _get_bash_capture_log
 
 
 class SaveLogsPeriodicallySidecar(object):
@@ -38,7 +38,7 @@ def _file_size(path):
             if new_sizes != sizes:
                 sizes = new_sizes
                 try:
-                    subprocess.call(BASH_SAVE_LOGS_ARGS)
+                    subprocess.call(_get_bash_capture_log()[0])
                 except:
                     pass
             time.sleep(update_delay(time.time() - start_time))

metaflow/plugins/argo/argo_workflows.py

@@ -45,7 +45,7 @@
     UI_URL,
 )
 from metaflow.metaflow_config_funcs import config_values
-from metaflow.mflog import BASH_SAVE_LOGS, bash_capture_logs, export_mflog_env_vars
+from metaflow.mflog import _get_bash_capture_log, bash_capture_logs, export_mflog_env_vars
 from metaflow.parameters import deploy_time_eval
 from metaflow.plugins.kubernetes.kubernetes import (
     parse_kube_keyvalue_list,
@@ -1194,7 +1194,7 @@ def _container_templates(self):
 
             cmd_str = "%s; c=$?; %s; exit $c" % (
                 " && ".join([init_cmds, bash_capture_logs(" && ".join(step_cmds))]),
-                BASH_SAVE_LOGS,
+                _get_bash_capture_log()[1],
             )
             cmds = shlex.split('bash -c "%s"' % cmd_str)
 

metaflow/plugins/aws/batch/batch.py

@@ -30,7 +30,7 @@
     export_mflog_env_vars,
     bash_capture_logs,
     tail_logs,
-    BASH_SAVE_LOGS,
+    _get_bash_capture_log,
 )
 
 from .batch_client import BatchClient
@@ -92,7 +92,7 @@ def _command(self, environment, code_package_url, step_name, step_cmds, task_spe
         # Note that if step_expr OOMs, this tail expression is never executed.
         # We lose the last logs in this scenario (although they are visible
         # still through AWS CloudWatch console).
-        cmd_str += "c=$?; %s; exit $c" % BASH_SAVE_LOGS
+        cmd_str += "c=$?; %s; exit $c" % _get_bash_capture_log()[1]
         return shlex.split('bash -c "%s"' % cmd_str)
 
     def _search_jobs(self, flow_name, run_id, user):

metaflow/plugins/kubernetes/kubernetes.py

@@ -41,7 +41,7 @@
 from metaflow.metaflow_config_funcs import config_values
 
 from metaflow.mflog import (
-    BASH_SAVE_LOGS,
+    _get_bash_capture_log,
     bash_capture_logs,
     export_mflog_env_vars,
     get_log_tailer,
@@ -130,7 +130,7 @@ def _command(
         # We lose the last logs in this scenario.
         #
         # TODO: Capture hard exit logs in Kubernetes.
-        cmd_str += "c=$?; %s; exit $c" % BASH_SAVE_LOGS
+        cmd_str += "c=$?; %s; exit $c" % _get_bash_capture_log()
         # For supporting sandboxes, ensure that a custom script is executed before
         # anything else is executed. The script is passed in as an env var.
         cmd_str = (

metaflow/plugins/pypi/conda_environment.py

@@ -331,19 +331,22 @@ def add_to_package(self):
             files.append((manifest, os.path.basename(manifest)))
         return files
 
-    def bootstrap_commands(self, step_name, datastore_type):
+    def bootstrap_commands(self, step_name, datastore_type, default_executable_path=None):
         # Bootstrap conda and execution environment for step
         step = next(step for step in self.flow if step.name == step_name)
         id_ = self.get_environment(step).get("id_")
         if id_:
+            python_executable_path = (
+                "python" if default_executable_path is None else default_executable_path
+            )
             return [
                 "echo 'Bootstrapping virtual environment...'",
                 # 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" linux-64'
-                % (self.flow.name, id_, self.datastore_type),
+                'DISABLE_TRACING=True %s -m metaflow.plugins.pypi.bootstrap "%s" %s "%s" linux-64'
+                % (python_executable_path, self.flow.name, id_, self.datastore_type),
                 "echo 'Environment bootstrapped.'",
                 "export PATH=$PATH:$(pwd)/micromamba",
             ]