Improve batch job timeout error message (#80)

Author: gregorybchris

hume/_batch/batch_job.py

@@ -14,6 +14,9 @@
 class BatchJob:
     """Batch job."""
 
+    TIMEOUT_MESSAGE = ("Connection to API has been terminated after {}s, but your job will continue to run. "
+                       "Get a reference to your job with `client.get_job('{}')` at any time.")
+
     def __init__(self, client: "HumeBatchClient", job_id: str):
         """Construct a BatchJob.
 
@@ -87,15 +90,15 @@ def await_complete(self, timeout: int = 300) -> BatchJobDetails:
         if timeout < 1:
             raise ValueError("timeout must be at least 1 second")
 
-        return self._await_complete(timeout=timeout)
+        # pylint: disable=unused-argument
+        @retry(timeout_message=self.TIMEOUT_MESSAGE.format(timeout, self.id))
+        def _await_complete(timeout: int = timeout) -> BatchJobDetails:
+            details = self._client.get_job_details(self.id)
+            if not BatchJobStatus.is_terminal(details.state.status):
+                raise RetryIterError
+            return details
 
-    # pylint: disable=unused-argument
-    @retry()
-    def _await_complete(self, timeout: int = 300) -> BatchJobDetails:
-        details = self._client.get_job_details(self.id)
-        if not BatchJobStatus.is_terminal(details.state.status):
-            raise RetryIterError
-        return details
+        return _await_complete(timeout=timeout)
 
     def __repr__(self) -> str:
         """Get the string representation of the `BatchJob`.

hume/_common/retry_utils.py

@@ -1,7 +1,7 @@
 """Function retry utilities."""
 import logging
 import time
-from typing import cast, Callable, Type, TypeVar
+from typing import Optional, cast, Callable, Type, TypeVar
 from typing_extensions import ParamSpec
 
 from hume.error.hume_client_exception import HumeClientException
@@ -24,6 +24,7 @@ def retry(
     max_delay: int = 300,
     backoff_factor: int = 2,
     error_type: Type[Exception] = RetryIterError,
+    timeout_message: Optional[str] = None,
 ) -> Callable[[Callable[P, R]], Callable[P, R]]:
     """Retry decorator for exponential backoff retry.
 
@@ -34,6 +35,8 @@ def retry(
         error_type (Type[Exception]): Class of exception to expect from decorated function when
             the function fails. Raise this exception type if the retry iteration has failed.
             Defaults to RetryIterError.
+        timeout_message (Optional[str]): A message that will be used when raising a
+            HumeClientException on timeout.
 
     Returns:
         Callable[[Callable[P, R]], Callable[P, R]]: Function decorator.
@@ -67,7 +70,10 @@ def func_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
 
                 retry_timeout = retry_kwargs["timeout"]
                 if total_await_time >= retry_timeout:
-                    raise HumeClientException(f"Request timed out after {retry_timeout}s")
+                    message = timeout_message
+                    if timeout_message is None:
+                        message = f"Request timed out after {retry_timeout}s"
+                    raise HumeClientException(message)
 
                 time.sleep(delay)
                 total_await_time += delay

pyproject.toml

@@ -25,7 +25,7 @@ license = "Proprietary"
 name = "hume"
 readme = "README.md"
 repository = "https://github.com/HumeAI/hume-python-sdk"
-version = "0.3.5"
+version = "0.3.6"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4"