Preserving Exception in the LazyEntity fetch (flyteorg#1412)

* Preserving Exception in the LazyEntity fetch Signed-off-by: Ketan Umare <[email protected]> * updated lint error Signed-off-by: Ketan Umare <[email protected]> * more tests Signed-off-by: Ketan Umare <[email protected]> Signed-off-by: Ketan Umare <[email protected]> Signed-off-by: byhsu <[email protected]>
ByronHsu · Feb 1, 2023 · 0cf7a35 · 0cf7a35
1 parent d3d258d
commit 0cf7a35
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 2 deletions.
diff --git a/flytekit/core/promise.py b/flytekit/core/promise.py
@@ -854,7 +854,8 @@ def create_and_link_node_from_remote(
     extra_inputs = used_inputs ^ set(kwargs.keys())
     if len(extra_inputs) > 0:
         raise _user_exceptions.FlyteAssertion(
-            "Too many inputs were specified for the interface.  Extra inputs were: {}".format(extra_inputs)
+            f"Too many inputs for [{entity.name}] Expected inputs: {typed_interface.inputs.keys()} "
+            f"- extra inputs: {extra_inputs}"
         )
 
     # Detect upstream nodes

diff --git a/flytekit/remote/lazy_entity.py b/flytekit/remote/lazy_entity.py
@@ -37,7 +37,12 @@ def entity(self) -> T:
         """
         with self._mutex:
             if self._entity is None:
-                self._entity = self._getter()
+                try:
+                    self._entity = self._getter()
+                except AttributeError as e:
+                    raise RuntimeError(
+                        f"Error downloading the entity {self._name}, (check original exception...)"
+                    ) from e
             return self._entity
 
     def __getattr__(self, item: str) -> typing.Any:

diff --git a/tests/flytekit/unit/remote/test_lazy_entity.py b/tests/flytekit/unit/remote/test_lazy_entity.py
@@ -63,3 +63,16 @@ def _getter():
     e.compile(ctx)
     assert e._entity is not None
     assert e.entity == dummy_task
+
+
+def test_lazy_loading_exception():
+    def _getter():
+        raise AttributeError("Error")
+
+    e = LazyEntity("x", _getter)
+    assert e.name == "x"
+    assert e._entity is None
+    with pytest.raises(RuntimeError) as exc:
+        assert e.blah
+
+    assert isinstance(exc.value.__cause__, AttributeError)