UDF error handling improvements #1464

dreadatour · 2025-11-17T15:41:34Z

Improve UDF error handling: return raw error from UDF to handle it properly in the parent process.

Copilot

Pull Request Overview

This PR improves UDF error handling by removing the intermediate process_safe wrapper and allowing raw exceptions from UDFs to propagate to the parent process, while adding proper stacktrace support for parallel execution.

Key Changes

Introduced UdfRunError exception class to wrap errors from parallel UDF execution with stacktrace information
Removed the process_safe method that was wrapping exceptions in DataChainError, allowing raw exceptions to propagate
Updated exception handling in parallel workers to capture and propagate stacktraces
Updated tests to expect the actual exception types raised by UDFs instead of DataChainError

Reviewed Changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
src/datachain/lib/udf.py	Introduced `UdfRunError` class, removed `process_safe` wrapper method and unused imports, replaced all `process_safe` calls with direct `process` calls
src/datachain/query/dispatch.py	Added stacktrace capture in worker error handling and `UdfRunError` wrapping for parallel execution failures
tests/unit/lib/test_checkpoints.py	Updated tests to expect `CustomMapperError` instead of `DataChainError`, added custom exception class for testing
tests/func/test_udf.py	Updated test to expect `RuntimeError` instead of `DataChainError` for UDF execution errors
tests/func/test_to_database.py	Updated test to expect `ValueError` instead of `DataChainError` for processing failures

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

src/datachain/lib/udf.py

codecov · 2025-11-17T15:50:45Z

Codecov Report

❌ Patch coverage is 90.32258% with 3 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/datachain/query/dispatch.py	40.00%	3 Missing ⚠️

📢 Thoughts on this report? Let us know!

dreadatour

Few highlights for the changes in this PR.

dreadatour · 2025-11-20T11:04:48Z

src/datachain/lib/udf.py

+        return self.__class__, (self.message,)
+
+
+class UdfRunError(Exception):


A new exception type that is caught by the parent process, capturing both the traceback and the UDF name for display in the UI.

dreadatour · 2025-11-20T11:06:53Z

src/datachain/lib/udf.py

-    def process_safe(self, obj_rows):
-        try:
-            result_objs = self.process(*obj_rows)
-        except Exception as e:  # noqa: BLE001
-            msg = f"============== Error in user code: '{self.name}' =============="
-            print(msg)
-            exc_type, exc_value, exc_traceback = sys.exc_info()
-            traceback.print_exception(exc_type, exc_value, exc_traceback.tb_next)
-            print("=" * len(msg))
-            raise DataChainError(
-                f"Error in user code in class '{self.name}': {e!s}"
-            ) from None
-        return result_objs


No need to catch and log this here — the parent process now catches the exception and handles logging and user-facing output.

dreadatour · 2025-11-20T11:07:20Z

src/datachain/query/dispatch.py

+                    {
+                        "status": FAILED_STATUS,
+                        "exception": e,
+                        "stacktrace": traceback.format_exc(),


Send also the traceback.

dreadatour · 2025-11-20T11:08:50Z

src/datachain/query/dispatch.py

-                        raise exc
+                        if isinstance(exc, KeyboardInterrupt):
+                            raise exc
+                        raise UdfRunError(exc, stacktrace=result.get("stacktrace"))


UdfRunError wrapper for exc exception is needed to preserve the traceback.

dreadatour · 2025-11-20T11:09:09Z

tests/func/test_to_database.py


    # The export should fail during processing, after table creation
-    with pytest.raises(dc.DataChainError, match="Processing failed on third item"):
+    with pytest.raises(ValueError, match="Processing failed on third item"):


Direct exception is now raised.

dreadatour · 2025-11-20T11:09:38Z

tests/unit/lib/test_checkpoints.py

+class CustomMapperError(Exception):
+    pass
+
+
+def mapper_fail(num: int) -> int:
+    raise CustomMapperError("Error")


This to check if we are raising and catching any exceptions as is.

Copilot

Pull Request Overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

.github/workflows/tests-studio.yml

Co-authored-by: Copilot <[email protected]>

shcheklein · 2025-11-20T20:39:56Z

.github/workflows/tests-studio.yml

          ref: ${{ env.STUDIO_BRANCH }}
          token: ${{ secrets.ITERATIVE_STUDIO_READ_ACCESS_TOKEN }}

+      - uses: astral-sh/setup-uv@v7


why this change?

We are using this trick to test datachain and studio related changes from the git branch in case both repos have branches with the same name.

Here, in datachain, we are already checking out changes from the desired branch.

Next, we do check if Studio have the branch with the same name — we will checkout the same branch from Studio.

But Studio itself have datachain as a dependency, and we need to make sure it will use the same datachain branch as in this PR.

Before Studio tests in datachain repo was broken sometimes because of this — rarely, only if API changes, but still. Same happens in this PR.

This change is not related directly to this PR, let me separate it 👍

dreadatour requested a review from Copilot November 17, 2025 15:41

dreadatour self-assigned this Nov 17, 2025

Copilot started reviewing on behalf of dreadatour November 17, 2025 15:41 View session

dreadatour marked this pull request as draft November 17, 2025 15:42

Copilot finished reviewing on behalf of dreadatour November 17, 2025 15:44

Copilot AI reviewed Nov 17, 2025

View reviewed changes

src/datachain/lib/udf.py Outdated Show resolved Hide resolved

dreadatour force-pushed the job-errors-observability branch from f166140 to 0f85aaa Compare November 18, 2025 01:53

dreadatour added 8 commits November 20, 2025 11:05

UDF error handling improvements

be5268a

Code review fix

4f76ebe

Fix distributed tests

ffabc70

Fix distributed tests

5e177ef

Fix distributed tests

be2592a

Fix distributed tests

e9359ad

Fix distributed tests

238bb27

Update github repos

ba760ea

dreadatour force-pushed the job-errors-observability branch from 6a51332 to ba760ea Compare November 20, 2025 04:05

dreadatour added 2 commits November 20, 2025 13:45

Add more tests

6a01c58

Invert unrelative changes

bce7498

dreadatour commented Nov 20, 2025

View reviewed changes

dreadatour marked this pull request as ready for review November 20, 2025 11:10

dreadatour requested review from a team, 0x2b3bfa0, amritghimire, Copilot, ilongin and shcheklein November 20, 2025 11:10

Copilot started reviewing on behalf of dreadatour November 20, 2025 11:11 View session

Copilot finished reviewing on behalf of dreadatour November 20, 2025 11:14

Copilot AI reviewed Nov 20, 2025

View reviewed changes

.github/workflows/tests-studio.yml Outdated Show resolved Hide resolved

Update .github/workflows/tests-studio.yml

34cb3ae

Co-authored-by: Copilot <[email protected]>

shcheklein reviewed Nov 20, 2025

View reviewed changes

		return self.__class__, (self.message,)


		class UdfRunError(Exception):

UDF error handling improvements #1464

Are you sure you want to change the base?

UDF error handling improvements #1464

Uh oh!

Conversation

dreadatour commented Nov 17, 2025

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Key Changes

Reviewed Changes

Uh oh!

Uh oh!

codecov bot commented Nov 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

dreadatour left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

codecov bot commented Nov 17, 2025 •

edited

Loading