From 40b7532f92dba0b4cf2a7a085b14d5d7b260c68d Mon Sep 17 00:00:00 2001 From: Anish Asthana Date: Tue, 22 Apr 2025 11:57:10 -0400 Subject: [PATCH] Add PodmanRunner This is almost entirely based on the existing docker runner implementation, with minor differences for podman specifics Signed-off-by: Anish Asthana --- sdk/CONTRIBUTING.md | 65 ++-- sdk/python/kfp/local/__init__.py | 2 + sdk/python/kfp/local/config.py | 20 +- sdk/python/kfp/local/podman_task_handler.py | 106 +++++++ .../kfp/local/podman_task_handler_test.py | 290 ++++++++++++++++++ sdk/python/kfp/local/task_dispatcher.py | 3 + sdk/python/requirements.in | 2 + sdk/python/requirements.txt | 14 +- 8 files changed, 469 insertions(+), 33 deletions(-) create mode 100644 sdk/python/kfp/local/podman_task_handler.py create mode 100644 sdk/python/kfp/local/podman_task_handler_test.py diff --git a/sdk/CONTRIBUTING.md b/sdk/CONTRIBUTING.md index ad995901238..5df77aa9d1b 100644 --- a/sdk/CONTRIBUTING.md +++ b/sdk/CONTRIBUTING.md @@ -1,4 +1,4 @@ -## Contributing to the `kfp` SDK +# Contributing to the `kfp` SDK For developing KFP v2 SDK, use the `master` branch. @@ -6,18 +6,18 @@ For developing KFP v1 SDK, use the [sdk/release-1.8 branch](https://github.com/k For general contribution guidelines including pull request conventions, see [pipelines/CONTRIBUTING.md](https://github.com/kubeflow/pipelines/blob/master/CONTRIBUTING.md). -### Pre-requisites +## Pre-requisites -Clone the repo: +Clone the repo: ```bash git clone https://github.com/kubeflow/pipelines.git && cd pipelines ``` -We suggest using a tool like [virtual env](https://docs.python.org/3/library/venv.html) or something similar for isolating -your environment and/or packages for you development environment. For this setup we'll stick with virtual env. +We suggest using a tool like [virtual env](https://docs.python.org/3/library/venv.html) or something similar for isolating +your environment and/or packages for you development environment. For this setup we'll stick with virtual env. -For supported python versions, see the sdk [setup.py](https://github.com/kubeflow/pipelines/blob/master/sdk/python/setup.py). +For supported python versions, see the sdk [setup.py](https://github.com/kubeflow/pipelines/blob/master/sdk/python/setup.py). ```bash # optional, replace with your tool of choice @@ -47,11 +47,12 @@ The SDK also relies on a couple other python packages also found within KFP. These consists of the [api proto package](https://github.com/kubeflow/pipelines/tree/master/api) and the kfp [kubernetes_platform](https://github.com/kubeflow/pipelines/tree/master/kubernetes_platform) package. For the proto code, we need protobuf-compiler to generate the python code. Read [here](../kubernetes_platform#dependencies) more about how to install this -dependency. +dependency. You can install both packages either in development mode if you are planning to do active development on these, or simply do a regular install. To install the proto package: + ```bash pushd api make generate python-dev # omit -dev for regular install @@ -59,77 +60,93 @@ popd ``` To install the kubernetes_platform package: + ```bash pushd kubernetes_platform make python-dev # omit -dev for regular install popd ``` -### Testing +## Testing + We suggest running unit tests using [`pytest`](https://docs.pytest.org/en/7.1.x/). From the project root, the following runs all KFP SDK unit tests: -```sh + +```bash pytest ``` To run tests in parallel for faster execution, you can run the tests using the `pytest-xdist` plugin: -```sh +```bash pytest -n auto ``` -### Code Style +## Code Style + Dependencies for code style checks/changes can be found in the kfp SDK [requirements-dev.txt](https://github.com/kubeflow/pipelines/blob/master/sdk/python/requirements-dev.txt). +### Style Guide [Required] -#### Style Guide [Required] The KFP SDK follows the [Google Python Style Guide](https://google.github.io/styleguide/pyguide.html). -#### Formatting [Required] +### Formatting [Required] + Please format your code using [yapf](https://github.com/google/yapf) according to the [`.style.yapf`](https://github.com/kubeflow/pipelines/blob/master/.style.yapf) file. From the project root, run the following code to format your code: -```sh + +```bash yapf --in-place --recursive ./sdk/python ``` -#### Docformatter [Required] +### Docformatter [Required] + We encourage you to lint your docstrings using [docformatter](https://github.com/PyCQA/docformatter). From the project root, run the following code to lint your docstrings: -```sh + +```bash docformatter --in-place --recursive ./sdk/python ``` -#### Formatting Imports [Required] +### Formatting Imports [Required] + Please organize your imports using [isort](https://pycqa.github.io/isort/index.html) according to the [`.isort.cfg`](https://github.com/kubeflow/pipelines/blob/master/.isort.cfg) file. From the project root, run the following code to format your code: -```sh + +```bash isort sdk/python ``` -#### Pylint [Encouraged] +### Pylint [Encouraged] + We encourage you to lint your code using [pylint](https://pylint.org/) according to the project [`.pylintrc`](https://github.com/kubeflow/pipelines/blob/master/.pylintrc) file. From the project root, run the following code to lint your code: -```sh + +```bash pylint ./sdk/python/kfp ``` Note: `kfp` is not currently fully pylint-compliant. Consider substituting the path argument with the files touched by your development. -#### Static Type Checking [Encouraged] +### Static Type Checking [Encouraged] + Please use [mypy](https://mypy.readthedocs.io/en/stable/) to check your type annotations. From the project root, run the following code to lint your docstrings: -```sh + +```bash mypy ./sdk/python/kfp/ ``` + Note: `kfp` is not currently fully mypy-compliant. Consider substituting the path argument with the files touched by your development. +### Pre-commit [Recommended] -#### Pre-commit [Recommended] Consider using [`pre-commit`](https://github.com/pre-commit/pre-commit) with the provided [`.pre-commit-config.yaml`](https://github.com/kubeflow/pipelines/blob/master/.pre-commit-config.yaml) to implement the above changes: -```sh + +```bash pre-commit install ``` diff --git a/sdk/python/kfp/local/__init__.py b/sdk/python/kfp/local/__init__.py index 6848df94003..d5c439814f0 100755 --- a/sdk/python/kfp/local/__init__.py +++ b/sdk/python/kfp/local/__init__.py @@ -16,10 +16,12 @@ from kfp.local.config import DockerRunner from kfp.local.config import init +from kfp.local.config import PodmanRunner from kfp.local.config import SubprocessRunner __all__ = [ 'init', 'SubprocessRunner', 'DockerRunner', + 'PodmanRunner', ] diff --git a/sdk/python/kfp/local/config.py b/sdk/python/kfp/local/config.py index 9ea01d18369..395fcbddef7 100755 --- a/sdk/python/kfp/local/config.py +++ b/sdk/python/kfp/local/config.py @@ -59,6 +59,20 @@ def __post_init__(self): ) from e +@dataclasses.dataclass +class PodmanRunner: + """Runner that indicates that local tasks should be run as a container + using Podman.""" + + def __post_init__(self): + try: + import podman # noqa + except ImportError as e: + raise ImportError( + f"Package 'podman' must be installed to use {PodmanRunner.__name__!r}. Install it using 'pip install podman'." + ) from e + + class LocalExecutionConfig: instance = None @@ -78,7 +92,7 @@ def __init__( pipeline_root: str, raise_on_error: bool, ) -> None: - permitted_runners = (SubprocessRunner, DockerRunner) + permitted_runners = (SubprocessRunner, DockerRunner, PodmanRunner) if not isinstance(runner, permitted_runners): raise ValueError( f'Got unknown runner {runner} of type {runner.__class__.__name__}. Runner should be one of the following types: {". ".join(prunner.__name__ for prunner in permitted_runners)}.' @@ -97,7 +111,7 @@ def validate(cls): def init( # annotate with subclasses, not parent class, for more helpful ref docs - runner: Union[SubprocessRunner, DockerRunner], + runner: Union[SubprocessRunner, DockerRunner, PodmanRunner], pipeline_root: str = './local_outputs', raise_on_error: bool = True, ) -> None: @@ -106,7 +120,7 @@ def init( Once called, components can be invoked locally outside of a pipeline definition. Args: - runner: The runner to use. Supported runners: kfp.local.SubprocessRunner and kfp.local.DockerRunner. + runner: The runner to use. Supported runners: kfp.local.SubprocessRunner, kfp.local.DockerRunner, and kfp.local.PodmanRunner. pipeline_root: Destination for task outputs. raise_on_error: If True, raises an exception when a local task execution fails. If False, fails gracefully and does not terminate the current program. """ diff --git a/sdk/python/kfp/local/podman_task_handler.py b/sdk/python/kfp/local/podman_task_handler.py new file mode 100644 index 00000000000..83f9f66c1e3 --- /dev/null +++ b/sdk/python/kfp/local/podman_task_handler.py @@ -0,0 +1,106 @@ +import os +from typing import Any, Dict, List + +from kfp.local import config +from kfp.local import status +from kfp.local import task_handler_interface + + +class PodmanTaskHandler(task_handler_interface.ITaskHandler): + """The task handler corresponding to PodmanRunner.""" + + def __init__( + self, + image: str, + full_command: List[str], + pipeline_root: str, + runner: config.PodmanRunner, + ) -> None: + self.image = image + self.full_command = full_command + self.pipeline_root = pipeline_root + self.runner = runner + + def get_volumes_to_mount(self) -> Dict[str, Any]: + """Gets the volume configuration to mount the pipeline root to the + container so that outputs can be obtained outside of the container.""" + if not os.path.isabs(self.pipeline_root): + # Defensive check that is enforced by upstream code. + # users should not hit this + raise ValueError( + "'pipeline_root' should be an absolute path to correctly construct the volume mount specification." + ) + return {self.pipeline_root: {'bind': self.pipeline_root, 'mode': 'rw'}} + + def run(self) -> status.Status: + """Runs the container and returns the status.""" + # nest podman import in case not available in user env so that + # this module is runnable, even if not using PodmanRunner + import podman + client = podman.PodmanClient() + try: + volumes = self.get_volumes_to_mount() + return_code = run_container( + client=client, + image=self.image, + command=self.full_command, + volumes=volumes, + ) + finally: + client.close() + return status.Status.SUCCESS if return_code == 0 else status.Status.FAILURE + + +def add_latest_tag_if_not_present(image: str) -> str: + """Adds the 'latest' tag if no tag is present in the image name.""" + if ':' not in image: + return f'{image}:latest' + return image + + +def run_container( + client: 'podman.PodmanClient', + image: str, + command: List[str], + volumes: Dict[str, Dict[str, str]], +) -> int: + """Runs a container using Podman. + + Args: + client: The Podman client instance. + image: The container image to use. + command: The command to run in the container. + volumes: Dictionary mapping host paths to volume configuration. + + Returns: + The exit code of the container. + """ + + image = add_latest_tag_if_not_present(image=image) + + image_exists = any( + image in existing_image.tags for existing_image in client.images.list()) + if image_exists: + print(f'Found image {image!r}\n') + else: + print(f'Pulling image {image!r}') + repository, tag = image.split(':') + client.images.pull(repository=repository, tag=tag) + print('Image pull complete\n') + + # Create and run container + container = client.containers.create( + image=image, + command=command, + volumes=volumes, + detach=True, + ) + + # Start container and stream logs + container.start() + for line in container.logs(stream=True): + print(line.decode(), end='') + + # Wait for container to finish and get exit code + container.wait() + return container.exit_code diff --git a/sdk/python/kfp/local/podman_task_handler_test.py b/sdk/python/kfp/local/podman_task_handler_test.py new file mode 100644 index 00000000000..d1acbb90e11 --- /dev/null +++ b/sdk/python/kfp/local/podman_task_handler_test.py @@ -0,0 +1,290 @@ +import os +from typing import Optional +import unittest +from unittest import mock + +from kfp import dsl +from kfp import local +from kfp.dsl import Artifact +from kfp.dsl import Output +from kfp.local import podman_task_handler +from kfp.local import testing_utilities + + +class PodmanMockTestCase(unittest.TestCase): + + def setUp(self): + super().setUp() + self.podman_mock = mock.Mock() + patcher = mock.patch('podman.PodmanClient') + self.mocked_podman_client = patcher.start().return_value + + mock_container = mock.Mock() + self.mocked_podman_client.containers.create.return_value = mock_container + # mock successful run + mock_container.logs.return_value = [ + 'fake'.encode('utf-8'), + 'container'.encode('utf-8'), + 'logs'.encode('utf-8'), + ] + mock_container.exit_code = 0 + + def tearDown(self): + super().tearDown() + self.podman_mock.reset_mock() + + +class TestRunContainer(PodmanMockTestCase): + + def test_no_volumes(self): + podman_task_handler.run_container( + client=self.mocked_podman_client, + image='alpine', + command=['echo', 'foo'], + volumes={}, + ) + + self.mocked_podman_client.containers.create.assert_called_once_with( + image='alpine:latest', + command=['echo', 'foo'], + volumes={}, + detach=True, + ) + + def test_cwd_volume(self): + current_test_dir = os.path.dirname(os.path.abspath(__file__)) + podman_task_handler.run_container( + client=self.mocked_podman_client, + image='alpine', + command=['cat', '/localdir/podman_task_handler_test.py'], + volumes={current_test_dir: { + 'bind': '/localdir', + 'mode': 'ro' + }}, + ) + self.mocked_podman_client.containers.create.assert_called_once_with( + image='alpine:latest', + command=['cat', '/localdir/podman_task_handler_test.py'], + volumes={current_test_dir: { + 'bind': '/localdir', + 'mode': 'ro' + }}, + detach=True, + ) + + +class TestPodmanTaskHandler(PodmanMockTestCase): + + def test_get_volumes_to_mount(self): + handler = podman_task_handler.PodmanTaskHandler( + image='alpine', + full_command=['echo', 'foo'], + pipeline_root=os.path.abspath('my_root'), + runner=local.PodmanRunner(), + ) + volumes = handler.get_volumes_to_mount() + self.assertEqual( + volumes, { + os.path.abspath('my_root'): { + 'bind': os.path.abspath('my_root'), + 'mode': 'rw' + } + }) + + def test_run(self): + handler = podman_task_handler.PodmanTaskHandler( + image='alpine', + full_command=['echo', 'foo'], + pipeline_root=os.path.abspath('my_root'), + runner=local.PodmanRunner(), + ) + + handler.run() + self.mocked_podman_client.containers.create.assert_called_once_with( + image='alpine:latest', + command=['echo', 'foo'], + volumes={ + os.path.abspath('my_root'): { + 'bind': os.path.abspath('my_root'), + 'mode': 'rw' + } + }, + detach=True, + ) + + def test_pipeline_root_relpath(self): + with self.assertRaisesRegex( + ValueError, + r"'pipeline_root' should be an absolute path to correctly construct the volume mount specification\." + ): + podman_task_handler.PodmanTaskHandler( + image='alpine', + full_command=['echo', 'foo'], + pipeline_root='my_relpath', + runner=local.PodmanRunner(), + ).run() + + +class TestE2E(PodmanMockTestCase, + testing_utilities.LocalRunnerEnvironmentTestCase): + + def setUp(self): + super().setUp() + local.init(runner=local.PodmanRunner()) + + def test_python(self): + + @dsl.component + def artifact_maker(x: str, a: Output[Artifact]): + with open(a.path, 'w') as f: + f.write(x) + + try: + artifact_maker(x='foo') + # cannot get outputs if they aren't created due to mock + except FileNotFoundError: + pass + + create_mock = self.mocked_podman_client.containers.create + create_mock.assert_called_once() + kwargs = create_mock.call_args[1] + self.assertEqual( + kwargs['image'], + 'python:3.9', + ) + self.assertTrue( + any('def artifact_maker' in c for c in kwargs['command'])) + self.assertTrue(kwargs['detach']) + root_vol_key = [ + key for key in kwargs['volumes'].keys() if 'local_outputs' in key + ][0] + self.assertEqual(kwargs['volumes'][root_vol_key]['bind'], root_vol_key) + self.assertEqual(kwargs['volumes'][root_vol_key]['mode'], 'rw') + + def test_empty_container_component(self): + + @dsl.container_component + def comp(): + return dsl.ContainerSpec(image='alpine') + + try: + comp() + # cannot get outputs if they aren't created due to mock + except FileNotFoundError: + pass + + create_mock = self.mocked_podman_client.containers.create + create_mock.assert_called_once() + kwargs = create_mock.call_args[1] + self.assertEqual( + kwargs['image'], + 'alpine:latest', + ) + self.assertEqual(kwargs['command'], []) + + def test_container_component(self): + + @dsl.container_component + def artifact_maker(x: str,): + return dsl.ContainerSpec( + image='alpine', + command=['sh', '-c', f'echo prefix-{x}'], + ) + + try: + artifact_maker(x='foo') + # cannot get outputs if they aren't created due to mock + except FileNotFoundError: + pass + + create_mock = self.mocked_podman_client.containers.create + create_mock.assert_called_once() + kwargs = create_mock.call_args[1] + self.assertEqual( + kwargs['image'], + 'alpine:latest', + ) + self.assertEqual(kwargs['command'], [ + 'sh', + '-c', + 'echo prefix-foo', + ]) + self.assertTrue(kwargs['detach']) + root_vol_key = [ + key for key in kwargs['volumes'].keys() if 'local_outputs' in key + ][0] + self.assertEqual(kwargs['volumes'][root_vol_key]['bind'], root_vol_key) + self.assertEqual(kwargs['volumes'][root_vol_key]['mode'], 'rw') + + def test_if_present_with_string_omitted(self): + + @dsl.container_component + def comp(x: Optional[str] = None): + return dsl.ContainerSpec( + image='alpine:3.19.0', + command=[ + dsl.IfPresentPlaceholder( + input_name='x', + then=['echo', x], + else_=['echo', 'No input provided!']) + ]) + + comp() + + create_mock = self.mocked_podman_client.containers.create + create_mock.assert_called_once() + kwargs = create_mock.call_args[1] + self.assertEqual( + kwargs['image'], + 'alpine:3.19.0', + ) + self.assertEqual(kwargs['command'], [ + 'echo', + 'No input provided!', + ]) + + def test_if_present_with_string_provided(self): + + @dsl.container_component + def comp(x: Optional[str] = None): + return dsl.ContainerSpec( + image='alpine:3.19.0', + command=[ + dsl.IfPresentPlaceholder( + input_name='x', + then=['echo', x], + else_=['echo', 'No artifact provided!']) + ]) + + comp(x='foo') + + create_mock = self.mocked_podman_client.containers.create + create_mock.assert_called_once() + kwargs = create_mock.call_args[1] + self.assertEqual( + kwargs['image'], + 'alpine:3.19.0', + ) + self.assertEqual(kwargs['command'], [ + 'echo', + 'foo', + ]) + + def test_concat_placeholder(self): + + @dsl.container_component + def comp(x: Optional[str] = None): + return dsl.ContainerSpec( + image='alpine', + command=[dsl.ConcatPlaceholder(['prefix-', x, '-suffix'])]) + + comp() + + create_mock = self.mocked_podman_client.containers.create + create_mock.assert_called_once() + kwargs = create_mock.call_args[1] + self.assertEqual( + kwargs['image'], + 'alpine:latest', + ) + self.assertEqual(kwargs['command'], ['prefix-null-suffix']) diff --git a/sdk/python/kfp/local/task_dispatcher.py b/sdk/python/kfp/local/task_dispatcher.py index 047bcd92b8b..2eb298f5e5f 100755 --- a/sdk/python/kfp/local/task_dispatcher.py +++ b/sdk/python/kfp/local/task_dispatcher.py @@ -22,6 +22,7 @@ from kfp.local import executor_output_utils from kfp.local import logging_utils from kfp.local import placeholder_utils +from kfp.local import podman_task_handler from kfp.local import status from kfp.local import subprocess_task_handler from kfp.local import task_handler_interface @@ -134,6 +135,8 @@ def run_single_task_implementation( subprocess_task_handler.SubprocessTaskHandler, local.DockerRunner: docker_task_handler.DockerTaskHandler, + local.PodmanRunner: + podman_task_handler.PodmanTaskHandler, } TaskHandler = task_handler_map[runner_type] diff --git a/sdk/python/requirements.in b/sdk/python/requirements.in index 8dc8682f52a..7e0aa576513 100644 --- a/sdk/python/requirements.in +++ b/sdk/python/requirements.in @@ -26,3 +26,5 @@ urllib3<3.0.0 ## standard library backports ## typing-extensions>=3.7.4,<5; python_version<"3.9" + +podman>=5.4.0.1 diff --git a/sdk/python/requirements.txt b/sdk/python/requirements.txt index 23ecd9647a4..0895dc05ddd 100644 --- a/sdk/python/requirements.txt +++ b/sdk/python/requirements.txt @@ -1,9 +1,5 @@ -# -# This file is autogenerated by pip-compile with Python 3.12 -# by the following command: -# -# pip-compile --no-emit-index-url requirements.in -# +# This file was autogenerated by uv via the following command: +# uv pip compile requirements.in -o requirements.txt cachetools==5.5.2 # via google-auth certifi==2025.1.31 @@ -53,6 +49,8 @@ oauthlib==3.2.2 # via # kubernetes # requests-oauthlib +podman==5.4.0.1 + # via -r requirements.in proto-plus==1.26.0 # via google-api-core protobuf==4.25.6 @@ -81,6 +79,7 @@ requests==2.32.3 # google-api-core # google-cloud-storage # kubernetes + # podman # requests-oauthlib # requests-toolbelt requests-oauthlib==2.0.0 @@ -96,11 +95,14 @@ six==1.17.0 # python-dateutil tabulate==0.9.0 # via -r requirements.in +tomli==2.2.1 + # via podman urllib3==2.4.0 # via # -r requirements.in # kfp-server-api # kubernetes + # podman # requests websocket-client==1.8.0 # via kubernetes