Infra: introduce retry func warpper

fsspec · Sep 12, 2024 · 44035aa · 44035aa
1 parent 904c982
commit 44035aa
Showing 1 changed file with 24 additions and 6 deletions.
diff --git a/tosfs/core.py b/tosfs/core.py
@@ -695,15 +695,33 @@ def put_file(
                         content=chunk,
                         content_type=content_type,
                     )
+                    retryable_func(
+                        lambda: self.tos_client.put_object(
+                            bucket,
+                            key,
+                            content=chunk,
+                            content_type=content_type,
+                        ),
+                        max_retry_num=self.max_retry_num,
+                    )
                 else:
-                    mpu = self.tos_client.create_multipart_upload(
-                        bucket, key, content_type=content_type
+                    mpu = retryable_func(
+                        lambda: self.tos_client.create_multipart_upload(
+                            bucket, key, content_type=content_type
+                        ),
+                        max_retry_num=self.max_retry_num,
                     )
-                    self.tos_client.upload_part_from_file(
-                        bucket, key, mpu.upload_id, file_path=lpath, part_number=1
+                    retryable_func(
+                        lambda: self.tos_client.upload_part_from_file(
+                            bucket, key, mpu.upload_id, file_path=lpath, part_number=1
+                        ),
+                        max_retry_num=self.max_retry_num,
                     )
-                    self.tos_client.complete_multipart_upload(
-                        bucket, key, mpu.upload_id, complete_all=True
+                    retryable_func(
+                        lambda: self.tos_client.complete_multipart_upload(
+                            bucket, key, mpu.upload_id, complete_all=True
+                        ),
+                        max_retry_num=self.max_retry_num,
                     )
         except (TosClientError, TosServerError) as e:
             raise e