feat: Support multiple pip index URLs in CustomTrainer #79

wassimbensalem · 2025-08-30T13:58:49Z

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

wassimbensalem · 2025-08-30T14:02:31Z

@andreyvelich @kramaranya Sorry I had to create a new PR because I had some issues with signing the commits ( I am new to this please bare with me :) ). I hope this PR includes all the feedback. If anything is missing let me know,

wassimbensalem · 2025-08-30T14:05:23Z

/assign

kramaranya · 2025-08-30T14:24:51Z

/ok-to-test

kramaranya

Thank you @wassimbensalem!
Could you please sign your commits? You can follow this guide https://github.com/kubeflow/community/tree/master/dco-signoff-hook#signing-off-commits

kramaranya · 2025-08-30T15:56:39Z

kubeflow/trainer/constants/constants.py

+# Handle environment variable for multiple URLs (comma-separated)
+DEFAULT_PIP_INDEX_URLS = os.getenv("DEFAULT_PIP_INDEX_URLS", DEFAULT_PYPI_URL).split(",")
+# Keep backward compatibility
+DEFAULT_PIP_INDEX_URL = DEFAULT_PYPI_URL


Why do we need this, since it's not exported anywhere?

kubeflow/trainer/types/types.py

kramaranya · 2025-08-30T23:39:48Z

@wassimbensalem you'll also need to update unit tests

sdk/kubeflow/trainer/backends/kubernetes/backend_test.py

