ob-metaflow 2.11.4.9__py2.py3-none-any.whl → 2.11.9.1__py2.py3-none-any.whl

metaflow/plugins/env_escape/client.py

@@ -34,8 +34,12 @@ from .communication.channel import Channel
 from .communication.socket_bytestream import SocketByteStream
 
 from .data_transferer import DataTransferer, ObjReference
-from .exception_transferer import load_exception
-from .override_decorators import LocalAttrOverride, LocalException, LocalOverride
+from .exception_transferer import ExceptionMetaClass, load_exception
+from .override_decorators import (
+    LocalAttrOverride,
+    LocalExceptionDeserializer,
+    LocalOverride,
+)
 from .stub import create_class
 from .utils import get_canonical_name
 
@@ -193,28 +197,41 @@ class Client(object):
         self._proxied_classes = {
             k: None
             for k in itertools.chain(
-                response[FIELD_CONTENT]["classes"], response[FIELD_CONTENT]["proxied"]
+                response[FIELD_CONTENT]["classes"],
+                response[FIELD_CONTENT]["proxied"],
+                (e[0] for e in response[FIELD_CONTENT]["exceptions"]),
             )
         }
 
+        self._exception_hierarchy = dict(response[FIELD_CONTENT]["exceptions"])
+        self._proxied_classnames = set(response[FIELD_CONTENT]["classes"]).union(
+            response[FIELD_CONTENT]["proxied"]
+        )
+        self._aliases = response[FIELD_CONTENT]["aliases"]
+
         # Determine all overrides
         self._overrides = {}
         self._getattr_overrides = {}
         self._setattr_overrides = {}
-        self._exception_overrides = {}
+        self._exception_deserializers = {}
         for override in override_values:
             if isinstance(override, (LocalOverride, LocalAttrOverride)):
                 for obj_name, obj_funcs in override.obj_mapping.items():
-                    if obj_name not in self._proxied_classes:
+                    canonical_name = get_canonical_name(obj_name, self._aliases)
+                    if canonical_name not in self._proxied_classes:
                         raise ValueError(
                             "%s does not refer to a proxied or override type" % obj_name
                         )
                     if isinstance(override, LocalOverride):
-                        override_dict = self._overrides.setdefault(obj_name, {})
+                        override_dict = self._overrides.setdefault(canonical_name, {})
                     elif override.is_setattr:
-                        override_dict = self._setattr_overrides.setdefault(obj_name, {})
+                        override_dict = self._setattr_overrides.setdefault(
+                            canonical_name, {}
+                        )
                     else:
-                        override_dict = self._getattr_overrides.setdefault(obj_name, {})
+                        override_dict = self._getattr_overrides.setdefault(
+                            canonical_name, {}
+                        )
                     if isinstance(obj_funcs, str):
                         obj_funcs = (obj_funcs,)
                     for name in obj_funcs:
@@ -223,11 +240,18 @@ class Client(object):
                                 "%s was already overridden for %s" % (name, obj_name)
                             )
                         override_dict[name] = override.func
-            if isinstance(override, LocalException):
-                cur_ex = self._exception_overrides.get(override.class_path, None)
-                if cur_ex is not None:
-                    raise ValueError("Exception %s redefined" % override.class_path)
-                self._exception_overrides[override.class_path] = override.wrapped_class
+            if isinstance(override, LocalExceptionDeserializer):
+                canonical_name = get_canonical_name(override.class_path, self._aliases)
+                if canonical_name not in self._exception_hierarchy:
+                    raise ValueError(
+                        "%s does not refer to an exception type" % override.class_path
+                    )
+                cur_des = self._exception_deserializers.get(canonical_name, None)
+                if cur_des is not None:
+                    raise ValueError(
+                        "Exception %s has multiple deserializers" % override.class_path
+                    )
+                self._exception_deserializers[canonical_name] = override.deserializer
 
         # Proxied standalone functions are functions that are proxied
         # as part of other objects like defaultdict for which we create a
@@ -243,8 +267,6 @@ class Client(object):
             "aliases": response[FIELD_CONTENT]["aliases"],
         }
 
-        self._aliases = response[FIELD_CONTENT]["aliases"]
-
     def __del__(self):
         self.cleanup()
 
@@ -288,8 +310,9 @@ class Client(object):
     def get_exports(self):
         return self._export_info
 
-    def get_local_exception_overrides(self):
-        return self._exception_overrides
+    def get_exception_deserializer(self, name):
+        cannonical_name = get_canonical_name(name, self._aliases)
+        return self._exception_deserializers.get(cannonical_name)
 
     def stub_request(self, stub, request_type, *args, **kwargs):
         # Encode the operation to send over the wire and wait for the response
@@ -313,7 +336,7 @@ class Client(object):
         if response_type == MSG_REPLY:
             return self.decode(response[FIELD_CONTENT])
         elif response_type == MSG_EXCEPTION:
-            raise load_exception(self._datatransferer, response[FIELD_CONTENT])
+            raise load_exception(self, response[FIELD_CONTENT])
         elif response_type == MSG_INTERNAL_ERROR:
             raise RuntimeError(
                 "Error in the server runtime:\n\n===== SERVER TRACEBACK =====\n%s"
@@ -334,10 +357,27 @@ class Client(object):
         # this connection will be converted to a local stub.
         return self._datatransferer.load(json_obj)
 
-    def get_local_class(self, name, obj_id=None):
+    def get_local_class(
+        self, name, obj_id=None, is_returned_exception=False, is_parent=False
+    ):
         # Gets (and creates if needed), the class mapping to the remote
         # class of name 'name'.
+
+        # We actually deal with four types of classes:
+        # - proxied functions
+        # - classes that are proxied regular classes AND proxied exceptions
+        # - classes that are proxied regular classes AND NOT proxied exceptions
+        # - classes that are NOT proxied regular classes AND are proxied exceptions
         name = get_canonical_name(name, self._aliases)
+
+        def name_to_parent_name(name):
+            return "parent:%s" % name
+
+        if is_parent:
+            lookup_name = name_to_parent_name(name)
+        else:
+            lookup_name = name
+
         if name == "function":
             # Special handling of pickled functions. We create a new class that
             # simply has a __call__ method that will forward things back to
@@ -346,17 +386,87 @@ class Client(object):
                 raise RuntimeError("Local function unpickling without an object ID")
             if obj_id not in self._proxied_standalone_functions:
                 self._proxied_standalone_functions[obj_id] = create_class(
-                    self, "__function_%s" % obj_id, {}, {}, {}, {"__call__": ""}
+                    self,
+                    "__main__.__function_%s" % obj_id,
+                    {},
+                    {},
+                    {},
+                    {"__call__": ""},
+                    [],
                 )
             return self._proxied_standalone_functions[obj_id]
+        local_class = self._proxied_classes.get(lookup_name, None)
+        if local_class is not None:
+            return local_class
+
+        is_proxied_exception = name in self._exception_hierarchy
+        is_proxied_non_exception = name in self._proxied_classnames
+
+        if not is_proxied_exception and not is_proxied_non_exception:
+            if is_returned_exception or is_parent:
+                # In this case, it may be a local exception that we need to
+                # recreate
+                try:
+                    ex_module, ex_name = name.rsplit(".", 1)
+                    __import__(ex_module, None, None, "*")
+                except Exception:
+                    pass
+                if ex_module in sys.modules and issubclass(
+                    getattr(sys.modules[ex_module], ex_name), BaseException
+                ):
+                    # This is a local exception that we can recreate
+                    local_exception = getattr(sys.modules[ex_module], ex_name)
+                    wrapped_exception = ExceptionMetaClass(
+                        ex_name,
+                        (local_exception,),
+                        dict(getattr(local_exception, "__dict__", {})),
+                    )
+                    wrapped_exception.__module__ = ex_module
+                    self._proxied_classes[lookup_name] = wrapped_exception
+                    return wrapped_exception
 
-        if name not in self._proxied_classes:
             raise ValueError("Class '%s' is not known" % name)
-        local_class = self._proxied_classes[name]
-        if local_class is None:
-            # We need to build up this class. To do so, we take everything that the
-            # remote class has and remove UNSUPPORTED things and overridden things
+
+        # At this stage:
+        # - we don't have a local_class for this
+        # - it is not an inbuilt exception so it is either a proxied exception, a
+        #   proxied class or a proxied object that is both an exception and a class.
+
+        parents = []
+        if is_proxied_exception:
+            # If exception, we need to get the parents from the exception
+            ex_parents = self._exception_hierarchy[name]
+            for parent in ex_parents:
+                # We always consider it to be an exception so that we wrap even non
+                # proxied builtins exceptions
+                parents.append(self.get_local_class(parent, is_parent=True))
+        # For regular classes, we get what it exposes from the server
+        if is_proxied_non_exception:
             remote_methods = self.stub_request(None, OP_GETMETHODS, name)
+        else:
+            remote_methods = {}
+
+        parent_local_class = None
+        local_class = None
+        if is_proxied_exception:
+            # If we are a proxied exception AND a proxied class, we create two classes:
+            # actually:
+            #  - the class itself (which is a stub)
+            #  - the class in the capacity of a parent class (to another exception
+            #    presumably). The reason for this is that if we have an exception/proxied
+            #    class A and another B and B inherits from A, the MRO order would be all
+            #    wrong since both A and B would also inherit from `Stub`. Here what we
+            #    do is:
+            #      - A_parent inherits from the actual parents of A (let's assume a
+            #        builtin exception)
+            #      - A inherits from (Stub, A_parent)
+            #      - B_parent inherits from A_parent and the builtin Exception
+            #      - B inherits from (Stub, B_parent)
+            ex_module, ex_name = name.rsplit(".", 1)
+            parent_local_class = ExceptionMetaClass(ex_name, (*parents,), {})
+            parent_local_class.__module__ = ex_module
+
+        if is_proxied_non_exception:
             local_class = create_class(
                 self,
                 name,
@@ -364,9 +474,26 @@ class Client(object):
                 self._getattr_overrides.get(name, {}),
                 self._setattr_overrides.get(name, {}),
                 remote_methods,
+                (parent_local_class,) if parent_local_class else None,
             )
+        if parent_local_class:
+            self._proxied_classes[name_to_parent_name(name)] = parent_local_class
+        if local_class:
             self._proxied_classes[name] = local_class
-        return local_class
+        else:
+            # This is for the case of pure proxied exceptions -- we want the lookup of
+            # foo.MyException to be the same class as looking of foo.MyException as a parent
+            # of another exception so `isinstance` works properly
+            self._proxied_classes[name] = parent_local_class
+
+        if is_parent:
+            # This should never happen but making sure
+            if not parent_local_class:
+                raise RuntimeError(
+                    "Exception parent class %s is not a proxied exception" % name
+                )
+            return parent_local_class
+        return self._proxied_classes[name]
 
     def can_pickle(self, obj):
         return getattr(obj, "___connection___", None) == self
@@ -395,7 +522,7 @@ class Client(object):
         obj_id = obj.identifier
         local_instance = self._proxied_objects.get(obj_id)
         if not local_instance:
-            local_class = self.get_local_class(remote_class_name, obj_id)
+            local_class = self.get_local_class(remote_class_name, obj_id=obj_id)
             local_instance = local_class(self, remote_class_name, obj_id)
         return local_instance
 

metaflow/plugins/env_escape/client_modules.py

@@ -7,7 +7,6 @@ import sys
 
 from .consts import OP_CALLFUNC, OP_GETVAL, OP_SETVAL
 from .client import Client
-from .override_decorators import LocalException
 from .utils import get_canonical_name
 
 
@@ -16,7 +15,7 @@ def _clean_client(client):
 
 
 class _WrappedModule(object):
-    def __init__(self, loader, prefix, exports, exception_classes, client):
+    def __init__(self, loader, prefix, exports, client):
         self._loader = loader
         self._prefix = prefix
         self._client = client
@@ -24,19 +23,20 @@ class _WrappedModule(object):
             r"^%s\.([a-zA-Z_][a-zA-Z0-9_]*)$" % prefix.replace(".", r"\.")  # noqa W605
         )
         self._exports = {}
