metaflow 2.12.8__py2.py3-none-any.whl → 2.12.9__py2.py3-none-any.whl

metaflow/runner/nbdeploy.py

@@ -0,0 +1,130 @@
+import os
+import tempfile
+from typing import Dict, Optional
+
+from metaflow import Deployer
+from metaflow.runner.utils import get_current_cell, format_flowfile
+
+DEFAULT_DIR = tempfile.gettempdir()
+
+
+class NBDeployerInitializationError(Exception):
+    """Custom exception for errors during NBDeployer initialization."""
+
+    pass
+
+
+class NBDeployer(object):
+    """
+    A  wrapper over `Deployer` for deploying flows defined in a Jupyter
+    notebook cell.
+
+    Instantiate this class on the last line of a notebook cell where
+    a `flow` is defined. In contrast to `Deployer`, this class is not
+    meant to be used in a context manager.
+
+    ```python
+    deployer = NBDeployer(FlowName)
+    ar = deployer.argo_workflows(name="madhur")
+    ar_obj = ar.create()
+    result = ar_obj.trigger(alpha=300)
+    print(result.status)
+    print(result.run)
+    result.terminate()
+    ```
+
+    Parameters
+    ----------
+    flow : FlowSpec
+        Flow defined in the same cell
+    show_output : bool, default True
+        Show the 'stdout' and 'stderr' to the console by default,
+    profile : Optional[str], default None
+        Metaflow profile to use to deploy this run. If not specified, the default
+        profile is used (or the one already set using `METAFLOW_PROFILE`)
+    env : Optional[Dict[str, str]], default None
+        Additional environment variables to set. This overrides the
+        environment set for this process.
+    base_dir : Optional[str], default None
+        The directory to run the subprocess in; if not specified, a temporary
+        directory is used.
+    **kwargs : Any
+        Additional arguments that you would pass to `python myflow.py` i.e. options
+        listed in `python myflow.py --help`
+
+    """
+
+    def __init__(
+        self,
+        flow,
+        show_output: bool = True,
+        profile: Optional[str] = None,
+        env: Optional[Dict] = None,
+        base_dir: str = DEFAULT_DIR,
+        **kwargs,
+    ):
+        try:
+            from IPython import get_ipython
+
+            ipython = get_ipython()
+        except ModuleNotFoundError:
+            raise NBDeployerInitializationError(
+                "'NBDeployer' requires an interactive Python environment (such as Jupyter)"
+            )
+
+        self.cell = get_current_cell(ipython)
+        self.flow = flow
+        self.show_output = show_output
+        self.profile = profile
+        self.env = env
+        self.cwd = base_dir
+        self.top_level_kwargs = kwargs
+
+        self.env_vars = os.environ.copy()
+        self.env_vars.update(env or {})
+        # clears the Jupyter parent process ID environment variable
+        # prevents server from interfering with Metaflow
+        self.env_vars.update({"JPY_PARENT_PID": ""})
+
+        if self.profile:
+            self.env_vars["METAFLOW_PROFILE"] = self.profile
+
+        if not self.cell:
+            raise ValueError("Couldn't find a cell.")
+
+        self.tmp_flow_file = tempfile.NamedTemporaryFile(
+            prefix=self.flow.__name__,
+            suffix=".py",
+            mode="w",
+            dir=self.cwd,
+            delete=False,
+        )
+
+        self.tmp_flow_file.write(format_flowfile(self.cell))
+        self.tmp_flow_file.flush()
+        self.tmp_flow_file.close()
+
+        self.flow_file = self.tmp_flow_file.name
+
+        self.deployer = Deployer(
+            flow_file=self.flow_file,
+            show_output=self.show_output,
+            profile=self.profile,
+            env=self.env_vars,
+            cwd=self.cwd,
+            **kwargs,
+        )
+
+        from metaflow.plugins import DEPLOYER_IMPL_PROVIDERS
+
+        for provider_class in DEPLOYER_IMPL_PROVIDERS:
+            method_name = provider_class.TYPE.replace("-", "_")
+            setattr(
+                NBDeployer, method_name, self.deployer.__make_function(provider_class)
+            )
+
+    def cleanup(self):
+        """
+        Delete any temporary files created during execution.
+        """
+        os.remove(self.flow_file)

