Skip to content

Conversation

Prateekbala
Copy link

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
Copy link

[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:

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
Copy link

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
Copy link

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

Copy link
Collaborator

@mprahl mprahl left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

/ok-to-test

@mprahl
Copy link
Collaborator

mprahl commented Oct 15, 2025

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

@Prateekbala Prateekbala force-pushed the feature/kubeflow-sdk-integrate branch from 8f7cf52 to 8c10e7b Compare October 15, 2025 20:40
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]>
@Prateekbala Prateekbala force-pushed the feature/kubeflow-sdk-integrate branch from 8c10e7b to f5c34d1 Compare October 15, 2025 21:02
@Prateekbala
Copy link
Author

@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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[feature] Include Kubeflow SDK in KFP SDK

2 participants