Repository: pytorch/torchx

Issue:
Docker build verbosity
## Description
Changing the docker image build to its low level implementation so it can be more verbose.

## Motivation/Background
Building the docker image can take quite some time, and for new users this makes it seem like the program is stuck (especially since the default base image that includes torchx is so big). Making it more verbose is not only a quality of life improvement for all users of the docker workspace, it also gives better visibility into the build process, potentially allowing optimization on the dockerfile.

On a side note, what is the rational for naming Dockerfile.torchx instead of just using a normal Dockerfile, is there a difference in the format?


## Detailed Proposal
Replacing the current docker build API call with its low level implementation. This would require instantiating a low-level client and processing the build event stream to show in real time to docker build commands. Also a processing function for the stream to be printing to screen correctly.


## Alternatives



## Additional context/links
https://github.com/pytorch/torchx/blob/19497eb1d2649f66cd12ca1eeed77353085f07e0/torchx/workspace/docker_workspace.py#L118

Hints from issue comments:
@d4l3k: https://docker-py.readthedocs.io/en/stable/api.html#docker.api.build.BuildApiMixin.build@ccharest93: Yeah the PR is written #815 and in draft mode, just waiting on #819 to merge before removing it from draft mode. @d4l3k: Ah nice, just saw that

I recently ran into this issue and well

On Wed, Feb 28, 2024, 11:12 Maatra ***@***.***> wrote:

> Yeah the PR is written #815 <https://github.com/pytorch/torchx/pull/815>
> and in draft mode, just waiting on #819
> <https://github.com/pytorch/torchx/pull/819> to merge before removing it
> from draft mode.
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/pytorch/torchx/issues/813#issuecomment-1969672534>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AAG56MCVKNT7DUHHWKK2JMLYV56TBAVCNFSM6AAAAABCTSMEUKVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSNRZGY3TENJTGQ>
> .
> You are receiving this because you commented.Message ID:
> ***@***.***>
>


Solution patch:
diff --git a/torchx/workspace/docker_workspace.py b/torchx/workspace/docker_workspace.py
index d93fb667..008b6e09 100644
--- a/torchx/workspace/docker_workspace.py
+++ b/torchx/workspace/docker_workspace.py
@@ -18,6 +18,7 @@ from typing import Dict, IO, Iterable, Mapping, Optional, TextIO, Tuple, TYPE_CH
 import fsspec
 
 import torchx
+from docker.errors import BuildError
 from torchx.specs import AppDef, CfgVal, Role, runopts
 from torchx.workspace.api import walk_workspace, WorkspaceMixin
 
@@ -93,6 +94,12 @@ class DockerWorkspaceMixin(WorkspaceMixin[Dict[str, Tuple[str, str]]]):
             type_=str,
             help="(remote jobs) the image repository to use when pushing patched images, must have push access. Ex: example.com/your/container",
         )
+        opts.add(
+            "quiet",
+            type_=bool,
+            default=False,
+            help="whether to suppress verbose output for image building. Defaults to ``False``.",
+        )
         return opts
 
     def build_workspace_and_update_role(
@@ -121,7 +128,7 @@ class DockerWorkspaceMixin(WorkspaceMixin[Dict[str, Tuple[str, str]]]):
                     f"failed to pull image {role.image}, falling back to local: {e}"
                 )
             log.info("Building workspace docker image (this may take a while)...")
-            image, _ = self._docker_client.images.build(
+            build_events = self._docker_client.api.build(
                 fileobj=context,
                 custom_context=True,
                 dockerfile=TORCHX_DOCKERFILE,
@@ -131,12 +138,26 @@ class DockerWorkspaceMixin(WorkspaceMixin[Dict[str, Tuple[str, str]]]):
                 },
                 pull=False,
                 rm=True,
+                decode=True,
                 labels={
                     self.LABEL_VERSION: torchx.__version__,
                 },
             )
