feat(sdk): auto-detect and install Kubeflow SDK for component functions #12027 #12368

Prateekbala · 2025-10-15T15:23:28Z

Description

This PR implements automatic Kubeflow SDK integration for KFP components as requested in #12027. It adds the ability to automatically detect and install the Kubeflow SDK when components use kubeflow imports, while also providing explicit extras for opt-in installation.

Changes Made

1. Added Kubeflow Extras to setup.py

Added kubeflow = ['kubeflow'] extras option
Included kubeflow in the 'all' extras bundle
Allows users to install with pip install kfp[kubeflow] or pip install kfp[all]

2. Implemented AST-based Auto-detection

Added _detect_kubeflow_imports_in_function() in component_factory.py
Uses Abstract Syntax Tree parsing (ast + inspect + textwrap.dedent) to detect kubeflow imports in component functions
Supports multiple import patterns:
- import kubeflow
- import kubeflow.<submodule>
- from kubeflow import <symbol>

3. Automatic Package Installation

Modified _get_packages_to_install_command(...) to auto-detect kubeflow usage
Automatically adds 'kubeflow' to packages_to_install when detected and not already specified
Respects explicit user-provided packages and version pins (no duplication)
Recognizes kubeflow when specified via VCS URLs
Fails closed: if source cannot be inspected or parsed, no auto-add occurs

4. Opt-out Control per Component

Extended @dsl.component with install_kubeflow_package: bool = True
- True (default): auto-add kubeflow if user code imports it
- False: never auto-add kubeflow for that component

5. Comprehensive Test Coverage

Added test cases covering:
- All supported kubeflow import patterns and negative cases
- Behavior when source inspection fails or syntax is invalid
- Package parsing (versions, extras, VCS URLs) and duplicate avoidance
- Decorator integration ensuring kubeflow is only installed when needed

Usage

Install via extras:

pip install kfp[kubeflow]
pip install kfp[all]

Default auto-detection (no user change needed):

@dsl.component
def my_comp(...):
    import kubeflow
    ...
# Kubeflow SDK is added to packages_to_install automatically

Opt-out:

@dsl.component(install_kubeflow_package=False)
def my_comp(...):
    import kubeflow
    ...
# Kubeflow SDK will not be auto-added

Related Issues

Fixes #12027

google-oss-prow · 2025-10-15T15:23:35Z

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by:
Once this PR has been reviewed and has the lgtm label, please assign chensun for approval. For more information see the Kubernetes Code Review Process.

The full list of commands accepted by this bot can be found here.

Needs approval from an approver in each of these files:

sdk/OWNERS

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

google-oss-prow · 2025-10-15T15:23:39Z

Hi @Prateekbala. Thanks for your PR.

I'm waiting for a kubeflow member to verify that this patch is reasonable to test. If it is, they should reply with /ok-to-test on its own line. Until that is done, I will not automatically test new commits in this PR, but the usual testing commands by org members will still work. Regular contributors should join the org to skip this step.

Once the patch is verified, the new status will be reflected by the ok-to-test label.

I understand the commands that are listed here.

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository.

github-actions · 2025-10-15T15:23:50Z

🚫 This command cannot be processed. Only organization members or owners can use the commands.

mprahl

/ok-to-test

mprahl · 2025-10-15T17:47:20Z

@Prateekbala could you please add a sign off to your commit?

This commit adds automatic detection and installation of kubeflow packages when they are imported in component functions. Key changes: - Added _detect_kubeflow_imports_in_function() to analyze AST for kubeflow imports - Added _parse_package_name() helper to extract clean package names - Enhanced _get_packages_to_install_command() to include kubeflow when detected - Added install_kubeflow_package parameter to @component decorator - Added kubeflow extras to setup.py for optional dependency management - Fixed type annotation syntax in component_factory_test.py - Fixed test function implementations to include actual kubeflow imports - Corrected InputPath/OutputPath type annotations with proper type parameters All tests now pass successfully. Signed-off-by: Prateek Bala <[email protected]>

…o feature/kubeflow-sdk-integrate

Prateekbala · 2025-10-16T06:49:30Z

@Prateekbala could you please add a sign off to your commit?
I’ve added the sign-off to my commit. Thank you for pointing that out.

mprahl · 2025-10-24T13:09:07Z

sdk/python/kfp/dsl/component_decorator.py

    pip_index_urls: Optional[List[str]] = None,
    output_component_file: Optional[str] = None,
    install_kfp_package: bool = True,
+    install_kubeflow_package: bool = True,