-        self._aliases = exports["aliases"]
+        self._aliases = exports.get("aliases", [])
         for k in ("classes", "functions", "values"):
             result = []
-            for item in exports[k]:
+            for item in exports.get(k, []):
                 m = is_match.match(item)
                 if m:
                     result.append(m.group(1))
             self._exports[k] = result
-        self._exception_classes = {}
-        for k, v in exception_classes.items():
-            m = is_match.match(k)
+        result = []
+        for item, _ in exports.get("exceptions", []):
+            m = is_match.match(item)
             if m:
-                self._exception_classes[m.group(1)] = v
+                result.append(m.group(1))
+        self._exports["exceptions"] = result
 
     def __getattr__(self, name):
         if name == "__loader__":
@@ -50,8 +50,8 @@ class _WrappedModule(object):
         name = get_canonical_name(self._prefix + "." + name, self._aliases)[
             len(self._prefix) + 1 :
         ]
-        if name in self._exports["classes"]:
-            # We load classes lazily
+        if name in self._exports["classes"] or name in self._exports["exceptions"]:
+            # We load classes and exceptions lazily
             return self._client.get_local_class("%s.%s" % (self._prefix, name))
         elif name in self._exports["functions"]:
             # TODO: Grab doc back from the remote side like in _make_method
@@ -67,8 +67,6 @@ class _WrappedModule(object):
             return self._client.stub_request(
                 None, OP_GETVAL, "%s.%s" % (self._prefix, name)
             )
