client: Raise error when returning result of failed task

Signed-off-by: Phoevos Kalemkeris <[email protected]>

Author: phoevos

client/cogstack_model_gateway_client/__init__.py

@@ -1,2 +1,3 @@
 from cogstack_model_gateway_client.client import GatewayClient as GatewayClient
 from cogstack_model_gateway_client.client import GatewayClientSync as GatewayClientSync
+from cogstack_model_gateway_client.exceptions import TaskFailedError as TaskFailedError

client/cogstack_model_gateway_client/client.py

@@ -5,7 +5,7 @@
 
 import httpx
 
-from cogstack_model_gateway_client.exceptions import retry_if_network_error
+from cogstack_model_gateway_client.exceptions import TaskFailedError, retry_if_network_error
 
 
 class GatewayClient:
@@ -103,7 +103,11 @@ async def submit_task(
         wait_for_completion: bool = False,
         return_result: bool = True,
     ):
-        """Submit a task to the Gateway and return the task info."""
+        """Submit a task to the Gateway and return the task info.
+
+        Raises:
+            TaskFailedError: If the task fails and wait_for_completion=True, return_result=True.
+        """
         model_name = model_name or self.default_model
         if not model_name:
             raise ValueError("Please provide a model name or set a default model for the client.")
@@ -118,7 +122,11 @@ async def submit_task(
             task_uuid = task_info["uuid"]
             task_info = await self.wait_for_task(task_uuid)
             if return_result:
-                return await self.get_task_result(task_uuid)
+                if task_info.get("status") == "succeeded":
+                    return await self.get_task_result(task_uuid)
+                else:
+                    error_message = task_info.get("error_message", "Unknown error")
+                    raise TaskFailedError(task_uuid, error_message, task_info)
         return task_info
 
     async def process(
@@ -128,7 +136,11 @@ async def process(
         wait_for_completion: bool = True,
         return_result: bool = True,
     ):
-        """Generate annotations for the provided text."""
+        """Generate annotations for the provided text.
+
+        Raises:
+            TaskFailedError: If the task fails and wait_for_completion=True, return_result=True.
+        """
         return await self.submit_task(
             model_name=model_name,
             task="process",
@@ -145,7 +157,11 @@ async def process_bulk(
         wait_for_completion: bool = True,
         return_result: bool = True,
     ):
-        """Generate annotations for a list of texts."""
+        """Generate annotations for a list of texts.
+
+        Raises:
+            TaskFailedError: If the task fails and wait_for_completion=True, return_result=True.
+        """
         return await self.submit_task(
             model_name=model_name,
             task="process_bulk",
@@ -166,7 +182,11 @@ async def redact(
         wait_for_completion: bool = True,
         return_result: bool = True,
     ):
-        """Redact sensitive information from the provided text."""
+        """Redact sensitive information from the provided text.
+
+        Raises:
+            TaskFailedError: If the task fails and wait_for_completion=True, return_result=True.
+        """
         params = {
             k: v
             for k, v in {
@@ -238,15 +258,20 @@ async def get_task_result(self, task_uuid: str, parse: bool = True):
     async def wait_for_task(
         self, task_uuid: str, detail: bool = True, raise_on_error: bool = False
     ):
-        """Poll Gateway until the task reaches a final state."""
+        """Poll Gateway until the task reaches a final state.
+
+        Raises:
+            TaskFailedError: If raise_on_error=True and the task fails.
+            TimeoutError: If timeout is reached before task completion.
+        """
         start = asyncio.get_event_loop().time()
         while True:
             task = await self.get_task(task_uuid, detail=detail)
             status = task.get("status")
             if status in ("succeeded", "failed"):
                 if status == "failed" and raise_on_error:
                     error_message = task.get("error_message", "Unknown error")
-                    raise RuntimeError(f"Task '{task_uuid}' failed: {error_message}")
+                    raise TaskFailedError(task_uuid, error_message, task)
                 return task
             if self.timeout is not None and asyncio.get_event_loop().time() - start > self.timeout:
                 raise TimeoutError(f"Timed out waiting for task '{task_uuid}' to complete")
@@ -365,7 +390,11 @@ def submit_task(
         wait_for_completion: bool = False,
         return_result: bool = True,
     ):
-        """Submit a task to the Gateway and return the task info."""
+        """Submit a task to the Gateway and return the task info.
+
+        Raises:
+            TaskFailedError: If the task fails and wait_for_completion=True, return_result=True.
+        """
         return asyncio.run(
             self._client.submit_task(
                 model_name=model_name,
@@ -387,7 +416,11 @@ def process(
         wait_for_completion: bool = True,
         return_result: bool = True,
     ):
-        """Generate annotations for the provided text."""
+        """Generate annotations for the provided text.
+
+        Raises:
+            TaskFailedError: If the task fails and wait_for_completion=True, return_result=True.
+        """
         return asyncio.run(
             self._client.process(
                 text=text,
@@ -404,7 +437,11 @@ def process_bulk(
         wait_for_completion: bool = True,
         return_result: bool = True,
     ):
-        """Generate annotations for a list of texts."""
+        """Generate annotations for a list of texts.
+
+        Raises:
+            TaskFailedError: If the task fails and wait_for_completion=True, return_result=True.
+        """
         return asyncio.run(
             self._client.process_bulk(
                 texts=texts,
@@ -425,7 +462,11 @@ def redact(
         wait_for_completion: bool = True,
         return_result: bool = True,
     ):
-        """Redact sensitive information from the provided text."""
+        """Redact sensitive information from the provided text.
+
+        Raises:
+            TaskFailedError: If the task fails and wait_for_completion=True, return_result=True.
+        """
         return asyncio.run(
             self._client.redact(
                 text=text,
@@ -452,7 +493,12 @@ def get_task_result(self, task_uuid: str, parse: bool = True):
         return asyncio.run(self._client.get_task_result(task_uuid=task_uuid, parse=parse))
 
     def wait_for_task(self, task_uuid: str, detail: bool = True, raise_on_error: bool = False):
-        """Poll Gateway until the task reaches a final state."""
+        """Poll Gateway until the task reaches a final state.
+
+        Raises:
+            TaskFailedError: If raise_on_error=True and the task fails.
+            TimeoutError: If timeout is reached before task completion.
+        """
         return asyncio.run(
             self._client.wait_for_task(
                 task_uuid=task_uuid, detail=detail, raise_on_error=raise_on_error

client/cogstack_model_gateway_client/exceptions.py

@@ -12,6 +12,22 @@
 log = logging.getLogger("cmg.client")
 
 
+class TaskFailedError(Exception):
+    """Raised when a Gateway task fails during execution.
+
+    Attributes:
+        task_uuid: The UUID of the failed task
+        error_message: The error message from the task
+        task_info: The full task information dict (optional)
+    """
+
+    def __init__(self, task_uuid: str, error_message: str, task_info: dict = None):
+        self.task_uuid = task_uuid
+        self.error_message = error_message
+        self.task_info = task_info
+        super().__init__(f"Task '{task_uuid}' failed: {error_message}")
+
+
 def is_network_error(exception: Exception):
     """Check if the exception is a network-related error."""
     return isinstance(

pyproject.toml

@@ -59,14 +59,14 @@ line-length = 100
 target-version = "py312"
 
 [tool.ruff.lint.isort]
-known-local-folder = ["common", "gateway", "scheduler", "client"]
+known-local-folder = ["common", "gateway", "scheduler", "cogstack_model_gateway_client"]
 
 [tool.ruff.lint]
 select = ["E", "F", "I", "UP"]
 
 [tool.pytest.ini_options]
 addopts = "-ra -s --disable-warnings --enable-cmg-logging"
-pythonpath = ["."]
+pythonpath = [".", "client"]
 testpaths = ["tests"]
 
 [build-system]

tests/unit/client/test_client.py

@@ -8,7 +8,8 @@
 import httpx
 import pytest
 
-from client.cogstack_model_gateway_client.client import GatewayClient, GatewayClientSync
+from cogstack_model_gateway_client.client import GatewayClient, GatewayClientSync
+from cogstack_model_gateway_client.exceptions import TaskFailedError
 
 
 @pytest.fixture
@@ -159,11 +160,11 @@ async def test_submit_task_wait_for_completion_and_return_result(mock_httpx_asyn
     mock_client_instance.request.return_value = mock_response
 
     mock_wait_for_task = mocker.patch(
-        "client.cogstack_model_gateway_client.client.GatewayClient.wait_for_task",
+        "cogstack_model_gateway_client.client.GatewayClient.wait_for_task",
         new=AsyncMock(return_value={"uuid": "task-123", "status": "succeeded"}),
     )
     mock_get_task_result = mocker.patch(
-        "client.cogstack_model_gateway_client.client.GatewayClient.get_task_result",
+        "cogstack_model_gateway_client.client.GatewayClient.get_task_result",
         new=AsyncMock(return_value="processed text"),
     )
 
@@ -180,6 +181,29 @@ async def test_submit_task_wait_for_completion_and_return_result(mock_httpx_asyn
     mock_wait_for_task.assert_awaited_once_with("task-123")
     mock_get_task_result.assert_awaited_once_with("task-123")
 
+    mock_wait_for_task = mocker.patch(
+        "cogstack_model_gateway_client.client.GatewayClient.wait_for_task",
+        new=AsyncMock(
+            return_value={
+                "uuid": "task-123",
+                "status": "failed",
+                "error_message": "Processing failed",
+            }
+        ),
+    )
+
+    async with GatewayClient(base_url="http://test-gateway.com") as client:
+        with pytest.raises(TaskFailedError, match="Task 'task-123' failed: Processing failed"):
+            await client.submit_task(
+                model_name="my_model",
+                task="process",
+                data="text",
+                wait_for_completion=True,
+                return_result=True,
+            )
+
+    mock_wait_for_task.assert_awaited_once_with("task-123")
+
 
 @pytest.mark.asyncio
 async def test_process_method(mocker):
@@ -335,9 +359,7 @@ async def test_wait_for_task_succeeded(mock_httpx_async_client, mocker):
         {"uuid": "task-polling", "status": "pending"},
         {"uuid": "task-polling", "status": "succeeded", "result": "done"},
     ]
-    mocker.patch(
-        "client.cogstack_model_gateway_client.client.GatewayClient.get_task", new=mock_get_task
-    )
+    mocker.patch("cogstack_model_gateway_client.client.GatewayClient.get_task", new=mock_get_task)
     mocker.patch("asyncio.sleep", new=AsyncMock())
 
     async with GatewayClient(base_url="http://test-gateway.com") as client:
@@ -352,9 +374,7 @@ async def test_wait_for_task_succeeded(mock_httpx_async_client, mocker):
 async def test_wait_for_task_timeout(mock_httpx_async_client, mocker):
     """Test wait_for_task raises TimeoutError."""
     mock_get_task = AsyncMock(return_value={"uuid": "task-polling", "status": "pending"})
-    mocker.patch(
-        "client.cogstack_model_gateway_client.client.GatewayClient.get_task", new=mock_get_task
-    )
+    mocker.patch("cogstack_model_gateway_client.client.GatewayClient.get_task", new=mock_get_task)
     mocker.patch("asyncio.sleep", new=AsyncMock())
 
     async with GatewayClient(base_url="http://test-gateway.com") as client:
@@ -383,9 +403,7 @@ async def mock_get_task_side_effect(*args, **kwargs):
             return {"uuid": "task-polling", "status": "succeeded"}
 
     mock_get_task = AsyncMock(side_effect=mock_get_task_side_effect)
-    mocker.patch(
-        "client.cogstack_model_gateway_client.client.GatewayClient.get_task", new=mock_get_task
-    )
+    mocker.patch("cogstack_model_gateway_client.client.GatewayClient.get_task", new=mock_get_task)
     mocker.patch("asyncio.sleep", new=AsyncMock())
 
     async with GatewayClient(base_url="http://test-gateway.com") as client:
@@ -400,19 +418,19 @@ async def mock_get_task_side_effect(*args, **kwargs):
 
 @pytest.mark.asyncio
 async def test_wait_for_task_failed_raise_on_error(mock_httpx_async_client, mocker):
-    """Test wait_for_task raises RuntimeError on task failure with raise_on_error."""
+    """Test wait_for_task raises TaskFailedError on task failure with raise_on_error."""
     mock_get_task = AsyncMock()
     mock_get_task.side_effect = [
         {"uuid": "task-polling", "status": "pending"},
         {"uuid": "task-polling", "status": "failed", "error_message": "Something went wrong"},
     ]
-    mocker.patch(
-        "client.cogstack_model_gateway_client.client.GatewayClient.get_task", new=mock_get_task
-    )
+    mocker.patch("cogstack_model_gateway_client.client.GatewayClient.get_task", new=mock_get_task)
     mocker.patch("asyncio.sleep", new=AsyncMock())
 
     async with GatewayClient(base_url="http://test-gateway.com") as client:
-        with pytest.raises(RuntimeError, match="Task 'task-polling' failed: Something went wrong"):
+        with pytest.raises(
+            TaskFailedError, match="Task 'task-polling' failed: Something went wrong"
+        ):
             await client.wait_for_task("task-polling", raise_on_error=True)
 
         assert mock_get_task.await_count == 2
@@ -426,9 +444,7 @@ async def test_wait_for_task_failed_no_raise_on_error(mock_httpx_async_client, m
         {"uuid": "task-polling", "status": "pending"},
         {"uuid": "task-polling", "status": "failed", "error_message": "Something went wrong"},
     ]
-    mocker.patch(
-        "client.cogstack_model_gateway_client.client.GatewayClient.get_task", new=mock_get_task
-    )
+    mocker.patch("cogstack_model_gateway_client.client.GatewayClient.get_task", new=mock_get_task)
     mocker.patch("asyncio.sleep", new=AsyncMock())
 
     async with GatewayClient(base_url="http://test-gateway.com") as client: