Repository: dwavesystems/dwave-cloud-client

Issue:
Upload part worker does not recover from failures to correctly retry upload
Individual part problem upload should be retried 2 times in case the upload fails due to any unhandled exception (e.g. SAPI returns `502 Bad Gateway` during overload). This is achieved with the `retried` decorator.

https://github.com/dwavesystems/dwave-cloud-client/blob/88e14e9b745957569d86ca8d8fcd1461ff7101b9/dwave/cloud/client.py#L1670-L1672

The problem is not all arguments are immutable. `part_stream` (uploaded part data binary stream) should always be at the beginning.

We can either always rewind the `part_stream`, or instead to pass in an immutable "chunk/part stream factory".

Discussion:

Patch:
diff --git a/dwave/cloud/client.py b/dwave/cloud/client.py
index 05bd128..8d17147 100644
--- a/dwave/cloud/client.py
+++ b/dwave/cloud/client.py
@@ -1668,7 +1668,7 @@ class Client(object):
 
     @staticmethod
     @retried(_UPLOAD_PART_RETRIES, backoff=_UPLOAD_RETRIES_BACKOFF)
-    def _upload_multipart_part(session, problem_id, part_id, part_stream,
+    def _upload_multipart_part(session, problem_id, part_id, part_generator,
                                uploaded_part_checksum=None):
         """Upload one problem part. Sync http request.
 
@@ -1679,8 +1679,9 @@ class Client(object):
                 Problem id.
             part_id (int):
                 Part number/id.
-            part_stream (:class:`io.BufferedIOBase`/binary-stream-like):
-                Problem part data container that supports `read` operation.
+            part_generator (generator of :class:`io.BufferedIOBase`/binary-stream-like):
+                Callable that produces problem part data container that supports
+                `read` and `seek` operations.
             uploaded_part_checksum (str/None):
                 Checksum of previously uploaded part. Optional, but if specified
                 checksum is verified, and part is uploaded only if checksums
@@ -1692,6 +1693,9 @@ class Client(object):
 
         logger.debug("Uploading part_id=%r of problem_id=%r", part_id, problem_id)
 
+        # generate the mutable part stream from immutable stream generator
+        part_stream = part_generator()
+
         # TODO: work-around to get a checksum of a binary stream (avoid 2x read)
         data = part_stream.read()
         digest = Client._digest(data)
@@ -1808,12 +1812,12 @@ class Client(object):
                 uploaded_parts[part_no] = checksum
         return uploaded_parts
 
-    def _upload_part_worker(self, problem_id, part_no, chunk_stream,
+    def _upload_part_worker(self, problem_id, part_no, chunk_generator,
                             uploaded_part_checksum=None):
 
         with self.create_session() as session:
             part_checksum = self._upload_multipart_part(
-                session, problem_id, part_id=part_no, part_stream=chunk_stream,
+                session, problem_id, part_id=part_no, part_generator=chunk_generator,
                 uploaded_part_checksum=uploaded_part_checksum)
 
             return part_no, part_checksum
@@ -1859,12 +1863,11 @@ class Client(object):
 
             # enqueue all parts, worker skips if checksum matches
             parts = {}
-            streams = collections.OrderedDict(enumerate(chunks))
-            for chunk_no, chunk_stream in streams.items():
+            for chunk_no, chunk_generator in enumerate(chunks.generators()):
                 part_no = chunk_no + 1
                 part_future = self._upload_part_executor.submit(
                     self._upload_part_worker,
-                    problem_id, part_no, chunk_stream,
+                    problem_id, part_no, chunk_generator,
                     uploaded_part_checksum=uploaded_parts.get(part_no))
 
                 parts[part_no] = part_future
diff --git a/dwave/cloud/upload.py b/dwave/cloud/upload.py
index 2b48bb8..85d603b 100644
--- a/dwave/cloud/upload.py
+++ b/dwave/cloud/upload.py
@@ -22,6 +22,7 @@ import threading
 
 from abc import abstractmethod
 from collections.abc import Sized
+from functools import partial
 
 __all__ = ['ChunkedData']
 
@@ -400,3 +401,9 @@ class ChunkedData(object):
     def __iter__(self):
         for idx in range(len(self)):
             yield self.chunk(idx)
+
+    def generators(self):
+        """Iterator of (immutable) chunk generators."""
+
+        for idx in range(len(self)):
+            yield partial(self.chunk, idx=idx)


Test patch:
diff --git a/tests/test_upload.py b/tests/test_upload.py
index baf12a2..6717d1e 100644
--- a/tests/test_upload.py
+++ b/tests/test_upload.py
@@ -358,6 +358,12 @@ class TestChunkedData(unittest.TestCase):
         chunks_expected = [self.data]
         self.verify_chunking(cd, chunks_expected)
 
+    def test_chunk_generators(self):
+        cd = ChunkedData(self.data, chunk_size=3)
+        chunks_expected = [b'012', b'345', b'678', b'9']
+        chunks_generated = [g().read() for g in cd.generators()]
+        self.assertListEqual(chunks_expected, chunks_generated)
+
 
 @unittest.skipUnless(config, "No live server configuration available.")
 class TestMultipartUpload(unittest.TestCase):
@@ -620,7 +626,6 @@ class TestMockedMultipartUpload(unittest.TestCase):
 
             def put(path, data, headers, seq=iter(range(Client._UPLOAD_PART_RETRIES+1))):
                 body = data.read()
-                data.seek(0)
                 headers = json.dumps(sorted(headers.items()))
                 keys = [
                     (