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")