From f55ebe9ee140fdb08c62d8ccda99ebb3b8ad330f Mon Sep 17 00:00:00 2001
From: Yael <fishel.yael@gmail.com>
Date: Wed, 13 Nov 2024 17:17:30 +0200
Subject: [PATCH] fix(sdk): Add error handling. Fixes #11164

Signed-off-by: Yael <fishel.yael@gmail.com>
---
 sdk/python/kfp/cli/component.py      | 11 ++++++++++-
 sdk/python/kfp/cli/component_test.py | 28 ++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/sdk/python/kfp/cli/component.py b/sdk/python/kfp/cli/component.py
index 459265fe393..008626f70c0 100644
--- a/sdk/python/kfp/cli/component.py
+++ b/sdk/python/kfp/cli/component.py
@@ -64,6 +64,10 @@
 _COMPONENT_ROOT_DIR = pathlib.Path('/usr/local/src/kfp/components')
 
 
+class ComponentBuildError(Exception):
+    pass
+
+
 @contextlib.contextmanager
 def _registered_modules():
     registered_modules = {}
@@ -313,8 +317,12 @@ def build_image(self, platform: str, push_image: bool):
             )
             for log in logs:
                 message = log.get('stream', '').rstrip('\n')
+                error = log.get('error', '').rstrip('\n')
                 if message:
                     logging.info(f'{docker_log_prefix}: {message}')
+                if error:
+                    raise ComponentBuildError(
+                        f'{docker_log_prefix} ERROR: {error}')
 
         except docker.errors.BuildError as e:
             for log in e.build_log:
@@ -339,7 +347,8 @@ def build_image(self, platform: str, push_image: bool):
                 if status:
                     logging.info(f'{docker_log_prefix}: {layer} {status}')
                 if error:
-                    logging.error(f'{docker_log_prefix} ERROR: {error}')
+                    raise ComponentBuildError(
+                        f'{docker_log_prefix} ERROR: {error}')
         except docker.errors.BuildError as e:
             logging.error(f'{docker_log_prefix}: {e}')
             raise e
diff --git a/sdk/python/kfp/cli/component_test.py b/sdk/python/kfp/cli/component_test.py
index d6240efdbe1..67c257ecb52 100644
--- a/sdk/python/kfp/cli/component_test.py
+++ b/sdk/python/kfp/cli/component_test.py
@@ -449,6 +449,34 @@ def test_docker_client_is_called_to_build_but_skips_pushing(self):
         self._docker_client.api.build.assert_called_once()
         self._docker_client.images.push.assert_not_called()
 
+    def test_docker_client_failed_to_build_img(self):
+        self._docker_client.api.build.return_value = [{'error': 'Error log'}]
+        created_component = _make_component(
+            func_name='train', target_image='custom-image')
+        _write_components('components.py', created_component)
+
+        result = self.runner.invoke(
+            self.cli,
+            ['build', str(self._working_dir), '--push-image'],
+        )
+
+        self.assertEqual(result.exit_code, 1)
+        self._docker_client.api.build.assert_called_once()
+        self._docker_client.images.push.assert_not_called()
+
+    def test_docker_client_failed_to_push_img(self):
+        self._docker_client.images.push.return_value = [{'error': 'Error log'}]
+        created_component = _make_component(
+            func_name='train', target_image='custom-image')
+        _write_components('components.py', created_component)
+
+        result = self.runner.invoke(
+            self.cli,
+            ['build', str(self._working_dir), '--push-image'],
+        )
+        self.assertEqual(result.exit_code, 1)
+        self._docker_client.images.push.assert_called_once()
+
     @mock.patch('kfp.__version__', '1.2.3')
     def test_docker_file_is_created_correctly(self):
         component = _make_component(