Feature/3506: Improve Type Conversion errors, use rich to prettify er…

…ror messages (#1582)

* improve input type conversion error

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* fix lint

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* fix

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* add tests

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* add tests

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* add rich

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* fix lint

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* remove prototyping script, update loggers

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* update __init__.py

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* update logger

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* update logger

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* fix GE and pandera tests

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

* fix lint

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

---------

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>

.gitignore

@@ -32,3 +32,4 @@ htmlcov
 *.ipynb
 *dat
 docs/source/_tags/
+.hypothesis

dev-requirements.in

@@ -2,6 +2,7 @@
 
 git+https://github.com/flyteorg/pytest-flyte@main#egg=pytest-flyte
 coverage[toml]
+hypothesis
 joblib
 mock
 pytest

flytekit/__init__.py

@@ -196,6 +196,8 @@
 import sys
 from typing import Generator
 
+from rich import traceback
+
 if sys.version_info < (3, 10):
  from importlib_metadata import entry_points
 else:
@@ -298,3 +300,6 @@ def load_implicit_plugins():
 
 # Load all implicit plugins
 load_implicit_plugins()
+
+# Pretty-print exception messages
+traceback.install(width=None, extra_lines=0)

flytekit/core/base_task.py

@@ -37,7 +37,7 @@
  translate_inputs_to_literals,
 )
 from flytekit.core.tracker import TrackedInstance
-from flytekit.core.type_engine import TypeEngine
+from flytekit.core.type_engine import TypeEngine, TypeTransformerFailedError
 from flytekit.deck.deck import Deck
 from flytekit.loggers import logger
 from flytekit.models import dynamic_job as _dynamic_job
@@ -239,12 +239,17 @@ def local_execute(self, ctx: FlyteContext, **kwargs) -> Union[Tuple[Promise], Pr
  # Promises as essentially inputs from previous task executions
  # native constants are just bound to this specific task (default values for a task input)
  # Also along with promises and constants, there could be dictionary or list of promises or constants
- kwargs = translate_inputs_to_literals(
- ctx,
- incoming_values=kwargs,
- flyte_interface_types=self.interface.inputs, # type: ignore
- native_types=self.get_input_types(),
- )
+ try:
+ kwargs = translate_inputs_to_literals(
+ ctx,
+ incoming_values=kwargs,
+ flyte_interface_types=self.interface.inputs,
+ native_types=self.get_input_types(), # type: ignore
+ )
+ except TypeTransformerFailedError as exc:
+ msg = f"Failed to convert inputs of task '{self.name}':\n {exc}"
+ logger.error(msg)
+ raise TypeError(msg) from exc
  input_literal_map = _literal_models.LiteralMap(literals=kwargs)
 
  # if metadata.cache is set, check memoized version
@@ -503,7 +508,14 @@ def dispatch_execute(
  ) as exec_ctx:
  # TODO We could support default values here too - but not part of the plan right now
  # Translate the input literals to Python native
- native_inputs = TypeEngine.literal_map_to_kwargs(exec_ctx, input_literal_map, self.python_interface.inputs)
+ try:
+ native_inputs = TypeEngine.literal_map_to_kwargs(
+ exec_ctx, input_literal_map, self.python_interface.inputs
+ )
+ except Exception as exc:
+ msg = f"Failed to convert inputs of task '{self.name}':\n {exc}"
+ logger.error(msg)
+ raise type(exc)(msg) from exc
 
  # TODO: Logger should auto inject the current context information to indicate if the task is running within
  # a workflow or a subworkflow etc
@@ -547,19 +559,20 @@ def dispatch_execute(
  # We manually construct a LiteralMap here because task inputs and outputs actually violate the assumption
  # built into the IDL that all the values of a literal map are of the same type.
  literals = {}
- for k, v in native_outputs_as_map.items():
+ for i, (k, v) in enumerate(native_outputs_as_map.items()):
  literal_type = self._outputs_interface[k].type
  py_type = self.get_type_for_output_var(k, v)
 
  if isinstance(v, tuple):
- raise TypeError(f"Output({k}) in task{self.name} received a tuple {v}, instead of {py_type}")
+ raise TypeError(f"Output({k}) in task '{self.name}' received a tuple {v}, instead of {py_type}")
  try:
  literals[k] = TypeEngine.to_literal(exec_ctx, v, py_type, literal_type)
  except Exception as e:
- logger.error(f"Failed to convert return value for var {k} with error {type(e)}: {e}")
- raise TypeError(
- f"Failed to convert return value for var {k} for function {self.name} with error {type(e)}: {e}"
- ) from e
+ # only show the name of output key if it's user-defined (by default Flyte names these as "o<n>")
+ key = k if k != f"o{i}" else i
+ msg = f"Failed to convert outputs of task '{self.name}' at position {key}:\n {e}"
+ logger.error(msg)
+ raise TypeError(msg) from e
 
  if self._disable_deck is False:
  INPUT = "input"

flytekit/core/promise.py

@@ -13,7 +13,7 @@
 from flytekit.core.context_manager import BranchEvalMode, ExecutionState, FlyteContext, FlyteContextManager
 from flytekit.core.interface import Interface
 from flytekit.core.node import Node
-from flytekit.core.type_engine import DictTransformer, ListTransformer, TypeEngine
+from flytekit.core.type_engine import DictTransformer, ListTransformer, TypeEngine, TypeTransformerFailedError
 from flytekit.exceptions import user as _user_exceptions
 from flytekit.models import interface as _interface_models
 from flytekit.models import literals as _literal_models
@@ -141,7 +141,10 @@ def extract_value(
  raise ValueError(f"Received unexpected keyword argument {k}")
  var = flyte_interface_types[k]
  t = native_types[k]
- result[k] = extract_value(ctx, v, t, var.type)
+ try:
+ result[k] = extract_value(ctx, v, t, var.type)
+ except TypeTransformerFailedError as exc:
+ raise TypeTransformerFailedError(f"Failed argument '{k}': {exc}") from exc
 
  return result
 
@@ -477,10 +480,14 @@ def create_native_named_tuple(
 
  if isinstance(promises, Promise):
  k, v = [(k, v) for k, v in entity_interface.outputs.items()][0] # get output native type
+ # only show the name of output key if it's user-defined (by default Flyte names these as "o<n>")
+ key = k if k != "o0" else 0
  try:
  return TypeEngine.to_python_value(ctx, promises.val, v)
  except Exception as e:
- raise AssertionError(f"Failed to convert value of output {k}, expected type {v}.") from e
+ raise TypeError(
+ f"Failed to convert output in position {key} of value {promises.val}, expected type {v}."
+ ) from e
 
  if len(promises) == 0:
  return None
@@ -490,7 +497,7 @@ def create_native_named_tuple(
  named_tuple_name = entity_interface.output_tuple_name
 
  outputs = {}
- for p in promises:
+ for i, p in enumerate(cast(Tuple[Promise], promises)):
  if not isinstance(p, Promise):
  raise AssertionError(
  "Workflow outputs can only be promises that are returned by tasks. Found a value of"
@@ -500,7 +507,9 @@ def create_native_named_tuple(
  try:
  outputs[p.var] = TypeEngine.to_python_value(ctx, p.val, t)
  except Exception as e:
- raise AssertionError(f"Failed to convert value of output {p.var}, expected type {t}.") from e
+ # only show the name of output key if it's user-defined (by default Flyte names these as "o<n>")
+ key = p.var if p.var != f"o{i}" else i
+ raise TypeError(f"Failed to convert output in position {key} of value {p.val}, expected type {t}.") from e
 
  # Should this class be part of the Interface?
  t = collections.namedtuple(named_tuple_name, list(outputs.keys()))
@@ -1055,7 +1064,7 @@ def flyte_entity_call_handler(entity: SupportsNodeCreation, *args, **kwargs):
  for k, v in kwargs.items():
  if k not in cast(SupportsNodeCreation, entity).python_interface.inputs:
  raise ValueError(
- f"Received unexpected keyword argument {k} in function {cast(SupportsNodeCreation, entity).name}"
+ f"Received unexpected keyword argument '{k}' in function '{cast(SupportsNodeCreation, entity).name}'"
  )
 
  ctx = FlyteContextManager.current_context()

flytekit/core/type_engine.py

@@ -88,7 +88,7 @@ def type_assertions_enabled(self) -> bool:
 
  def assert_type(self, t: Type[T], v: T):
  if not hasattr(t, "__origin__") and not isinstance(v, t):
- raise TypeTransformerFailedError(f"Type of Val '{v}' is not an instance of {t}")
+ raise TypeTransformerFailedError(f"Expected value of type {t} but got '{v}' of type {type(v)}")
 
  @abstractmethod
  def get_literal_type(self, t: Type[T]) -> LiteralType:
@@ -166,7 +166,9 @@ def get_literal_type(self, t: Type[T] = None) -> LiteralType:
 
  def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType) -> Literal:
  if type(python_val) != self._type:
- raise TypeTransformerFailedError(f"Expected value of type {self._type} but got type {type(python_val)}")
+ raise TypeTransformerFailedError(
+ f"Expected value of type {self._type} but got '{python_val}' of type {type(python_val)}"
+ )
  return self._to_literal_transformer(python_val)
 
  def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> T:
@@ -185,7 +187,7 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type:
  return res
  except AttributeError:
  # Assume that this is because a property on `lv` was None
- raise TypeTransformerFailedError(f"Cannot convert literal {lv}")
+ raise TypeTransformerFailedError(f"Cannot convert literal {lv} to {self._type}")
 
  def guess_python_type(self, literal_type: LiteralType) -> Type[T]:
  if literal_type.simple is not None and literal_type.simple == self._lt.simple:
@@ -852,7 +854,13 @@ def literal_map_to_kwargs(
  raise ValueError(
  f"Received more input values {len(lm.literals)}" f" than allowed by the input spec {len(python_types)}"
  )
- return {k: TypeEngine.to_python_value(ctx, lm.literals[k], python_types[k]) for k, v in lm.literals.items()}
+ kwargs = {}
+ for i, k in enumerate(lm.literals):
+ try:
+ kwargs[k] = TypeEngine.to_python_value(ctx, lm.literals[k], python_types[k])
+ except TypeTransformerFailedError as exc:
+ raise TypeTransformerFailedError(f"Error converting input '{k}' at position {i}:\n {exc}") from exc
+ return kwargs
 
  @classmethod
  def dict_to_literal_map(

flytekit/core/workflow.py

@@ -32,7 +32,7 @@
 from flytekit.core.python_auto_container import PythonAutoContainerTask
 from flytekit.core.reference_entity import ReferenceEntity, WorkflowReference
 from flytekit.core.tracker import extract_task_module
-from flytekit.core.type_engine import TypeEngine
+from flytekit.core.type_engine import TypeEngine, TypeTransformerFailedError
 from flytekit.exceptions import scopes as exception_scopes
 from flytekit.exceptions.user import FlyteValidationException, FlyteValueException
 from flytekit.loggers import logger
@@ -258,7 +258,11 @@ def __call__(self, *args, **kwargs):
  input_kwargs = self.python_interface.default_inputs_as_kwargs
  input_kwargs.update(kwargs)
  self.compile()
- return flyte_entity_call_handler(self, *args, **input_kwargs)
+ try:
+ return flyte_entity_call_handler(self, *args, **input_kwargs)
+ except Exception as exc:
+ exc.args = (f"Encountered error while executing workflow '{self.name}':\n {exc}", *exc.args[1:])
+ raise exc
 
  def execute(self, **kwargs):
  raise Exception("Should not be called")
@@ -272,7 +276,12 @@ def local_execute(self, ctx: FlyteContext, **kwargs) -> Union[Tuple[Promise], Pr
  for k, v in kwargs.items():
  if not isinstance(v, Promise):
  t = self.python_interface.inputs[k]
- kwargs[k] = Promise(var=k, val=TypeEngine.to_literal(ctx, v, t, self.interface.inputs[k].type))
+ try:
+ kwargs[k] = Promise(var=k, val=TypeEngine.to_literal(ctx, v, t, self.interface.inputs[k].type))
+ except TypeTransformerFailedError as exc:
+ raise TypeError(
+ f"Failed to convert input argument '{k}' of workflow '{self.name}':\n {exc}"
+ ) from exc
 
  # The output of this will always be a combination of Python native values and Promises containing Flyte
  # Literals.

flytekit/exceptions/scopes.py

@@ -194,10 +194,13 @@ def user_entry_point(wrapped, instance, args, kwargs):
  _CONTEXT_STACK.append(_USER_CONTEXT)
  if _is_base_context():
  # See comment at this location for system_entry_point
+ fn_name = wrapped.__name__
  try:
  return wrapped(*args, **kwargs)
- except FlyteScopedException as ex:
- raise ex.value
+ except FlyteScopedException as exc:
+ raise exc.type(f"Error encountered while executing '{fn_name}':\n {exc.value}") from exc
+ except Exception as exc:
+ raise type(exc)(f"Error encountered while executing '{fn_name}':\n {exc}") from exc
  else:
  try:
  return wrapped(*args, **kwargs)

flytekit/loggers.py

@@ -2,6 +2,8 @@
 import os
 
 from pythonjsonlogger import jsonlogger
+from rich.console import Console
+from rich.logging import RichHandler
 
 # Note:
 # The environment variable controls exposed to affect the individual loggers should be considered to be beta.
@@ -10,6 +12,7 @@
 # For now, assume this is the environment variable whose usage will remain unchanged and controls output for all
 # loggers defined in this file.
 LOGGING_ENV_VAR = "FLYTE_SDK_LOGGING_LEVEL"
+LOGGING_FMT_ENV_VAR = "FLYTE_SDK_LOGGING_FORMAT"
 
 # By default, the root flytekit logger to debug so everything is logged, but enable fine-tuning
 logger = logging.getLogger("flytekit")
@@ -33,8 +36,18 @@
 user_space_logger = child_loggers["user_space"]
 
 # create console handler
-ch = logging.StreamHandler()
-ch.setLevel(logging.DEBUG)
+try:
+ handler = RichHandler(
+ rich_tracebacks=True,
+ omit_repeated_times=False,
+ keywords=["[flytekit]"],
+ log_time_format="%Y-%m-%d %H:%M:%S,%f",
+ console=Console(width=os.get_terminal_size().columns),
+ )
+except OSError:
+ handler = logging.StreamHandler()
+
+handler.setLevel(logging.DEBUG)
 
 # Root logger control
 # Don't want to import the configuration library since that will cause all sorts of circular imports, let's
@@ -63,10 +76,14 @@
  child_logger.setLevel(logging.WARNING)
 
 # create formatter
-formatter = jsonlogger.JsonFormatter(fmt="%(asctime)s %(name)s %(levelname)s %(message)s")
+logging_fmt = os.environ.get(LOGGING_FMT_ENV_VAR, "json")
+if logging_fmt == "json":
+ formatter = jsonlogger.JsonFormatter(fmt="%(asctime)s %(name)s %(levelname)s %(message)s")
+else:
+ formatter = logging.Formatter(fmt="[%(name)s] %(message)s")
 
-# add formatter to ch
-ch.setFormatter(formatter)
+# add formatter to the handler
+handler.setFormatter(formatter)
 
 # add ch to logger
-logger.addHandler(ch)
+logger.addHandler(handler)

flytekit/models/common.py

@@ -1,5 +1,6 @@
 import abc as _abc
 import json as _json
+import re
 
 from flyteidl.admin import common_pb2 as _common_pb2
 from google.protobuf import json_format as _json_format
@@ -57,7 +58,8 @@ def short_string(self):
  """
  :rtype: Text
  """
- return str(self.to_flyte_idl())
+ literal_str = re.sub(r"\s+", " ", str(self.to_flyte_idl())).strip()
+ return f"<FlyteLiteral {literal_str}>"
 
  def verbose_string(self):
  """

plugins/flytekit-pandera/tests/test_plugin.py

@@ -55,7 +55,13 @@ def invalid_wf() -> pandera.typing.DataFrame[OutSchema]:
  def wf_with_df_input(df: pandera.typing.DataFrame[InSchema]) -> pandera.typing.DataFrame[OutSchema]:
  return transform2(df=transform1(df=df))
 
- with pytest.raises(pandera.errors.SchemaError, match="^expected series 'col2' to have type float64, got object"):
+ with pytest.raises(
+ pandera.errors.SchemaError,
+ match=(
+ "^Encountered error while executing workflow 'test_plugin.wf_with_df_input':\n"
+ " expected series 'col2' to have type float64, got object"
+ ),
+ ):
  wf_with_df_input(df=invalid_df)
 
  # raise error when executing workflow with invalid output
@@ -67,7 +73,14 @@ def transform2_noop(df: pandera.typing.DataFrame[IntermediateSchema]) -> pandera
  def wf_invalid_output(df: pandera.typing.DataFrame[InSchema]) -> pandera.typing.DataFrame[OutSchema]:
  return transform2_noop(df=transform1(df=df))
 
- with pytest.raises(TypeError, match="^Failed to convert return value"):
+ with pytest.raises(
+ TypeError,
+ match=(
+ "^Encountered error while executing workflow 'test_plugin.wf_invalid_output':\n"
+ " Error encountered while executing 'wf_invalid_output':\n"
+ " Failed to convert outputs of task"
+ ),
+ ):
  wf_invalid_output(df=valid_df)
 
 

setup.py

@@ -77,6 +77,7 @@
  "numpy<1.24.0",
  "gitpython",
  "kubernetes>=12.0.1",
+ "rich",
  ],
  extras_require=extras_require,
  scripts=[