added test for create_experiment in katib_client

Signed-off-by: tariq-hasan <mmtariquehsn@gmail.com>

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.

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}"

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")