From acfa7cc971dea25aeaaf0503507044df592e0fbf Mon Sep 17 00:00:00 2001 From: Ketan Umare <16888709+kumare3@users.noreply.github.com> Date: Fri, 20 Jan 2023 11:48:30 -0800 Subject: [PATCH] Preserving Exception in the LazyEntity fetch (#1412) * Preserving Exception in the LazyEntity fetch Signed-off-by: Ketan Umare * updated lint error Signed-off-by: Ketan Umare * more tests Signed-off-by: Ketan Umare Signed-off-by: Ketan Umare --- flytekit/core/promise.py | 3 ++- flytekit/remote/lazy_entity.py | 7 ++++++- tests/flytekit/unit/remote/test_lazy_entity.py | 13 +++++++++++++ 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/flytekit/core/promise.py b/flytekit/core/promise.py index 53048cb03f..bef86cc9ed 100644 --- 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 index b40c6e3ff7..4755aad99d 100644 --- 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 index 1ed191aea4..5328a2caf0 100644 --- 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)