Lines 228 to 239 in 06d77f9

    
           return models.TrainerV1alpha1Trainer( 
        
               command=[ 
        
                   "bash", 
        
                   "-c", 
        
                   '\nif ! [ -x "$(command -v pip)" ]; then\n    python -m ensurepip ' 
        
                   "|| python -m ensurepip --user || apt-get install python-pip" 
        
                   "\nfi\n\nPIP_DISABLE_PIP_VERSION_CHECK=1 python -m pip install --quiet" 
        
                   "         --no-warn-script-location --index-url https://pypi.org/simple " 
        
                   "torch numpy \n\nread -r -d '' SCRIPT << EOM\n\nfunc=lambda: " 
        
                   'print("Hello World"),\n\n<lambda>(' 
        
                   "{'learning_rate': 0.001, 'batch_size': 32})\n\nEOM\nprintf \"%s\" " 
        
                   '"$SCRIPT" > "backend_test.py"\ntorchrun "backend_test.py"',

And it would be great to add a new test case for multiple pip index urls!

wassimbensalem · 2025-09-01T18:48:43Z

@kramaranya Is there anything still missing?

andreyvelich · 2025-09-01T18:55:49Z

@wassimbensalem Please can you rebase your PR ?

- Add pip_index_urls parameter to CustomTrainer (list of URLs) - Update constants to support multiple URLs with environment variable - Modify utils to handle multiple index URLs in pip installation - Update tests to use new parameter - Fix mutable default issue with dataclasses.field(default_factory=...) Closes kubeflow#72 Signed-off-by: wassimbensalem <bswassim@gmail.com>

Signed-off-by: wassimbensalem <bswassim@gmail.com>

…ation Signed-off-by: wassimbensalem <bswassim@gmail.com>

Signed-off-by: wassimbensalem <bswassim@gmail.com>

andreyvelich · 2025-09-01T21:52:38Z

kubeflow/trainer/backends/kubernetes/backend_test.py

+    if packages_to_install is None:
+        packages_to_install = ["torch", "numpy"]


Let's just explicitly set the packages from the input args here.

Suggested change

if packages_to_install is None:

packages_to_install = ["torch", "numpy"]

andreyvelich · 2025-09-01T22:08:44Z

kubeflow/trainer/backends/kubernetes/backend_test.py

+
+def test_get_script_for_python_packages_multiple_urls():
+    """Test get_script_for_python_packages with multiple pip index URLs."""
+    from kubeflow.trainer.utils import utils
+
+    # Test with multiple pip index URLs
+    packages = ["torch", "numpy", "custom-package"]
+    pip_index_urls = [
+        "https://pypi.org/simple",
+        "https://private.repo.com/simple",
+        "https://internal.company.com/simple"
+    ]
+
+    script = utils.get_script_for_python_packages(
+        packages_to_install=packages,
+        pip_index_urls=pip_index_urls,


Can you move those test cases to the utils_test.py, and refactor them similar to other tests, e.g.

@pytest.mark.parametrize( "test_case", [ TestCase( name="....", config={....}, expected_output={....} ), ], )

andreyvelich · 2025-09-01T22:10:19Z

kubeflow/trainer/backends/kubernetes/backend_test.py

            expected_error=ValueError,
        ),
+        TestCase(
+            name="valid flow with multiple pip index URLs",


You can remove this test case, since we verify it in the utils_test.py.

@wassimbensalem Did you address this comment ?

andreyvelich · 2025-09-01T22:11:21Z

kubeflow/trainer/constants/constants.py


-# The default PIP index URL to download Python packages.
-DEFAULT_PIP_INDEX_URL = os.getenv("DEFAULT_PIP_INDEX_URL", "https://pypi.org/simple")
+# Handle environment variable for multiple URLs (comma-separated)


Suggested change

# Handle environment variable for multiple URLs (comma-separated)

# Handle environment variable for multiple URLs (comma-separated).

# The first URL will be the index-url, and remaining ones are extra-index-urls.

andreyvelich · 2025-09-01T22:12:49Z

kubeflow/trainer/utils/utils.py

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


Let's not set the defaults here, since it should be controlled by the public APIs.

Suggested change

pip_index_urls: list[str] = constants.DEFAULT_PIP_INDEX_URLS,

is_mpi: bool = False,

pip_index_urls: list[str],

is_mpi: bool,

andreyvelich · 2025-09-01T22:13:10Z

kubeflow/trainer/utils/utils.py

 import textwrap
 import threading
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Dict, Optional


Suggested change

from typing import Any, Callable, Dict, Optional

from typing import Any, Callable, Optional

andreyvelich · 2025-09-01T22:13:19Z

kubeflow/trainer/utils/utils.py

    train_func: Callable,
-    train_func_parameters: Optional[dict[str, Any]],
-    pip_index_url: str,
+    train_func_parameters: Optional[Dict[str, Any]],


Suggested change

train_func_parameters: Optional[Dict[str, Any]],

train_func_parameters: Optional[dict[str, Any]],

andreyvelich · 2025-09-01T22:13:37Z

kubeflow/trainer/utils/utils.py

-    train_func_parameters: Optional[dict[str, Any]],
-    pip_index_url: str,
+    train_func_parameters: Optional[Dict[str, Any]],
+    pip_index_urls: list[str] = constants.DEFAULT_PIP_INDEX_URLS,


Same comment.

Suggested change

pip_index_urls: list[str] = constants.DEFAULT_PIP_INDEX_URLS,

pip_index_urls: list[str],

kubeflow/trainer/utils/utils.py

coveralls · 2025-09-01T22:17:42Z

Pull Request Test Coverage Report for Build 17472441226

Warning: This coverage report may be inaccurate.

This pull request's base commit is no longer the HEAD commit of its target branch. This means it includes changes from outside the original pull request, including, potentially, unrelated coverage changes.

For more information on this, see Tracking coverage changes with pull request builds.
To avoid this issue with future PRs, see these Recommended CI Configurations.
For a quick fix, rebase this PR at GitHub. Your next report should be accurate.

Details

4 of 5 (80.0%) changed or added relevant lines in 1 file are covered.
73 unchanged lines in 2 files lost coverage.
Overall coverage increased (+5.2%) to 70.192%

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
kubeflow/trainer/utils/utils.py	4	5	80.0%

Files with Coverage Reduction	New Missed Lines	%
kubeflow/trainer/backends/kubernetes/backend.py	17	78.1%
kubeflow/trainer/utils/utils.py	56	62.14%

Totals
Change from base Build 17382047367:	5.2%
Covered Lines:	292
Relevant Lines:	416

💛 - Coveralls

kramaranya · 2025-09-01T22:28:22Z

/milestone v0.1

…d format Signed-off-by: wassimbensalem <bswassim@gmail.com>

andreyvelich · 2025-09-02T14:49:49Z

kubeflow/trainer/backends/kubernetes/backend_test.py

            expected_error=ValueError,
        ),
+        TestCase(
+            name="valid flow with multiple pip index URLs",


@wassimbensalem Did you address this comment ?

andreyvelich · 2025-09-02T14:50:23Z

kubeflow/trainer/backends/kubernetes/backend_test.py

+    pip_index_urls: Optional[list[str]] = None,
+    packages_to_install: list[str] = ["torch", "numpy"],
 ) -> models.TrainerV1alpha1Trainer:
    """
    Get the custom trainer for the TrainJob.
    """
+    if pip_index_urls is None:
+        pip_index_urls = constants.DEFAULT_PIP_INDEX_URLS


Suggested change

pip_index_urls: Optional[list[str]] = None,

packages_to_install: list[str] = ["torch", "numpy"],

) -> models.TrainerV1alpha1Trainer:

"""

Get the custom trainer for the TrainJob.

"""

if pip_index_urls is None:

pip_index_urls = constants.DEFAULT_PIP_INDEX_URLS

pip_index_urls: Optional[list[str]] = constants.DEFAULT_PIP_INDEX_URLS,

packages_to_install: list[str] = ["torch", "numpy"],

) -> models.TrainerV1alpha1Trainer:

