diff --git a/sdk/python/kfp/cli/component.py b/sdk/python/kfp/cli/component.py index 459265fe3931..008626f70c08 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 d6240efdbe11..67c257ecb52e 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(