-        elif name in self._exception_classes:
-            return self._exception_classes[name]
         else:
             # Try to see if this is a submodule that we can load
             m = None
@@ -173,7 +171,6 @@ class ModuleImporter(object):
 
             # Get information about overrides and what the server knows about
             exports = self._client.get_exports()
-            ex_overrides = self._client.get_local_exception_overrides()
 
             prefixes = set()
             export_classes = exports.get("classes", [])
@@ -182,42 +179,13 @@ class ModuleImporter(object):
             export_exceptions = exports.get("exceptions", [])
             self._aliases = exports.get("aliases", {})
             for name in itertools.chain(
-                export_classes, export_functions, export_values
+                export_classes,
+                export_functions,
+                export_values,
+                (e[0] for e in export_exceptions),
             ):
                 splits = name.rsplit(".", 1)
                 prefixes.add(splits[0])
-
-            # Now look at the exceptions coming from the server
-            formed_exception_classes = {}
-            for ex_name, ex_parents in export_exceptions:
-                # Exception is a tuple (name, (parents,))
-                # Exceptions are also given in order of instantiation (ie: the
-                # server already topologically sorted them)
-                ex_class_dict = ex_overrides.get(ex_name, None)
-                if ex_class_dict is None:
-                    ex_class_dict = {}
-                else:
-                    ex_class_dict = dict(ex_class_dict.__dict__)
-                parents = []
-                for fake_base in ex_parents:
-                    if fake_base.startswith("builtins."):
-                        # This is something we know of here
-                        parents.append(eval(fake_base[9:]))
-                    else:
-                        # It's in formed_classes
-                        parents.append(formed_exception_classes[fake_base])
-                splits = ex_name.rsplit(".", 1)
-                ex_class_dict["__user_defined__"] = set(ex_class_dict.keys())
-                new_class = type(splits[1], tuple(parents), ex_class_dict)
-                new_class.__module__ = splits[0]
-                new_class.__name__ = splits[1]
-                formed_exception_classes[ex_name] = new_class
-
-            # Now update prefixes as needed
-            for name in formed_exception_classes:
-                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
@@ -235,7 +203,7 @@ class ModuleImporter(object):
             self._handled_modules = {}
             for prefix in prefixes:
                 self._handled_modules[prefix] = _WrappedModule(
-                    self, prefix, exports, formed_exception_classes, self._client
+                    self, prefix, exports, self._client
                 )
         canonical_fullname = get_canonical_name(fullname, self._aliases)
         # Modules are created canonically but we need to return something for any

metaflow/plugins/env_escape/configurations/emulate_test_lib/overrides.py

@@ -5,56 +5,44 @@ from metaflow.plugins.env_escape.override_decorators import (
     remote_override,
     remote_getattr_override,
     remote_setattr_override,
-    local_exception,
+    local_exception_deserialize,
     remote_exception_serialize,
 )
 
 
 @local_override({"test_lib.TestClass1": "print_value"})
 def local_print_value(stub, func):