metaflow/runner/nbrun.py

@@ -1,9 +1,9 @@
-import ast
 import os
 import tempfile
 from typing import Dict, Optional
 
 from metaflow import Runner
+from metaflow.runner.utils import get_current_cell, format_flowfile
 
 DEFAULT_DIR = tempfile.gettempdir()
 
@@ -14,32 +14,6 @@ class NBRunnerInitializationError(Exception):
     pass
 
 
-def get_current_cell(ipython):
-    if ipython:
-        return ipython.history_manager.input_hist_raw[-1]
-    return None
-
-
-def format_flowfile(cell):
-    """
-    Formats the given cell content to create a valid Python script that can be executed as a Metaflow flow.
-    """
-    flowspec = [
-        x
-        for x in ast.parse(cell).body
-        if isinstance(x, ast.ClassDef) and any(b.id == "FlowSpec" for b in x.bases)
-    ]
-
-    if not flowspec:
-        raise ModuleNotFoundError(
-            "The cell doesn't contain any class that inherits from 'FlowSpec'"
-        )
-
-    lines = cell.splitlines()[: flowspec[0].end_lineno]
-    lines += ["if __name__ == '__main__':", f"    {flowspec[0].name}()"]
-    return "\n".join(lines)
-
-
 class NBRunner(object):
     """
     A  wrapper over `Runner` for executing flows defined in a Jupyter
@@ -47,7 +21,7 @@ class NBRunner(object):
 
     Instantiate this class on the last line of a notebook cell where
     a `flow` is defined. In contrast to `Runner`, this class is not
-    meant to be used a context manager. Instead, use a blocking helper
+    meant to be used in a context manager. Instead, use a blocking helper
     function like `nbrun` (which calls `cleanup()` internally) or call
     `cleanup()` explictly when using non-blocking APIs.
 
@@ -101,6 +75,8 @@ class NBRunner(object):
 
         self.env_vars = os.environ.copy()
         self.env_vars.update(env or {})
+        # clears the Jupyter parent process ID environment variable
+        # prevents server from interfering with Metaflow
         self.env_vars.update({"JPY_PARENT_PID": ""})
         if profile:
             self.env_vars["METAFLOW_PROFILE"] = profile

metaflow/runner/utils.py

@@ -0,0 +1,49 @@
+import os
+import ast
+import time
+from typing import Dict
+
+
+def get_current_cell(ipython):
+    if ipython:
+        return ipython.history_manager.input_hist_raw[-1]
+    return None
+
+
+def format_flowfile(cell):
+    """
+    Formats the given cell content to create a valid Python script that can be executed as a Metaflow flow.
+    """
+    flowspec = [
+        x
+        for x in ast.parse(cell).body
+        if isinstance(x, ast.ClassDef) and any(b.id == "FlowSpec" for b in x.bases)
+    ]
+
+    if not flowspec:
+        raise ModuleNotFoundError(
+            "The cell doesn't contain any class that inherits from 'FlowSpec'"
+        )
+
+    lines = cell.splitlines()[: flowspec[0].end_lineno]
+    lines += ["if __name__ == '__main__':", f"    {flowspec[0].name}()"]
+    return "\n".join(lines)
+
+
+def clear_and_set_os_environ(env: Dict):
+    os.environ.clear()
+    os.environ.update(env)
+
+
+def read_from_file_when_ready(file_path: str, timeout: float = 5):
+    start_time = time.time()
+    with open(file_path, "r", encoding="utf-8") as file_pointer:
+        content = file_pointer.read()
+        while not content:
+            if time.time() - start_time > timeout:
+                raise TimeoutError(
+                    "Timeout while waiting for file content from '%s'" % file_path
+                )
+            time.sleep(0.1)
+            content = file_pointer.read()
+        return content