fix(sdk): Add error handling. Fixes #11164

Signed-off-by: Yael <fishel.yael@gmail.com>
kubeflow · Nov 13, 2024 · f55ebe9 · f55ebe9
1 parent 60a8865
commit f55ebe9
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 1 deletion.
diff --git 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
@@ -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(