Refactor function arguments and address comments

Author: talsperre

metaflow/mflog/__init__.py

@@ -42,22 +42,36 @@
 )
 
 
-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)
+def _get_bash_capture_log(default_python_executable=None):
+    python_executable = default_python_executable or "python"
+    bash_save_log_args = [python_executable, "-m", "metaflow.mflog.save_logs"]
+    return 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, default_python_executable_path="python"):
+def bash_capture_logs(
+    bash_expr, var_transform=None, default_python_executable_path=None
+):
+    python_executable = default_python_executable_path or "python"
     if var_transform is None:
         var_transform = lambda s: "$%s" % s
 
     cmd = "%s -m metaflow.mflog.tee %s %s"
     parts = (
         bash_expr,
-        cmd % (default_python_executable_path, TASK_LOG_SOURCE, var_transform("MFLOG_STDOUT")),
-        cmd % (default_python_executable_path, TASK_LOG_SOURCE, var_transform("MFLOG_STDERR")),
+        cmd
+        % (
+            python_executable,
+            TASK_LOG_SOURCE,
+            var_transform("MFLOG_STDOUT"),
+        ),
+        cmd
+        % (
+            python_executable,
+            TASK_LOG_SOURCE,
+            var_transform("MFLOG_STDERR"),
+        ),
     )
     return "(%s) 1>> >(%s) 2>> >(%s >&2)" % parts
 

metaflow/mflog/save_logs_periodically.py

@@ -38,7 +38,7 @@ def _file_size(path):
             if new_sizes != sizes:
                 sizes = new_sizes
                 try:
-                    subprocess.call(_get_bash_capture_log()[0])
+                    subprocess.call(_get_bash_capture_log())
                 except:
                     pass
             time.sleep(update_delay(time.time() - start_time))

metaflow/plugins/argo/argo_workflows.py

@@ -45,7 +45,11 @@
     UI_URL,
 )
 from metaflow.metaflow_config_funcs import config_values
-from metaflow.mflog import _get_bash_capture_log, 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 +1198,7 @@ def _container_templates(self):
 
             cmd_str = "%s; c=$?; %s; exit $c" % (
                 " && ".join([init_cmds, bash_capture_logs(" && ".join(step_cmds))]),
-                _get_bash_capture_log()[1],
+                " ".join(_get_bash_capture_log()),
             )
             cmds = shlex.split('bash -c "%s"' % cmd_str)
 

metaflow/plugins/aws/batch/batch.py

@@ -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" % _get_bash_capture_log()[1]
+        cmd_str += "c=$?; %s; exit $c" % " ".join(_get_bash_capture_log())
         return shlex.split('bash -c "%s"' % cmd_str)
 
     def _search_jobs(self, flow_name, run_id, user):

metaflow/plugins/kubernetes/kubernetes.py

@@ -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" % _get_bash_capture_log()
+        cmd_str += "c=$?; %s; exit $c" % " ".join(_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 = (