Add new exit code for aborted queries (#1587)

Author: dreadatour

src/datachain/catalog/__init__.py

@@ -1,5 +1,6 @@
 from .catalog import (
     QUERY_DATASET_PREFIX,
+    QUERY_SCRIPT_ABORTED_EXIT_CODE,
     QUERY_SCRIPT_CANCELED_EXIT_CODE,
     Catalog,
     is_namespace_local,
@@ -8,6 +9,7 @@
 
 __all__ = [
     "QUERY_DATASET_PREFIX",
+    "QUERY_SCRIPT_ABORTED_EXIT_CODE",
     "QUERY_SCRIPT_CANCELED_EXIT_CODE",
     "Catalog",
     "get_catalog",

src/datachain/catalog/catalog.py

@@ -78,6 +78,8 @@
 DATASET_INTERNAL_ERROR_MESSAGE = "Internal error on creating dataset"
 # exit code we use if query script was canceled
 QUERY_SCRIPT_CANCELED_EXIT_CODE = 11
+# exit code we use if the job is already in a terminal state (failed/canceled elsewhere)
+QUERY_SCRIPT_ABORTED_EXIT_CODE = 12
 QUERY_SCRIPT_SIGTERM_EXIT_CODE = -15  # if query script was terminated by SIGTERM
 
 # dataset pull

src/datachain/error.py

@@ -85,7 +85,17 @@ def __init__(self, message: str, return_code: int = 0):
 
 
 class QueryScriptCancelError(QueryScriptRunError):
-    pass
+    """Raised when a running script is canceled by the user."""
+
+
+class QueryScriptAbortError(QueryScriptRunError):
+    """Raised when execution should stop because the job is already in a terminal state.
+
+    Unlike QueryScriptCancelError (user-initiated cancellation), this signals that
+    the job has already finished (e.g., failed elsewhere) and continuing execution
+    would be wasteful. The distinction allows handlers to differentiate between
+    "user clicked Cancel" and "job already dead, stop working".
+    """
 
 
 class ClientError(RuntimeError):

src/datachain/query/dataset.py

@@ -43,6 +43,7 @@
 from datachain.dataset import DatasetDependency, DatasetStatus, RowDict
 from datachain.error import (
     DatasetNotFoundError,
+    QueryScriptAbortError,
     QueryScriptCancelError,
     TableMissingError,
 )
@@ -592,7 +593,7 @@ def create_result_query(
         to select
         """
 
-    def populate_udf_output_table(
+    def populate_udf_output_table(  # noqa: PLR0915
         self,
         udf_table: "Table",
         query: Select,
@@ -612,7 +613,10 @@ def populate_udf_output_table(
             )
             return
 
-        from datachain.catalog import QUERY_SCRIPT_CANCELED_EXIT_CODE
+        from datachain.catalog import (
+            QUERY_SCRIPT_ABORTED_EXIT_CODE,
+            QUERY_SCRIPT_CANCELED_EXIT_CODE,
+        )
         from datachain.catalog.loader import (
             DISTRIBUTED_IMPORT_PATH,
             get_udf_distributor_class,
@@ -726,6 +730,14 @@ def populate_udf_output_table(
                                 "UDF execution was canceled by the user."
                             ) from None
                         if retval := process.poll():
+                            if retval == QUERY_SCRIPT_CANCELED_EXIT_CODE:
+                                raise QueryScriptCancelError(
+                                    "UDF execution was canceled by the user."
+                                )
+                            if retval == QUERY_SCRIPT_ABORTED_EXIT_CODE:
+                                raise QueryScriptAbortError(
+                                    "UDF execution aborted: job already terminated."
+                                )
                             raise RuntimeError(
                                 f"UDF Execution Failed! Exit code: {retval}"
                             )
@@ -761,6 +773,9 @@ def populate_udf_output_table(
                         processed_cb.close()
                         generated_cb.close()
 
+            except QueryScriptAbortError:
+                catalog.warehouse.close()
+                sys.exit(QUERY_SCRIPT_ABORTED_EXIT_CODE)
             except QueryScriptCancelError:
                 catalog.warehouse.close()
                 sys.exit(QUERY_SCRIPT_CANCELED_EXIT_CODE)

tests/unit/test_udf_subprocess_exit_codes.py

@@ -0,0 +1,130 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+import datachain as dc
+from datachain.catalog import (
+    QUERY_SCRIPT_ABORTED_EXIT_CODE,
+    QUERY_SCRIPT_CANCELED_EXIT_CODE,
+)
+
+
+def _make_mock_popen(poll_return_value=0, communicate_side_effect=None):
+    mock_process = MagicMock()
+    mock_process.__enter__ = MagicMock(return_value=mock_process)
+    mock_process.__exit__ = MagicMock(return_value=False)
+    if communicate_side_effect:
+        mock_process.communicate.side_effect = communicate_side_effect
+    else:
+        mock_process.communicate.return_value = (b"", b"")
+    mock_process.poll.return_value = poll_return_value
+    return mock_process
+
+
+def _build_parallel_chain(session):
+    def identity(x: int) -> int:
+        return x
+
+    return (
+        dc.read_values(x=list(range(5)), session=session)
+        .settings(parallel=2)
+        .map(identity, output={"result": int})
+    )
+
+
+def test_aborted_exit_code_causes_sys_exit(test_session_tmpfile, monkeypatch):
+    monkeypatch.delenv("DATACHAIN_DISTRIBUTED", raising=False)
+    chain = _build_parallel_chain(test_session_tmpfile)
+
+    mock_process = _make_mock_popen(poll_return_value=QUERY_SCRIPT_ABORTED_EXIT_CODE)
+
+    with patch("datachain.query.dataset.subprocess.Popen", return_value=mock_process):
+        with pytest.raises(SystemExit) as exc_info:
+            chain.to_list()
+
+    assert exc_info.value.code == QUERY_SCRIPT_ABORTED_EXIT_CODE
+
+
+def test_canceled_exit_code_causes_sys_exit(test_session_tmpfile, monkeypatch):
+    monkeypatch.delenv("DATACHAIN_DISTRIBUTED", raising=False)
+    chain = _build_parallel_chain(test_session_tmpfile)
+
+    mock_process = _make_mock_popen(poll_return_value=QUERY_SCRIPT_CANCELED_EXIT_CODE)
+
+    with patch("datachain.query.dataset.subprocess.Popen", return_value=mock_process):
+        with pytest.raises(SystemExit) as exc_info:
+            chain.to_list()
+
+    assert exc_info.value.code == QUERY_SCRIPT_CANCELED_EXIT_CODE
+
+
+def test_unknown_exit_code_raises_runtime_error(test_session_tmpfile, monkeypatch):
+    monkeypatch.delenv("DATACHAIN_DISTRIBUTED", raising=False)
+    chain = _build_parallel_chain(test_session_tmpfile)
+
+    mock_process = _make_mock_popen(poll_return_value=42)
+
+    with patch("datachain.query.dataset.subprocess.Popen", return_value=mock_process):
+        with pytest.raises(RuntimeError, match="UDF Execution Failed! Exit code: 42"):
+            chain.to_list()
+
+
+def test_keyboard_interrupt_causes_cancel_sys_exit(test_session_tmpfile, monkeypatch):
+    monkeypatch.delenv("DATACHAIN_DISTRIBUTED", raising=False)
+    chain = _build_parallel_chain(test_session_tmpfile)
+
+    mock_process = _make_mock_popen(
+        communicate_side_effect=KeyboardInterrupt("simulated ctrl-c")
+    )
+
+    with patch("datachain.query.dataset.subprocess.Popen", return_value=mock_process):
+        with pytest.raises(SystemExit) as exc_info:
+            chain.to_list()
+
+    assert exc_info.value.code == QUERY_SCRIPT_CANCELED_EXIT_CODE
+
+
+def test_aborted_exit_code_closes_warehouse(test_session_tmpfile, monkeypatch):
+    monkeypatch.delenv("DATACHAIN_DISTRIBUTED", raising=False)
+    chain = _build_parallel_chain(test_session_tmpfile)
+
+    mock_process = _make_mock_popen(poll_return_value=QUERY_SCRIPT_ABORTED_EXIT_CODE)
+    warehouse = test_session_tmpfile.catalog.warehouse
+    original_close = warehouse.close
+    close_called = False
+
+    def tracking_close():
+        nonlocal close_called
+        close_called = True
+        original_close()
+
+    monkeypatch.setattr(warehouse, "close", tracking_close)
+
+    with patch("datachain.query.dataset.subprocess.Popen", return_value=mock_process):
+        with pytest.raises(SystemExit):
+            chain.to_list()
+
+    assert close_called
+
+
+def test_canceled_exit_code_closes_warehouse(test_session_tmpfile, monkeypatch):
+    monkeypatch.delenv("DATACHAIN_DISTRIBUTED", raising=False)
+    chain = _build_parallel_chain(test_session_tmpfile)
+
+    mock_process = _make_mock_popen(poll_return_value=QUERY_SCRIPT_CANCELED_EXIT_CODE)
+    warehouse = test_session_tmpfile.catalog.warehouse
+    original_close = warehouse.close
+    close_called = False
+
+    def tracking_close():
+        nonlocal close_called
+        close_called = True
+        original_close()
+
+    monkeypatch.setattr(warehouse, "close", tracking_close)
+
+    with patch("datachain.query.dataset.subprocess.Popen", return_value=mock_process):
+        with pytest.raises(SystemExit):
+            chain.to_list()
+
+    assert close_called