🎨Image pulling progress: pass current/total size (Part 3) (ITISFounda…

…tion#5685)

packages/models-library/src/models_library/progress_bar.py

@@ -0,0 +1,36 @@
+from typing import Any, ClassVar, Literal, TypeAlias
+
+from pydantic import BaseModel
+
+# NOTE: keep a list of possible unit, and please use correct official unit names
+ProgressUnit: TypeAlias = Literal["Byte"]
+
+
+class ProgressReport(BaseModel):
+    actual_value: float
+    total: float
+    unit: ProgressUnit | None = None
+
+    @property
+    def percent_value(self) -> float:
+        if self.total != 0:
+            return max(min(self.actual_value / self.total, 1.0), 0.0)
+        return 0
+
+    class Config:
+        frozen = True
+        schema_extra: ClassVar[dict[str, Any]] = {
+            "examples": [
+                # typical percent progress (no units)
+                {
+                    "actual_value": 0.3,
+                    "total": 1.0,
+                },
+                # typical byte progress
+                {
+                    "actual_value": 128.5,
+                    "total": 1024.0,
+                    "unit": "Byte",
+                },
+            ]
+        }

packages/models-library/src/models_library/rabbitmq_messages.py

@@ -6,10 +6,10 @@
 from typing import Any, Literal, TypeAlias
 
 import arrow
-from models_library.products import ProductName
 from pydantic import BaseModel, Field
-from pydantic.types import NonNegativeFloat
 
+from .products import ProductName
+from .progress_bar import ProgressReport
 from .projects import ProjectID
 from .projects_nodes_io import NodeID
 from .projects_state import RunningState
@@ -99,7 +99,7 @@ class ProgressMessageMixin(RabbitMessageBase):
     progress_type: ProgressType = (
         ProgressType.COMPUTATION_RUNNING
     )  # NOTE: backwards compatible
-    progress: NonNegativeFloat
+    report: ProgressReport
 
     def routing_key(self) -> str | None:
         return None

packages/models-library/tests/test_rabbit_messages.py

@@ -2,6 +2,7 @@
 
 import pytest
 from faker import Faker
+from models_library.progress_bar import ProgressReport
 from models_library.rabbitmq_messages import (
     ProgressRabbitMessageNode,
     ProgressRabbitMessageProject,
@@ -21,7 +22,7 @@
                 user_id=faker.uuid4(cast_to=None),
                 node_id=faker.uuid4(cast_to=None),
                 progress_type=ProgressType.SERVICE_OUTPUTS_PULLING,
-                progress=0.4,
+                report=ProgressReport(actual_value=0.4, total=1),
             ).json(),
             ProgressRabbitMessageNode,
             id="node_progress",
@@ -31,7 +32,7 @@
                 project_id=faker.uuid4(cast_to=None),
                 user_id=faker.uuid4(cast_to=None),
                 progress_type=ProgressType.PROJECT_CLOSING,
-                progress=0.4,
+                report=ProgressReport(actual_value=0.4, total=1),
             ).json(),
             ProgressRabbitMessageProject,
             id="project_progress",

packages/service-library/src/servicelib/docker_utils.py

@@ -204,7 +204,9 @@ async def pull_image(
                     assert parsed_progress.id  # nosec
                     layer_id_to_size.setdefault(
                         parsed_progress.id, _PulledStatus(0)
-                    ).downloaded = layer_id_to_size[parsed_progress.id].size
+                    ).downloaded = layer_id_to_size.setdefault(
+                        parsed_progress.id, _PulledStatus(0)
+                    ).size
                 case "extracting":
                     assert parsed_progress.id  # nosec
                     assert parsed_progress.progress_detail  # nosec

packages/service-library/src/servicelib/fastapi/docker_utils.py

