ob-metaflow 2.11.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/env_escape/server.py

@@ -264,6 +264,7 @@ class Server(object):
     def serve(self, path=None, port=None):
         # Open up a connection
         if path is not None:
+            # Keep the print line to facilitate debugging
             # print("SERVER: Starting at %s" % path)
             sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
             __try_op__("bind", sock.bind, BIND_RETRY, path)
@@ -503,6 +504,12 @@ class Server(object):
         class_type = self._known_classes.get(class_name)
         if class_type is None:
             raise ValueError("Unknown class %s" % class_name)
+        # Check if __init__ is overridden
+        override_mapping = self._overrides.get(class_type)
+        if override_mapping:
+            override_func = override_mapping.get("__init__")
+            if override_func:
+                return override_func(None, class_type, *args, **kwargs)
         return class_type(*args, **kwargs)
 
     def _handle_subclasscheck(self, target, class_name, otherclass_name, reverse=False):

metaflow/plugins/env_escape/stub.py

@@ -146,7 +146,7 @@ class Stub(with_metaclass(StubMetaClass, object)):
             return object.__getattribute__(self, name)
 
     def __getattr__(self, name):
-        if name in DELETED_ATTRS:
+        if name in DELETED_ATTRS or self.___is_returned_exception___:
             raise AttributeError()
         return fwd_request(self, OP_GETATTR, name)
 
@@ -166,6 +166,8 @@ class Stub(with_metaclass(StubMetaClass, object)):
         ):
             object.__setattr__(self, name, value)
         else:
+            if self.___is_returned_exception___:
+                raise AttributeError()
             fwd_request(self, OP_SETATTR, name, value)
 
     def __dir__(self):
@@ -274,9 +276,22 @@ class MetaWithConnection(StubMetaClass):
         if len(args) > 0 and id(args[0]) == id(cls.___class_connection___):
             return super(MetaWithConnection, cls).__call__(*args, **kwargs)
         else:
-            return cls.___class_connection___.stub_request(
-                None, OP_INIT, cls.___class_remote_class_name___, *args, **kwargs
-            )
+            if hasattr(cls, "__overriden_init__"):
+                return cls.__overriden_init__(
+                    None,
+                    functools.partial(
+                        cls.___class_connection___.stub_request,
+                        None,
+                        OP_INIT,
+                        cls.___class_remote_class_name___,
+                    ),
+                    *args,
+                    **kwargs
+                )
+            else:
+                return cls.___class_connection___.stub_request(
+                    None, OP_INIT, cls.___class_remote_class_name___, *args, **kwargs
+                )
 
     def __subclasscheck__(cls, subclass):
         subclass_name = "%s.%s" % (subclass.__module__, subclass.__name__)
@@ -379,7 +394,10 @@ def create_class(
             name = name[7:]
             method_type = CLASS_METHOD
         if name in overriden_methods:
-            if method_type == NORMAL_METHOD:
+            if name == "__init__":
+                class_dict["__overriden_init__"] = overriden_methods["__init__"]
+
+            elif method_type == NORMAL_METHOD:
                 class_dict[name] = (
                     lambda override, orig_method: lambda obj, *args, **kwargs: override(
                         obj, functools.partial(orig_method, obj), *args, **kwargs
@@ -410,6 +428,7 @@ def create_class(
             class_dict[name] = _make_method(
                 method_type, connection, class_name, name, doc
             )
+
     # Check for any getattr/setattr overrides
     special_attributes = set(getattr_overrides.keys())
     special_attributes.update(set(setattr_overrides.keys()))