Repository: PrefectHQ/prefect

Issue:
Docker tasks need better logging

Discussion:

Patch:
diff --git a/changes/pr3672.yaml b/changes/pr3672.yaml
new file mode 100644
index 0000000000..7309e9d060
--- /dev/null
+++ b/changes/pr3672.yaml
@@ -0,0 +1,2 @@
+enhancement:
+  - "Add debug logging for Docker tasks PullImage and BuildImage" - [#3672](https://github.com/PrefectHQ/prefect/pull/3672)"
\ No newline at end of file
diff --git a/src/prefect/tasks/docker/images.py b/src/prefect/tasks/docker/images.py
index 683de9a0c0..d347ceb10c 100644
--- a/src/prefect/tasks/docker/images.py
+++ b/src/prefect/tasks/docker/images.py
@@ -109,6 +109,7 @@ class PullImage(Task):
         - extra_docker_kwargs (dict, optional): Extra keyword arguments to pass through to the
             Docker call (cf. method `pull`). See
             https://docker-py.readthedocs.io/en/stable/api.html for more details
+        - stream_logs (bool, optional): Adds debug logs with pull status info from Docker
         - **kwargs (dict, optional): Additional keyword arguments to pass to the Task
             constructor
     """
@@ -119,17 +120,19 @@ class PullImage(Task):
         tag: str = None,
         docker_server_url: str = "unix:///var/run/docker.sock",
         extra_docker_kwargs: dict = None,
+        stream_logs: bool = False,
         **kwargs: Any,
     ):
         self.repository = repository
         self.tag = tag
         self.docker_server_url = docker_server_url
         self.extra_docker_kwargs = extra_docker_kwargs
+        self.stream_logs = stream_logs
 
         super().__init__(**kwargs)
 
     @defaults_from_attrs(
-        "repository", "tag", "docker_server_url", "extra_docker_kwargs"
+        "repository", "tag", "docker_server_url", "extra_docker_kwargs", "stream_logs"
     )
     def run(
         self,
@@ -137,6 +140,7 @@ class PullImage(Task):
         tag: str = None,
         docker_server_url: str = "unix:///var/run/docker.sock",
         extra_docker_kwargs: dict = None,
+        stream_logs: bool = False,
     ) -> str:
         """
         Task run method.
@@ -151,12 +155,14 @@ class PullImage(Task):
             - extra_docker_kwargs (dict, optional): Extra keyword arguments to pass through to the
                 Docker call (cf. method `pull`). See
                 https://docker-py.readthedocs.io/en/stable/api.html for more details
+            - stream_logs (bool, optional): Adds debug logs with pull status info from Docker
 
         Returns:
             - str: The output from Docker for pulling the image
 
         Raises:
             - ValueError: if `repository` is `None`
+            - HTTPError: if image doesn't exist or provided repository/tag is not correct
         """
         if not repository:
             raise ValueError("A repository to pull the image from must be specified.")
@@ -164,15 +170,33 @@ class PullImage(Task):
         # 'import docker' is expensive time-wise, we should do this just-in-time to keep
         # the 'import prefect' time low
         import docker
+        from requests.exceptions import HTTPError
 
         client = docker.APIClient(base_url=docker_server_url, version="auto")
         self.logger.debug(f"Pulling image {repository}:{tag}")
-        api_result = client.pull(
-            repository=repository, tag=tag, **(extra_docker_kwargs or dict())
-        )
-
-        self.logger.debug(f"Pulled image {repository}:{tag}")
-        return api_result
+        lines = []
+        try:
+            api_result = client.pull(
+                repository=repository,
+                tag=tag,
+                stream=stream_logs,
+                decode=True,
+                **(extra_docker_kwargs or dict()),
+            )
+            if isinstance(api_result, str):
+                return "".join(line for line in api_result.split("\r"))
+            for line in api_result:
+                status_line = line.get("status")
+                if status_line and stream_logs:
+                    self.logger.debug(status_line)
+                lines.append(str(line))
+            self.logger.debug(f"Pulled image {repository}:{tag}")
+            return "\n".join(lines)
+        except HTTPError as exc:
+            self.logger.exception(exc)
+            raise HTTPError(
+                "Can't pull the image, check a repository name or image tag."
+            ) from exc
 
 
 class PushImage(Task):
@@ -444,6 +468,7 @@ class BuildImage(Task):
         - extra_docker_kwargs (dict, optional): Extra keyword arguments to pass through to the
             Docker call (cf. method `build`). See
             https://docker-py.readthedocs.io/en/stable/api.html for more details
+        - stream_logs (bool, optional): Adds debug logs with build image status info from Docker
         - **kwargs (dict, optional): Additional keyword arguments to pass to the Task
             constructor
     """
@@ -457,6 +482,7 @@ class BuildImage(Task):
         forcerm: bool = False,
         docker_server_url: str = "unix:///var/run/docker.sock",
         extra_docker_kwargs: dict = None,
+        stream_logs: bool = False,
         **kwargs: Any,
     ):
         self.path = path