-    print("Encoding before sending to server")
     v = func()
-    print("Adding 5")
     return v + 5
 
 
 @remote_override({"test_lib.TestClass1": "print_value"})
 def remote_print_value(obj, func):
-    print("Decoding from client")
     v = func()
-    print("Encoding for client")
-    return v
+    return v + 3
 
 
 @local_getattr_override({"test_lib.TestClass1": "override_value"})
 def local_get_value2(stub, name, func):
-    print("In local getattr override for %s" % name)
     r = func()
-    print("In local getattr override, got %s" % r)
-    return r
-
-
-@local_setattr_override({"test_lib.TestClass1": "override_value"})
-def local_set_value2(stub, name, func, v):
-    print("In local setattr override for %s" % name)
-    r = func(v)
-    print("In local setattr override, got %s" % r)
-    return r
+    return r + 5
 
 
 @remote_getattr_override({"test_lib.TestClass1": "override_value"})
 def remote_get_value2(obj, name):
-    print("In remote getattr override for %s" % name)
     r = getattr(obj, name)
-    print("In remote getattr override, got %s" % r)
+    return r + 3
+
+
+@local_setattr_override({"test_lib.TestClass1": "override_value"})
+def local_set_value2(stub, name, func, v):
+    r = func(v + 5)
     return r
 
 
 @remote_setattr_override({"test_lib.TestClass1": "override_value"})
 def remote_set_value2(obj, name, v):
-    print("In remote setattr override for %s" % name)
-    r = setattr(obj, name, v)
-    print("In remote setattr override, got %s" % r)
+    r = setattr(obj, name, v + 3)
     return r
 
 
@@ -63,30 +51,31 @@ def unsupported_method(stub, func, *args, **kwargs):
     return NotImplementedError("Just because")
 
 
-@local_override({"test_lib.package.TestClass3": "thirdfunction"})
-def iamthelocalthird(stub, func, val):
-    print("Locally the Third")
-    v = func(val)
-    return v
+@local_exception_deserialize("test_lib.SomeException")
+def some_exception_deserialize(ex, json_obj):
+    ex.user_value = json_obj
 
 
-@remote_override({"test_lib.package.TestClass3": "thirdfunction"})
-def iamtheremotethird(obj, func, val):
-    print("Remotely the Third")
-    v = func(val)
-    return v
+@remote_exception_serialize("test_lib.SomeException")
+def some_exception_serialize(ex):
+    return 42
 
 
-@local_exception("test_lib.SomeException")
-class SomeException:
-    def __str__(self):
-        parent_val = super(self.__realclass__, self).__str__()
-        return parent_val + " In SomeException str override: %s" % self.user_value
+@local_exception_deserialize("test_lib.ExceptionAndClass")
+def exception_and_class_deserialize(ex, json_obj):
+    ex.user_value = json_obj
 
-    def _deserialize_user(self, json_obj):
-        self.user_value = json_obj
 
+@remote_exception_serialize("test_lib.ExceptionAndClass")
+def exception_and_class_serialize(ex):
+    return 43
 
-@remote_exception_serialize("test_lib.SomeException")
-def some_exception_serialize(ex):
-    return 42
+
+@local_exception_deserialize("test_lib.ExceptionAndClassChild")
+def exception_and_class_child_deserialize(ex, json_obj):
+    ex.user_value = json_obj
+
+
+@remote_exception_serialize("test_lib.ExceptionAndClassChild")
+def exception_and_class_child_serialize(ex):
+    return 44

metaflow/plugins/env_escape/configurations/emulate_test_lib/server_mappings.py

