convert S3 progress to progress bar data

Author: sanderegg

services/storage/src/simcore_service_storage/simcore_s3_dsm.py

@@ -35,9 +35,9 @@
 )
 from models_library.users import UserID
 from pydantic import AnyUrl, ByteSize, NonNegativeInt, TypeAdapter, ValidationError
-from servicelib.aiohttp.long_running_tasks.server import TaskProgress
 from servicelib.fastapi.client_session import get_client_session
 from servicelib.logging_utils import log_context
+from servicelib.progress_bar import ProgressBarData
 from servicelib.utils import ensure_ends_with, limited_gather
 from simcore_postgres_database.utils_repos import transaction_context
 from sqlalchemy.ext.asyncio import AsyncEngine
@@ -77,7 +77,7 @@
 from .modules.db.projects import ProjectRepository
 from .modules.db.tokens import TokenRepository
 from .modules.s3 import get_s3_client
-from .utils.s3_utils import S3TransferDataCB, update_task_progress
+from .utils.s3_utils import S3TransferDataCB
 from .utils.simcore_s3_dsm_utils import (
     compute_file_id_prefix,
     expand_directory,
@@ -751,7 +751,7 @@ async def deep_copy_project_simcore_s3(
         src_project: dict[str, Any],
         dst_project: dict[str, Any],
         node_mapping: dict[NodeID, NodeID],
-        task_progress: TaskProgress | None = None,
+        task_progress: ProgressBarData,
     ) -> None:
         src_project_uuid: ProjectID = ProjectID(src_project["uuid"])
         dst_project_uuid: ProjectID = ProjectID(dst_project["uuid"])
@@ -761,7 +761,7 @@ async def deep_copy_project_simcore_s3(
             msg=f"{src_project_uuid} -> {dst_project_uuid}: "
             "Step 1: check access rights (read of src and write of dst)",
         ):
-            update_task_progress(task_progress, "Checking study access rights...")
+            task_progress.description = "Checking study access rights..."
 
             for prj_uuid in [src_project_uuid, dst_project_uuid]:
                 if not await ProjectRepository.instance(
@@ -789,8 +789,8 @@ async def deep_copy_project_simcore_s3(
             msg=f"{src_project_uuid} -> {dst_project_uuid}:"
             " Step 2: collect what to copy",
         ):
-            update_task_progress(
-                task_progress, f"Collecting files of '{src_project['name']}'..."
+            task_progress.description = (
+                f"Collecting files of '{src_project['name']}'..."
             )
 
             src_project_files = await FileMetaDataRepository.instance(
@@ -814,68 +814,70 @@ async def deep_copy_project_simcore_s3(
             src_project_total_data_size: ByteSize = TypeAdapter(
                 ByteSize
             ).validate_python(sum(n for n, _ in sizes_and_num_files))
-        with log_context(
-            _logger,
-            logging.INFO,
-            msg=f"{src_project_uuid} -> {dst_project_uuid}:"
-            " Step 3.1: prepare copy tasks for files referenced from simcore",
-        ):
-            copy_tasks = []
-            s3_transfered_data_cb = S3TransferDataCB(
-                task_progress,
-                src_project_total_data_size,
-                task_progress_message_prefix=f"Copying {total_num_of_files} files to '{dst_project['name']}'",
-            )
-            for src_fmd in src_project_files:
-                if not src_fmd.node_id or (src_fmd.location_id != self.location_id):
-                    msg = (
-                        "This is not foreseen, stem from old decisions, and needs to "
-                        f"be implemented if needed. Faulty metadata: {src_fmd=}"
-                    )
-                    raise NotImplementedError(msg)
-
-                if new_node_id := node_mapping.get(src_fmd.node_id):
-                    copy_tasks.append(
-                        self._copy_path_s3_s3(
-                            user_id,
-                            src_fmd=src_fmd,
-                            dst_file_id=TypeAdapter(SimcoreS3FileID).validate_python(
-                                f"{dst_project_uuid}/{new_node_id}/{src_fmd.object_name.split('/', maxsplit=2)[-1]}"
-                            ),
-                            bytes_transfered_cb=s3_transfered_data_cb.copy_transfer_cb,
+
+        async with S3TransferDataCB(
+            task_progress,
+            src_project_total_data_size,
+            task_progress_message_prefix=f"Copying {total_num_of_files} files to '{dst_project['name']}'",
+        ) as s3_transfered_data_cb:
+            with log_context(
+                _logger,
+                logging.INFO,
+                msg=f"{src_project_uuid} -> {dst_project_uuid}:"
+                " Step 3.1: prepare copy tasks for files referenced from simcore",
+            ):
+                copy_tasks = []
+                for src_fmd in src_project_files:
+                    if not src_fmd.node_id or (src_fmd.location_id != self.location_id):
+                        msg = (
+                            "This is not foreseen, stem from old decisions, and needs to "
+                            f"be implemented if needed. Faulty metadata: {src_fmd=}"
                         )
-                    )
-        with log_context(
-            _logger,
-            logging.INFO,
-            msg=f"{src_project_uuid} -> {dst_project_uuid}:"
-            " Step 3.1: prepare copy tasks for files referenced from DAT-CORE",
-        ):
-            for node_id, node in dst_project.get("workbench", {}).items():
-                copy_tasks.extend(
-                    [
-                        self._copy_file_datcore_s3(
-                            user_id=user_id,
-                            source_uuid=output["path"],
-                            dest_project_id=dst_project_uuid,
-                            dest_node_id=NodeID(node_id),
-                            file_storage_link=output,
-                            bytes_transfered_cb=s3_transfered_data_cb.upload_transfer_cb,
+                        raise NotImplementedError(msg)
+
+                    if new_node_id := node_mapping.get(src_fmd.node_id):
+                        copy_tasks.append(
+                            self._copy_path_s3_s3(
+                                user_id,
+                                src_fmd=src_fmd,
+                                dst_file_id=TypeAdapter(
+                                    SimcoreS3FileID
+                                ).validate_python(
+                                    f"{dst_project_uuid}/{new_node_id}/{src_fmd.object_name.split('/', maxsplit=2)[-1]}"
+                                ),
+                                bytes_transfered_cb=s3_transfered_data_cb.copy_transfer_cb,
+                            )
                         )
-                        for output in node.get("outputs", {}).values()
-                        if isinstance(output, dict)
-                        and (int(output.get("store", self.location_id)) == DATCORE_ID)
-                    ]
-                )
-        with log_context(
-            _logger,
-            logging.INFO,
-            msg=f"{src_project_uuid} -> {dst_project_uuid}: Step 3.3: effective copying {len(copy_tasks)} files",
-        ):
-            await limited_gather(*copy_tasks, limit=MAX_CONCURRENT_S3_TASKS)
-
-        # ensure the full size is reported
-        s3_transfered_data_cb.finalize_transfer()
+            with log_context(
+                _logger,
+                logging.INFO,
+                msg=f"{src_project_uuid} -> {dst_project_uuid}:"
+                " Step 3.1: prepare copy tasks for files referenced from DAT-CORE",
+            ):
+                for node_id, node in dst_project.get("workbench", {}).items():
+                    copy_tasks.extend(
+                        [
+                            self._copy_file_datcore_s3(
+                                user_id=user_id,
+                                source_uuid=output["path"],
+                                dest_project_id=dst_project_uuid,
+                                dest_node_id=NodeID(node_id),
+                                file_storage_link=output,
+                                bytes_transfered_cb=s3_transfered_data_cb.upload_transfer_cb,
+                            )
+                            for output in node.get("outputs", {}).values()
+                            if isinstance(output, dict)
+                            and (
+                                int(output.get("store", self.location_id)) == DATCORE_ID
+                            )
+                        ]
+                    )
+            with log_context(
+                _logger,
+                logging.INFO,
+                msg=f"{src_project_uuid} -> {dst_project_uuid}: Step 3.3: effective copying {len(copy_tasks)} files",
+            ):
+                await limited_gather(*copy_tasks, limit=MAX_CONCURRENT_S3_TASKS)
 
     async def _get_size_and_num_files(
         self, fmd: FileMetaDataAtDB

services/storage/src/simcore_service_storage/utils/s3_utils.py

@@ -1,51 +1,61 @@
+import asyncio
+import datetime
 import logging
 from collections import defaultdict
 from dataclasses import dataclass, field
 
 from pydantic import ByteSize, TypeAdapter
-from servicelib.aiohttp.long_running_tasks.server import (
-    ProgressMessage,
-    ProgressPercent,
-    TaskProgress,
-)
+from servicelib.async_utils import cancel_wait_task
+from servicelib.background_task import create_periodic_task
+from servicelib.progress_bar import ProgressBarData
 
 _logger = logging.getLogger(__name__)
 
 
-def update_task_progress(
-    task_progress: TaskProgress | None,
-    message: ProgressMessage | None = None,
-    progress: ProgressPercent | None = None,
-) -> None:
-    _logger.debug("%s [%s]", message or "", progress or "n/a")
-    if task_progress:
-        task_progress.update(message=message, percent=progress)
-
-
 @dataclass
 class S3TransferDataCB:
-    task_progress: TaskProgress | None
+    task_progress: ProgressBarData
     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)
     )
+    _update_task_event: asyncio.Event = field(default_factory=asyncio.Event)
+    _async_update_periodic_task: asyncio.Task | None = None
 
     def __post_init__(self) -> None:
+        self._async_update_periodic_task = create_periodic_task(
+            self._async_update,
+            interval=datetime.timedelta(seconds=1),
+            task_name="s3_transfer_cb_update",
+        )
         self._update()
 
-    def _update(self) -> None:
-        update_task_progress(
-            self.task_progress,
+    async def __aenter__(self) -> "S3TransferDataCB":
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback) -> None:
+        self.finalize_transfer()
+        await asyncio.sleep(0)
+        assert self._async_update_periodic_task  # nosec
+        await cancel_wait_task(self._async_update_periodic_task)
+
+    async def _async_update(self) -> None:
+        await self._update_task_event.wait()
+        self._update_task_event.clear()
+        self.task_progress.description = (
             f"{self.task_progress_message_prefix} - "
-            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)
-            ),
+            f"{self.total_bytes_to_transfer.human_readable()}"
+        )
+        await self.task_progress.update(
+            min(self._total_bytes_copied, self.total_bytes_to_transfer)
+            / (self.total_bytes_to_transfer or 1)
         )
 
+    def _update(self) -> None:
+        self._update_task_event.set()
+
     def finalize_transfer(self) -> None:
         self._total_bytes_copied = (
             self.total_bytes_to_transfer - self._total_bytes_copied