ob-metaflow 2.15.13.1__py2.py3-none-any.whl → 2.19.7.1rc0__py2.py3-none-any.whl

metaflow/plugins/env_escape/client_modules.py

@@ -1,5 +1,6 @@
 import atexit
 import importlib
+import importlib.util
 import itertools
 import pickle
 import re
@@ -41,6 +42,8 @@ class _WrappedModule(object):
     def __getattr__(self, name):
         if name == "__loader__":
             return self._loader
+        if name == "__spec__":
+            return importlib.util.spec_from_loader(self._prefix, self._loader)
         if name in ("__name__", "__package__"):
             return self._prefix
         if name in ("__file__", "__path__"):
@@ -71,7 +74,8 @@ class _WrappedModule(object):
             # Try to see if this is a submodule that we can load
             m = None
             try:
-                m = self._loader.load_module(".".join([self._prefix, name]))
+                submodule_name = ".".join([self._prefix, name])
+                m = importlib.import_module(submodule_name)
             except ImportError:
                 pass
             if m is None:
@@ -117,7 +121,28 @@ class _WrappedModule(object):
 
 
 class ModuleImporter(object):
-    # This ModuleImporter implements the Importer Protocol defined in PEP 302
+    """
+    A custom import hook that proxies module imports to a different Python environment.
+
+    This class implements the MetaPathFinder and Loader protocols (PEP 451) to enable
+    "environment escape" - allowing the current Python process to import and use modules
+    from a different Python interpreter with potentially different versions or packages.
+
+    When a module is imported through this importer:
+    1. A client spawns a server process in the target Python environment
+    2. The module is loaded in the remote environment
+    3. A _WrappedModule proxy is returned that forwards all operations (function calls,
+       attribute access, etc.) to the remote environment via RPC
+    4. Data is serialized/deserialized using pickle for cross-environment communication
+
+    Args:
+        python_executable: Path to the Python interpreter for the remote environment
+        pythonpath: Python path to use in the remote environment
+        max_pickle_version: Maximum pickle protocol version supported by remote interpreter
+        config_dir: Directory containing configuration for the environment escape
+        module_prefixes: List of module name prefixes to handle
+    """
+
     def __init__(
         self,
         python_executable,
@@ -135,84 +160,89 @@ class ModuleImporter(object):
         self._handled_modules = None
         self._aliases = {}
 
-    def find_module(self, fullname, path=None):
+    def find_spec(self, fullname, path=None, target=None):
         if self._handled_modules is not None:
             if get_canonical_name(fullname, self._aliases) in self._handled_modules:
-                return self
+                return importlib.util.spec_from_loader(fullname, self)
             return None
         if any([fullname.startswith(prefix) for prefix in self._module_prefixes]):
             # We potentially handle this
-            return self
+            return importlib.util.spec_from_loader(fullname, self)
         return None
 
-    def load_module(self, fullname):
-        if fullname in sys.modules:
-            return sys.modules[fullname]
-        if self._client is None:
-            if sys.version_info[0] < 3:
-                raise NotImplementedError(
-                    "Environment escape imports are not supported in Python 2"
-                )
-            # We initialize a client and query the modules we handle
-            # The max_pickle_version is the pickle version that the server (so
-            # the underlying interpreter we call into) supports; we determine
-            # what version the current environment support and take the minimum
-            # of those two
-            max_pickle_version = min(self._max_pickle_version, pickle.HIGHEST_PROTOCOL)
-
-            self._client = Client(
-                self._module_prefixes,
-                self._python_executable,
-                self._pythonpath,
-                max_pickle_version,
-                self._config_dir,
-            )
-            atexit.register(_clean_client, self._client)
-
-            # Get information about overrides and what the server knows about
-            exports = self._client.get_exports()
-
-            prefixes = set()
-            export_classes = exports.get("classes", [])
-            export_functions = exports.get("functions", [])
-            export_values = exports.get("values", [])
-            export_exceptions = exports.get("exceptions", [])
-            self._aliases = exports.get("aliases", {})
-            for name in itertools.chain(
-                export_classes,
-                export_functions,
-                export_values,
-                (e[0] for e in export_exceptions),
-            ):
-                splits = name.rsplit(".", 1)
-                prefixes.add(splits[0])
-            # We will make sure that we create modules even for "empty" prefixes
-            # because packages are always loaded hierarchically so if we have
-            # something in `a.b.c` but nothing directly in `a`, we still need to
-            # create a module named `a`. There is probably a better way of doing this
-            all_prefixes = list(prefixes)
-            for prefix in all_prefixes:
-                parts = prefix.split(".")
-                cur = parts[0]
-                for i in range(1, len(parts)):
-                    prefixes.add(cur)
-                    cur = ".".join([cur, parts[i]])
-
-            # We now know all the modules that we can handle. We update
-            # handled_module and return the module if we have it or raise ImportError
-            self._handled_modules = {}
-            for prefix in prefixes:
-                self._handled_modules[prefix] = _WrappedModule(
-                    self, prefix, exports, self._client
-                )
+    def create_module(self, spec):
+        # Return the pre-created wrapped module for this spec
+        self._initialize_client()
+
+        fullname = spec.name
         canonical_fullname = get_canonical_name(fullname, self._aliases)
-        # Modules are created canonically but we need to return something for any
-        # of the aliases.
-        module = self._handled_modules.get(canonical_fullname)
-        if module is None:
-            raise ImportError
-        sys.modules[fullname] = module
-        return module
+        # Modules are created canonically but we need to handle any of the aliases.
+        wrapped_module = self._handled_modules.get(canonical_fullname)
+        if wrapped_module is None:
+            raise ImportError(f"No module named '{fullname}'")
+        return wrapped_module
+
+    def exec_module(self, module):
+        # No initialization needed since the wrapped module returned by
+        # create_module() is fully initialized
+        pass
+
+    def _initialize_client(self):
+        if self._client is not None:
+            return
+
+        # We initialize a client and query the modules we handle
+        # The max_pickle_version is the pickle version that the server (so
+        # the underlying interpreter we call into) supports; we determine
+        # what version the current environment support and take the minimum
+        # of those two
+        max_pickle_version = min(self._max_pickle_version, pickle.HIGHEST_PROTOCOL)
+
+        self._client = Client(
+            self._module_prefixes,
+            self._python_executable,
+            self._pythonpath,
+            max_pickle_version,
+            self._config_dir,
+        )
+        atexit.register(_clean_client, self._client)
+
+        # Get information about overrides and what the server knows about
+        exports = self._client.get_exports()
+
+        prefixes = set()
+        export_classes = exports.get("classes", [])
+        export_functions = exports.get("functions", [])
+        export_values = exports.get("values", [])
+        export_exceptions = exports.get("exceptions", [])
+        self._aliases = exports.get("aliases", {})
+        for name in itertools.chain(
+            export_classes,
+            export_functions,
+            export_values,
+            (e[0] for e in export_exceptions),
+        ):
+            splits = name.rsplit(".", 1)
+            prefixes.add(splits[0])
+        # We will make sure that we create modules even for "empty" prefixes
+        # because packages are always loaded hierarchically so if we have
+        # something in `a.b.c` but nothing directly in `a`, we still need to
+        # create a module named `a`. There is probably a better way of doing this
+        all_prefixes = list(prefixes)
+        for prefix in all_prefixes:
+            parts = prefix.split(".")
+            cur = parts[0]
+            for i in range(1, len(parts)):
+                prefixes.add(cur)
+                cur = ".".join([cur, parts[i]])
+
+        # We now know all the modules that we can handle. We update
+        # handled_module and return the module if we have it or raise ImportError
+        self._handled_modules = {}
+        for prefix in prefixes:
+            self._handled_modules[prefix] = _WrappedModule(
+                self, prefix, exports, self._client
+            )
 
 
 def create_modules(python_executable, pythonpath, max_pickle_version, path, prefixes):

metaflow/plugins/events_decorator.py

@@ -70,7 +70,26 @@ class TriggerDecorator(FlowDecorator):
         "options": {},
     }
 
