Raise JobTimeoutError instead of bare TimeoutError

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Changed
+
+- `Beaker.(job|experiment).(wait_for|as_completed)()` methods now raises a more specific `JobTimeoutError` (which inherits from `TimeoutError`) instead of a generic `TimeoutError`.
+
 ## [v1.5.0](https://github.com/allenai/beaker-py/releases/tag/v1.5.0) - 2022-06-16
 
 ### Added

beaker/exceptions.py

@@ -175,3 +175,7 @@ class ChecksumFailedError(BeakerError):
 
 class JobFailedError(BeakerError):
     pass
+
+
+class JobTimeoutError(BeakerError, TimeoutError):
+    pass

beaker/services/experiment.py

@@ -347,7 +347,7 @@ def wait_for(
             :class:`~beaker.exceptions.JobFailedError` will be raised for non-zero exit codes.
 
         :raises ExperimentNotFound: If any experiment can't be found.
-        :raises TimeoutError: If the ``timeout`` expires.
+        :raises JobTimeoutError: If the ``timeout`` expires.
         :raises DuplicateExperimentError: If the same experiment is given as an argument more than once.
         :raises JobFailedError: If ``strict=True`` and any job finishes with a non-zero exit code.
         :raises BeakerError: Any other :class:`~beaker.exceptions.BeakerError` type that can occur.
@@ -401,7 +401,7 @@ def as_completed(
             :class:`~beaker.exceptions.JobFailedError` will be raised for non-zero exit codes.
 
         :raises ExperimentNotFound: If any experiment can't be found.
-        :raises TimeoutError: If the ``timeout`` expires.
+        :raises JobTimeoutError: If the ``timeout`` expires.
         :raises DuplicateExperimentError: If the same experiment is given as an argument more than once.
         :raises JobFailedError: If ``strict=True`` and any job finishes with a non-zero exit code.
         :raises BeakerError: Any other :class:`~beaker.exceptions.BeakerError` type that can occur.
@@ -488,7 +488,7 @@ def complete_experiment(exp_id: str) -> Experiment:
                 # Check for timeout.
                 elapsed = time.time() - start
                 if timeout is not None and elapsed >= timeout:
-                    raise TimeoutError
+                    raise JobTimeoutError
 
                 if incomplete_jobs:
                     # Wait for current stack of incomplete jobs to finalize.

beaker/services/job.py

@@ -264,7 +264,7 @@ def wait_for(
             :class:`~beaker.exceptions.JobFailedError` will be raised for non-zero exit codes.
 
         :raises JobNotFound: If any job can't be found.
-        :raises TimeoutError: If the ``timeout`` expires.
+        :raises JobTimeoutError: If the ``timeout`` expires.
         :raises DuplicateJobError: If the same job is given as an argument more than once.
         :raises JobFailedError: If ``strict=True`` and any job finishes with a non-zero exit code.
         :raises BeakerError: Any other :class:`~beaker.exceptions.BeakerError` type that can occur.
@@ -317,7 +317,7 @@ def as_completed(
             :class:`~beaker.exceptions.JobFailedError` will be raised for non-zero exit codes.
 
         :raises JobNotFound: If any job can't be found.
-        :raises TimeoutError: If the ``timeout`` expires.
+        :raises JobTimeoutError: If the ``timeout`` expires.
         :raises DuplicateJobError: If the same job is given as an argument more than once.
         :raises JobFailedError: If ``strict=True`` and any job finishes with a non-zero exit code.
         :raises BeakerError: Any other :class:`~beaker.exceptions.BeakerError` type that can occur.
@@ -412,7 +412,7 @@ def display_name(j: Job) -> str:
 
                 elapsed = time.time() - start
                 if timeout is not None and elapsed >= timeout:
-                    raise TimeoutError
+                    raise JobTimeoutError
                 time.sleep(poll_interval)
         finally:
             if owned_progress: