From 26a216ebafe2cb1acad384982a9deb501172e473 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 10 Jan 2023 15:59:55 -0800 Subject: [PATCH 1/4] AWS batch return error code once it fails Signed-off-by: Kevin Su --- flytekit/bin/entrypoint.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/flytekit/bin/entrypoint.py b/flytekit/bin/entrypoint.py index 3d5017675e..c976c58806 100644 --- a/flytekit/bin/entrypoint.py +++ b/flytekit/bin/entrypoint.py @@ -161,6 +161,11 @@ def _dispatch_execute( logger.info(f"Engine folder written successfully to the output prefix {output_prefix}") logger.debug("Finished _dispatch_execute") + if task_def.task_type == "aws-batch" and _constants.ERROR_FILE_NAME in output_file_dict: + # AWS batch job get the status from the exit code, so once we catch the error, + # we should return the error code here + exit(1) + def get_one_of(*args) -> str: """ From 3de0912e855647896ba95327b39f1cba79f814c8 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 10 Jan 2023 17:15:59 -0800 Subject: [PATCH 2/4] AWS batch return error code once it fails Signed-off-by: Kevin Su --- flytekit/bin/entrypoint.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/flytekit/bin/entrypoint.py b/flytekit/bin/entrypoint.py index c976c58806..fe38f946f9 100644 --- a/flytekit/bin/entrypoint.py +++ b/flytekit/bin/entrypoint.py @@ -161,7 +161,8 @@ def _dispatch_execute( logger.info(f"Engine folder written successfully to the output prefix {output_prefix}") logger.debug("Finished _dispatch_execute") - if task_def.task_type == "aws-batch" and _constants.ERROR_FILE_NAME in output_file_dict: + if os.environ.get("FLYTE_FAIL_ON_ERROR", "").lower() == "true" and _constants.ERROR_FILE_NAME in output_file_dict: + # This env is set by the flytepropeller # AWS batch job get the status from the exit code, so once we catch the error, # we should return the error code here exit(1) From 57781e3bd90f5506efab1410e132a9f1344c31fe Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 11 Jan 2023 02:52:54 -0800 Subject: [PATCH 3/4] update tests Signed-off-by: Kevin Su --- tests/flytekit/unit/bin/test_python_entrypoint.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/flytekit/unit/bin/test_python_entrypoint.py b/tests/flytekit/unit/bin/test_python_entrypoint.py index 479ad9e7bd..2df40c8b30 100644 --- a/tests/flytekit/unit/bin/test_python_entrypoint.py +++ b/tests/flytekit/unit/bin/test_python_entrypoint.py @@ -3,6 +3,7 @@ from collections import OrderedDict import mock +import pytest from flyteidl.core.errors_pb2 import ErrorDocument from flytekit.bin.entrypoint import _dispatch_execute, normalize_inputs, setup_execution @@ -109,6 +110,11 @@ def verify_output(*args, **kwargs): _dispatch_execute(ctx, python_task, "inputs path", "outputs prefix") assert mock_write_to_file.call_count == 1 + with pytest.raises(SystemExit) as cm: + os.environ["FLYTE_FAIL_ON_ERROR"] = "true" + _dispatch_execute(ctx, python_task, "inputs path", "outputs prefix") + pytest.assertEqual(cm.value.code, 1) + # This function collects outputs instead of writing them to a file. # See flytekit.core.utils.write_proto_to_file for the original From ce0938a0c691194268ce84078b7231bd8fbb2dd9 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 11 Jan 2023 13:38:59 -0800 Subject: [PATCH 4/4] Update tests Signed-off-by: Kevin Su --- .../unit/bin/test_python_entrypoint.py | 28 ++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/tests/flytekit/unit/bin/test_python_entrypoint.py b/tests/flytekit/unit/bin/test_python_entrypoint.py index 2df40c8b30..6a8b8c430e 100644 --- a/tests/flytekit/unit/bin/test_python_entrypoint.py +++ b/tests/flytekit/unit/bin/test_python_entrypoint.py @@ -110,8 +110,34 @@ def verify_output(*args, **kwargs): _dispatch_execute(ctx, python_task, "inputs path", "outputs prefix") assert mock_write_to_file.call_count == 1 + +@mock.patch.dict(os.environ, {"FLYTE_FAIL_ON_ERROR": "True"}) +@mock.patch("flytekit.core.utils.load_proto_from_file") +@mock.patch("flytekit.core.data_persistence.FileAccessProvider.get_data") +@mock.patch("flytekit.core.data_persistence.FileAccessProvider.put_data") +@mock.patch("flytekit.core.utils.write_proto_to_file") +def test_dispatch_execute_return_error_code(mock_write_to_file, mock_upload_dir, mock_get_data, mock_load_proto): + 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: + python_task = mock.MagicMock() + python_task.dispatch_execute.side_effect = Exception("random") + + empty_literal_map = _literal_models.LiteralMap({}).to_flyte_idl() + mock_load_proto.return_value = empty_literal_map + + def verify_output(*args, **kwargs): + assert isinstance(args[0], ErrorDocument) + + mock_write_to_file.side_effect = verify_output + with pytest.raises(SystemExit) as cm: - os.environ["FLYTE_FAIL_ON_ERROR"] = "true" _dispatch_execute(ctx, python_task, "inputs path", "outputs prefix") pytest.assertEqual(cm.value.code, 1)