From 703eb2bf53922afafaf2717f2ef02b0b32c240c2 Mon Sep 17 00:00:00 2001 From: tariq-hasan Date: Sun, 5 May 2024 08:46:16 -0400 Subject: [PATCH] added test for create_experiment in katib_client Signed-off-by: tariq-hasan --- Makefile | 3 + .../kubeflow/katib/api/katib_client.py | 5 +- .../kubeflow/katib/api/katib_client_test.py | 286 ++++++++++++++++++ 3 files changed, 290 insertions(+), 4 deletions(-) create mode 100644 sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py diff --git a/Makefile b/Makefile index 061df399f8c..ebc42a7ddfa 100755 --- a/Makefile +++ b/Makefile @@ -178,6 +178,9 @@ pytest: prepare-pytest prepare-pytest-testdata PYTHONPATH=$(PYTHONPATH) pytest ./test/unit/v1beta1/suggestion --ignore=./test/unit/v1beta1/suggestion/test_skopt_service.py PYTHONPATH=$(PYTHONPATH) pytest ./test/unit/v1beta1/earlystopping PYTHONPATH=$(PYTHONPATH) pytest ./test/unit/v1beta1/metricscollector + cp ./pkg/apis/manager/v1beta1/python/api_pb2.py ./sdk/python/v1beta1/kubeflow/katib/katib_api_pb2.py + PYTHONPATH=$(PYTHONPATH) pytest ./sdk/python/v1beta1/kubeflow/katib + rm ./sdk/python/v1beta1/kubeflow/katib/katib_api_pb2.py # The skopt service doesn't work appropriately with Python 3.11. # So, we need to run the test with Python 3.9. diff --git a/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py b/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py index 351eda8fb7e..966cc5ac5af 100644 --- a/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py +++ b/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py @@ -120,16 +120,13 @@ def create_experiment( raise ValueError("Experiment must have a name or generateName") try: - outputs = self.custom_api.create_namespaced_custom_object( + self.custom_api.create_namespaced_custom_object( constants.KUBEFLOW_GROUP, constants.KATIB_VERSION, namespace, constants.EXPERIMENT_PLURAL, experiment, ) - experiment_name = outputs["metadata"][ - "name" - ] # if "generate_name" is used, "name" gets a prefix from server except multiprocessing.TimeoutError: raise TimeoutError( f"Timeout to create Katib Experiment: {namespace}/{experiment_name}" diff --git a/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py b/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py new file mode 100644 index 00000000000..dd9fc9bf68b --- /dev/null +++ b/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py @@ -0,0 +1,286 @@ +import multiprocessing +from typing import List, Optional +from unittest.mock import patch, Mock + +import pytest +from kubernetes.client import V1ObjectMeta + +from kubeflow.katib import KatibClient +from kubeflow.katib import V1beta1AlgorithmSpec +from kubeflow.katib import V1beta1Experiment +from kubeflow.katib import V1beta1ExperimentSpec +from kubeflow.katib import V1beta1FeasibleSpace +from kubeflow.katib import V1beta1ObjectiveSpec +from kubeflow.katib import V1beta1ParameterSpec +from kubeflow.katib import V1beta1TrialParameterSpec +from kubeflow.katib import V1beta1TrialTemplate +from kubeflow.katib.constants import constants + + +class ConflictException(Exception): + def __init__(self): + self.status = 409 + + +def create_namespaced_custom_object_response(*args, **kwargs): + if args[2] == "timeout": + raise multiprocessing.TimeoutError() + elif args[2] == "conflict": + raise ConflictException() + elif args[2] == "runtime": + raise Exception() + + +def generate_trial_template() -> V1beta1TrialTemplate: + trial_spec={ + "apiVersion": "batch/v1", + "kind": "Job", + "spec": { + "template": { + "metadata": { + "annotations": { + "sidecar.istio.io/inject": "false" + } + }, + "spec": { + "containers": [ + { + "name": "training-container", + "image": "docker.io/kubeflowkatib/pytorch-mnist-cpu:v0.14.0", + "command": [ + "python3", + "/opt/pytorch-mnist/mnist.py", + "--epochs=1", + "--batch-size=64", + "--lr=${trialParameters.learningRate}", + "--momentum=${trialParameters.momentum}", + ] + } + ], + "restartPolicy": "Never" + } + } + } + } + + return V1beta1TrialTemplate( + primary_container_name="training-container", + trial_parameters=[ + V1beta1TrialParameterSpec( + name="learningRate", + description="Learning rate for the training model", + reference="lr" + ), + V1beta1TrialParameterSpec( + name="momentum", + description="Momentum for the training model", + reference="momentum" + ), + ], + trial_spec=trial_spec + ) + + +def generate_experiment( + metadata: V1ObjectMeta, + algorithm_spec: V1beta1AlgorithmSpec, + objective_spec: V1beta1ObjectiveSpec, + parameters: List[V1beta1ParameterSpec], + trial_template: V1beta1TrialTemplate, +) -> V1beta1Experiment: + return V1beta1Experiment( + api_version=constants.API_VERSION, + kind=constants.EXPERIMENT_KIND, + metadata=metadata, + spec=V1beta1ExperimentSpec( + max_trial_count=3, + parallel_trial_count=2, + max_failed_trial_count=1, + algorithm=algorithm_spec, + objective=objective_spec, + parameters=parameters, + trial_template=trial_template, + ) + ) + + +def create_experiment( + name: Optional[str] = None, + generate_name: Optional[str] = None +) -> V1beta1Experiment: + experiment_namespace = "test" + + if name is not None: + metadata = V1ObjectMeta(name=name, namespace=experiment_namespace) + elif generate_name is not None: + metadata = V1ObjectMeta(generate_name=generate_name, namespace=experiment_namespace) + else: + metadata = V1ObjectMeta(namespace=experiment_namespace) + + algorithm_spec=V1beta1AlgorithmSpec( + algorithm_name="random" + ) + + objective_spec=V1beta1ObjectiveSpec( + type="minimize", + goal= 0.001, + objective_metric_name="loss", + ) + + parameters=[ + V1beta1ParameterSpec( + name="lr", + parameter_type="double", + feasible_space=V1beta1FeasibleSpace( + min="0.01", + max="0.06" + ), + ), + V1beta1ParameterSpec( + name="momentum", + parameter_type="double", + feasible_space=V1beta1FeasibleSpace( + min="0.5", + max="0.9" + ), + ), + ] + + trial_template = generate_trial_template() + + experiment = generate_experiment( + metadata, + algorithm_spec, + objective_spec, + parameters, + trial_template + ) + return experiment + + +test_create_experiment_data = [ + ( + "experiment name and generate_name missing", + {"experiment": create_experiment()}, + ValueError, + ), + ( + "create_namespaced_custom_object timeout error", + { + "experiment": create_experiment(name="experiment-mnist-ci-test"), + "namespace": "timeout", + }, + TimeoutError, + ), + ( + "create_namespaced_custom_object conflict error", + { + "experiment": create_experiment(name="experiment-mnist-ci-test"), + "namespace": "conflict", + }, + Exception, + ), + ( + "create_namespaced_custom_object runtime error", + { + "experiment": create_experiment(name="experiment-mnist-ci-test"), + "namespace": "runtime", + }, + RuntimeError, + ), + ( + "valid flow with experiment type V1beta1Experiment and name", + { + "experiment": create_experiment(name="experiment-mnist-ci-test"), + "namespace": "test", + }, + "success", + ), + ( + "valid flow with experiment type V1beta1Experiment and generate_name", + { + "experiment": create_experiment(generate_name="experiment-mnist-ci-test"), + "namespace": "test", + }, + "success", + ), + ( + "valid flow with experiment type V1beta1Experiment and name and generate_name", + { + "experiment": create_experiment( + name="experiment-mnist-ci-test", + generate_name="experiment-mnist-ci-test", + ), + "namespace": "test" + }, + "success", + ), + ( + "valid flow with experiment JSON and name", + { + "experiment": { + "metadata": { + "name": "experiment-mnist-ci-test", + "namespace": "test", + } + } + }, + "success", + ), + ( + "valid flow with experiment JSON and generate_name", + { + "experiment": { + "metadata": { + "generate_name": "experiment-mnist-ci-test", + "namespace": "test", + } + } + }, + "success", + ), + ( + "valid flow with experiment JSON and name and generate_name", + { + "experiment": { + "metadata": { + "name": "experiment-mnist-ci-test", + "generate_name": "experiment-mnist-ci-test", + "namespace": "test", + } + } + }, + "success", + ), +] + + +@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 + ) + ), + ), patch( + "kubernetes.config.load_kube_config", + return_value=Mock() + ): + client = KatibClient() + yield client + + +@pytest.mark.parametrize("test_name,kwargs,expected_output", test_create_experiment_data) +def test_create_experiment(katib_client, test_name, kwargs, expected_output): + """ + test create_experiment function of katib client + """ + print("Executing test:", test_name) + try: + katib_client.create_experiment(**kwargs) + assert expected_output == "success" + except Exception as e: + assert type(e) is expected_output + print("test execution complete")