Fix: Always propagate pytorch task worker process exception timestamp…

… to task exception

Signed-off-by: Fabio Grätz <[email protected]>

flytekit/exceptions/user.py

@@ -15,14 +15,15 @@ class FlyteUserException(_FlyteException):
 class FlyteUserRuntimeException(_FlyteException):
     _ERROR_CODE = "USER:RuntimeError"
 
-    def __init__(self, exc_value: Exception):
+    def __init__(self, exc_value: Exception, timestamp: typing.Optional[float] = None):
         """
         FlyteUserRuntimeException is thrown when a user code raises an exception.
 
         :param exc_value: The exception that was raised from user code.
+        :param timestamp: The timestamp as fractional seconds since epoch when the exception was raised.
         """
         self._exc_value = exc_value
-        super().__init__(str(exc_value))
+        super().__init__(str(exc_value), timestamp=timestamp)
 
     @property
     def value(self):

plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py

@@ -18,7 +18,7 @@
 from flytekit.core.context_manager import FlyteContextManager, OutputMetadata
 from flytekit.core.pod_template import PodTemplate
 from flytekit.core.resources import convert_resources_to_resource_model
-from flytekit.exceptions.user import FlyteRecoverableException
+from flytekit.exceptions.user import FlyteRecoverableException, FlyteUserRuntimeException
 from flytekit.extend import IgnoreOutputs, TaskPlugins
 from flytekit.loggers import logger
 
@@ -475,7 +475,7 @@ def fn_partial():
                 # the automatically assigned timestamp based on exception creation time
                 raise FlyteRecoverableException(e.format_msg(), timestamp=first_failure.timestamp)
             else:
-                raise RuntimeError(e.format_msg())
+                raise FlyteUserRuntimeException(e, timestamp=first_failure.timestamp)
         except SignalException as e:
             logger.exception(f"Elastic launch agent process terminating: {e}")
             raise IgnoreOutputs()

plugins/flytekit-kf-pytorch/tests/test_elastic_task.py

@@ -276,3 +276,37 @@ def test_task_omp_set():
         assert os.environ["OMP_NUM_THREADS"] == "42"
 
     test_task_omp_set()
+
+
+def test_exception_timestamp() -> None:
+    """Test that the timestamp of the worker process exception is propagated to the task exception."""
+    @task(
+        task_config=Elastic(
+            nnodes=1,
+            nproc_per_node=2,
+        )
+    )
+    def test_task():
+        raise Exception("Test exception")
+
+    with pytest.raises(Exception) as e:
+        test_task()
+
+    assert e.value.timestamp is not None
+
+
+def test_recoverable_exception_timestamp() -> None:
+    """Test that the timestamp of the worker process exception is propagated to the task exception."""
+    @task(
+        task_config=Elastic(
+            nnodes=1,
+            nproc_per_node=2,
+        )
+    )
+    def test_task():
+        raise FlyteRecoverableException("Recoverable test exception")
+
+    with pytest.raises(Exception) as e:
+        test_task()
+
+    assert e.value.timestamp is not None