-    def process_event_name(self, event):
+    def process_event(self, event):
+        """
+        Process a single event and return a dictionary if static trigger and a function
+        if deploy-time trigger.
+
+        Parameters
+        ----------
+        event : Union[str, Dict[str, Any], Callable]
+            Event to process
+
+        Returns
+        -------
+        Union[Dict[str, Union[str, Callable]], Callable]
+            Processed event
+
+        Raises
+        ------
+        MetaflowException
+            If the event is not in the correct format
+        """
         if is_stringish(event):
             return {"name": str(event)}
         elif isinstance(event, dict):
@@ -82,12 +101,26 @@ class TriggerDecorator(FlowDecorator):
                 event["name"], DeployTimeField
             ):
                 event["name"] = DeployTimeField(
-                    "event_name", str, None, event["name"], False
+                    "event_name",
+                    str,
+                    None,
+                    event["name"],
+                    False,
+                    print_representation=str(event["name"]),
                 )
-            event["parameters"] = self.process_parameters(event.get("parameters", {}))
+            event["parameters"] = self.process_parameters(
+                event.get("parameters", {}), event["name"]
+            )
             return event
         elif callable(event) and not isinstance(event, DeployTimeField):
-            return DeployTimeField("event", [str, dict], None, event, False)
+            return DeployTimeField(
+                "event",
+                [str, dict],
+                None,
+                event,
+                False,
+                print_representation=str(event),
+            )
         else:
             raise MetaflowException(
                 "Incorrect format for *event* attribute in *@trigger* decorator. "
@@ -96,47 +129,67 @@ class TriggerDecorator(FlowDecorator):
                 "'parameters': {'alpha': 'beta'}})"
             )
 
