Repository: inductiva/task-runner
Pull request #161: Retry Output Uploads on HTTP Errors
Author: login: HenriquePreto
Created: 2025-05-09T13:29:56Z | Merged: 2025-05-13T10:39:28Z

Closes https://github.com/inductiva/tasks/issues/1089.

Adds retry logic for output uploads using [tenacity](https://github.com/jd/tenacity), which has a nicer and more flexible interface than our current `@utils.retry`. Might be worth replacing our custom retry with this across the codebase later on.

**Tests**:

1) To test the retry mechanism, I generated a signed URL for downloading, which triggers a 403 error when used for uploading.

```python
import uuid
from task_runner import event_logger, file_manager
from task_runner.utils import files

URL = "<download-signed-url>"

class MockEventLogger(event_logger.BaseEventLogger):
    def log(self, event):
        print(event)

stream_process = files.get_seven_zip_stream_process(
    "../benchmark-upload-of-task-runner-outputs/myfolder")
data = stream_process.stdout

response = file_manager.WebApiFileManager.retry_upload(
    method="PUT",
    url=URL,
    data=data,
    task_id="id",
    task_runner_uuid=uuid.uuid4(),
    event_logger=MockEventLogger(),
)

print(response)
```

2) For safety, I also ran a benchmark, but as expected, this had no impact on performance.

```python
import uuid
from task_runner import event_logger, file_manager, utils
from task_runner.utils import files


URL = "<upload-signed-url>"

@utils.execution_time_with_result
def upload_no_retry(url, data):
    return file_manager.WebApiFileManager.upload("PUT", url, data)


class MockEventLogger(event_logger.BaseEventLogger):
    def log(self, event):
        print(event)


if __name__ == "__main__":
    local_path = "<10GB-output-folder>"

    stream_process = files.get_seven_zip_stream_process(local_path)
    r1, t1 = upload_no_retry(URL, stream_process.stdout)

    assert r1.status_code == 200

    print(r1, t1)
    # <Response [200]> 220.21145510673523

    stream_process = files.get_seven_zip_stream_process(local_path)
    r2, t2 = file_manager.WebApiFileManager.retry_upload(
        method="PUT",
        url=URL,
        data=stream_process.stdout,
        task_id="id",
        task_runner_uuid=uuid.uuid4(),
        event_logger=MockEventLogger(),
    )

    assert r2.status_code == 200

    print(r2, t2)
    # <Response [200]> 198.83276176452637
```

Note: During benchmarking, the upload raised an exception that wasn't among the expected HTTPError types (it was a `requests.exceptions.SSLError`). So, I decided to catch **all** exceptions during upload (check this [commit](https://github.com/inductiva/task-runner/pull/161/commits/c2eb36c6cd5f1088ce41cabdcaa527cb4879c322))



Diff:
diff --git a/task-runner/requirements.txt b/task-runner/requirements.txt
index 3ef7fc09..2c705591 100644
--- a/task-runner/requirements.txt
+++ b/task-runner/requirements.txt
@@ -6,4 +6,5 @@ fsspec
 gcsfs
 stream-zip==0.0.81
 pysocks
-gputil
\ No newline at end of file
+gputil
+tenacity
\ No newline at end of file
diff --git a/task-runner/task_runner/file_manager.py b/task-runner/task_runner/file_manager.py
index c27e421b..cdf4d39a 100644
--- a/task-runner/task_runner/file_manager.py
+++ b/task-runner/task_runner/file_manager.py
@@ -1,12 +1,14 @@
 import abc
 import os
 import pathlib
-import time
+import traceback
 import urllib
 import urllib.request
 import uuid
 
 import requests
+import tenacity
+from inductiva_api import events
 from typing_extensions import override
 
 import task_runner
@@ -32,7 +34,9 @@ def upload_output(
         task_id: str,
         task_dir_remote: str,
         local_path: str,
+        task_runner_uuid: uuid.UUID,
         operations_logger: OperationsLogger,
+        event_logger: task_runner.BaseEventLogger,
         stream_zip: bool = True,
         compress_with: str = "AUTO",
     ):
