def test_dispatch_execute_user_error_non_recov(mock_write_to_file, mock_upload_dir, mock_get_data, mock_load_proto):
# Just leave these here, mock them out so nothing happens
mock_get_data.return_value = True
mock_upload_dir.return_value = True
@task
def t1(a: int) -> str:
# Should be interpreted as a non-recoverable user error
raise ValueError(f"some exception {a}")
ctx = context_manager.FlyteContext.current_context()
with context_manager.FlyteContextManager.with_context(
ctx.with_execution_state(
ctx.execution_state.with_params(mode=context_manager.ExecutionState.Mode.TASK_EXECUTION)
)
) as ctx:
input_literal_map = TypeEngine.dict_to_literal_map(ctx, {"a": 5})
mock_load_proto.return_value = input_literal_map.to_flyte_idl()
files = OrderedDict()
mock_write_to_file.side_effect = get_output_collector(files)
# See comment in test_dispatch_execute_ignore for why we need to decorate
system_entry_point(_dispatch_execute)(ctx, t1, "inputs path", "outputs prefix")
assert len(files) == 1
# Exception should've caused an error file
k = list(files.keys())[0]
assert "error.pb" in k
v = list(files.values())[0]
ed = error_models.ErrorDocument.from_flyte_idl(v)
assert ed.error.kind == error_models.ContainerError.Kind.NON_RECOVERABLE
assert "some exception 5" in ed.error.message
assert ed.error.origin == execution_models.ExecutionError.ErrorKind.USER
def test_dispatch_execute_system_error(mock_write_to_file, mock_upload_dir, mock_get_data, mock_load_proto):
# Just leave these here, mock them out so nothing happens
mock_get_data.return_value = True
mock_upload_dir.return_value = True
ctx = context_manager.FlyteContext.current_context()
with context_manager.FlyteContextManager.with_context(
ctx.with_execution_state(
ctx.execution_state.with_params(mode=context_manager.ExecutionState.Mode.TASK_EXECUTION)
)
) as ctx:
input_literal_map = TypeEngine.dict_to_literal_map(ctx, {"a": 5})
mock_load_proto.return_value = input_literal_map.to_flyte_idl()
python_task = mock.MagicMock()
python_task.dispatch_execute.side_effect = Exception("some system exception")
files = OrderedDict()
mock_write_to_file.side_effect = get_output_collector(files)
# See comment in test_dispatch_execute_ignore for why we need to decorate
system_entry_point(_dispatch_execute)(ctx, python_task, "inputs path", "outputs prefix")
assert len(files) == 1
# Exception should've caused an error file
k = list(files.keys())[0]
assert "error.pb" in k
v = list(files.values())[0]
ed = error_models.ErrorDocument.from_flyte_idl(v)
# System errors default to recoverable
assert ed.error.kind == error_models.ContainerError.Kind.RECOVERABLE
assert "some system exception" in ed.error.message
assert ed.error.origin == execution_models.ExecutionError.ErrorKind.SYSTEM
def test_dispatch_execute_normal(mock_write_to_file, mock_upload_dir, mock_get_data, mock_load_proto):
# Just leave these here, mock them out so nothing happens
mock_get_data.return_value = True
mock_upload_dir.return_value = True
@task
def t1(a: int) -> str:
return f"string is: {a}"
ctx = context_manager.FlyteContext.current_context()
with context_manager.FlyteContextManager.with_context(
ctx.with_execution_state(
ctx.execution_state.with_params(mode=context_manager.ExecutionState.Mode.TASK_EXECUTION)
)
) as ctx:
input_literal_map = TypeEngine.dict_to_literal_map(ctx, {"a": 5})
mock_load_proto.return_value = input_literal_map.to_flyte_idl()
files = OrderedDict()
mock_write_to_file.side_effect = get_output_collector(files)
# See comment in test_dispatch_execute_ignore for why we need to decorate
system_entry_point(_dispatch_execute)(ctx, t1, "inputs path", "outputs prefix")
assert len(files) == 1
# A successful run should've written an outputs file.
k = list(files.keys())[0]
assert "outputs.pb" in k
v = list(files.values())[0]
lm = _literal_models.LiteralMap.from_flyte_idl(v)
assert lm.literals["o0"].scalar.primitive.string_value == "string is: 5"
def test_dispatch_execute_ignore(mock_write_to_file, mock_upload_dir, mock_get_data, mock_load_proto):
# Just leave these here, mock them out so nothing happens
mock_get_data.return_value = True
mock_upload_dir.return_value = True
ctx = context_manager.FlyteContext.current_context()
# IgnoreOutputs
with context_manager.FlyteContextManager.with_context(
ctx.with_execution_state(
ctx.execution_state.with_params(mode=context_manager.ExecutionState.Mode.TASK_EXECUTION)
)
) as ctx:
python_task = mock.MagicMock()
python_task.dispatch_execute.side_effect = IgnoreOutputs()
empty_literal_map = _literal_models.LiteralMap({}).to_flyte_idl()
mock_load_proto.return_value = empty_literal_map
# The system_entry_point decorator does different thing based on whether or not it's the
# first time it's called. Using it here to mimic the fact that _dispatch_execute is
# called by _execute_task, which also has a system_entry_point
system_entry_point(_dispatch_execute)(ctx, python_task, "inputs path", "outputs prefix")
assert mock_write_to_file.call_count == 0
def _dispatch_execute(
ctx: FlyteContext,
task_def: PythonTask,
inputs_path: str,
output_prefix: str,
):
"""
Dispatches execute to PythonTask
Step1: Download inputs and load into a literal map
Step2: Invoke task - dispatch_execute
Step3:
a: [Optional] Record outputs to output_prefix
b: OR if IgnoreOutputs is raised, then ignore uploading outputs
c: OR if an unhandled exception is retrieved - record it as an errors.pb
"""
output_file_dict = {}
logger.debug(f"Starting _dispatch_execute for {task_def.name}")
try:
# Step1
local_inputs_file = os.path.join(ctx.execution_state.working_dir, "inputs.pb")
ctx.file_access.get_data(inputs_path, local_inputs_file)
input_proto = utils.load_proto_from_file(_literals_pb2.LiteralMap, local_inputs_file)
idl_input_literals = _literal_models.LiteralMap.from_flyte_idl(input_proto)
# Step2
# Decorate the dispatch execute function before calling it, this wraps all exceptions into one
# of the FlyteScopedExceptions
outputs = _scoped_exceptions.system_entry_point(task_def.dispatch_execute)(ctx, idl_input_literals)
# Step3a
if isinstance(outputs, VoidPromise):
logger.warning("Task produces no outputs")
output_file_dict = {_constants.OUTPUT_FILE_NAME: _literal_models.LiteralMap(literals={})}
elif isinstance(outputs, _literal_models.LiteralMap):
output_file_dict = {_constants.OUTPUT_FILE_NAME: outputs}
elif isinstance(outputs, _dynamic_job.DynamicJobSpec):
output_file_dict = {_constants.FUTURES_FILE_NAME: outputs}
else:
logger.error(f"SystemError: received unknown outputs from task {outputs}")
output_file_dict[_constants.ERROR_FILE_NAME] = _error_models.ErrorDocument(
_error_models.ContainerError(
"UNKNOWN_OUTPUT",
f"Type of output received not handled {type(outputs)} outputs: {outputs}",
_error_models.ContainerError.Kind.RECOVERABLE,
_execution_models.ExecutionError.ErrorKind.SYSTEM,
)
)
# Handle user-scoped errors
except _scoped_exceptions.FlyteScopedUserException as e:
if isinstance(e.value, IgnoreOutputs):
logger.warning(f"User-scoped IgnoreOutputs received! Outputs.pb will not be uploaded. reason {e}!!")
return
output_file_dict[_constants.ERROR_FILE_NAME] = _error_models.ErrorDocument(
_error_models.ContainerError(
e.error_code, e.verbose_message, e.kind, _execution_models.ExecutionError.ErrorKind.USER
)
)
logger.error("!! Begin User Error Captured by Flyte !!")
logger.error(e.verbose_message)
logger.error("!! End Error Captured by Flyte !!")
# Handle system-scoped errors
except _scoped_exceptions.FlyteScopedSystemException as e:
if isinstance(e.value, IgnoreOutputs):
logger.warning(f"System-scoped IgnoreOutputs received! Outputs.pb will not be uploaded. reason {e}!!")
return
output_file_dict[_constants.ERROR_FILE_NAME] = _error_models.ErrorDocument(
_error_models.ContainerError(
e.error_code, e.verbose_message, e.kind, _execution_models.ExecutionError.ErrorKind.SYSTEM
)
)
logger.error("!! Begin System Error Captured by Flyte !!")
logger.error(e.verbose_message)
logger.error("!! End Error Captured by Flyte !!")
# Interpret all other exceptions (some of which may be caused by the code in the try block outside of
# dispatch_execute) as recoverable system exceptions.
except Exception as e:
# Step 3c
exc_str = _traceback.format_exc()
output_file_dict[_constants.ERROR_FILE_NAME] = _error_models.ErrorDocument(
_error_models.ContainerError(
"SYSTEM:Unknown",
exc_str,
_error_models.ContainerError.Kind.RECOVERABLE,
_execution_models.ExecutionError.ErrorKind.SYSTEM,
)
)
logger.error(f"Exception when executing task {task_def.name or task_def.id.name}, reason {str(e)}")
logger.error("!! Begin Unknown System Error Captured by Flyte !!")
logger.error(exc_str)
logger.error("!! End Error Captured by Flyte !!")
for k, v in output_file_dict.items():
utils.write_proto_to_file(v.to_flyte_idl(), os.path.join(ctx.execution_state.engine_dir, k))
ctx.file_access.put_data(ctx.execution_state.engine_dir, output_prefix, is_multipart=True)
logger.info(f"Engine folder written successfully to the output prefix {output_prefix}")
logger.debug("Finished _dispatch_execute")