+            image_id = None
+            for event in build_events:
+                if message := event.get("stream"):
+                    if not cfg.get("quiet", False):
+                        message = message.strip("\r\n").strip("\n")
+                        if message:
+                            log.info(message)
+                if aux := event.get("aux"):
+                    image_id = aux["ID"]
+                if error := event.get("error"):
+                    raise BuildError(reason=error, build_log=None)
             if len(old_imgs) == 0 or role.image not in old_imgs:
-                role.image = image.id
+                assert image_id, "image id was not found"
+                role.image = image_id
+
         finally:
             context.close()
 


Test patch:
diff --git a/torchx/schedulers/test/kubernetes_scheduler_test.py b/torchx/schedulers/test/kubernetes_scheduler_test.py
index 8e6b356a..8ad0b0e8 100644
--- a/torchx/schedulers/test/kubernetes_scheduler_test.py
+++ b/torchx/schedulers/test/kubernetes_scheduler_test.py
@@ -700,6 +700,7 @@ spec:
         self.assertEqual(
             set(runopts._opts.keys()),
             {
+                "quiet",
                 "queue",
                 "namespace",
                 "image_repo",
diff --git a/torchx/workspace/test/docker_workspace_test.py b/torchx/workspace/test/docker_workspace_test.py
index fa9a4335..d8418d35 100644
--- a/torchx/workspace/test/docker_workspace_test.py
+++ b/torchx/workspace/test/docker_workspace_test.py
@@ -14,9 +14,10 @@ import tempfile
 import unittest
 from io import StringIO
 from pathlib import Path
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import fsspec
+from docker.errors import BuildError
 
 from torchx.specs import AppDef, Role
 from torchx.workspace.docker_workspace import (
@@ -64,7 +65,7 @@ class DockerWorkspaceMockTest(unittest.TestCase):
     def test_runopts(self) -> None:
         self.assertCountEqual(
             DockerWorkspaceMixin().workspace_opts()._opts.keys(),
-            {"image_repo"},
+            {"image_repo", "quiet"},
         )
 
     def test_update_app_images(self) -> None:
@@ -130,6 +131,49 @@ class DockerWorkspaceMockTest(unittest.TestCase):
         self.assertEqual(img.tag.call_count, 2)
         self.assertEqual(client.images.push.call_count, 2)
 
+    @patch("torchx.workspace.docker_workspace.log")
+    def test_build_workspace_and_update_role(self, mock_log: MagicMock) -> None:
+        client = MagicMock()
+        img = MagicMock()
+        client.images.get.return_value = img
+        image_id = "some_id"
+        client.api.build.return_value = [
+            {"stream": "some message"},
+            {"aux": {"ID": image_id}},
+        ]
+        workspace = DockerWorkspaceMixin(docker_client=client)
+        role = Role(name="b", image="example.com/repo:hasha")
+        workspace.build_workspace_and_update_role(
+            role=role,
+            workspace="bogus",
+            cfg={"image_repo": "example.com/repo", "quiet": "true"},
+        )
+        self.assertEqual(client.api.build.call_count, 1)
+        self.assertEqual(role.image, image_id)
+        self.assertEqual(mock_log.info.call_count, 2)
+
+    @patch("torchx.workspace.docker_workspace.log")
+    def test_build_workspace_and_update_raises_error(self, mock_log: MagicMock) -> None:
+        client = MagicMock()
+        img = MagicMock()
+        client.images.get.return_value = img
+        error = "some error"
+        message = "some message"
+        client.api.build.return_value = [
+            {"stream": message},
+            {"error": error},
+        ]
+        workspace = DockerWorkspaceMixin(docker_client=client)
+        role = Role(name="b", image="example.com/repo:hasha")
+        with self.assertRaisesRegex(BuildError, error):
+            workspace.build_workspace_and_update_role(
+                role=role,
+                workspace="bogus",
+                cfg={"image_repo": "example.com/repo"},
+            )
+        self.assertEqual(mock_log.info.call_count, 3)
+        self.assertEqual(mock_log.info.call_args_list[-1][0][0], message)
+
     def test_push_images_empty(self) -> None:
         workspace = DockerWorkspaceMixin()
         workspace.push_images({})