@@ -78,13 +82,67 @@ def download_input(
         url = self._api_client.get_download_input_url(storage_dir)
         urllib.request.urlretrieve(url, dest_path)
 
+    @staticmethod
+    def upload(method: str, url: str, data) -> requests.Response:
+        response = requests.request(
+            method=method,
+            url=url,
+            data=data,
+            timeout=WebApiFileManager.REQUEST_TIMEOUT_S,
+            headers={"Content-Type": "application/octet-stream"},
+        )
+        response.raise_for_status()
+        return response
+
+    @staticmethod
+    def make_fail_upload_hook(task_id: str, task_runner_uuid: uuid.UUID,
+                              event_logger: task_runner.BaseEventLogger):
+
+        def fail_upload_hook(retry_state: tenacity.RetryCallState):
+            error = retry_state.outcome.exception()
+            if error is None:
+                return
+
+            message = utils.get_exception_root_cause_message(error)
+
+            event_logger.log(
+                events.TaskOutputUploadFailed(
+                    id=task_id,
+                    machine_id=task_runner_uuid,
+                    error_message=message,
+                    traceback=traceback.format_exc(),
+                ))
+
+        return fail_upload_hook
+
+    @staticmethod
+    @utils.execution_time_with_result
+    def retry_upload(
+            method: str, url: str, data, task_id: str,
+            task_runner_uuid: uuid.UUID,
+            event_logger: task_runner.BaseEventLogger) -> requests.Response:
+        for attempt in tenacity.Retrying(
+                stop=tenacity.stop_after_attempt(5),
+                wait=tenacity.wait_exponential(multiplier=10),
+                before_sleep=WebApiFileManager.make_fail_upload_hook(
+                    task_id=task_id,
+                    task_runner_uuid=task_runner_uuid,
+                    event_logger=event_logger,
+                ),
+                reraise=True,
+        ):
+            with attempt:
+                return WebApiFileManager.upload(method, url, data)
+
     @override
     def upload_output(
         self,
         task_id: str,
         task_dir_remote: str,
         local_path: str,
+        task_runner_uuid: uuid.UUID,
         operations_logger: OperationsLogger,
+        event_logger: task_runner.BaseEventLogger,
         stream_zip: bool = True,
         compress_with: str = "AUTO",
     ):
@@ -112,18 +170,14 @@ def upload_output(
 
         operation = operations_logger.start_operation(
             OperationName.UPLOAD_OUTPUT, task_id)
-        start_time = time.time()
-        resp = requests.request(
+        resp, upload_time = WebApiFileManager.retry_upload(
             method=upload_info.method,
             url=upload_info.url,
             data=data,
-            timeout=self.REQUEST_TIMEOUT_S,
-            headers={
-                "Content-Type": "application/octet-stream",
-            },
+            task_id=task_id,
+            task_runner_uuid=task_runner_uuid,
+            event_logger=event_logger,
         )
-        upload_time = time.time() - start_time
-        resp.raise_for_status()
 
         operation.end(attributes={"execution_time_s": upload_time})
 
diff --git a/task-runner/task_runner/task_request_handler.py b/task-runner/task_runner/task_request_handler.py
index 4796b90d..e00dc3c8 100644
--- a/task-runner/task_runner/task_request_handler.py
+++ b/task-runner/task_runner/task_request_handler.py
@@ -393,7 +393,7 @@ def __call__(self, request: dict[str, str]) -> None:
                 self._publish_event(
                     events.TaskOutputUploadFailed(
                         id=self.task_id,
-                        machine_id=self.executer_uuid,
+                        machine_id=self.task_runner_uuid,
                         error_message=message,
                         traceback=traceback.format_exc(),
                     ))
@@ -582,7 +582,7 @@ def _pack_output(self) -> int:
         if output_total_files is not None:
             self._post_task_metric(utils.OUTPUT_TOTAL_FILES, output_total_files)
 
-        output_zipped_bytes, zip_duration, upload_duration = (
+        output_zipped_bytes, zip_duration, upload_duration = \
             self.file_manager.upload_output(
                 self.task_id,
                 self.task_dir_remote,
@@ -590,7 +590,9 @@ def _pack_output(self) -> int:
                 stream_zip=self.stream_zip,
                 compress_with=self.compress_with,
                 operations_logger=self._operations_logger,
-            ))
+                event_logger=self.event_logger,
+                task_runner_uuid=self.task_runner_uuid,
+            )
 
         logging.info("Output zipped in: %s seconds", zip_duration)
 


Reviews:
user: login: hpenedones
state: COMMENTED
author_association: NONE
submitted_at: 2025-05-09T15:18:19Z
user: login: hpenedones
state: COMMENTED
author_association: NONE
submitted_at: 2025-05-09T15:18:37Z
user: login: HenriquePreto
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-12T09:26:50Z
user: login: hpenedones
state: COMMENTED
author_association: NONE
submitted_at: 2025-05-12T15:08:14Z
user: login: bgregoinductiva
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T10:30:43Z
user: login: rcvalerio
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T10:34:41Z

Review comments:
diff_hunk: @@ -78,13 +82,67 @@ def download_input(
         url = self._api_client.get_download_input_url(storage_dir)
         urllib.request.urlretrieve(url, dest_path)
 
+    @staticmethod
+    def upload(method: str, url: str, data) -> requests.Response:
+        response = requests.request(
+            method=method,
+            url=url,
+            data=data,
+            timeout=WebApiFileManager.REQUEST_TIMEOUT_S,
+            headers={"Content-Type": "application/octet-stream"},
+        )
+        response.raise_for_status()
+        return response
+
+    @staticmethod
+    def make_fail_upload_hook(task_id: str, task_runner_uuid: uuid.UUID,
+                              event_logger: task_runner.BaseEventLogger):
+
+        def fail_upload_hook(retry_state: tenacity.RetryCallState):
+            error = retry_state.outcome.exception()
+            if error is None:
+                return
+
+            message = utils.get_exception_root_cause_message(error)
+
+            event_logger.log(
+                events.TaskOutputUploadFailed(
+                    id=task_id,
+                    machine_id=task_runner_uuid,
+                    error_message=message,
+                    traceback=traceback.format_exc(),
+                ))
+
+        return fail_upload_hook
+
+    @staticmethod
+    @utils.execution_time_with_result
+    def retry_upload(
+            method: str, url: str, data, task_id: str,
+            task_runner_uuid: uuid.UUID,
+            event_logger: task_runner.BaseEventLogger) -> requests.Response:
+        for attempt in tenacity.Retrying(
+                stop=tenacity.stop_never,
path: task-runner/task_runner/file_manager.py
user: login: hpenedones
body: I think that stop_never might be too much tenacity -- because this keeps the VM running while it is trying, which incurs costs. I think that after a fair amount of trials this might be wasteful (especially in the cases where the simulation might have been "fast" -- so it is not too expensive to re-run). Maybe this is not "trivial" to implement and we can do it later (now we could just retry, say 3 or 5 times), but we could consider indexing how much we try, based on the amount of compute the task took... Say, at most I will spend 50% of total task duration retrying to upload, but beyond that it starts to be more wasteful than losing the data...
created_at: 2025-05-09T15:18:18Z
author_association: NONE
diff_hunk: @@ -1,12 +1,14 @@
 import abc
 import os
 import pathlib
-import time
+import traceback
 import urllib
 import urllib.request
 import uuid
 
 import requests
+import tenacity
path: task-runner/task_runner/file_manager.py
user: login: hpenedones
body: Super cool! Nice finding!
created_at: 2025-05-09T15:18:36Z
author_association: NONE
diff_hunk: @@ -78,13 +82,67 @@ def download_input(
         url = self._api_client.get_download_input_url(storage_dir)
         urllib.request.urlretrieve(url, dest_path)
 
+    @staticmethod
+    def upload(method: str, url: str, data) -> requests.Response:
+        response = requests.request(
+            method=method,
+            url=url,
+            data=data,
+            timeout=WebApiFileManager.REQUEST_TIMEOUT_S,
+            headers={"Content-Type": "application/octet-stream"},
+        )
+        response.raise_for_status()
+        return response
+
+    @staticmethod
+    def make_fail_upload_hook(task_id: str, task_runner_uuid: uuid.UUID,
+                              event_logger: task_runner.BaseEventLogger):
+
+        def fail_upload_hook(retry_state: tenacity.RetryCallState):
+            error = retry_state.outcome.exception()
+            if error is None:
+                return
+
+            message = utils.get_exception_root_cause_message(error)
+
+            event_logger.log(
+                events.TaskOutputUploadFailed(
+                    id=task_id,
+                    machine_id=task_runner_uuid,
+                    error_message=message,
+                    traceback=traceback.format_exc(),
+                ))
+
+        return fail_upload_hook
+
+    @staticmethod
+    @utils.execution_time_with_result
+    def retry_upload(
+            method: str, url: str, data, task_id: str,
+            task_runner_uuid: uuid.UUID,
+            event_logger: task_runner.BaseEventLogger) -> requests.Response:
+        for attempt in tenacity.Retrying(
+                stop=tenacity.stop_never,
path: task-runner/task_runner/file_manager.py
user: login: HenriquePreto
body: Yes, I see your point -- that sounds good. My goal was to give us the opportunity to act manually in case something unexpected happens. But I agree it's a bit too generic. I think it's a good idea to leave it at 5 attempts for now, and if the issue happens again, we can reevaluate this.

created_at: 2025-05-12T09:26:50Z
author_association: CONTRIBUTOR
path: task-runner/tests/unit/test_file_manager.py
user: login: hpenedones
body: seems like this file is empty...
created_at: 2025-05-12T15:08:14Z
author_association: NONE

PR comments:
user: login: hpenedones
created_at: 2025-05-12T15:01:09Z
body: So, if there is a consistent error what is the maximum time this will be retrying (how long do the 5 tries take, with the exponential decay)?
author_association: NONE
user: login: HenriquePreto
created_at: 2025-05-12T15:03:23Z
body: > So, if there is a consistent error what is the maximum time this will be retrying (how long do the 5 tries take, with the exponential decay)?

It's like: 1s, 2s, 4s, 8s, 16s ...
author_association: CONTRIBUTOR
user: login: hpenedones
created_at: 2025-05-12T15:09:27Z
body: Ah, ok, thanks -- so the first retry is just after 1s, and therefore the other retries are also relatively fast...
author_association: NONE
user: login: hpenedones
created_at: 2025-05-12T15:10:11Z
body: Or is it 1 min (i.e. 60 seconds)?
https://github.com/inductiva/task-runner/blob/23739ffe60e429e2850ebe121d99277b39471361/task-runner/task_runner/file_manager.py#L53


author_association: NONE
user: login: HenriquePreto
created_at: 2025-05-12T15:15:16Z
body: > Ah, ok, thanks -- so the first retry is just after 1s, and therefore the other retries are also relatively fast...

Yes.
author_association: CONTRIBUTOR
user: login: HenriquePreto
created_at: 2025-05-12T15:18:35Z
body: > Or is it 1 min (i.e. 60 seconds)?

No, that's only the timeout argument that represents how long the client will wait for a server to send a response. If the server does not respond within this time, a `requests.exceptions.Timeout` will be raised. But there are other exceptions that might happen besides this one.

We can increase the time between retries by tweaking the arguments passed to `tenacity.wait_exponential()`.
author_association: CONTRIBUTOR
user: login: hpenedones
created_at: 2025-05-12T15:24:04Z
body: Ah, ok -- thanks for clarifying! yeah, maybe let's then increase the initial time between retries to be maybe 10 seconds -- this way the total time of 5 retries, instead of ~30 seconds, will be about 5 minutes which would be more reasonable to have some chances of some network problem going away... What do you think?
author_association: NONE
user: login: HenriquePreto
created_at: 2025-05-12T15:28:33Z
body: > Ah, ok -- thanks for clarifying! yeah, maybe let's then increase the initial time between retries to be maybe 10 seconds -- this way the total time of 5 retries, instead of ~30 seconds, will be about 5 minutes which would be more reasonable to have some chances of some network problem going away... What do you think?

Sounds good!
author_association: CONTRIBUTOR