@@ -99,7 +99,7 @@ async def retrieve_image_layer_information(
     return None
 
 
-_DEFAULT_MIN_IMAGE_SIZE: Final[ByteSize] = parse_obj_as(ByteSize, "50MiB")
+_DEFAULT_MIN_IMAGE_SIZE: Final[ByteSize] = parse_obj_as(ByteSize, "200MiB")
 
 
 async def pull_images(
@@ -122,6 +122,7 @@ async def pull_images(
     async with ProgressBarData(
         num_steps=images_total_size,
         progress_report_cb=progress_cb,
+        progress_unit="Byte",
     ) as pbar:
 
         await asyncio.gather(

packages/service-library/src/servicelib/progress_bar.py

@@ -4,6 +4,9 @@
 from inspect import isawaitable
 from typing import Final, Optional, Protocol, runtime_checkable
 
+from models_library.progress_bar import ProgressReport, ProgressUnit
+from pydantic import parse_obj_as
+
 from .logging_utils import log_catch
 
 _logger = logging.getLogger(__name__)
@@ -14,13 +17,13 @@
 
 @runtime_checkable
 class AsyncReportCB(Protocol):
-    async def __call__(self, progress_value: float) -> None:
+    async def __call__(self, report: ProgressReport) -> None:
         ...
 
 
 @runtime_checkable
 class ReportCB(Protocol):
-    def __call__(self, progress_value: float) -> None:
+    def __call__(self, report: ProgressReport) -> None:
         ...
 
 
@@ -73,6 +76,7 @@ async def main_fct():
             "description": "Optionally defines the step relative weight (defaults to steps of equal weights)"
         },
     )
+    progress_unit: ProgressUnit | None = None
     progress_report_cb: AsyncReportCB | ReportCB | None = None
     _current_steps: float = _INITIAL_VALUE
     _children: list = field(default_factory=list)
@@ -81,6 +85,8 @@ async def main_fct():
     _last_report_value: float = _INITIAL_VALUE
 
     def __post_init__(self) -> None:
+        if self.progress_unit is not None:
+            parse_obj_as(ProgressUnit, self.progress_unit)
         self._continuous_value_lock = asyncio.Lock()
         self.num_steps = max(1, self.num_steps)
         if self.step_weights:
@@ -112,7 +118,14 @@ async def _report_external(self, value: float, *, force: bool = False) -> None:
                 or ((value - self._last_report_value) > _MIN_PROGRESS_UPDATE_PERCENT)
                 or value == _FINAL_VALUE
             ):
-                call = self.progress_report_cb(value)
+                call = self.progress_report_cb(
+                    ProgressReport(
+                        # NOTE: here we convert back to actual value since this is possibly weighted
+                        actual_value=value * self.num_steps,
+                        total=self.num_steps,
+                        unit=self.progress_unit,
+                    ),
+                )
                 if isawaitable(call):
                     await call
                 self._last_report_value = value

packages/service-library/tests/aiohttp/test_docker_utils.py

@@ -1,10 +1,16 @@
 # pylint: disable=protected-access
+# pylint: disable=redefined-outer-name
+# pylint: disable=too-many-arguments
+# pylint: disable=unused-argument
+# pylint: disable=unused-variable
 from collections.abc import Awaitable, Callable
 from typing import Any
+from unittest import mock
 from unittest.mock import call
 
 import pytest
 from models_library.docker import DockerGenericTag
+from models_library.progress_bar import ProgressReport
 from pydantic import parse_obj_as
 from pytest_mock import MockerFixture
 from servicelib import progress_bar
@@ -68,6 +74,33 @@ async def test_retrieve_image_layer_information_from_external_registry(
     assert layer_information
 
 
+@pytest.fixture
+async def mocked_log_cb(mocker: MockerFixture) -> mock.AsyncMock:
+    async def _log_cb(*args, **kwargs) -> None:
+        print(f"received log: {args}, {kwargs}")
+
+    return mocker.AsyncMock(side_effect=_log_cb)
+
+
+@pytest.fixture
+async def mocked_progress_cb(mocker: MockerFixture) -> mock.AsyncMock:
+    async def _progress_cb(*args, **kwargs) -> None:
+        print(f"received progress: {args}, {kwargs}")
+
+    return mocker.AsyncMock(side_effect=_progress_cb)
+
+
+def _assert_progress_report_values(
+    mocked_progress_cb: mock.AsyncMock, *, total: float
+) -> None:
+    assert mocked_progress_cb.call_args_list[0] == call(
+        ProgressReport(actual_value=0, total=total)
+    )
+    assert mocked_progress_cb.call_args_list[-1] == call(
+        ProgressReport(actual_value=total, total=total)
+    )
+
+
 @pytest.mark.parametrize(
     "image",
     ["itisfoundation/sleeper:1.0.0", "nginx:latest"],
@@ -76,55 +109,61 @@ async def test_pull_image(
     remove_images_from_host: Callable[[list[str]], Awaitable[None]],
     image: DockerGenericTag,
     registry_settings: RegistrySettings,
-    mocker: MockerFixture,
+    mocked_log_cb: mock.AsyncMock,
+    mocked_progress_cb: mock.AsyncMock,
     caplog: pytest.LogCaptureFixture,
 ):
     # clean first
     await remove_images_from_host([image])
     layer_information = await retrieve_image_layer_information(image, registry_settings)
     assert layer_information
 
-    async def _log_cb(*args, **kwargs) -> None:
-        print(f"received log: {args}, {kwargs}")
-
-    fake_progress_report_cb = mocker.AsyncMock()
     async with progress_bar.ProgressBarData(
         num_steps=layer_information.layers_total_size,
-        progress_report_cb=fake_progress_report_cb,
+        progress_report_cb=mocked_progress_cb,
     ) as main_progress_bar:
-        fake_log_cb = mocker.AsyncMock(side_effect=_log_cb)
         await pull_image(
-            image, registry_settings, main_progress_bar, fake_log_cb, layer_information
+            image,
+            registry_settings,
+            main_progress_bar,
+            mocked_log_cb,
+            layer_information,
         )
-        fake_log_cb.assert_called()
+        mocked_log_cb.assert_called()
         assert (
             main_progress_bar._current_steps  # noqa: SLF001
             == layer_information.layers_total_size
         )
-    assert fake_progress_report_cb.call_args_list[0] == call(0.0)
-    fake_progress_report_cb.assert_called_with(1.0)
+    _assert_progress_report_values(
+        mocked_progress_cb, total=layer_information.layers_total_size
+    )
+
+    mocked_progress_cb.reset_mock()
+    mocked_log_cb.reset_mock()
 
     # check there were no warnings popping up from the docker pull
     # NOTE: this would pop up in case docker changes its pulling statuses
-    for record in caplog.records:
-        assert record.levelname != "WARNING", record.message
+    assert not [r.message for r in caplog.records if r.levelname == "WARNING"]
 
     # pull a second time should, the image is already there
     async with progress_bar.ProgressBarData(
         num_steps=layer_information.layers_total_size,
-        progress_report_cb=fake_progress_report_cb,
+        progress_report_cb=mocked_progress_cb,
     ) as main_progress_bar:
-        fake_log_cb = mocker.AsyncMock(side_effect=_log_cb)
         await pull_image(
-            image, registry_settings, main_progress_bar, fake_log_cb, layer_information
+            image,
+            registry_settings,
+            main_progress_bar,
+            mocked_log_cb,
+            layer_information,
         )
-        fake_log_cb.assert_called()
+        mocked_log_cb.assert_called()
         assert (
             main_progress_bar._current_steps  # noqa: SLF001
             == layer_information.layers_total_size
         )
-    assert fake_progress_report_cb.call_args_list[0] == call(0.0)
-    fake_progress_report_cb.assert_called_with(1.0)
+    _assert_progress_report_values(
+        mocked_progress_cb, total=layer_information.layers_total_size
+    )
     # check there were no warnings
-    for record in caplog.records:
-        assert record.levelname != "WARNING", record.message
+    assert not [r.message for r in caplog.records if r.levelname == "WARNING"]