@@ -10,17 +10,22 @@ sys.path.append(
 import test_lib as lib
 
 EXPORTED_CLASSES = {
-    "test_lib": {
+    ("test_lib", "test_lib.alias"): {
         "TestClass1": lib.TestClass1,
         "TestClass2": lib.TestClass2,
-        "package.TestClass3": lib.TestClass3,
+        "BaseClass": lib.BaseClass,
+        "ChildClass": lib.ChildClass,
+        "ExceptionAndClass": lib.ExceptionAndClass,
+        "ExceptionAndClassChild": lib.ExceptionAndClassChild,
     }
 }
 
 EXPORTED_EXCEPTIONS = {
-    "test_lib": {
+    ("test_lib", "test_lib.alias"): {
         "SomeException": lib.SomeException,
         "MyBaseException": lib.MyBaseException,
+        "ExceptionAndClass": lib.ExceptionAndClass,
+        "ExceptionAndClassChild": lib.ExceptionAndClassChild,
     }
 }
 

metaflow/plugins/env_escape/configurations/test_lib_impl/test_lib.py

@@ -1,4 +1,5 @@
 import functools
+from html.parser import HTMLParser
 
 
 class MyBaseException(Exception):
@@ -9,8 +10,48 @@ class SomeException(MyBaseException):
     pass
 
 
-class TestClass1(object):
+class ExceptionAndClass(MyBaseException):
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def method_on_exception(self):
+        return "method_on_exception"
+
+    def __str__(self):
+        return "ExceptionAndClass Str: %s" % super().__str__()
+
+
+class ExceptionAndClassChild(ExceptionAndClass):
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def method_on_child_exception(self):
+        return "method_on_child_exception"
+
+    def __str__(self):
+        return "ExceptionAndClassChild Str: %s" % super().__str__()
+
 
+class BaseClass(HTMLParser):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._output = []
+
+    def handle_starttag(self, tag, attrs):
+        self._output.append(tag)
+        return super().handle_starttag(tag, attrs)
+
+    def get_output(self):
+        return self._output
+
+
+class ChildClass(BaseClass):
+    def handle_endtag(self, tag):
+        self._output.append(tag)
+        return super().handle_endtag(tag)
+
+
+class TestClass1(object):
     cls_object = 25
 
     def __init__(self, value):
@@ -41,11 +82,11 @@ class TestClass1(object):
         return TestClass2(self._value, stride, count)
 
     @staticmethod
-    def somethingstatic(val):
+    def static_method(val):
         return val + 42
 
     @classmethod
-    def somethingclass(cls):
+    def class_method(cls):
         return cls.cls_object
 
     @property
@@ -56,13 +97,42 @@ class TestClass1(object):
     def override_value(self, value):
         self._value2 = value
 
+    def __hidden(self, name, value):
+        setattr(self, name, value)
+
+    def weird_indirection(self, name):
+        return functools.partial(self.__hidden, name)
+
+    def returnChild(self):
+        return ChildClass()
+
+    def raiseOrReturnValueError(self, doRaise=False):
+        if doRaise:
+            raise ValueError("I raised")
+        return ValueError("I returned")
+
+    def raiseOrReturnSomeException(self, doRaise=False):
+        if doRaise:
+            raise SomeException("I raised")
+        return SomeException("I returned")
+
+    def raiseOrReturnExceptionAndClass(self, doRaise=False):
+        if doRaise:
+            raise ExceptionAndClass("I raised")
+        return ExceptionAndClass("I returned")
+
+    def raiseOrReturnExceptionAndClassChild(self, doRaise=False):
+        if doRaise:
+            raise ExceptionAndClassChild("I raised")
+        return ExceptionAndClassChild("I returned")
+
 
 class TestClass2(object):
     def __init__(self, value, stride, count):
         self._mylist = [value + stride * i for i in range(count)]
 
     def something(self, val):
-        return "In Test2 with %s" % val
+        return "Test2:Something:%s" % val
 
     def __iter__(self):
         self._pos = 0
@@ -75,24 +145,6 @@ class TestClass2(object):
         raise StopIteration
 
 
-class TestClass3(object):
-    def __init__(self):
-        print("I am Class3")
-
-    def thirdfunction(self, val):
-        print("Got value: %s" % val)
-        # raise AttributeError("Some weird error")
-
-    def raiseSomething(self):
-        raise SomeException("Something went wrong")
-
-    def __hidden(self, name, value):
-        setattr(self, name, value)
-
-    def weird_indirection(self, name):
-        return functools.partial(self.__hidden, name)
-
-
 def test_func(*args, **kwargs):
     return "In test func"