"""

Get the custom trainer for the TrainJob.

"""

andreyvelich · 2025-09-02T14:55:15Z

kubeflow/trainer/utils/utils_test.py

+    for expected_text in test_case.expected_output["contains"]:
+        assert expected_text in script, (
+            f"Expected '{expected_text}' to be in script"
+        )
+
+    # Verify the script does not contain unexpected elements
+    for unexpected_text in test_case.expected_output["not_contains"]:
+        assert unexpected_text not in script, (
+            f"Expected '{unexpected_text}' to NOT be in script"
+        )


Instead of having contains and not_contains check, can you assert the entire script that get_script_for_python_packages returns?
e.g.

assert test_case.expected_output == script

@andreyvelich All comments should be addressed now. Please let me know if anything is missing.

…ontains/not_contains checks - Changed TestCase expected_output from Dict[str, Any] to str - Replaced contains/not_contains logic with direct string equality - Updated all test cases to specify complete expected script output - Tests now assert the entire script returned by get_script_for_python_packages Signed-off-by: wassimbensalem <bswassim@gmail.com>

- Add support for multiple pip index URLs in CustomTrainer - Update constants to handle comma-separated environment variables - Add integration test for multiple pip index URLs flow - Update utils to generate correct pip install commands with --index-url and --extra-index-url - Ensure backward compatibility with single URL usage Signed-off-by: wassimbensalem <bswassim@gmail.com>

- Remove redundant test case since comprehensive tests already exist in utils test - Keep integration tests focused on core functionality rather than implementation details - Maintain test coverage while reducing duplication Signed-off-by: wassimbensalem <bswassim@gmail.com>

…backend" This reverts commit 0ec003d. Signed-off-by: wassimbensalem <bswassim@gmail.com>

Signed-off-by: wassimbensalem <bswassim@gmail.com>

andreyvelich

Thank you for this great contribution @wassimbensalem!
/lgtm
/assign @kubeflow/kubeflow-sdk-team

astefanutti · 2025-09-05T07:36:46Z

kubeflow/trainer/backends/kubernetes/backend_test.py

