feat: Support multiple pip index URLs in CustomTrainer #74

wassimbensalem · 2025-08-28T20:59:54Z

What this PR does

Implements support for multiple pip index URLs in CustomTrainer as requested in #72.

Changes Made

Types: Changed pip_index_url to pip_index_urls: list[str] in CustomTrainer
Constants: Added DEFAULT_PIP_INDEX_URLS with environment variable support
Utils: Implemented logic for --index-url and --extra-index-url in pip install commands
Tests: Updated test cases to use new field names
Backward Compatibility: Maintained DEFAULT_PIP_INDEX_URL for existing code

Usage Examples

# Single URL (backward compatible)
trainer = CustomTrainer(
    func=training_function,
    pip_index_urls=["https://pypi.org/simple"]
)

# Multiple URLs
trainer = CustomTrainer(
    func=training_function,
    pip_index_urls=[
        "https://pypi.org/simple",
        "https://private.repo.com/simple"
    ]
)

Environment Variable Support

export DEFAULT_PIP_INDEX_URLS="https://pypi.org/simple,https://private.repo.com/simple"

Technical Details

Uses field(default_factory=...) to prevent mutable default issues
First URL becomes --index-url, subsequent URLs become --extra-index-url
Automatically includes PyPI if not in the provided URLs
Follows the same pattern as KFP for consistency

Closes #72

- Change pip_index_url to pip_index_urls: list[str] in CustomTrainer - Add DEFAULT_PIP_INDEX_URLS constant with environment variable support - Implement --index-url and --extra-index-url logic in pip install commands - Maintain backward compatibility with DEFAULT_PIP_INDEX_URL - Update tests to use new field names - Fix mutable default issue using default_factory Closes kubeflow#72

google-oss-prow · 2025-08-28T20:59:59Z

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

OWNERS

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

andreyvelich

Thank you for this contribution @wassimbensalem!

andreyvelich · 2025-08-28T23:12:28Z

kubeflow/trainer/types/types.py

        packages_to_install (`Optional[List[str]]`):
            A list of Python packages to install before running the function.
-        pip_index_url (`Optional[str]`): The PyPI URL from which to install Python packages.
+        pip_index_urls (`Optional[list[str]]`): The PyPI URLs from which to install Python packages.


Please add the comment that the first URL will be used as the default index, and the rest will be used as additional indexes.

andreyvelich · 2025-08-28T23:12:56Z

kubeflow/trainer/types/types.py

    func_args: Optional[Dict] = None
    packages_to_install: Optional[list[str]] = None
-    pip_index_url: str = constants.DEFAULT_PIP_INDEX_URL
+    pip_index_urls: list[str] = field(default_factory=lambda: constants.DEFAULT_PIP_INDEX_URLS)


Why do you want to use field here ?

If we use simply the constants.DEFAULT_PIP_INDEX_URLS, in any instance that we create if someone modifies constants.DEFAULT_PIP_INDEX_ULRLS it will affect all created instance. Sothe field with default factory will assure that we have a fresh copy of it each time. This is optional, I can remove it as well !

I see, we probably need to do this to return new copy all the time object is created ?

pip_index_urls: list[str] = field(default_factory=lambda: list(constants.DEFAULT_PIP_INDEX_URLS))

andreyvelich · 2025-08-28T23:15:16Z

kubeflow/trainer/utils/utils.py

-    is_mpi: bool,
+    pip_index_urls: list[str] = constants.DEFAULT_PIP_INDEX_URLS,
+    is_mpi: bool = False,
+    include_pypi: bool = True


Why do you want to introduce this flag ?
I would let user explicitly sets the PyPI index in the list if they need it:

pip_index_urls=["pypi.custom.com/simple", "https://pypi.org/simple"]

I Just was thinking if the user forgot to set the default pypi. For instance he chose pip_index_urls = ["repoA", "repoB"], but the package wasnt there, so we have a fallback solution which is the default pypi. This is also optional and I can remove it :)

We can have it in the future, if we get any complains from users.
I would just define single constant that can be overridden by env:

DEFAULT_PIP_INDEX_URLS = os.getenv( "DEFAULT_PIP_INDEX_URLS", "https://pypi.org/simple" ).split(",")

Agree with @andreyvelich

andreyvelich · 2025-08-28T23:18:48Z

kubeflow/trainer/utils/utils.py

+        --no-warn-script-location {} {}
        """.format(
-            pip_index_url,
+            options_str,


You can directly set it here.

Suggested change

options_str,

" ".join(options),

astefanutti · 2025-08-29T08:44:52Z

@wassimbensalem could you please sign the commits?

astefanutti · 2025-08-29T08:45:04Z

/ok-to-test

kramaranya

Thank you @wassimbensalem!
Could you also sign your commits and fix the code formatting -- https://github.com/kubeflow/sdk/blob/main/CONTRIBUTING.md#development-workflow?

kramaranya · 2025-08-29T14:24:55Z

kubeflow/trainer/constants/constants.py

+DEFAULT_PIP_INDEX_URLS = (
+    os.getenv("DEFAULT_PIP_INDEX_URLS", DEFAULT_PYPI_URL).split(",") 
+    if os.getenv("DEFAULT_PIP_INDEX_URLS") 
+    else [DEFAULT_PYPI_URL]
+)


We could simplify this to not call getenv twice

kramaranya · 2025-08-29T14:27:37Z

kubeflow/trainer/utils/utils.py

-    is_mpi: bool,
+    pip_index_urls: list[str] = constants.DEFAULT_PIP_INDEX_URLS,
+    is_mpi: bool = False,
+    include_pypi: bool = True


Agree with @andreyvelich

- Simplify DEFAULT_PIP_INDEX_URLS logic - Remove include_pypi parameter and automatic PyPI fallback - Clean up function signatures

google-oss-prow bot requested review from Electronic-Waste, andreyvelich and astefanutti August 28, 2025 21:00

google-oss-prow bot added the size/M label Aug 28, 2025

andreyvelich reviewed Aug 28, 2025

View reviewed changes

wassimbensalem added 2 commits August 29, 2025 09:46

fix: join list inline

ce4333d

docs: add comment for the PIP_INDEX_URLS

6e1bd92

google-oss-prow bot added the ok-to-test label Aug 29, 2025

kramaranya reviewed Aug 29, 2025

View reviewed changes

google-oss-prow bot added size/L and removed size/M labels Aug 30, 2025

refactor: Simplify constants and remove unused parameters

872285b

- Simplify DEFAULT_PIP_INDEX_URLS logic - Remove include_pypi parameter and automatic PyPI fallback - Clean up function signatures

wassimbensalem force-pushed the feature/support-multiple-pip-index-urls branch from 0f794c2 to 872285b Compare August 30, 2025 13:35

wassimbensalem closed this Aug 30, 2025

kramaranya mentioned this pull request Aug 30, 2025

feat: Support multiple pip index URLs in CustomTrainer #79

Merged

Uh oh!

feat: Support multiple pip index URLs in CustomTrainer #74

feat: Support multiple pip index URLs in CustomTrainer #74

Uh oh!

Conversation

wassimbensalem commented Aug 28, 2025

What this PR does

Changes Made

Usage Examples

Environment Variable Support

Technical Details

Uh oh!

google-oss-prow bot commented Aug 28, 2025

Uh oh!

andreyvelich 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!

andreyvelich Aug 29, 2025 • 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!

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!

astefanutti commented Aug 29, 2025

Uh oh!

astefanutti commented Aug 29, 2025

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

andreyvelich Aug 29, 2025 •

edited

Loading