return Run objects

Author: madhur-ob

metaflow/cli.py

@@ -686,6 +686,13 @@ def common_run_options(func):
         type=str,
         help="Write the ID of this run to the file specified.",
     )
+    @click.option(
+        "--flow-name-file",
+        default=None,
+        show_default=True,
+        type=str,
+        help="Write the flow name of this run to the file specified.",
+    )
     @wraps(func)
     def wrapper(*args, **kwargs):
         return func(*args, **kwargs)
@@ -736,6 +743,7 @@ def resume(
     max_log_size=None,
     decospecs=None,
     run_id_file=None,
+    flow_name_file=None,
 ):
     before_run(obj, tags, decospecs + obj.environment.decospecs())
 
@@ -790,7 +798,9 @@ def resume(
         max_num_splits=max_num_splits,
         max_log_size=max_log_size * 1024 * 1024,
     )
-    write_run_id(run_id_file, runtime.run_id)
+    write_file(run_id_file, runtime.run_id)
+    write_file(flow_name_file, obj.flow.name)
+
     runtime.persist_constants()
     runtime.execute()
 
@@ -819,6 +829,7 @@ def run(
     max_log_size=None,
     decospecs=None,
     run_id_file=None,
+    flow_name_file=None,
     user_namespace=None,
     **kwargs
 ):
@@ -842,17 +853,19 @@ def run(
         max_log_size=max_log_size * 1024 * 1024,
     )
     write_latest_run_id(obj, runtime.run_id)
-    write_run_id(run_id_file, runtime.run_id)
+    write_file(run_id_file, runtime.run_id)
+    write_file(flow_name_file, obj.flow.name)
 
     obj.flow._set_constants(obj.graph, kwargs)
     runtime.persist_constants()
     runtime.execute()
 
 
-def write_run_id(run_id_file, run_id):
-    if run_id_file is not None:
-        with open(run_id_file, "w") as f:
-            f.write(str(run_id))
+def write_file(file_path, content):
+    if file_path is not None:
+        with open(file_path, "w") as f:
+            f.write(str(content))
+            f.close()
 
 
 def before_run(obj, tags, decospecs):

metaflow/metaflow_runner.py

@@ -1,63 +1,117 @@
+import os
+import sys
+import time
 import asyncio
+import tempfile
+import subprocess
+from metaflow import Run
+from typing import List, Dict
 from concurrent.futures import ProcessPoolExecutor
-import subprocess, sys, os
-from typing import List, Dict, Optional
 
 
-def convert_params_to_cli_args(params: List[Dict]):
-    converted_params = [item for pair in params.items() for item in pair]
-    converted_params = [
-        str(val) if idx % 2 else f"--{val}" for idx, val in enumerate(converted_params)
-    ]
-    return converted_params
+def cli_runner(flow_file: str, command: str, args: List[str], env_vars: Dict):
+    process = subprocess.Popen(
+        [sys.executable, flow_file, command, *args],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        env=env_vars,
+    )
+    return process
 
 
-def cli_runner(flow_file: str, command: str, args: List[str], env_vars: Dict):
-    result = subprocess.run([sys.executable, flow_file, command, *args], env=env_vars)
-    return result.returncode
+def read_from_file_when_ready(file_pointer):
+    content = file_pointer.read().decode()
+    while not content:
+        time.sleep(0.1)
+        content = file_pointer.read().decode()
+    return content
 
 
 class Pool(object):
     def __init__(
         self,
         flow_file: str,
         num_processes: int,
-        profile: Optional[str] = None,
-        with_context: Optional[str] = None,
     ):
         self.flow_file = flow_file
         self.num_processes = num_processes
-        self.profile = profile
-        self.with_context = with_context
+        self.runner = Runner(self.flow_file)
 
-    def map(self, params: List[Dict], command: str = "run"):
-        return asyncio.run(self._map(params, command))
+    def __enter__(self):
+        return self
 
-    async def _map(self, params: List[Dict], command: str = "run"):
-        self.command = command
+    def map(self, paramater_space: List[Dict], blocking: bool = False):
+        return asyncio.run(self._map(paramater_space, blocking))
+
+    async def _map(self, paramater_space: List[Dict], blocking: bool = False):
         loop = asyncio.get_running_loop()
         tasks, runs = [], []
         with ProcessPoolExecutor(max_workers=self.num_processes) as executor:
-            for each_param in params:
-                tasks.append(loop.run_in_executor(executor, self._execute, each_param))
+            for set_of_params in paramater_space:
+                tasks.append(
+                    loop.run_in_executor(
+                        executor, self.runner.run, set_of_params, blocking
+                    )
+                )
             for done in asyncio.as_completed(tasks):
                 runs.append(await done)
 
         return runs
 
-    def _execute(self, params: List[Dict]):
-        env_vars = os.environ.copy()
-        if self.profile:
-            env_vars.update({"METAFLOW_PROFILE": self.profile})
-        params = convert_params_to_cli_args(params)
-        print(f"Starting for {params}")
-        result = cli_runner(self.flow_file, self.command, params, env_vars)
-        print(f"Result ready for {params}")
-        return {" ".join(params): result}
+    def __exit__(self, exc_type, exc_value, traceback):
+        pass
+
+
+# consider more args in constructor,
+# list of them is available using:
+# `python ../try.py --help`
+class Runner(object):
+    def __init__(
+        self,
+        flow_file: str,
+    ):
+        self.flow_file = flow_file
 
     def __enter__(self):
-        print("Start....")
         return self
 
+    # consider more args for run method,
+    # list of them is available using:
+    # `python ../try.py run --help`
+    def run(
+        self,
+        params: Dict,  # eventually, parse the file for parameters? contains stuff like {'alpha': 30}
+        blocking: bool = False,
+    ):
+        env_vars = os.environ.copy()
+
+        params_as_cli_args = []
+        for (k, v) in params.items():
+            params_as_cli_args.extend(["--" + str(k), str(v)])
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            tfp_flow_name = tempfile.NamedTemporaryFile(dir=temp_dir, delete=False)
+            tfp_run_id = tempfile.NamedTemporaryFile(dir=temp_dir, delete=False)
+
+            params_as_cli_args.extend(["--run-id-file", tfp_run_id.name])
+            params_as_cli_args.extend(["--flow-name-file", tfp_flow_name.name])
+
+            process = cli_runner(
+                self.flow_file,
+                command="run",
+                args=params_as_cli_args,
+                env_vars=env_vars,
+            )
+            if blocking:
+                process.wait()
+
+            flow_name = read_from_file_when_ready(tfp_flow_name)
+            run_id = read_from_file_when_ready(tfp_run_id)
+
+            pathspec_components = (flow_name, run_id)
+            run_object = Run("/".join(pathspec_components), _namespace_check=False)
+
+            return run_object
+
     def __exit__(self, exc_type, exc_value, traceback):
-        print("End...")
+        pass