Could we make this an enum value instead? I'm thinking of the three options or similar:

auto

install

skip

Then we can default to auto but allow the user to override (force an installation) in cases where auto doesn't work. I suppose the user could always force it with adding a value to packages_to_install but this makes the behavior clearer.

mprahl · 2025-10-24T13:16:09Z

sdk/python/kfp/dsl/component_factory.py

+        if detected_kubeflow:
+            # Parse existing packages to check for kubeflow
+            existing_package_names = [
+                _parse_package_name(pkg) for pkg in packages_to_install


My preference is to use a regex to specifically detect the presence of kubeflow rather than having generic package name detection.

mprahl · 2025-10-24T13:16:53Z

sdk/python/kfp/dsl/component_factory_test.py

+    def test_detect_kubeflow_imports_simple_import(self):
+        """Test detection of 'import kubeflow'."""
+
+        def component_with_kubeflow_import():


These test components don't seem to actually import kubeflow.

- Replace Union[bool, KubeflowPackageInstallMode] with enum-only approach - Implement specific regex patterns for kubeflow package detection - Remove _parse_package_name() and add _is_kubeflow_package() function - Update test implementations with proper mocking and import statements - Export KubeflowPackageInstallMode enum in __init__.py Signed-off-by: Prateek Bala <[email protected]>

Prateekbala · 2025-11-02T20:18:29Z

I've implemented the requested changes

mprahl · 2025-11-05T15:58:50Z

Please rebase so that there aren't any merge commits.

Prateekbala · 2025-11-06T14:19:41Z

Please rebase so that there aren't any merge commits.

Done

chensun · 2025-11-13T21:03:47Z

It adds the ability to automatically detect and install the Kubeflow SDK when components use kubeflow imports, while also providing explicit extras for opt-in installation.

@Prateekbala @mprahl
I generally dislike over-smart behaviors like this. Kubeflow imports inside users' component code are no different from any other library imports, and it should be users' responsibilities specifying the dependencies they want to install.

We should stick to standard Python dependency management. Inferring dependencies from imports is problematic and rarely the right approach.

In practice, as the sdk evolves, there could be breaking changes, the version you install may not actually work with the user code, and there could be transient dependencies incompatible with user-specified dependencies--one of the reason we opted to install kfp with no dependencies.

Prateekbala · 2025-11-15T17:50:10Z

It adds the ability to automatically detect and install the Kubeflow SDK when components use kubeflow imports, while also providing explicit extras for opt-in installation.

@Prateekbala @mprahl I generally dislike over-smart behaviors like this. Kubeflow imports inside users' component code are no different from any other library imports, and it should be users' responsibilities specifying the dependencies they want to install.

Hi @chensun, I understand the concern around inferring dependencies. In the earlier discussion on #12027 and in the Kubeflow SDK design thread, @mprahl proposed using an AST-based check to detect kubeflow imports by providing a kubeflow extras option. The direction in that conversation was to implement this AST-based detection together with the extras entry as an initial solution, with the expectation that it could be revised when the broader Kubeflow SDK integration work evolves.

andreyvelich

Thanks for this great contribution @Prateekbala!
/assign @kubeflow/kubeflow-sdk-team @franciscojavierarceo @Fiona-Waters @abhijeet-dhumal it would be great if you could help with review to include Kubeflow SDK into KFP!

kramaranya

Thank you @Prateekbala!
I left a few initial comments

kramaranya · 2025-11-17T20:11:16Z

sdk/python/kfp/dsl/component_decorator.py

+class KubeflowPackageInstallMode(str, enum.Enum):
+    """Installation mode for the Kubeflow SDK package in components.
+
+    Attributes:
+        AUTO: Automatically detect and install if kubeflow is imported in the component function.
+        INSTALL: Always install the kubeflow package regardless of usage.
+        SKIP: Never install the kubeflow package, even if detected in the component function.
+    """
+    AUTO = 'auto'
+    INSTALL = 'install'
+    SKIP = 'skip'
+
+


Why do we duplicate this definition from https://github.com/Prateekbala/pipelines/blob/feature/kubeflow-sdk-integrate/sdk/python/kfp/dsl/component_factory.py#L52-L62? We can just import it

kramaranya · 2025-11-17T20:12:05Z

sdk/python/kfp/dsl/component_decorator.py

+    install_kubeflow_package:
+    KubeflowPackageInstallMode = KubeflowPackageInstallMode.AUTO,


Suggested change

install_kubeflow_package:

KubeflowPackageInstallMode = KubeflowPackageInstallMode.AUTO,

install_kubeflow_package: KubeflowPackageInstallMode = KubeflowPackageInstallMode.AUTO,

kramaranya · 2025-11-17T20:25:25Z

sdk/python/kfp/dsl/component_factory.py

    packages_to_install = packages_to_install or []
+
+    # Handle kubeflow installation based on the mode
+    if install_kubeflow_package == KubeflowPackageInstallMode.INSTALL:


This won't work if user uses KubeflowPackageInstallMode from other definition, we should keep enum in one place only so it's always the same type

kramaranya · 2025-11-17T20:28:11Z

sdk/python/kfp/dsl/component_factory.py

+
+    Detects these import patterns:
+        - import kubeflow
+        - import kubeflow.training (any submodule)


There is no training submodule in sdk https://github.com/kubeflow/sdk/tree/main/kubeflow/trainer

Suggested change

- import kubeflow.training (any submodule)

- import kubeflow.trainer (any submodule)

kramaranya · 2025-11-17T20:45:11Z

sdk/python/kfp/dsl/component_decorator.py

+        install_kubeflow_package: Controls whether the Kubeflow SDK is installed.
+            Can be KubeflowPackageInstallMode.AUTO (default), KubeflowPackageInstallMode.INSTALL,
+            or KubeflowPackageInstallMode.SKIP.
+            - AUTO: Detects kubeflow imports in the component function via AST parsing
+              and automatically adds 'kubeflow' to packages_to_install if detected.
+            - INSTALL: Always installs the kubeflow package regardless of whether it's used.
+            - SKIP: Never installs kubeflow, even if detected (useful if pre-installed in base_image).


install_kubeflow_package reads like a boolean rather than a mode.
Would kubeflow_sdk_install_mode be clearer? cc @kubeflow/kubeflow-sdk-team @kubeflow/wg-pipeline-leads

Shall we just follow the same signature as install_kfp_package ?
If the default value is AUTO, why would user change it to INSTALL ?

kramaranya · 2025-11-17T20:59:47Z

I'm wondering if we should revisit this given @chensun's concerns? Any other concerns from @kubeflow/wg-pipeline-leads about this approach or are we happy to move forward?

Fiona-Waters

I think that there could be some complications with this implementation around version clashes etc. If we do go ahead we would need clear user documentation around the install mode.

mprahl · 2025-11-18T21:19:25Z

It adds the ability to automatically detect and install the Kubeflow SDK when components use kubeflow imports, while also providing explicit extras for opt-in installation.

@Prateekbala @mprahl I generally dislike over-smart behaviors like this. Kubeflow imports inside users' component code are no different from any other library imports, and it should be users' responsibilities specifying the dependencies they want to install.

We should stick to standard Python dependency management. Inferring dependencies from imports is problematic and rarely the right approach.

In practice, as the sdk evolves, there could be breaking changes, the version you install may not actually work with the user code, and there could be transient dependencies incompatible with user-specified dependencies--one of the reason we opted to install kfp with no dependencies.

@chensun Thanks for the review. You raise valid points regarding the risks of "magic" behavior and dependency management. I definitely agree that we want to avoid inferring dependencies for general libraries.

However, I believe the Kubeflow SDK warrants a special exception separate from standard third-party packages. Our goal is to encourage seamless usage of the wider Kubeflow ecosystem within components. Conceptually, this somewhat aligns with how we already pre-install the kfp SDK into the container at runtime; we are simply extending that convenience to the broader Kubeflow namespace.

To address your concerns about version drift and user control, I propose the following safeguards for this PR:

Respecting Opt-Outs: If the user explicitly disables kfp package installation in the component decorator, we should also disable the automatic detection/installation of the Kubeflow SDK.
Strict Version Pinning: We could pin the Kubeflow SDK version to an x.y version in setup of the KFP SDK. This ensures the injected SDK is always aligned with the runtime environment, mitigating the risk of breaking changes or incompatibility.
Best-Effort Detection: We could keep the autodetect logic for kubeflow imports, but it will only trigger the installation of the pinned, compatible version defined above.
CI Verification: Let's add a pipeline run in CI to explicitly test this autodetection and installation path, ensuring the pinned version installs successfully without conflict.
Visibility: We should log in the executor that Kubeflow SDK usage was detected and it will be automatically installed. We can also provide instructions to disable it.

Prateekbala · 2025-11-19T12:57:22Z

To address your concerns about version drift and user control, I propose the following safeguards for this PR:

Respecting Opt-Outs: If the user explicitly disables kfp package installation in the component decorator, we should also disable the automatic detection/installation of the Kubeflow SDK.

Strict Version Pinning: We could pin the Kubeflow SDK version to an x.y version in setup of the KFP SDK. This ensures the injected SDK is always aligned with the runtime environment, mitigating the risk of breaking changes or incompatibility.

Best-Effort Detection: We could keep the autodetect logic for kubeflow imports, but it will only trigger the installation of the pinned, compatible version defined above.

CI Verification: Let's add a pipeline run in CI to explicitly test this autodetection and installation path, ensuring the pinned version installs successfully without conflict.

Visibility: We should log in the executor that Kubeflow SDK usage was detected and it will be automatically installed. We can also provide instructions to disable it.

Hey @mprahl , I've addressed the review feedback and added some safeguards based on the concerns. A few things I need clarification on:

Kubeflow SDK version - What version should we actually pin to is it 1.9.0 ? Should it match a specific KFP version or use a range like >=1.9.0,<2.0?
Logging - Where should we log the auto-detection? At compile time when building the spec, or at runtime in the executor?
Default mode - Currently AUTO is the default. Given the concerns , should we start with SKIP as default instead and let users opt-in?

mprahl · 2025-11-19T14:42:49Z

To address your concerns about version drift and user control, I propose the following safeguards for this PR:

Respecting Opt-Outs: If the user explicitly disables kfp package installation in the component decorator, we should also disable the automatic detection/installation of the Kubeflow SDK.

Strict Version Pinning: We could pin the Kubeflow SDK version to an x.y version in setup of the KFP SDK. This ensures the injected SDK is always aligned with the runtime environment, mitigating the risk of breaking changes or incompatibility.

Best-Effort Detection: We could keep the autodetect logic for kubeflow imports, but it will only trigger the installation of the pinned, compatible version defined above.

CI Verification: Let's add a pipeline run in CI to explicitly test this autodetection and installation path, ensuring the pinned version installs successfully without conflict.

Visibility: We should log in the executor that Kubeflow SDK usage was detected and it will be automatically installed. We can also provide instructions to disable it.

Hey @mprahl , I've addressed the review feedback and added some safeguards based on the concerns. A few things I need clarification on:

Kubeflow SDK version - What version should we actually pin to is it 1.9.0 ? Should it match a specific KFP version or use a range like >=1.9.0,<2.0?

I think we can keep the range >=0.2.0,<0.3.0 for now and bump it every KFP release. @andreyvelich @kramaranya do you agree with that?

Logging - Where should we log the auto-detection? At compile time when building the spec, or at runtime in the executor?

I was thinking the executor so the user can have debug information about which Kubeflow SDK was used at runtime.

Default mode - Currently AUTO is the default. Given the concerns , should we start with SKIP as default instead and let users opt-in?

I think if the install_kfp_package argument is True default to auto. Otherwise let the user explicitly specify the behavior.

kramaranya · 2025-11-19T16:12:15Z

I think we can keep the range >=0.2.0,<0.3.0 for now and bump it every KFP release. @andreyvelich @kramaranya do you agree with that?

I think we can just keep >=0.2.0 and it should install the latest version, so we don't need to bump it every time

andreyvelich · 2025-11-19T18:19:11Z

I think we can just keep >=0.2.0 and it should install the latest version, so we don't need to bump it every time

@kramaranya If we are planning to introduce breaking changes between minor releases, it might be better to stitch with <0.3 for now to not break users' pipelines after upgrade.

mprahl · 2025-11-19T20:16:08Z

Thank you very much for your efforts @Prateekbala!

We discussed this PR during today's KFP community call. The maintainers have decided to put a pause on this effort due to a versioning concern:

The Concern: The KFP SDK version used to compile a pipeline does not necessarily match the active KFP deployment.
The Risk: If we pin the Kubeflow SDK at the KFP SDK level, we risk installing a Kubeflow SDK that is incompatible with the actual Kubeflow deployment.

Alternative Proposal

Instead, we believe this should be handled as an API Server concern, not a KFP SDK concern. The system should ideally resolve the associated Kubeflow SDK version (using major.minor/x.y alignment) based on the deployed KFP version. This ensures we always use the SDKs compatible with what is actually installed.

Since this architectural change has significant repercussions, we need a member of the KFP community to propose a KEP (Kubeflow Enhancement Proposal) to define this path forward.

Prateekbala · 2025-11-19T20:35:15Z

Thank you very much for your efforts @Prateekbala!

We discussed this PR during today's KFP community call. The maintainers have decided to put a pause on this effort due to a versioning concern:

Thanks for the update @mprahl! I’ll pause the PR for now. Happy to help and contribute more.

google-oss-prow bot requested review from VaniHaripriya and droctothorpe October 15, 2025 15:23

google-oss-prow bot added the needs-ok-to-test label Oct 15, 2025

google-oss-prow bot added the size/L label Oct 15, 2025

mprahl reviewed Oct 15, 2025

View reviewed changes

google-oss-prow bot added ok-to-test and removed needs-ok-to-test labels Oct 15, 2025

Prateekbala force-pushed the feature/kubeflow-sdk-integrate branch from 8f7cf52 to 8c10e7b Compare October 15, 2025 20:40

Prateekbala added 2 commits October 16, 2025 02:31

Merge branch 'master' of https://github.com/Prateekbala/pipelines int…

f5c34d1

…o feature/kubeflow-sdk-integrate

Prateekbala force-pushed the feature/kubeflow-sdk-integrate branch from 8c10e7b to f5c34d1 Compare October 15, 2025 21:02

mprahl reviewed Oct 24, 2025

View reviewed changes

andreyvelich mentioned this pull request Oct 27, 2025

Integrate the Kubeflow Pipelines with Kubeflow SDK kubeflow/sdk#125

Open

Prateekbala added 2 commits November 2, 2025 19:28

Merge branch 'kubeflow:master' into feature/kubeflow-sdk-integrate

afc6674

google-oss-prow bot added size/XL and removed size/L labels Nov 2, 2025

Merge branch 'master' into feature/kubeflow-sdk-integrate

2eb0f2d

Prateekbala added 2 commits November 8, 2025 19:15

Merge branch 'master' into feature/kubeflow-sdk-integrate

7996f9c

Merge branch 'master' into feature/kubeflow-sdk-integrate

fefffba

Prateekbala added 2 commits November 14, 2025 12:30

Merge branch 'master' into feature/kubeflow-sdk-integrate

6382c72

Merge branch 'master' into feature/kubeflow-sdk-integrate

2b4ffc0

andreyvelich reviewed Nov 17, 2025

View reviewed changes

kramaranya reviewed Nov 17, 2025

View reviewed changes

Fiona-Waters reviewed Nov 18, 2025

View reviewed changes

Merge branch 'kubeflow:master' into feature/kubeflow-sdk-integrate

acc8fad

		install_kubeflow_package:
		KubeflowPackageInstallMode = KubeflowPackageInstallMode.AUTO,

	install_kubeflow_package:
	KubeflowPackageInstallMode = KubeflowPackageInstallMode.AUTO,
	install_kubeflow_package: KubeflowPackageInstallMode = KubeflowPackageInstallMode.AUTO,

	- import kubeflow.training (any submodule)
	- import kubeflow.trainer (any submodule)

feat(sdk): auto-detect and install Kubeflow SDK for component functions #12027 #12368

Are you sure you want to change the base?

feat(sdk): auto-detect and install Kubeflow SDK for component functions #12027 #12368

Uh oh!

Conversation

Prateekbala commented Oct 15, 2025

Description

Changes Made

1. Added Kubeflow Extras to setup.py

2. Implemented AST-based Auto-detection

3. Automatic Package Installation

4. Opt-out Control per Component

5. Comprehensive Test Coverage

Usage

Install via extras:

Default auto-detection (no user change needed):

Opt-out:

Related Issues

Uh oh!

google-oss-prow bot commented Oct 15, 2025

Uh oh!

google-oss-prow bot commented Oct 15, 2025

Uh oh!

github-actions bot commented Oct 15, 2025

Uh oh!

mprahl left a comment

Choose a reason for hiding this comment

Uh oh!

mprahl commented Oct 15, 2025

Uh oh!

Prateekbala commented Oct 16, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Prateekbala commented Nov 2, 2025

Uh oh!

mprahl commented Nov 5, 2025

Uh oh!

Prateekbala commented Nov 6, 2025

Uh oh!

chensun commented Nov 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Prateekbala commented Nov 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

andreyvelich left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kramaranya left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kramaranya commented Nov 17, 2025

Uh oh!

Fiona-Waters left a comment

Choose a reason for hiding this comment

Uh oh!

mprahl commented Nov 18, 2025

Uh oh!

Prateekbala commented Nov 19, 2025

Uh oh!

mprahl commented Nov 19, 2025

Uh oh!

chensun commented Nov 13, 2025 •

edited

Loading

Prateekbala commented Nov 15, 2025 •

edited

Loading

andreyvelich left a comment •

edited

Loading