FlyteDownloadDataException

Signed-off-by: Kevin Su <[email protected]>

flytekit/bin/entrypoint.py

@@ -159,19 +159,7 @@ def _dispatch_execute(
         logger.error(exc_str)
         logger.error("!! End Error Captured by Flyte !!")
 
-    # All the Non-user errors are captured here, and are considered system errors
-    except Exception as e:
-        exc_str = get_traceback_str(e)
-        output_file_dict[_constants.ERROR_FILE_NAME] = _error_models.ErrorDocument(
-            _error_models.ContainerError(
-                "SYSTEM",
-                exc_str,
-                _error_models.ContainerError.Kind.RECOVERABLE,
-                _execution_models.ExecutionError.ErrorKind.USER,
-            )
-        )
-
-    # All the Non-user errors are captured here, and are considered system errors
+    # All the Non-user runtime errors are captured here, and are considered system errors
     except Exception as e:
         kind = _error_models.ContainerError.Kind.RECOVERABLE
         err = e
@@ -186,7 +174,7 @@ def _dispatch_execute(
                 "SYSTEM",
                 exc_str,
                 kind,
-                _execution_models.ExecutionError.ErrorKind.SYSTEM,
+                _execution_models.ExecutionError.ErrorKind.USER,
             )
         )
 

flytekit/core/base_task.py

@@ -71,7 +71,11 @@
 from flytekit.core.type_engine import TypeEngine, TypeTransformerFailedError
 from flytekit.core.utils import timeit
 from flytekit.deck import DeckField
-from flytekit.exceptions.system import FlyteNonRecoverableSystemException
+from flytekit.exceptions.system import (
+    FlyteDownloadDataException,
+    FlyteNonRecoverableSystemException,
+    FlyteUploadDataException,
+)
 from flytekit.exceptions.user import FlyteUserRuntimeException
 from flytekit.loggers import logger
 from flytekit.models import dynamic_job as _dynamic_job
@@ -729,6 +733,8 @@ def dispatch_execute(
             # Translate the input literals to Python native
             try:
                 native_inputs = self._literal_map_to_python_input(input_literal_map, exec_ctx)
+            except FlyteUploadDataException:
+                raise
             except Exception as exc:
                 raise FlyteNonRecoverableSystemException(exc) from exc
 
@@ -785,6 +791,8 @@ def dispatch_execute(
             try:
                 literals_map, native_outputs_as_map = self._output_to_literal_map(native_outputs, exec_ctx)
                 self._write_decks(native_inputs, native_outputs_as_map, ctx, new_user_params)
+            except FlyteDownloadDataException:
+                raise
             except Exception as exc:
                 raise FlyteNonRecoverableSystemException(exc) from exc
             # After the execute has been successfully completed

flytekit/core/data_persistence.py

@@ -36,6 +36,7 @@
 from flytekit.configuration import DataConfig
 from flytekit.core.local_fsspec import FlyteLocalFileSystem
 from flytekit.core.utils import timeit
+from flytekit.exceptions.system import FlyteDownloadDataException, FlyteUploadDataException
 from flytekit.exceptions.user import FlyteAssertion, FlyteDataNotFoundException
 from flytekit.interfaces.random import random
 from flytekit.loggers import logger
@@ -561,7 +562,7 @@ def get_data(self, remote_path: str, local_path: str, is_multipart: bool = False
         except FlyteDataNotFoundException:
             raise
         except Exception as ex:
-            raise FlyteAssertion(
+            raise FlyteDownloadDataException(
                 f"Failed to get data from {remote_path} to {local_path} (recursive={is_multipart}).\n\n"
                 f"Original exception: {str(ex)}"
             )
@@ -589,7 +590,7 @@ def put_data(
                     return put_result
                 return remote_path
         except Exception as ex:
-            raise FlyteAssertion(
+            raise FlyteUploadDataException(
                 f"Failed to put data from {local_path} to {remote_path} (recursive={is_multipart}).\n\n"
                 f"Original exception: {str(ex)}"
             ) from ex

flytekit/exceptions/system.py

@@ -65,3 +65,11 @@ def __init__(self, exc_value: Exception):
     @property
     def value(self):
         return self._exc_value
+
+
+class FlyteDownloadDataException(FlyteSystemException):
+    _ERROR_CODE = "SYSTEM:DownloadDataError"
+
+
+class FlyteUploadDataException(FlyteSystemException):
+    _ERROR_CODE = "SYSTEM:UploadDataError"