feat(sdk): Support multiple pip index URLs. (#2566)

* feat: add extra index urls to the katib tune

Signed-off-by: Andrey Velichkevich <[email protected]>

* Remove test with None index

Signed-off-by: Andrey Velichkevich <[email protected]>

---------

Signed-off-by: Andrey Velichkevich <[email protected]>
Co-authored-by: Andrey Velichkevich <[email protected]>

Author: wassimbensalem

sdk/python/v1beta1/kubeflow/katib/api/katib_client.py

@@ -215,7 +215,7 @@ def tune(
         ] = None,
         retain_trials: bool = False,
         packages_to_install: List[str] = None,
-        pip_index_url: str = "https://pypi.org/simple",
+        pip_index_urls: Optional[List[str]] = ["https://pypi.org/simple"],
         metrics_collector_config: Dict[str, Any] = {"kind": "StdOut"},
         trial_active_deadline_seconds: Optional[int] = None,
     ):
@@ -353,7 +353,7 @@ class name in this argument.
             packages_to_install: List of Python packages to install in addition
                 to the base image packages. These packages are installed before
                 executing the objective function.
-            pip_index_url: The PyPI url from which to install Python packages.
+            pip_index_urls: List of PyPI urls from which to install Python packages.
             metrics_collector_config: Specify the config of metrics collector,
                 for example, `metrics_collector_config = {"kind": "Push"}`.
                 Currently, we only support `StdOut` and `Push` metrics collector.
@@ -469,7 +469,7 @@ class name in this argument.
                 entrypoint,
                 input_params,
                 packages_to_install,
-                pip_index_url,
+                pip_index_urls,
             )
 
             # Generate container spec for PyTorchJob or Job.

sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py

@@ -569,6 +569,20 @@ def create_experiment(
         },
         TEST_RESULT_SUCCESS,
     ),
+    (
+        "valid flow with pip_index_urls",
+        {
+            "name": "tune_test",
+            "objective": lambda x: print(f"a={x}"),
+            "parameters": {"a": katib.search.int(min=10, max=100)},
+            "packages_to_install": ["pandas", "numpy"],
+            "pip_index_urls": [
+                "https://pypi.org/simple",
+                "https://private-repo.com/simple",
+            ],
+        },
+        TEST_RESULT_SUCCESS,
+    ),
 ]
 
 
@@ -682,32 +696,39 @@ def read_namespaced_pod_log_response(*args, **kwargs):
 
 @pytest.fixture
 def katib_client():