@@ -466,6 +492,7 @@ class BuildImage(Task):
         self.forcerm = forcerm
         self.docker_server_url = docker_server_url
         self.extra_docker_kwargs = extra_docker_kwargs
+        self.stream_logs = stream_logs
 
         super().__init__(**kwargs)
 
@@ -477,6 +504,7 @@ class BuildImage(Task):
         "forcerm",
         "docker_server_url",
         "extra_docker_kwargs",
+        "stream_logs",
     )
     def run(
         self,
@@ -487,6 +515,7 @@ class BuildImage(Task):
         forcerm: bool = False,
         docker_server_url: str = "unix:///var/run/docker.sock",
         extra_docker_kwargs: dict = None,
+        stream_logs: bool = False,
     ) -> None:
         """
         Task run method.
@@ -504,6 +533,7 @@ class BuildImage(Task):
             - extra_docker_kwargs (dict, optional): Extra keyword arguments to pass through to the
                 Docker call (cf. method `build`). See
                 https://docker-py.readthedocs.io/en/stable/api.html for more details
+            - stream_logs (bool, optional): Adds debug logs with build image status info from Docker
 
         Returns:
             - List[dict]: a cleaned dictionary of the output of `client.build`
@@ -542,4 +572,9 @@ class BuildImage(Task):
             for line in resp.split(b"\r\n")
             if line
         ]
+        if self.stream_logs:
+            for line in output:
+                stream_line = line.get("stream")
+                if stream_line and stream_line != "\n":
+                    self.logger.debug(stream_line)
         return output


Test patch:
diff --git a/tests/tasks/docker/test_images.py b/tests/tasks/docker/test_images.py
index 89835e9dc2..6cd69927a9 100644
--- a/tests/tasks/docker/test_images.py
+++ b/tests/tasks/docker/test_images.py
@@ -198,6 +198,28 @@ class TestPullImageTask(DockerLoggingTestingUtilityMixin):
             result = task.run(repository="test")
             assert result == expected_docker_output
 
+    def test_returns_pull_value_with_stream_logs(self, monkeypatch, caplog):
+
+        task = PullImage()
+
+        api = MagicMock()
+        expected_docker_output = [
+            {"status": "Pulling"},
+            {"status": "Digest"},
+            {"status": "Test"},
+        ]
+        expected_result = (
+            "{'status': 'Pulling'}\n{'status': 'Digest'}\n{'status': 'Test'}"
+        )
+
+        api.return_value.pull = MagicMock(return_value=expected_docker_output)
+        monkeypatch.setattr("docker.APIClient", api)
+
+        with caplog.at_level(logging.DEBUG, logger=task.logger.name):
+
+            result = task.run(repository="original", stream_logs=True)
+            assert result == expected_result
+
     def test_logs_twice_on_success(self, monkeypatch, caplog):
         tag = "A very specific tag for an image"
         repository = "An even more specific repository"