🐛Storage: Copying returns wrong size (#6272)

Author: sanderegg

packages/aws-library/src/aws_library/s3/__init__.py

@@ -1,4 +1,8 @@
-from ._client import SimcoreS3API
+from ._client import (
+    CopiedBytesTransferredCallback,
+    SimcoreS3API,
+    UploadedBytesTransferredCallback,
+)
 from ._constants import PRESIGNED_LINK_MAX_SIZE, S3_MAX_FILE_SIZE
 from ._errors import (
     S3AccessError,
@@ -18,20 +22,22 @@
 )
 
 __all__: tuple[str, ...] = (
-    "SimcoreS3API",
+    "CopiedBytesTransferredCallback",
+    "MultiPartUploadLinks",
     "PRESIGNED_LINK_MAX_SIZE",
     "S3_MAX_FILE_SIZE",
     "S3AccessError",
     "S3BucketInvalidError",
     "S3DestinationNotEmptyError",
+    "S3DirectoryMetaData",
     "S3KeyNotFoundError",
+    "S3MetaData",
     "S3NotConnectedError",
+    "S3ObjectKey",
     "S3RuntimeError",
     "S3UploadNotFoundError",
-    "S3DirectoryMetaData",
-    "S3MetaData",
-    "S3ObjectKey",
-    "MultiPartUploadLinks",
+    "SimcoreS3API",
+    "UploadedBytesTransferredCallback",
     "UploadID",
 )
 

packages/aws-library/src/aws_library/s3/_client.py

@@ -1,11 +1,12 @@
 import asyncio
 import contextlib
+import functools
 import logging
 import urllib.parse
-from collections.abc import AsyncGenerator, Callable, Sequence
+from collections.abc import AsyncGenerator, Sequence
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Final, cast
+from typing import Any, Final, Protocol, cast
 
 import aioboto3
 from aiobotocore.session import ClientCreatorContext
@@ -43,6 +44,16 @@
 _AWS_MAX_ITEMS_PER_PAGE: Final[int] = 1000
 
 
+class UploadedBytesTransferredCallback(Protocol):
+    def __call__(self, bytes_transferred: int, *, file_name: str) -> None:
+        ...
+
+
+class CopiedBytesTransferredCallback(Protocol):
+    def __call__(self, total_bytes_copied: int, *, file_name: str) -> None:
+        ...
+
+
 @dataclass(frozen=True)
 class SimcoreS3API:  # pylint: disable=too-many-public-methods
     _client: S3Client
@@ -372,7 +383,7 @@ async def upload_file(
         bucket: S3BucketName,
         file: Path,
         object_key: S3ObjectKey,
-        bytes_transfered_cb: Callable[[int], None] | None,
+        bytes_transfered_cb: UploadedBytesTransferredCallback | None,
     ) -> None:
         """upload a file using aioboto3 transfer manager (e.g. works >5Gb and creates multiple threads)"""
         upload_options: dict[str, Any] = {
@@ -381,7 +392,11 @@ async def upload_file(
             "Config": TransferConfig(max_concurrency=self.transfer_max_concurrency),
         }
         if bytes_transfered_cb:
-            upload_options |= {"Callback": bytes_transfered_cb}
+            upload_options |= {
+                "Callback": functools.partial(
+                    bytes_transfered_cb, file_name=f"{object_key}"
+                )
+            }
         await self._client.upload_file(f"{file}", **upload_options)
 
     @s3_exception_handler(_logger)
@@ -391,7 +406,7 @@ async def copy_object(
         bucket: S3BucketName,
         src_object_key: S3ObjectKey,
         dst_object_key: S3ObjectKey,
-        bytes_transfered_cb: Callable[[int], None] | None,
+        bytes_transfered_cb: CopiedBytesTransferredCallback | None,
     ) -> None:
         """copy a file in S3 using aioboto3 transfer manager (e.g. works >5Gb and creates multiple threads)"""
         copy_options: dict[str, Any] = {
@@ -404,7 +419,11 @@ async def copy_object(
             ),
         }
         if bytes_transfered_cb:
-            copy_options |= {"Callback": bytes_transfered_cb}
+            copy_options |= {
+                "Callback": functools.partial(
+                    bytes_transfered_cb, file_name=f"{dst_object_key}"
+                )
+            }
         await self._client.copy(**copy_options)
 
     @s3_exception_handler(_logger)
@@ -414,7 +433,7 @@ async def copy_objects_recursively(
         bucket: S3BucketName,
         src_prefix: str,
         dst_prefix: str,
-        bytes_transfered_cb: Callable[[int], None] | None,
+        bytes_transfered_cb: CopiedBytesTransferredCallback | None,
     ) -> None:
         """copy from 1 location in S3 to another recreating the same structure"""
         dst_metadata = await self.get_directory_metadata(

packages/aws-library/tests/test_s3_client.py

@@ -257,17 +257,34 @@ async def _uploader(
 
 
 @dataclass
-class _ProgressCallback:
+class _UploadProgressCallback:
     file_size: int
     action: str
     logger: logging.Logger
     _total_bytes_transfered: int = 0
 
-    def __call__(self, bytes_transferred: int) -> None:
+    def __call__(self, bytes_transferred: int, *, file_name: str) -> None:
         self._total_bytes_transfered += bytes_transferred
-        self.logger.debug(
+        assert self._total_bytes_transfered <= self.file_size
+        self.logger.info(
             "progress: %s",
-            f"{self.action} {self._total_bytes_transfered} / {self.file_size} bytes",
+            f"{self.action} {file_name=} {self._total_bytes_transfered} / {self.file_size} bytes",
+        )
+
+
+@dataclass
+class _CopyProgressCallback:
+    file_size: int
+    action: str
+    logger: logging.Logger
+    _total_bytes_transfered: int = 0
+
+    def __call__(self, total_bytes_copied: int, *, file_name: str) -> None:
+        self._total_bytes_transfered = total_bytes_copied
+        assert self._total_bytes_transfered <= self.file_size
+        self.logger.info(
+            "progress: %s",
+            f"{self.action} {file_name=} {self._total_bytes_transfered} / {self.file_size} bytes",
         )
 
 
@@ -287,7 +304,7 @@ async def _uploader(file: Path, base_path: Path | None = None) -> UploadedFile:
         with log_context(
             logging.INFO, msg=f"uploading {file} to {with_s3_bucket}/{object_key}"
         ) as ctx:
-            progress_cb = _ProgressCallback(
+            progress_cb = _UploadProgressCallback(
                 file_size=file.stat().st_size, action="uploaded", logger=ctx.logger
             )
             response = await simcore_s3_api.upload_file(
@@ -365,7 +382,7 @@ async def _copier(src_key: S3ObjectKey, dst_key: S3ObjectKey) -> S3ObjectKey:
             bucket=with_s3_bucket, object_key=src_key
         )
         with log_context(logging.INFO, msg=f"copying {src_key} to {dst_key}") as ctx:
-            progress_cb = _ProgressCallback(
+            progress_cb = _CopyProgressCallback(
                 file_size=file_metadata.size, action="copied", logger=ctx.logger
             )
             await simcore_s3_api.copy_object(
@@ -397,7 +414,7 @@ async def _copier(src_prefix: str, dst_prefix: str) -> str:
             logging.INFO,
             msg=f"copying {src_prefix} [{ByteSize(src_directory_metadata.size).human_readable()}] to {dst_prefix}",
         ) as ctx:
-            progress_cb = _ProgressCallback(
+            progress_cb = _CopyProgressCallback(
                 file_size=src_directory_metadata.size,
                 action="copied",
                 logger=ctx.logger,

services/storage/src/simcore_service_storage/s3_utils.py

@@ -1,5 +1,6 @@
 import logging
-from dataclasses import dataclass
+from collections import defaultdict
+from dataclasses import dataclass, field
 
 from pydantic import ByteSize, parse_obj_as
 from servicelib.aiohttp.long_running_tasks.server import (
@@ -27,23 +28,44 @@ class S3TransferDataCB:
     total_bytes_to_transfer: ByteSize
     task_progress_message_prefix: str = ""
     _total_bytes_copied: int = 0
+    _file_total_bytes_copied: dict[str, int] = field(
+        default_factory=lambda: defaultdict(int)
+    )
 
-    def __post_init__(self):
-        self.copy_transfer_cb(0)
+    def __post_init__(self) -> None:
+        self._update()
 
-    def finalize_transfer(self):
-        self.copy_transfer_cb(self.total_bytes_to_transfer - self._total_bytes_copied)
+    def _update(self) -> None:
+        update_task_progress(
+            self.task_progress,
+            f"{self.task_progress_message_prefix} - "
+            f"{parse_obj_as(ByteSize,self._total_bytes_copied).human_readable()}"
+            f"/{self.total_bytes_to_transfer.human_readable()}]",
+            ProgressPercent(
+                min(self._total_bytes_copied, self.total_bytes_to_transfer)
+                / (self.total_bytes_to_transfer or 1)
+            ),
+        )
 
-    def copy_transfer_cb(self, copied_bytes: int):
-        self._total_bytes_copied += copied_bytes
+    def finalize_transfer(self) -> None:
+        self._total_bytes_copied = (
+            self.total_bytes_to_transfer - self._total_bytes_copied
+        )
+        self._update()
+
+    def copy_transfer_cb(self, total_bytes_copied: int, *, file_name: str) -> None:
+        logger.debug(
+            "Copied %s of %s",
+            parse_obj_as(ByteSize, total_bytes_copied).human_readable(),
+            file_name,
+        )
+        self._file_total_bytes_copied[file_name] = total_bytes_copied
+        self._total_bytes_copied = sum(self._file_total_bytes_copied.values())
+        if self.total_bytes_to_transfer != 0:
+            self._update()
+
+    def upload_transfer_cb(self, bytes_transferred: int, *, file_name: str) -> None:
+        self._file_total_bytes_copied[file_name] += bytes_transferred
+        self._total_bytes_copied = sum(self._file_total_bytes_copied.values())
         if self.total_bytes_to_transfer != 0:
-            update_task_progress(
-                self.task_progress,
-                f"{self.task_progress_message_prefix} - "
-                f"{parse_obj_as(ByteSize,self._total_bytes_copied).human_readable()}"
-                f"/{self.total_bytes_to_transfer.human_readable()}]",
-                ProgressPercent(
-                    max(self._total_bytes_copied, self.total_bytes_to_transfer)
-                    / self.total_bytes_to_transfer
-                ),
-            )
+            self._update()

services/storage/src/simcore_service_storage/simcore_s3_dsm.py

@@ -3,7 +3,7 @@
 import logging
 import tempfile
 import urllib.parse
-from collections.abc import Callable, Coroutine
+from collections.abc import Coroutine
 from contextlib import suppress
 from dataclasses import dataclass
 from pathlib import Path
@@ -13,7 +13,13 @@
 from aiohttp import web
 from aiopg.sa import Engine
 from aiopg.sa.connection import SAConnection
-from aws_library.s3 import S3DirectoryMetaData, S3KeyNotFoundError, S3MetaData
+from aws_library.s3 import (
+    CopiedBytesTransferredCallback,
+    S3DirectoryMetaData,
+    S3KeyNotFoundError,
+    S3MetaData,
+    UploadedBytesTransferredCallback,
+)
 from models_library.api_schemas_storage import LinkType, S3BucketName, UploadedPart
 from models_library.basic_types import SHA256Str
 from models_library.projects import ProjectID
@@ -668,7 +674,7 @@ async def deep_copy_project_simcore_s3(
                             dest_project_id=dst_project_uuid,
                             dest_node_id=NodeID(node_id),
                             file_storage_link=output,
-                            bytes_transfered_cb=s3_transfered_data_cb.copy_transfer_cb,
+                            bytes_transfered_cb=s3_transfered_data_cb.upload_transfer_cb,
                         )
                         for output in node.get("outputs", {}).values()
                         if isinstance(output, dict)
@@ -953,7 +959,7 @@ async def _copy_file_datcore_s3(
         dest_project_id: ProjectID,
         dest_node_id: NodeID,
         file_storage_link: dict[str, Any],
-        bytes_transfered_cb: Callable[[int], None],
+        bytes_transfered_cb: UploadedBytesTransferredCallback,
     ) -> FileMetaData:
         session = get_client_session(self.app)
         # 2 steps: Get download link for local copy, then upload to S3
@@ -1004,7 +1010,7 @@ async def _copy_path_s3_s3(
         *,
         src_fmd: FileMetaDataAtDB,
         dst_file_id: SimcoreS3FileID,
-        bytes_transfered_cb: Callable[[int], None],
+        bytes_transfered_cb: CopiedBytesTransferredCallback,
     ) -> FileMetaData:
         with log_context(
             _logger,