-    with patch(
-        "kubernetes.client.CustomObjectsApi",
-        return_value=Mock(
-            create_namespaced_custom_object=Mock(
-                side_effect=create_namespaced_custom_object_response
-            ),
-            get_namespaced_custom_object=Mock(
-                side_effect=get_namespaced_custom_object_response
+    with (
+        patch(
+            "kubernetes.client.CustomObjectsApi",
+            return_value=Mock(
+                create_namespaced_custom_object=Mock(
+                    side_effect=create_namespaced_custom_object_response
+                ),
+                get_namespaced_custom_object=Mock(
+                    side_effect=get_namespaced_custom_object_response
+                ),
             ),
         ),
-    ), patch("kubernetes.config.load_kube_config", return_value=Mock()), patch(
-        "kubeflow.katib.katib_api_pb2_grpc.DBManagerStub",
-        return_value=Mock(
-            GetObservationLog=Mock(side_effect=get_observation_log_response)
-        ),
-    ), patch(
-        "kubernetes.client.CoreV1Api",
-        return_value=Mock(
-            create_namespaced_persistent_volume_claim=Mock(
-                side_effect=create_namespaced_persistent_volume_claim_response
+        patch("kubernetes.config.load_kube_config", return_value=Mock()),
+        patch(
+            "kubeflow.katib.katib_api_pb2_grpc.DBManagerStub",
+            return_value=Mock(
+                GetObservationLog=Mock(side_effect=get_observation_log_response)
             ),
-            list_namespaced_persistent_volume_claim=Mock(
-                side_effect=list_namespaced_persistent_volume_claim_response
+        ),
+        patch(
+            "kubernetes.client.CoreV1Api",
+            return_value=Mock(
+                create_namespaced_persistent_volume_claim=Mock(
+                    side_effect=create_namespaced_persistent_volume_claim_response
+                ),
+                list_namespaced_persistent_volume_claim=Mock(
+                    side_effect=list_namespaced_persistent_volume_claim_response
+                ),
+                list_namespaced_pod=Mock(side_effect=list_namespaced_pod_response),
+                read_namespaced_pod_log=Mock(
+                    side_effect=read_namespaced_pod_log_response
+                ),
             ),
-            list_namespaced_pod=Mock(side_effect=list_namespaced_pod_response),
-            read_namespaced_pod_log=Mock(side_effect=read_namespaced_pod_log_response),
         ),
     ):
         client = KatibClient()
@@ -860,6 +881,18 @@ def test_tune(katib_client, test_name, kwargs, expected_output):
                         additional_metric_names=[],
                     )
 
+                elif test_name == "valid flow with pip_index_urls":
+                    # Verify pip install command in container args.
+                    args_content = "".join(
+                        experiment.spec.trial_template.trial_spec.spec.template.spec.containers[
+                            0
+                        ].args
+                    )
+                    assert (
+                        "--index-url https://pypi.org/simple --extra-index-url "
+                        "https://private-repo.com/simple pandas numpy" in args_content
+                    )
+
         except Exception as e:
             assert type(e) is expected_output
         print("test execution complete")

sdk/python/v1beta1/kubeflow/katib/utils/utils.py

@@ -115,7 +115,18 @@ def validate_objective_function(objective: Callable):
         )
 
 
-def get_script_for_python_packages(packages_to_install, pip_index_url):
+def format_pip_index_urls(
+    pip_index_urls: List[str] = ["https://pypi.org/simple"],
+) -> str:
+    index_url = f"--index-url {pip_index_urls[0]}"
+    for url in pip_index_urls[1:]:
+        index_url += f" --extra-index-url {url}"
+    return index_url
+
+
+def get_script_for_python_packages(
+    packages_to_install, pip_index_urls=["https://pypi.org/simple"]
+):
     packages_str = " ".join([str(package) for package in packages_to_install])
 
     script_for_python_packages = textwrap.dedent(
@@ -125,7 +136,7 @@ def get_script_for_python_packages(packages_to_install, pip_index_url):
         fi
 
         PIP_DISABLE_PIP_VERSION_CHECK=1 python3 -m pip install --prefer-binary --quiet \
-        --no-warn-script-location --index-url {pip_index_url} {packages_str}
+        --no-warn-script-location {format_pip_index_urls(pip_index_urls)} {packages_str}
         """
     )
 
@@ -233,7 +244,7 @@ def get_exec_script_from_objective(
     entrypoint: str,
     input_params: Dict[str, Any],
     packages_to_install: Optional[List[str]] = None,
-    pip_index_url: str = "https://pypi.org/simple",
+    pip_index_urls: Optional[List[str]] = ["https://pypi.org/simple"],
 ) -> str:
     """
     Get executable script for container args from the given objective function and parameters.
@@ -277,7 +288,7 @@ def get_exec_script_from_objective(
     # Install Python packages if that is required.
     if packages_to_install is not None:
         exec_script = (
-            get_script_for_python_packages(packages_to_install, pip_index_url)
+            get_script_for_python_packages(packages_to_install, pip_index_urls)
             + exec_script
         )
 

sdk/python/v1beta1/kubeflow/katib/utils/utils_test.py

@@ -0,0 +1,25 @@
+import pytest
+from kubeflow.katib.utils import utils
+
+
+@pytest.mark.parametrize(
+    "pip_index_urls, expected",
+    [
+        (["https://pypi.org/simple"], "--index-url https://pypi.org/simple"),
+        (
+            ["https://pypi.org/simple", "https://private-repo.com/simple"],
+            "--index-url https://pypi.org/simple --extra-index-url https://private-repo.com/simple",
+        ),
+        (
+            [
+                "https://pypi.org/simple",
+                "https://private-repo.com/simple",
+                "https://another-repo.com/simple",
+            ],
+            "--index-url https://pypi.org/simple --extra-index-url https://private-repo.com/simple "
+            "--extra-index-url https://another-repo.com/simple",
+        ),
+    ],
+)
+def test_format_pip_index_urls(pip_index_urls, expected):
+    assert utils.format_pip_index_urls(pip_index_urls) == expected