-    def process_parameters(self, parameters):
+    def process_parameters(self, parameters, event_name):
+        """
+        Process the parameters for an event and return a dictionary of parameter mappings if
+        parameters was statically defined or a function if deploy-time trigger.
+
+        Parameters
+        ----------
+        Parameters : Union[Dict[str, str], List[Union[str, Tuple[str, str]]], Callable]
+            Parameters to process
+
+        event_name : Union[str, callable]
+            Name of the event
+
+        Returns
+        -------
+        Union[Dict[str, str], Callable]
+            Processed parameters
+
+        Raises
+        ------
+        MetaflowException
+            If the parameters are not in the correct format
+        """
         new_param_values = {}
-        if isinstance(parameters, (list, tuple)):
+        if isinstance(parameters, list):
             for mapping in parameters:
                 if is_stringish(mapping):
+                    # param_name
                     new_param_values[mapping] = mapping
-                elif callable(mapping) and not isinstance(mapping, DeployTimeField):
-                    mapping = DeployTimeField(
-                        "parameter_val", str, None, mapping, False
-                    )
-                    new_param_values[mapping] = mapping
-                elif isinstance(mapping, (list, tuple)) and len(mapping) == 2:
-                    if callable(mapping[0]) and not isinstance(
-                        mapping[0], DeployTimeField
-                    ):
-                        mapping[0] = DeployTimeField(
-                            "parameter_val", str, None, mapping[0], False
-                        )
-                    if callable(mapping[1]) and not isinstance(
-                        mapping[1], DeployTimeField
-                    ):
-                        mapping[1] = DeployTimeField(
-                            "parameter_val", str, None, mapping[1], False
+                elif isinstance(mapping, tuple) and len(mapping) == 2:
+                    # (param_name, field_name)
+                    param_name, field_name = mapping
+                    if not is_stringish(param_name) or not is_stringish(field_name):
+                        raise MetaflowException(
+                            f"The *parameters* attribute for event {event_name} is invalid. "
+                            "It should be a list/tuple of strings and lists/tuples of size 2."
                         )
-                    new_param_values[mapping[0]] = mapping[1]
+                    new_param_values[param_name] = field_name
                 else:
                     raise MetaflowException(
                         "The *parameters* attribute for event is invalid. "
                         "It should be a list/tuple of strings and lists/tuples of size 2"
                     )
-        elif callable(parameters) and not isinstance(parameters, DeployTimeField):
-            return DeployTimeField(
-                "parameters", [list, dict, tuple], None, parameters, False
-            )
         elif isinstance(parameters, dict):
             for key, value in parameters.items():
-                if callable(key) and not isinstance(key, DeployTimeField):
-                    key = DeployTimeField("flow_parameter", str, None, key, False)
-                if callable(value) and not isinstance(value, DeployTimeField):
-                    value = DeployTimeField("signal_parameter", str, None, value, False)
+                if not is_stringish(key) or not is_stringish(value):
+                    raise MetaflowException(
+                        f"The *parameters* attribute for event {event_name} is invalid. "
+                        "It should be a dictionary of string keys and string values."
+                    )
                 new_param_values[key] = value
+        elif callable(parameters) and not isinstance(parameters, DeployTimeField):
+            # func
+            return DeployTimeField(
+                "parameters",
+                [list, dict, tuple],
+                None,
+                parameters,
+                False,
+                print_representation=str(parameters),
+            )
         return new_param_values
 
     def flow_init(
@@ -158,7 +211,7 @@ class TriggerDecorator(FlowDecorator):
             )
         elif self.attributes["event"]:
             event = self.attributes["event"]
-            processed_event = self.process_event_name(event)
+            processed_event = self.process_event(event)
             self.triggers.append(processed_event)
         elif self.attributes["events"]:
             # events attribute supports the following formats -
@@ -169,13 +222,18 @@ class TriggerDecorator(FlowDecorator):
             if isinstance(self.attributes["events"], list):
                 # process every event in events
                 for event in self.attributes["events"]:
-                    processed_event = self.process_event_name(event)
+                    processed_event = self.process_event(event)
                     self.triggers.append(processed_event)
             elif callable(self.attributes["events"]) and not isinstance(
                 self.attributes["events"], DeployTimeField
             ):
                 trig = DeployTimeField(
-                    "events", list, None, self.attributes["events"], False
+                    "events",
+                    list,
+                    None,
+                    self.attributes["events"],
+                    False,
+                    print_representation=str(self.attributes["events"]),
                 )
                 self.triggers.append(trig)
             else:
@@ -208,101 +266,40 @@ class TriggerDecorator(FlowDecorator):
 
     def format_deploytime_value(self):
         new_triggers = []
+
+        # First pass to evaluate DeployTimeFields
         for trigger in self.triggers:
             # Case where trigger is a function that returns a list of events
             # Need to do this bc we need to iterate over list later
             if isinstance(trigger, DeployTimeField):
                 evaluated_trigger = deploy_time_eval(trigger)
-                if isinstance(evaluated_trigger, dict):
-                    trigger = evaluated_trigger
-                elif isinstance(evaluated_trigger, str):
-                    trigger = {"name": evaluated_trigger}
                 if isinstance(evaluated_trigger, list):
-                    for trig in evaluated_trigger:
-                        if is_stringish(trig):
-                            new_triggers.append({"name": trig})
-                        else:  # dict or another deploytimefield
-                            new_triggers.append(trig)
+                    for event in evaluated_trigger:
+                        new_triggers.append(self.process_event(event))
                 else:
-                    new_triggers.append(trigger)
+                    new_triggers.append(self.process_event(evaluated_trigger))
             else:
                 new_triggers.append(trigger)
 
+        # Second pass to evaluate names
+        for trigger in new_triggers:
+            name = trigger.get("name")
+            if isinstance(name, DeployTimeField):
+                trigger["name"] = deploy_time_eval(name)
+                if not is_stringish(trigger["name"]):
+                    raise MetaflowException(
+                        f"The *name* attribute for event {trigger} is not a valid string"
+                    )
+
+        # third pass to evaluate parameters
+        for trigger in new_triggers:
+            parameters = trigger.get("parameters", {})
+            if isinstance(parameters, DeployTimeField):
+                parameters_eval = deploy_time_eval(parameters)
+                parameters = self.process_parameters(parameters_eval, trigger["name"])
+                trigger["parameters"] = parameters
+
         self.triggers = new_triggers
-        for trigger in self.triggers:
-            old_trigger = trigger
-            trigger_params = trigger.get("parameters", {})
-            # Case where param is a function (can return list or dict)
-            if isinstance(trigger_params, DeployTimeField):
-                trigger_params = deploy_time_eval(trigger_params)
-            # If params is a list of strings, convert to dict with same key and value
-            if isinstance(trigger_params, (list, tuple)):
-                new_trigger_params = {}
-                for mapping in trigger_params:
-                    if is_stringish(mapping) or callable(mapping):
-                        new_trigger_params[mapping] = mapping
-                    elif callable(mapping) and not isinstance(mapping, DeployTimeField):
-                        mapping = DeployTimeField(
-                            "parameter_val", str, None, mapping, False
-                        )
-                        new_trigger_params[mapping] = mapping
-                    elif isinstance(mapping, (list, tuple)) and len(mapping) == 2:
-                        if callable(mapping[0]) and not isinstance(
-                            mapping[0], DeployTimeField
-                        ):
-                            mapping[0] = DeployTimeField(
-                                "parameter_val",
-                                str,
-                                None,
-                                mapping[1],
-                                False,
-                            )
-                        if callable(mapping[1]) and not isinstance(
-                            mapping[1], DeployTimeField
-                        ):
-                            mapping[1] = DeployTimeField(
-                                "parameter_val",
-                                str,
-                                None,
-                                mapping[1],
-                                False,
-                            )
-
-                        new_trigger_params[mapping[0]] = mapping[1]
-                    else:
-                        raise MetaflowException(
-                            "The *parameters* attribute for event '%s' is invalid. "
-                            "It should be a list/tuple of strings and lists/tuples "
-                            "of size 2" % self.attributes["event"]["name"]
-                        )
-                trigger_params = new_trigger_params
-            trigger["parameters"] = trigger_params
-
-            trigger_name = trigger.get("name")
-            # Case where just the name is a function (always a str)
-            if isinstance(trigger_name, DeployTimeField):
-                trigger_name = deploy_time_eval(trigger_name)
-                trigger["name"] = trigger_name
-
-            # Third layer
-            # {name:, parameters:[func, ..., ...]}
-            # {name:, parameters:{func : func2}}
-            for trigger in self.triggers:
-                old_trigger = trigger
-                trigger_params = trigger.get("parameters", {})
-                new_trigger_params = {}
-                for key, value in trigger_params.items():
-                    if isinstance(value, DeployTimeField) and key is value:
-                        evaluated_param = deploy_time_eval(value)
-                        new_trigger_params[evaluated_param] = evaluated_param
-                    elif isinstance(value, DeployTimeField):
-                        new_trigger_params[key] = deploy_time_eval(value)
-                    elif isinstance(key, DeployTimeField):
-                        new_trigger_params[deploy_time_eval(key)] = value
-                    else:
-                        new_trigger_params[key] = value
-                trigger["parameters"] = new_trigger_params
-            self.triggers[self.triggers.index(old_trigger)] = trigger
 
 
 class TriggerOnFinishDecorator(FlowDecorator):
@@ -402,7 +399,14 @@ class TriggerOnFinishDecorator(FlowDecorator):
             if callable(flow) and not isinstance(
                 self.attributes["flow"], DeployTimeField
             ):
-                trig = DeployTimeField("fq_name", [str, dict], None, flow, False)
+                trig = DeployTimeField(
+                    "fq_name",
+                    [str, dict],
+                    None,
+                    flow,
+                    False,
+                    print_representation=str(flow),
+                )
                 self.triggers.append(trig)
             else:
                 self.triggers.extend(self._parse_static_triggers([flow]))
@@ -411,7 +415,9 @@ class TriggerOnFinishDecorator(FlowDecorator):
             #     1. flows=['FooFlow', 'BarFlow']
             flows = self.attributes["flows"]
             if callable(flows) and not isinstance(flows, DeployTimeField):
-                trig = DeployTimeField("flows", list, None, flows, False)
+                trig = DeployTimeField(
+                    "flows", list, None, flows, False, print_representation=str(flows)
+                )
                 self.triggers.append(trig)
             elif isinstance(flows, list):
                 self.triggers.extend(self._parse_static_triggers(flows))

metaflow/plugins/exit_hook/exit_hook_decorator.py

@@ -0,0 +1,46 @@
+from metaflow.decorators import FlowDecorator
+from metaflow.exception import MetaflowException
+
+
+class ExitHookDecorator(FlowDecorator):
+    name = "exit_hook"
+    allow_multiple = True
+
+    defaults = {
+        "on_success": [],
+        "on_error": [],
+        "options": {},
+    }
+
+    def flow_init(
+        self, flow, graph, environment, flow_datastore, metadata, logger, echo, options
+    ):
+        on_success = self.attributes["on_success"]
+        on_error = self.attributes["on_error"]
+
+        if not on_success and not on_error:
+            raise MetaflowException(
+                "Choose at least one of the options on_success/on_error"
+            )
+
+        self.success_hooks = []
+        self.error_hooks = []
+        for success_fn in on_success:
+            if isinstance(success_fn, str):
+                self.success_hooks.append(success_fn)
+            elif callable(success_fn):
+                self.success_hooks.append(success_fn.__name__)
+            else:
+                raise ValueError(
+                    "Exit hooks inside 'on_success' must be a function or a string referring to the function"
+                )
+
+        for error_fn in on_error:
+            if isinstance(error_fn, str):
+                self.error_hooks.append(error_fn)
+            elif callable(error_fn):
+                self.error_hooks.append(error_fn.__name__)
+            else:
+                raise ValueError(
+                    "Exit hooks inside 'on_error' must be a function or a string referring to the function"
+                )

metaflow/plugins/exit_hook/exit_hook_script.py

@@ -0,0 +1,52 @@
+import os
+import inspect
+import importlib
+import sys
+
+
+def main(flow_file, fn_name_or_path, run_pathspec):
+    hook_fn = None
+
+    try:
+        module_path, function_name = fn_name_or_path.rsplit(".", 1)
+        module = importlib.import_module(module_path)
+        hook_fn = getattr(module, function_name)
+    except (ImportError, AttributeError, ValueError):
+        try:
+            module_name = os.path.splitext(os.path.basename(flow_file))[0]
+            spec = importlib.util.spec_from_file_location(module_name, flow_file)
+            module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(module)
+            hook_fn = getattr(module, fn_name_or_path)
+        except (AttributeError, IOError) as e:
+            print(
+                f"[exit_hook] Could not load function '{fn_name_or_path}' "
+                f"as an import path or from '{flow_file}': {e}"
+            )
+            sys.exit(1)
+
+    argspec = inspect.getfullargspec(hook_fn)
+
+    # Check if fn expects a run object as an arg.
+    if "run" in argspec.args or argspec.varkw is not None:
+        from metaflow import Run
+
+        try:
+            _run = Run(run_pathspec, _namespace_check=False)
+        except Exception as ex:
+            print(ex)
+            _run = None
+
+        hook_fn(run=_run)
+    else:
+        hook_fn()
+
+
+if __name__ == "__main__":
+    try:
+        flow_file, fn_name, run_pathspec = sys.argv[1:4]
+    except Exception:
+        print("Usage: exit_hook_script.py <flow_file> <function_name> <run_pathspec>")
+        sys.exit(1)
+
+    main(flow_file, fn_name, run_pathspec)