-                train_job_trainer=get_custom_trainer(),
+                train_job_trainer=get_custom_trainer(
+                    pip_index_urls=constants.DEFAULT_PIP_INDEX_URLS,
+                    packages_to_install=["torch", "numpy"],


Small nit: could be other packages since those are already in the runtime image.

astefanutti · 2025-09-05T07:37:38Z

/lgtm

Thanks @wassimbensalem!

kramaranya

Thank you for this contribution @wassimbensalem!
/lgtm

astefanutti · 2025-09-05T08:25:36Z

/approve

Great work!

google-oss-prow · 2025-09-05T08:25:44Z

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: astefanutti

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

The pull request process is described here

Needs approval from an approver in each of these files:

~~OWNERS~~ [astefanutti]

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

* feat: Support multiple pip index URLs in CustomTrainer - Add pip_index_urls parameter to CustomTrainer (list of URLs) - Update constants to support multiple URLs with environment variable - Modify utils to handle multiple index URLs in pip installation - Update tests to use new parameter - Fix mutable default issue with dataclasses.field(default_factory=...) Closes kubeflow#72 Signed-off-by: wassimbensalem <bswassim@gmail.com> * test: add comprehensive test cases for multiple pip index URLs Signed-off-by: wassimbensalem <bswassim@gmail.com> * refactor: simplify constants for pip index URLs Signed-off-by: wassimbensalem <bswassim@gmail.com> * fix: remove Optional from pip_index_urls docstring to match implementation Signed-off-by: wassimbensalem <bswassim@gmail.com> * test: adjust unit test for get_custom_trainer Signed-off-by: wassimbensalem <bswassim@gmail.com> * fix: add missing Dict import in utils.py Signed-off-by: wassimbensalem <bswassim@gmail.com> * refactor: move test cases to utils_test.py and convert to parametrized format Signed-off-by: wassimbensalem <bswassim@gmail.com> * refactor: update utils test to use exact string matching instead of contains/not_contains checks - Changed TestCase expected_output from Dict[str, Any] to str - Replaced contains/not_contains logic with direct string equality - Updated all test cases to specify complete expected script output - Tests now assert the entire script returned by get_script_for_python_packages Signed-off-by: wassimbensalem <bswassim@gmail.com> * feat: add multiple pip index URLs support with comprehensive tests - Add support for multiple pip index URLs in CustomTrainer - Update constants to handle comma-separated environment variables - Add integration test for multiple pip index URLs flow - Update utils to generate correct pip install commands with --index-url and --extra-index-url - Ensure backward compatibility with single URL usage Signed-off-by: wassimbensalem <bswassim@gmail.com> * refactor: remove redundant multiple pip index URLs test from backend - Remove redundant test case since comprehensive tests already exist in utils test - Keep integration tests focused on core functionality rather than implementation details - Maintain test coverage while reducing duplication Signed-off-by: wassimbensalem <bswassim@gmail.com> * Revert "refactor: remove redundant multiple pip index URLs test from backend" This reverts commit 0ec003d. Signed-off-by: wassimbensalem <bswassim@gmail.com> * refactor: remove duplicated test case from the backend. Signed-off-by: wassimbensalem <bswassim@gmail.com> --------- Signed-off-by: wassimbensalem <bswassim@gmail.com>

google-oss-prow bot requested review from Electronic-Waste, andreyvelich and astefanutti August 30, 2025 13:58

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

google-oss-prow bot assigned wassimbensalem Aug 30, 2025

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

kramaranya reviewed Aug 30, 2025

View reviewed changes

kubeflow/trainer/types/types.py Outdated Show resolved Hide resolved

wassimbensalem force-pushed the feature/support-multiple-pip-index-urls-clean branch 2 times, most recently from 4e08634 to 00daf6c Compare August 31, 2025 12:42

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

wassimbensalem added 6 commits September 1, 2025 21:58

test: add comprehensive test cases for multiple pip index URLs

2ffc235

Signed-off-by: wassimbensalem <bswassim@gmail.com>

refactor: simplify constants for pip index URLs

f8f1e78

Signed-off-by: wassimbensalem <bswassim@gmail.com>

fix: remove Optional from pip_index_urls docstring to match implement…

75084bb

…ation Signed-off-by: wassimbensalem <bswassim@gmail.com>

test: adjust unit test for get_custom_trainer

c4f6152

Signed-off-by: wassimbensalem <bswassim@gmail.com>

fix: add missing Dict import in utils.py

5eb42b1

Signed-off-by: wassimbensalem <bswassim@gmail.com>

wassimbensalem force-pushed the feature/support-multiple-pip-index-urls-clean branch from 4aab345 to 5eb42b1 Compare September 1, 2025 20:05

andreyvelich reviewed Sep 1, 2025

View reviewed changes

google-oss-prow bot added this to the v0.1 milestone Sep 1, 2025

refactor: move test cases to utils_test.py and convert to parametrize…

f87e9c4

…d format Signed-off-by: wassimbensalem <bswassim@gmail.com>

andreyvelich reviewed Sep 2, 2025

View reviewed changes

andreyvelich mentioned this pull request Sep 3, 2025

feat(trainer): Add support for param unpacking in the training function call #62

Merged

wassimbensalem added 5 commits September 4, 2025 19:45

Revert "refactor: remove redundant multiple pip index URLs test from …

39ca82d

…backend" This reverts commit 0ec003d. Signed-off-by: wassimbensalem <bswassim@gmail.com>

refactor: remove duplicated test case from the backend.

c084980

Signed-off-by: wassimbensalem <bswassim@gmail.com>

andreyvelich reviewed Sep 4, 2025

View reviewed changes

google-oss-prow bot assigned andreyvelich Sep 4, 2025

google-oss-prow bot added the lgtm label Sep 4, 2025

astefanutti reviewed Sep 5, 2025

View reviewed changes

google-oss-prow bot assigned astefanutti Sep 5, 2025

kramaranya reviewed Sep 5, 2025

View reviewed changes

google-oss-prow bot assigned kramaranya Sep 5, 2025

google-oss-prow bot added the approved label Sep 5, 2025

google-oss-prow bot merged commit 427b35d into kubeflow:main Sep 5, 2025
10 checks passed

		if packages_to_install is None:
		packages_to_install = ["torch", "numpy"]

	# Handle environment variable for multiple URLs (comma-separated)
	# Handle environment variable for multiple URLs (comma-separated).
	# The first URL will be the index-url, and remaining ones are extra-index-urls.

		pip_index_urls: list[str] = constants.DEFAULT_PIP_INDEX_URLS,
		is_mpi: bool = False,

	from typing import Any, Callable, Dict, Optional
	from typing import Any, Callable, Optional

	train_func_parameters: Optional[Dict[str, Any]],
	train_func_parameters: Optional[dict[str, Any]],

	pip_index_urls: list[str] = constants.DEFAULT_PIP_INDEX_URLS,
	pip_index_urls: list[str],

Uh oh!

feat: Support multiple pip index URLs in CustomTrainer #79

feat: Support multiple pip index URLs in CustomTrainer #79

Uh oh!

Conversation

wassimbensalem commented Aug 30, 2025

What this PR does

Changes Made

Usage Examples

Single URL (backward compatible)

Multiple URLs

Environment Variable Support

Technical Details

Uh oh!

wassimbensalem commented Aug 30, 2025

Uh oh!

wassimbensalem commented Aug 30, 2025

Uh oh!

kramaranya commented Aug 30, 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!

Uh oh!

kramaranya commented Aug 30, 2025

Uh oh!

wassimbensalem commented Sep 1, 2025

Uh oh!

andreyvelich commented Sep 1, 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!

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!

Uh oh!

coveralls commented Sep 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pull Request Test Coverage Report for Build 17472441226

Warning: This coverage report may be inaccurate.

Details

💛 - Coveralls

Uh oh!

kramaranya commented Sep 1, 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!

Choose a reason for hiding this comment

Uh oh!

andreyvelich left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

astefanutti commented Sep 5, 2025

Uh oh!

kramaranya left a comment

Choose a reason for hiding this comment

Uh oh!

astefanutti commented Sep 5, 2025

Uh oh!

coveralls commented Sep 1, 2025 •

edited

Loading