🐛 retry pulling image layer when they timeout (#7051)

Co-authored-by: Andrei Neagu <[email protected]>

Author: GitHK

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

@@ -2,14 +2,12 @@
 
 from pydantic import BaseModel, ConfigDict
 
-from .basic_types import IDStr
-
 # NOTE: keep a list of possible unit, and please use correct official unit names
 ProgressUnit: TypeAlias = Literal["Byte"]
 
 
 class ProgressStructuredMessage(BaseModel):
-    description: IDStr
+    description: str
     current: float
     total: int
     unit: str | None = None
@@ -51,6 +49,7 @@ class ProgressStructuredMessage(BaseModel):
 class ProgressReport(BaseModel):
     actual_value: float
     total: float = 1.0
+    attempt: int = 0
     unit: ProgressUnit | None = UNITLESS
     message: ProgressStructuredMessage | None = None
 

packages/service-library/src/servicelib/archiving_utils/_interface_7zip.py

@@ -10,7 +10,6 @@
 from typing import Final
 
 import tqdm
-from models_library.basic_types import IDStr
 from pydantic import NonNegativeInt
 from servicelib.logging_utils import log_catch
 from tqdm.contrib.logging import tqdm_logging_redirect
@@ -199,7 +198,7 @@ async def archive_dir(
 ) -> None:
     if progress_bar is None:
         progress_bar = ProgressBarData(
-            num_steps=1, description=IDStr(f"compressing {dir_to_compress.name}")
+            num_steps=1, description=f"compressing {dir_to_compress.name}"
         )
 
     options = " ".join(
@@ -223,7 +222,7 @@ async def archive_dir(
 
     async with AsyncExitStack() as exit_stack:
         sub_progress = await exit_stack.enter_async_context(
-            progress_bar.sub_progress(folder_size_bytes, description=IDStr("..."))
+            progress_bar.sub_progress(folder_size_bytes, description="...")
         )
 
         tqdm_progress = exit_stack.enter_context(
@@ -290,7 +289,7 @@ async def unarchive_dir(
 ) -> set[Path]:
     if progress_bar is None:
         progress_bar = ProgressBarData(
-            num_steps=1, description=IDStr(f"extracting {archive_to_extract.name}")
+            num_steps=1, description=f"extracting {archive_to_extract.name}"
         )
 
     # get archive information
@@ -304,7 +303,7 @@ async def unarchive_dir(
 
     async with AsyncExitStack() as exit_stack:
         sub_prog = await exit_stack.enter_async_context(
-            progress_bar.sub_progress(steps=total_bytes, description=IDStr("..."))
+            progress_bar.sub_progress(steps=total_bytes, description="...")
         )
 
         tqdm_progress = exit_stack.enter_context(

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

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 from collections.abc import Awaitable, Callable
 from contextlib import AsyncExitStack
@@ -11,8 +12,21 @@
 from models_library.docker import DockerGenericTag
 from models_library.generated_models.docker_rest_api import ProgressDetail
 from models_library.utils.change_case import snake_to_camel
-from pydantic import BaseModel, ByteSize, ConfigDict, TypeAdapter, ValidationError
+from pydantic import (
+    BaseModel,
+    ByteSize,
+    ConfigDict,
+    NonNegativeInt,
+    TypeAdapter,
+    ValidationError,
+)
 from settings_library.docker_registry import RegistrySettings
+from tenacity import (
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_random_exponential,
+)
 from yarl import URL
 
 from .logging_utils import LogLevelInt
@@ -209,6 +223,8 @@ async def pull_image(
     progress_bar: ProgressBarData,
     log_cb: LogCB,
     image_information: DockerImageManifestsV2 | None,
+    *,
+    retry_upon_error_count: NonNegativeInt = 10,
 ) -> None:
     """pull a docker image to the host machine.
 
@@ -219,7 +235,9 @@ async def pull_image(
         progress_bar -- the current progress bar
         log_cb -- a callback function to send logs to
         image_information -- the image layer information. If this is None, then no fine progress will be retrieved.
+        retry_upon_error_count -- number of tries if there is a TimeoutError. Usually cased by networking issues.
     """
+
     registry_auth = None
     if registry_settings.REGISTRY_URL and registry_settings.REGISTRY_URL in image:
         registry_auth = {
@@ -245,39 +263,64 @@ async def pull_image(
 
         client = await exit_stack.enter_async_context(aiodocker.Docker())
 
-        reported_progress = 0.0
-        async for pull_progress in client.images.pull(
-            image, stream=True, auth=registry_auth
-        ):
-            try:
-                parsed_progress = TypeAdapter(_DockerPullImage).validate_python(
-                    pull_progress
+        def _reset_progress_from_previous_attempt() -> None:
+            for pulled_status in layer_id_to_size.values():
+                pulled_status.downloaded = 0
+                pulled_status.extracted = 0
+
+        attempt: NonNegativeInt = 1
+
+        @retry(
+            wait=wait_random_exponential(),
+            stop=stop_after_attempt(retry_upon_error_count),
+            reraise=True,
+            retry=retry_if_exception_type(asyncio.TimeoutError),
+        )
+        async def _pull_image_with_retry() -> None:
+            nonlocal attempt
+            if attempt > 1:
+                # for each attempt rest the progress
+                progress_bar.reset()
+                _reset_progress_from_previous_attempt()
+            attempt += 1
+
+            _logger.info("attempt '%s' trying to pull image='%s'", attempt, image)
+
+            reported_progress = 0.0
+            async for pull_progress in client.images.pull(
+                image, stream=True, auth=registry_auth
+            ):
+                try:
+                    parsed_progress = TypeAdapter(_DockerPullImage).validate_python(
+                        pull_progress
+                    )
+                except ValidationError:
+                    _logger.exception(
+                        "Unexpected error while validating '%s'. "
+                        "TIP: This is probably an unforeseen pull status text that shall be added to the code. "
+                        "The pulling process will still continue.",
+                        f"{pull_progress=}",
+                    )
+                else:
+                    await _parse_pull_information(
+                        parsed_progress, layer_id_to_size=layer_id_to_size
+                    )
+
+                # compute total progress
+                total_downloaded_size = sum(
+                    layer.downloaded for layer in layer_id_to_size.values()
                 )
-            except ValidationError:
-                _logger.exception(
-                    "Unexpected error while validating '%s'. "
-                    "TIP: This is probably an unforeseen pull status text that shall be added to the code. "
-                    "The pulling process will still continue.",
-                    f"{pull_progress=}",
+                total_extracted_size = sum(
+                    layer.extracted for layer in layer_id_to_size.values()
                 )
-            else:
-                await _parse_pull_information(
-                    parsed_progress, layer_id_to_size=layer_id_to_size
+                total_progress = (total_downloaded_size + total_extracted_size) / 2.0
+                progress_to_report = total_progress - reported_progress
+                await progress_bar.update(progress_to_report)
+                reported_progress = total_progress
+
+                await log_cb(
+                    f"pulling {image_short_name}: {pull_progress}...",
+                    logging.DEBUG,
                 )
 
-            # compute total progress
-            total_downloaded_size = sum(
-                layer.downloaded for layer in layer_id_to_size.values()
-            )
-            total_extracted_size = sum(
-                layer.extracted for layer in layer_id_to_size.values()
-            )
-            total_progress = (total_downloaded_size + total_extracted_size) / 2.0
-            progress_to_report = total_progress - reported_progress
-            await progress_bar.update(progress_to_report)
-            reported_progress = total_progress
-
-            await log_cb(
-                f"pulling {image_short_name}: {pull_progress}...",
-                logging.DEBUG,
-            )
+        await _pull_image_with_retry()

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

@@ -3,7 +3,6 @@
 from typing import Final
 
 import httpx
-from models_library.basic_types import IDStr
 from models_library.docker import DockerGenericTag
 from pydantic import ByteSize, TypeAdapter, ValidationError
 from settings_library.docker_registry import RegistrySettings
@@ -129,9 +128,7 @@ async def pull_images(
         num_steps=images_total_size,
         progress_report_cb=progress_cb,
         progress_unit="Byte",
-        description=IDStr(
-            f"pulling {len(images)} images",
-        ),
+        description=f"pulling {len(images)} images",
     ) as pbar:
 
         await asyncio.gather(

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

@@ -4,7 +4,6 @@
 from inspect import isawaitable
 from typing import Final, Optional, Protocol, runtime_checkable
 
-from models_library.basic_types import IDStr
 from models_library.progress_bar import (
     ProgressReport,
     ProgressStructuredMessage,
@@ -18,6 +17,7 @@
 _MIN_PROGRESS_UPDATE_PERCENT: Final[float] = 0.01
 _INITIAL_VALUE: Final[float] = -1.0
 _FINAL_VALUE: Final[float] = 1.0
+_PROGRESS_ALREADY_REACGED_MAXIMUM: Final[str] = "Progress already reached maximum of"
 
 
 @runtime_checkable
@@ -84,10 +84,11 @@ async def main_fct():
             "description": "Optionally defines the step relative weight (defaults to steps of equal weights)"
         },
     )
-    description: IDStr = field(metadata={"description": "define the progress name"})
+    description: str = field(metadata={"description": "define the progress name"})
     progress_unit: ProgressUnit | None = None
     progress_report_cb: AsyncReportCB | ReportCB | None = None
     _current_steps: float = _INITIAL_VALUE
+    _currnet_attempt: int = 0
     _children: list["ProgressBarData"] = field(default_factory=list)
     _parent: Optional["ProgressBarData"] = None
     _continuous_value_lock: asyncio.Lock = field(init=False)
@@ -147,6 +148,7 @@ async def _report_external(self, value: float) -> None:
                         # NOTE: here we convert back to actual value since this is possibly weighted
                         actual_value=value * self.num_steps,
                         total=self.num_steps,
+                        attempt=self._currnet_attempt,
                         unit=self.progress_unit,
                         message=self.compute_report_message_stuct(),
                     ),
@@ -176,7 +178,7 @@ async def update(self, steps: float = 1) -> None:
             if new_steps_value > self.num_steps:
                 _logger.warning(
                     "%s",
-                    f"Progress already reached maximum of {self.num_steps=}, "
+                    f"{_PROGRESS_ALREADY_REACGED_MAXIMUM} {self.num_steps=}, "
                     f"cause: {self._current_steps=} is updated by {steps=}"
                     "TIP: sub progresses are not created correctly please check the stack trace",
                     stack_info=True,
@@ -197,6 +199,11 @@ async def update(self, steps: float = 1) -> None:
             await self._update_parent(parent_update_value)
         await self._report_external(new_progress_value)
 
+    def reset(self) -> None:
+        self._currnet_attempt += 1
+        self._current_steps = _INITIAL_VALUE
+        self._last_report_value = _INITIAL_VALUE
+
     async def set_(self, new_value: float) -> None:
         await self.update(new_value - self._current_steps)
 
@@ -207,7 +214,7 @@ async def finish(self) -> None:
     def sub_progress(
         self,
         steps: int,
-        description: IDStr,
+        description: str,
         step_weights: list[float] | None = None,
         progress_unit: ProgressUnit | None = None,
     ) -> "ProgressBarData":

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

@@ -95,15 +95,15 @@ def _assert_progress_report_values(
     # NOTE: we exclude the message part here as this is already tested in servicelib
     # check first progress
     assert mocked_progress_cb.call_args_list[0].args[0].dict(
-        exclude={"message"}
+        exclude={"message", "attempt"}
     ) == ProgressReport(actual_value=0, total=total, unit="Byte").model_dump(
-        exclude={"message"}
+        exclude={"message", "attempt"}
     )
     # check last progress
     assert mocked_progress_cb.call_args_list[-1].args[0].dict(
-        exclude={"message"}
+        exclude={"message", "attempt"}
     ) == ProgressReport(actual_value=total, total=total, unit="Byte").model_dump(
-        exclude={"message"}
+        exclude={"message", "attempt"}
     )
 
 

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

@@ -101,15 +101,15 @@ def _assert_progress_report_values(
     # NOTE: we exclude the message part here as this is already tested in servicelib
     # check first progress
     assert mocked_progress_cb.call_args_list[0].args[0].dict(
-        exclude={"message"}
+        exclude={"message", "attempt"}
     ) == ProgressReport(actual_value=0, total=total, unit="Byte").model_dump(
-        exclude={"message"}
+        exclude={"message", "attempt"}
     )
     # check last progress
     assert mocked_progress_cb.call_args_list[-1].args[0].dict(
-        exclude={"message"}
+        exclude={"message", "attempt"}
     ) == ProgressReport(actual_value=total, total=total, unit="Byte").model_dump(
-        exclude={"message"}
+        exclude={"message", "attempt"}
     )
 
 

packages/service-library/tests/test_progress_bar.py

@@ -15,6 +15,7 @@
 from servicelib.progress_bar import (
     _INITIAL_VALUE,
     _MIN_PROGRESS_UPDATE_PERCENT,
+    _PROGRESS_ALREADY_REACGED_MAXIMUM,
     ProgressBarData,
 )
 
@@ -258,10 +259,43 @@ async def test_set_progress(caplog: pytest.LogCaptureFixture, faker: Faker):
         assert root._current_steps == pytest.approx(34)  # noqa: SLF001
         await root.set_(58)
         assert root._current_steps == pytest.approx(50)  # noqa: SLF001
-        assert "already reached maximum" in caplog.messages[0]
+        assert "WARNING" in caplog.text
+        assert _PROGRESS_ALREADY_REACGED_MAXIMUM in caplog.messages[0]
         assert "TIP:" in caplog.messages[0]
 
 
+async def test_reset_progress(caplog: pytest.LogCaptureFixture, faker: Faker):
+    async with ProgressBarData(num_steps=50, description=faker.pystr()) as root:
+        assert root._current_steps == pytest.approx(0)  # noqa: SLF001
+        assert root.num_steps == 50
+        assert root.step_weights is None
+        await root.set_(50)
+        assert root._current_steps == pytest.approx(50)  # noqa: SLF001
+        assert "WARNING" not in caplog.text
+        assert _PROGRESS_ALREADY_REACGED_MAXIMUM not in caplog.text
+        await root.set_(51)
+        assert root._current_steps == pytest.approx(50)  # noqa: SLF001
+        assert "WARNING" in caplog.text
+        assert _PROGRESS_ALREADY_REACGED_MAXIMUM in caplog.text
+
+        caplog.clear()
+        root.reset()
+
+        assert root._current_steps == pytest.approx(-1)  # noqa: SLF001
+        assert "WARNING" not in caplog.text
+        assert _PROGRESS_ALREADY_REACGED_MAXIMUM not in caplog.text
+
+        await root.set_(12)
+        assert root._current_steps == pytest.approx(12)  # noqa: SLF001
+        assert "WARNING" not in caplog.text
+        assert _PROGRESS_ALREADY_REACGED_MAXIMUM not in caplog.text
+
+        await root.set_(51)
+        assert root._current_steps == pytest.approx(50)  # noqa: SLF001
+        assert "WARNING" in caplog.text
+        assert _PROGRESS_ALREADY_REACGED_MAXIMUM in caplog.text
+
+
 async def test_concurrent_progress_bar(faker: Faker):
     async def do_something(root: ProgressBarData):
         async with root.sub_progress(steps=50, description=faker.pystr()) as sub:
@@ -304,7 +338,7 @@ async def test_too_many_updates_does_not_raise_but_show_warning_with_stack(
         await root.update()
         await root.update()
         await root.update()
-        assert "already reached maximum" in caplog.messages[0]
+        assert _PROGRESS_ALREADY_REACGED_MAXIMUM in caplog.messages[0]
         assert "TIP:" in caplog.messages[0]