diff --git a/alab_management/resource_manager/resource_requester.py b/alab_management/resource_manager/resource_requester.py index 2988b5c5..496adaaa 100644 --- a/alab_management/resource_manager/resource_requester.py +++ b/alab_management/resource_manager/resource_requester.py @@ -2,7 +2,7 @@ TaskLauncher is the core module of the system, which actually executes the tasks. """ - +import concurrent import time from concurrent.futures import Future from datetime import datetime @@ -28,6 +28,14 @@ ] # the raw request sent by task process +class CombinedTimeoutError(TimeoutError, concurrent.futures.TimeoutError): + """ + Combined TimeoutError. + + If you catch either TimeoutError or concurrent.futures.TimeoutError, this will catch both. + """ + + class ResourcesRequest(BaseModel): """ This class is used to validate the resource request. Each request should have a format of @@ -231,7 +239,12 @@ def request_resources( ) # DB_ACCESS_OUTSIDE_VIEW _id: ObjectId = cast(ObjectId, result.inserted_id) self._waiting[_id] = {"f": f, "device_str_to_request": device_str_to_request} - result = f.result(timeout=timeout) + try: + result = f.result(timeout=timeout) + except concurrent.futures.TimeoutError as e: + raise CombinedTimeoutError( + f"Request {result.inserted_id} timed out after {timeout} seconds." + ) from e return { **self._post_process_requested_resource( devices=result["devices"],