From bf2673e3a52524a8c20524a2031c9bdc7b9ec859 Mon Sep 17 00:00:00 2001
From: Electronic-Waste <2690692950@qq.com>
Date: Fri, 26 Jul 2024 13:43:06 +0000
Subject: [PATCH 01/15] fix: fix bugs in report_metrics.

Signed-off-by: Electronic-Waste <2690692950@qq.com>
---
 .../kubeflow/katib/api/report_metrics.py      | 50 +++++++++----------
 1 file changed, 24 insertions(+), 26 deletions(-)

diff --git a/sdk/python/v1beta1/kubeflow/katib/api/report_metrics.py b/sdk/python/v1beta1/kubeflow/katib/api/report_metrics.py
index 97aadd236ff..d5726bacd99 100644
--- a/sdk/python/v1beta1/kubeflow/katib/api/report_metrics.py
+++ b/sdk/python/v1beta1/kubeflow/katib/api/report_metrics.py
@@ -20,6 +20,7 @@
 import grpc
 from kubeflow.katib.constants import constants
 import kubeflow.katib.katib_api_pb2 as katib_api_pb2
+import kubeflow.katib.katib_api_pb2_grpc as katib_api_pb2_grpc
 from kubeflow.katib.utils import utils
 
 
@@ -53,35 +54,32 @@ def report_metrics(
         )
 
     # Get channel for grpc call to db manager
-    db_manager_address = db_manager_address.split(":")
-    channel = grpc.beta.implementations.insecure_channel(
-        db_manager_address[0], int(db_manager_address[1])
-    )
+    channel = grpc.insecure_channel(db_manager_address)
 
     # Validate metrics value in dict
     for value in metrics.values():
         utils.validate_metrics_value(value)
     
     # Dial katib db manager to report metrics
-    with katib_api_pb2.beta_create_DBManager_stub(channel) as client:
-        try:
-            timestamp = datetime.now(timezone.utc).strftime(constants.RFC3339_FORMAT)
-            client.ReportObservationLog(
-                request=katib_api_pb2.ReportObservationLogRequest(
-                    trial_name=name,
-                    observation_logs=katib_api_pb2.ObservationLog(
-                        metric_logs=[
-                            katib_api_pb2.MetricLog(
-                                time_stamp=timestamp,
-                                metric=katib_api_pb2.Metric(name=name,value=str(value))   
-                            )
-                            for name, value in metrics.items()
-                        ]
-                    )
-                ),
-                timeout=timeout,
-            )
-        except Exception as e:
-            raise RuntimeError(
-                f"Unable to push metrics to Katib DB for Trial {namespace}/{name}. Exception: {e}"
-            )
+    client = katib_api_pb2_grpc.DBManagerStub(channel)
+    try:
+        timestamp = datetime.now(timezone.utc).strftime(constants.RFC3339_FORMAT)
+        client.ReportObservationLog(
+            request=katib_api_pb2.ReportObservationLogRequest(
+                trial_name=name,
+                observation_log=katib_api_pb2.ObservationLog(
+                    metric_logs=[
+                        katib_api_pb2.MetricLog(
+                            time_stamp=timestamp,
+                            metric=katib_api_pb2.Metric(name=name,value=str(value))   
+                        )
+                        for name, value in metrics.items()
+                    ]
+                )
+            ),
+            timeout=timeout,
+        )
+    except Exception as e:
+        raise RuntimeError(
+            f"Unable to push metrics to Katib DB for Trial {namespace}/{name}. Exception: {e}"
+        )

From 69e3d6272f0977f1d6b94b7c319fc0c64c01d312 Mon Sep 17 00:00:00 2001
From: Electronic-Waste <2690692950@qq.com>
Date: Fri, 26 Jul 2024 13:44:08 +0000
Subject: [PATCH 02/15] fix: fix bugs in tune.

Signed-off-by: Electronic-Waste <2690692950@qq.com>
---
 sdk/python/v1beta1/kubeflow/katib/api/katib_client.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py b/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py
index b18a81cad81..ceb3be2ce77 100644
--- a/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py
+++ b/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py
@@ -386,7 +386,7 @@ def tune(
 
         # Add metrics collector to the Katib Experiment.
         # Up to now, We only support parameter `kind`, of which default value is `StdOut`, to specify the kind of metrics collector. 
-        experiment.spec.metrics_collector = models.V1beta1MetricsCollectorSpec(
+        experiment.spec.metrics_collector_spec = models.V1beta1MetricsCollectorSpec(
             collector=models.V1beta1CollectorSpec(kind=metrics_collector_config["kind"])
         )
 

From a3d2e9f4c6b4751637cec27ece004e76cb1feaa4 Mon Sep 17 00:00:00 2001
From: Electronic-Waste <2690692950@qq.com>
Date: Fri, 26 Jul 2024 13:47:58 +0000
Subject: [PATCH 03/15] fix: fix bugs in get_trial_metrics.

Signed-off-by: Electronic-Waste <2690692950@qq.com>
---
 .../kubeflow/katib/api/katib_client.py        | 30 +++++++++----------
 1 file changed, 14 insertions(+), 16 deletions(-)

diff --git a/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py b/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py
index ceb3be2ce77..b7c64fb6080 100644
--- a/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py
+++ b/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py
@@ -24,6 +24,7 @@
 from kubeflow.katib.api_client import ApiClient
 from kubeflow.katib.constants import constants
 import kubeflow.katib.katib_api_pb2 as katib_api_pb2
+import kubeflow.katib.katib_api_pb2_grpc as katib_api_pb2_grpc
 from kubeflow.katib.utils import utils
 from kubernetes import client
 from kubernetes import config
@@ -1280,21 +1281,18 @@ def get_trial_metrics(
 
         namespace = namespace or self.namespace
 
-        db_manager_address = db_manager_address.split(":")
-        channel = grpc.beta.implementations.insecure_channel(
-            db_manager_address[0], int(db_manager_address[1])
-        )
+        channel = grpc.insecure_channel(db_manager_address)
 
-        with katib_api_pb2.beta_create_DBManager_stub(channel) as client:
-            try:
-                # When metric name is empty, we select all logs from the Katib DB.
-                observation_logs = client.GetObservationLog(
-                    katib_api_pb2.GetObservationLogRequest(trial_name=name),
-                    timeout=timeout,
-                )
-            except Exception as e:
-                raise RuntimeError(
-                    f"Unable to get metrics for Trial {namespace}/{name}. Exception: {e}"
-                )
+        client = katib_api_pb2_grpc.DBManagerStub(channel)
+        try:
+            # When metric name is empty, we select all logs from the Katib DB.
+            observation_logs = client.GetObservationLog(
+                katib_api_pb2.GetObservationLogRequest(trial_name=name),
+                timeout=timeout,
+            )
+        except Exception as e:
+            raise RuntimeError(
+                f"Unable to get metrics for Trial {namespace}/{name}. Exception: {e}"
+            )
 
-            return observation_logs.observation_log.metric_logs
+        return observation_logs.observation_log.metric_logs

From 22fe26ada091d321e1da927472e1b5ac568e4fb8 Mon Sep 17 00:00:00 2001
From: Electronic-Waste <2690692950@qq.com>
Date: Fri, 26 Jul 2024 13:50:55 +0000
Subject: [PATCH 04/15] fix: update .gitignore and setup.py.

Signed-off-by: Electronic-Waste <2690692950@qq.com>
---
 sdk/python/v1beta1/.gitignore |  1 +
 sdk/python/v1beta1/setup.py   | 14 ++++++++++++++
 2 files changed, 15 insertions(+)

diff --git a/sdk/python/v1beta1/.gitignore b/sdk/python/v1beta1/.gitignore
index 81f90cfca9b..64aa80409b8 100644
--- a/sdk/python/v1beta1/.gitignore
+++ b/sdk/python/v1beta1/.gitignore
@@ -3,3 +3,4 @@ dist/
 
 # Katib gRPC APIs
 kubeflow/katib/katib_api_pb2.py
+kubeflow/katib/katib_api_pb2_grpc.py
diff --git a/sdk/python/v1beta1/setup.py b/sdk/python/v1beta1/setup.py
index 6b9c152f2d2..c329c4a4c95 100644
--- a/sdk/python/v1beta1/setup.py
+++ b/sdk/python/v1beta1/setup.py
@@ -28,6 +28,7 @@
 ]
 
 katib_grpc_api_file = "../../../pkg/apis/manager/v1beta1/python/api_pb2.py"
+katib_grpc_svc_file = "../../../pkg/apis/manager/v1beta1/python/api_pb2_grpc.py"
 
 # Copy Katib gRPC Python APIs to use it in the Katib SDK Client.
 # We need to always copy this file only on the SDK building stage, not on SDK installation stage.
@@ -37,6 +38,19 @@
         "kubeflow/katib/katib_api_pb2.py",
     )
 
+if os.path.exists(katib_grpc_svc_file):
+    shutil.copy(
+        katib_grpc_svc_file,
+        "kubeflow/katib/katib_api_pb2_grpc.py",
+    )
+
+    with open("kubeflow/katib/katib_api_pb2_grpc.py", 'r+') as file:
+        content = file.read()
+        new_content = content.replace("api_pb2", "kubeflow.katib.katib_api_pb2", 1)
+        file.seek(0)
+        file.write(new_content)
+        file.truncate()
+
 setuptools.setup(
     name="kubeflow-katib",
     version="0.17.0",

From 570584c8a94f6ce5b96a10c54250faf26d1549f4 Mon Sep 17 00:00:00 2001
From: Electronic-Waste <2690692950@qq.com>
Date: Fri, 26 Jul 2024 14:37:43 +0000
Subject: [PATCH 05/15] fix: update Makefile.

Signed-off-by: Electronic-Waste <2690692950@qq.com>
---
 Makefile                    | 4 +++-
 sdk/python/v1beta1/setup.py | 2 +-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/Makefile b/Makefile
index 3edfc9ece13..80073202e4c 100755
--- a/Makefile
+++ b/Makefile
@@ -171,8 +171,10 @@ pytest: prepare-pytest prepare-pytest-testdata
 	pytest ./test/unit/v1beta1/earlystopping
 	pytest ./test/unit/v1beta1/metricscollector
 	cp ./pkg/apis/manager/v1beta1/python/api_pb2.py ./sdk/python/v1beta1/kubeflow/katib/katib_api_pb2.py
+	cp ./pkg/apis/manager/v1beta1/python/api_pb2_grpc.py ./sdk/python/v1beta1/kubeflow/katib/katib_api_pb2_grpc.py
+	sed -i "s/api_pb2/kubeflow\.katib\.katib_api_pb2/g" ./sdk/python/v1beta1/kubeflow/katib/katib_api_pb2_grpc.py
 	pytest ./sdk/python/v1beta1/kubeflow/katib
-	rm ./sdk/python/v1beta1/kubeflow/katib/katib_api_pb2.py
+	rm ./sdk/python/v1beta1/kubeflow/katib/katib_api_pb2.py ./sdk/python/v1beta1/kubeflow/katib/katib_api_pb2_grpc.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/setup.py b/sdk/python/v1beta1/setup.py
index c329c4a4c95..32538052baf 100644
--- a/sdk/python/v1beta1/setup.py
+++ b/sdk/python/v1beta1/setup.py
@@ -46,7 +46,7 @@
 
     with open("kubeflow/katib/katib_api_pb2_grpc.py", 'r+') as file:
         content = file.read()
-        new_content = content.replace("api_pb2", "kubeflow.katib.katib_api_pb2", 1)
+        new_content = content.replace("api_pb2", "kubeflow.katib.katib_api_pb2")
         file.seek(0)
         file.write(new_content)
         file.truncate()

From bfbb5b7e7b9067247a6f9157a0a11152d6059389 Mon Sep 17 00:00:00 2001
From: Electronic-Waste <2690692950@qq.com>
Date: Sun, 28 Jul 2024 15:11:11 +0000
Subject: [PATCH 06/15] feat: add report_metrics_test.py.

Signed-off-by: Electronic-Waste <2690692950@qq.com>
---
 .../kubeflow/katib/api/report_metrics.py      |   4 +-
 .../kubeflow/katib/api/report_metrics_test.py | 134 ++++++++++++++++++
 2 files changed, 136 insertions(+), 2 deletions(-)
 create mode 100644 sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py

diff --git a/sdk/python/v1beta1/kubeflow/katib/api/report_metrics.py b/sdk/python/v1beta1/kubeflow/katib/api/report_metrics.py
index d5726bacd99..6250052310a 100644
--- a/sdk/python/v1beta1/kubeflow/katib/api/report_metrics.py
+++ b/sdk/python/v1beta1/kubeflow/katib/api/report_metrics.py
@@ -40,9 +40,9 @@ def report_metrics(
         timeout: Optional, gRPC API Server timeout in seconds to report metrics.
     
     Raises:
-        ValueError: The Trial name is not passed to environment variables.
-        RuntimeError: Unable to push Trial metrics to Katib DB or
+        ValueError: The Trial name is not passed to environment variables or
             metrics value has incorrect format (cannot be converted to type `float`).
+        RuntimeError: Unable to push Trial metrics to Katib DB.
     """
 
     # Get Trial's namespace and name
diff --git a/sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py b/sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py
new file mode 100644
index 00000000000..e5476091e83
--- /dev/null
+++ b/sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py
@@ -0,0 +1,134 @@
+from unittest.mock import patch
+
+from kubeflow.katib.constants import constants
+from kubeflow.katib import report_metrics
+import pytest
+
+TEST_RESULT_SUCCESS = "success"
+ENV_VARIABLE_EMPTY = True
+ENV_VARIABLE_NOT_EMPTY = False
+
+
+def report_observation_log_response(*args, **kwargs):
+    if kwargs.get("timeout") == 0:
+        raise TimeoutError
+
+
+test_report_metrics_data = [
+    (
+        "valid metrics with float type",
+        {
+            "metrics": {
+                "result": 0.99
+            },
+            "db_manager_address": constants.DEFAULT_DB_MANAGER_ADDRESS,
+            "timeout": constants.DEFAULT_TIMEOUT
+
+        },
+        TEST_RESULT_SUCCESS,
+        ENV_VARIABLE_NOT_EMPTY
+    ),
+    (
+        "valid metrics with string type",
+        {
+            "metrics": {
+                "result": "0.99"
+            },
+            "db_manager_address": constants.DEFAULT_DB_MANAGER_ADDRESS,
+            "timeout": constants.DEFAULT_TIMEOUT
+        },
+        TEST_RESULT_SUCCESS,
+        ENV_VARIABLE_NOT_EMPTY
+    ),
+    (
+        "valid metrics with int type",
+        {
+            "metrics": {
+                "result": 1
+            },
+            "db_manager_address": constants.DEFAULT_DB_MANAGER_ADDRESS,
+            "timeout": constants.DEFAULT_TIMEOUT
+        },
+        TEST_RESULT_SUCCESS,
+        ENV_VARIABLE_NOT_EMPTY
+    ),
+    (
+        "ReportObservationLog timeout error",
+        {
+            "metrics": {
+                "result": 0.99
+            },
+            "db_manager_address": constants.DEFAULT_DB_MANAGER_ADDRESS,
+            "timeout": 0
+        },
+        RuntimeError,
+        ENV_VARIABLE_NOT_EMPTY
+    ),
+    (
+        "invalid metrics with type string",
+        {
+            "metrics": {
+                "result": "abc"
+            },
+            "db_manager_address": constants.DEFAULT_DB_MANAGER_ADDRESS,
+            "timeout": constants.DEFAULT_TIMEOUT
+        },
+        ValueError,
+        ENV_VARIABLE_NOT_EMPTY
+    ),
+    (
+        "Trial name is not passed to env variables",
+        {
+            "metrics": {
+                "result": 0.99
+            },
+            "db_manager_address": constants.DEFAULT_DB_MANAGER_ADDRESS,
+            "timeout": constants.DEFAULT_TIMEOUT
+        },
+        ValueError,
+        ENV_VARIABLE_EMPTY
+    )
+]
+
+
+@pytest.fixture
+def mock_getenv(request):
+    with patch("os.getenv") as mock:
+        if request.param is ENV_VARIABLE_EMPTY:
+            mock.side_effect = ValueError
+        else:
+            mock.return_value = "example"
+        yield mock
+
+
+@pytest.fixture
+def mock_get_current_k8s_namespace():
+    with patch("kubeflow.katib.utils.utils.get_current_k8s_namespace") as mock:
+        mock.return_value = "test"
+        yield mock
+
+
+@pytest.fixture
+def mock_report_observation_log():
+    with patch("kubeflow.katib.katib_api_pb2_grpc.DBManagerStub") as mock:
+        mock_instance = mock.return_value
+        mock_instance.ReportObservationLog.side_effect = report_observation_log_response
+        yield mock_instance
+
+
+@pytest.mark.parametrize(
+    "test_name, kwargs, expected_output, mock_getenv",
+    test_report_metrics_data,
+    indirect=["mock_getenv"]
+)
+def test_report_metrics(test_name, kwargs, expected_output, mock_getenv, mock_get_current_k8s_namespace, mock_report_observation_log):
+    """
+    test report_metrics function
+    """
+    print("\n\nExecuting test:", test_name)
+    try:
+        report_metrics(**kwargs)
+        assert expected_output == TEST_RESULT_SUCCESS
+    except Exception as e:
+        assert type(e) is expected_output
+    print("test execution complete")
\ No newline at end of file

From 6f0e9edfac08b5eb8c355b141094fadd9ebfa820 Mon Sep 17 00:00:00 2001
From: Electronic-Waste <2690692950@qq.com>
Date: Sun, 28 Jul 2024 15:14:41 +0000
Subject: [PATCH 07/15] fix: fix lint error.

Signed-off-by: Electronic-Waste <2690692950@qq.com>
---
 sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py b/sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py
index e5476091e83..a1ec51b4103 100644
--- a/sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py
+++ b/sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py
@@ -1,7 +1,7 @@
 from unittest.mock import patch
 
-from kubeflow.katib.constants import constants
 from kubeflow.katib import report_metrics
+from kubeflow.katib.constants import constants
 import pytest
 
 TEST_RESULT_SUCCESS = "success"

From 462347fa38df0821637ad274d075687d3ed8bb5d Mon Sep 17 00:00:00 2001
From: Electronic-Waste <2690692950@qq.com>
Date: Mon, 29 Jul 2024 10:28:26 +0000
Subject: [PATCH 08/15] feat: add UTs for get_trial_metrics.

Signed-off-by: Electronic-Waste <2690692950@qq.com>
---
 .../kubeflow/katib/api/katib_client_test.py   | 96 ++++++++++++++++++-
 .../kubeflow/katib/api/report_metrics_test.py | 10 +-
 2 files changed, 95 insertions(+), 11 deletions(-)

diff --git a/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py b/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py
index 7f9f5e3fdba..19b82939c7f 100644
--- a/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py
+++ b/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py
@@ -13,6 +13,7 @@
 from kubeflow.katib import V1beta1TrialParameterSpec
 from kubeflow.katib import V1beta1TrialTemplate
 from kubeflow.katib.constants import constants
+import kubeflow.katib.katib_api_pb2 as katib_api_pb2
 from kubernetes.client import V1ObjectMeta
 import pytest
 
@@ -238,7 +239,7 @@ def create_experiment(
 
 
 @pytest.fixture
-def katib_client():
+def katib_client_create_experiment():
     with patch(
         "kubernetes.client.CustomObjectsApi",
         return_value=Mock(
@@ -255,14 +256,103 @@ def katib_client():
 
 
 @pytest.mark.parametrize("test_name,kwargs,expected_output", test_create_experiment_data)
-def test_create_experiment(katib_client, test_name, kwargs, expected_output):
+def test_create_experiment(katib_client_create_experiment, test_name, kwargs, expected_output):
     """
     test create_experiment function of katib client
     """
     print("\n\nExecuting test:", test_name)
     try:
-        katib_client.create_experiment(**kwargs)
+        katib_client_create_experiment.create_experiment(**kwargs)
         assert expected_output == TEST_RESULT_SUCCESS
     except Exception as e:
         assert type(e) is expected_output
     print("test execution complete")
+
+
+def get_observation_log_response(*args, **kwargs):
+    if kwargs.get("timeout") == 0:
+        raise TimeoutError
+    elif args[0].trial_name == "invalid":
+        raise RuntimeError
+    else:
+        return katib_api_pb2.GetObservationLogReply(
+            observation_log=katib_api_pb2.ObservationLog(
+                metric_logs=[
+                    katib_api_pb2.MetricLog(
+                        time_stamp="2024-07-29T15:09:08Z",
+                        metric=katib_api_pb2.Metric(name="result",value="0.99")
+                    )
+                ]
+            )
+        )
+
+test_get_trial_metrics_data = [
+    (
+        "valid trial name",
+        {
+            "name": "example",
+            "namespace": "valid",
+            "timeout": constants.DEFAULT_TIMEOUT
+        },
+        [
+            katib_api_pb2.MetricLog(
+                time_stamp="2024-07-29T15:09:08Z",
+                metric=katib_api_pb2.Metric(name="result",value="0.99")
+            )
+        ]
+    ),
+    (
+        "invalid trial name",
+        {
+            "name": "invalid",
+            "namespace": "invalid",
+            "timeout": constants.DEFAULT_TIMEOUT
+        },
+        RuntimeError
+    ),
+    (
+        "GetObservationLog timeout error",
+        {
+            "name": "example",
+            "namespace": "valid",
+            "timeout": 0
+        },
+        RuntimeError
+    )
+]
+
+
+@pytest.fixture
+def katib_client_get_trial_metrics():
+    with patch(
+        "kubernetes.client.CustomObjectsApi",
+        return_value=Mock(),
+    ), patch(
+        "kubernetes.config.load_kube_config",
+        return_value=Mock()
+    ):
+        client = KatibClient()
+        yield client
+
+
+@pytest.fixture
+def mock_get_observation_log():
+    with patch("kubeflow.katib.katib_api_pb2_grpc.DBManagerStub") as mock:
+        mock_instance = mock.return_value
+        mock_instance.GetObservationLog.side_effect = get_observation_log_response
+        yield mock_instance
+
+
+@pytest.mark.parametrize("test_name,kwargs,expected_output", test_get_trial_metrics_data)
+def test_get_trial_metrics(test_name, kwargs, expected_output, katib_client_get_trial_metrics, mock_get_observation_log):
+    """
+    test get_trial_metrics function of katib client
+    """
+    print("\n\nExecuting test:", test_name)
+    try:
+        metrics = katib_client_get_trial_metrics.get_trial_metrics(**kwargs)
+        for i in range(len(metrics)):
+            assert metrics[i] == expected_output[i]
+    except Exception as e:
+        assert type(e) is expected_output
+    print("test execution complete")
diff --git a/sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py b/sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py
index a1ec51b4103..129afcf4577 100644
--- a/sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py
+++ b/sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py
@@ -21,7 +21,6 @@ def report_observation_log_response(*args, **kwargs):
             "metrics": {
                 "result": 0.99
             },
-            "db_manager_address": constants.DEFAULT_DB_MANAGER_ADDRESS,
             "timeout": constants.DEFAULT_TIMEOUT
 
         },
@@ -34,7 +33,6 @@ def report_observation_log_response(*args, **kwargs):
             "metrics": {
                 "result": "0.99"
             },
-            "db_manager_address": constants.DEFAULT_DB_MANAGER_ADDRESS,
             "timeout": constants.DEFAULT_TIMEOUT
         },
         TEST_RESULT_SUCCESS,
@@ -46,7 +44,6 @@ def report_observation_log_response(*args, **kwargs):
             "metrics": {
                 "result": 1
             },
-            "db_manager_address": constants.DEFAULT_DB_MANAGER_ADDRESS,
             "timeout": constants.DEFAULT_TIMEOUT
         },
         TEST_RESULT_SUCCESS,
@@ -58,7 +55,6 @@ def report_observation_log_response(*args, **kwargs):
             "metrics": {
                 "result": 0.99
             },
-            "db_manager_address": constants.DEFAULT_DB_MANAGER_ADDRESS,
             "timeout": 0
         },
         RuntimeError,
@@ -70,7 +66,6 @@ def report_observation_log_response(*args, **kwargs):
             "metrics": {
                 "result": "abc"
             },
-            "db_manager_address": constants.DEFAULT_DB_MANAGER_ADDRESS,
             "timeout": constants.DEFAULT_TIMEOUT
         },
         ValueError,
@@ -82,7 +77,6 @@ def report_observation_log_response(*args, **kwargs):
             "metrics": {
                 "result": 0.99
             },
-            "db_manager_address": constants.DEFAULT_DB_MANAGER_ADDRESS,
             "timeout": constants.DEFAULT_TIMEOUT
         },
         ValueError,
@@ -117,7 +111,7 @@ def mock_report_observation_log():
 
 
 @pytest.mark.parametrize(
-    "test_name, kwargs, expected_output, mock_getenv",
+    "test_name,kwargs,expected_output,mock_getenv",
     test_report_metrics_data,
     indirect=["mock_getenv"]
 )
@@ -131,4 +125,4 @@ def test_report_metrics(test_name, kwargs, expected_output, mock_getenv, mock_ge
         assert expected_output == TEST_RESULT_SUCCESS
     except Exception as e:
         assert type(e) is expected_output
-    print("test execution complete")
\ No newline at end of file
+    print("test execution complete")

From a7a845fd6739be4e89cc9b842cab32dc732cfd5e Mon Sep 17 00:00:00 2001
From: Electronic-Waste <2690692950@qq.com>
Date: Tue, 30 Jul 2024 06:31:23 +0000
Subject: [PATCH 09/15] fix: update post_gen.py.

Signed-off-by: Electronic-Waste <2690692950@qq.com>
---
 hack/gen-python-sdk/post_gen.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/hack/gen-python-sdk/post_gen.py b/hack/gen-python-sdk/post_gen.py
index a0d1649d205..b7bda1f251f 100644
--- a/hack/gen-python-sdk/post_gen.py
+++ b/hack/gen-python-sdk/post_gen.py
@@ -41,8 +41,8 @@ def _rewrite_helper(input_file, output_file, rewrite_rules):
     if (output_file == "sdk/python/v1beta1/kubeflow/katib/__init__.py"):
         lines.append("# Import Katib API client.\n")
         lines.append("from kubeflow.katib.api.katib_client import KatibClient\n")
-        lines.append("# Import Katib report metrics functions")
-        lines.append("from kubeflow.katib.api.report_metrics import report_metrics")
+        lines.append("# Import Katib report metrics functions\n")
+        lines.append("from kubeflow.katib.api.report_metrics import report_metrics\n")
         lines.append("# Import Katib helper functions.\n")
         lines.append("import kubeflow.katib.api.search as search\n")
         lines.append("# Import Katib helper constants.\n")

From 8b095add55d7dd8983eba3d3b4bab34d51dc2f35 Mon Sep 17 00:00:00 2001
From: Electronic-Waste <2690692950@qq.com>
Date: Tue, 30 Jul 2024 06:45:31 +0000
Subject: [PATCH 10/15] refactor: rebase to master.

Signed-off-by: Electronic-Waste <2690692950@qq.com>
---
 sdk/python/v1beta1/kubeflow/katib/__init__.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/sdk/python/v1beta1/kubeflow/katib/__init__.py b/sdk/python/v1beta1/kubeflow/katib/__init__.py
index 0ea206ea7ca..7aef4c9897d 100644
--- a/sdk/python/v1beta1/kubeflow/katib/__init__.py
+++ b/sdk/python/v1beta1/kubeflow/katib/__init__.py
@@ -71,7 +71,9 @@
 
 # Import Katib API client.
 from kubeflow.katib.api.katib_client import KatibClient
-# Import Katib report metrics functionsfrom kubeflow.katib.api.report_metrics import report_metrics# Import Katib helper functions.
+# Import Katib report metrics functions
+from kubeflow.katib.api.report_metrics import report_metrics
+# Import Katib helper functions.
 import kubeflow.katib.api.search as search
 # Import Katib helper constants.
 from kubeflow.katib.constants.constants import BASE_IMAGE_TENSORFLOW

From f6997f77d528c335bcf221008a3d6d651f39d254 Mon Sep 17 00:00:00 2001
From: Electronic-Waste <2690692950@qq.com>
Date: Fri, 9 Aug 2024 12:27:34 +0000
Subject: [PATCH 11/15] test(sdk): use single katib_client.

Signed-off-by: Electronic-Waste <2690692950@qq.com>
---
 .../kubeflow/katib/api/katib_client_test.py   | 103 ++++++++----------
 1 file changed, 45 insertions(+), 58 deletions(-)

diff --git a/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py b/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py
index 19b82939c7f..ead5ca90a9a 100644
--- a/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py
+++ b/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py
@@ -38,6 +38,24 @@ def create_namespaced_custom_object_response(*args, **kwargs):
         return {"metadata": {"name": "12345-experiment-mnist-ci-test"}}
 
 
+def get_observation_log_response(*args, **kwargs):
+    if kwargs.get("timeout") == 0:
+        raise TimeoutError
+    elif args[0].trial_name == "invalid":
+        raise RuntimeError
+    else:
+        return katib_api_pb2.GetObservationLogReply(
+            observation_log=katib_api_pb2.ObservationLog(
+                metric_logs=[
+                    katib_api_pb2.MetricLog(
+                        time_stamp="2024-07-29T15:09:08Z",
+                        metric=katib_api_pb2.Metric(name="result",value="0.99")
+                    )
+                ]
+            )
+        )
+
+
 def generate_trial_template() -> V1beta1TrialTemplate:
     trial_spec={
         "apiVersion": "batch/v1",
@@ -238,54 +256,6 @@ def create_experiment(
 ]
 
 
-@pytest.fixture
-def katib_client_create_experiment():
-    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_create_experiment, test_name, kwargs, expected_output):
-    """
-    test create_experiment function of katib client
-    """
-    print("\n\nExecuting test:", test_name)
-    try:
-        katib_client_create_experiment.create_experiment(**kwargs)
-        assert expected_output == TEST_RESULT_SUCCESS
-    except Exception as e:
-        assert type(e) is expected_output
-    print("test execution complete")
-
-
-def get_observation_log_response(*args, **kwargs):
-    if kwargs.get("timeout") == 0:
-        raise TimeoutError
-    elif args[0].trial_name == "invalid":
-        raise RuntimeError
-    else:
-        return katib_api_pb2.GetObservationLogReply(
-            observation_log=katib_api_pb2.ObservationLog(
-                metric_logs=[
-                    katib_api_pb2.MetricLog(
-                        time_stamp="2024-07-29T15:09:08Z",
-                        metric=katib_api_pb2.Metric(name="result",value="0.99")
-                    )
-                ]
-            )
-        )
-
 test_get_trial_metrics_data = [
     (
         "valid trial name",
@@ -323,34 +293,51 @@ def get_observation_log_response(*args, **kwargs):
 
 
 @pytest.fixture
-def katib_client_get_trial_metrics():
+def katib_client():
     with patch(
         "kubernetes.client.CustomObjectsApi",
-        return_value=Mock(),
+        return_value=Mock(
+            create_namespaced_custom_object=Mock(
+                side_effect=create_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
+            )
+        )
     ):
         client = KatibClient()
         yield client
 
 
-@pytest.fixture
-def mock_get_observation_log():
-    with patch("kubeflow.katib.katib_api_pb2_grpc.DBManagerStub") as mock:
-        mock_instance = mock.return_value
-        mock_instance.GetObservationLog.side_effect = get_observation_log_response
-        yield mock_instance
+@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("\n\nExecuting test:", test_name)
+    try:
+        katib_client.create_experiment(**kwargs)
+        assert expected_output == TEST_RESULT_SUCCESS
+    except Exception as e:
+        assert type(e) is expected_output
+    print("test execution complete")
 
 
 @pytest.mark.parametrize("test_name,kwargs,expected_output", test_get_trial_metrics_data)
-def test_get_trial_metrics(test_name, kwargs, expected_output, katib_client_get_trial_metrics, mock_get_observation_log):
+def test_get_trial_metrics(katib_client, test_name, kwargs, expected_output):
     """
     test get_trial_metrics function of katib client
     """
     print("\n\nExecuting test:", test_name)
     try:
-        metrics = katib_client_get_trial_metrics.get_trial_metrics(**kwargs)
+        metrics = katib_client.get_trial_metrics(**kwargs)
         for i in range(len(metrics)):
             assert metrics[i] == expected_output[i]
     except Exception as e:

From a0902cb8f79122ee20ffd211d8eb97bdc221d003 Mon Sep 17 00:00:00 2001
From: Electronic-Waste <2690692950@qq.com>
Date: Thu, 22 Aug 2024 17:27:01 +0000
Subject: [PATCH 12/15] fix(sdk): add TODO for import rewrite.

Signed-off-by: Electronic-Waste <2690692950@qq.com>
---
 Makefile                    | 2 ++
 sdk/python/v1beta1/setup.py | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/Makefile b/Makefile
index 80073202e4c..a6708de7f5b 100755
--- a/Makefile
+++ b/Makefile
@@ -166,6 +166,8 @@ ifeq ("$(wildcard $(TEST_TENSORFLOW_EVENT_FILE_PATH))", "")
 	python examples/v1beta1/trial-images/tf-mnist-with-summaries/mnist.py --epochs 5 --batch-size 200 --log-path $(TEST_TENSORFLOW_EVENT_FILE_PATH)
 endif
 
+# TODO(Electronic-Waste): Remove the import rewrite when protobuf supports `python_package` option.
+# REF: https://github.com/protocolbuffers/protobuf/issues/7061
 pytest: prepare-pytest prepare-pytest-testdata
 	pytest ./test/unit/v1beta1/suggestion --ignore=./test/unit/v1beta1/suggestion/test_skopt_service.py
 	pytest ./test/unit/v1beta1/earlystopping
diff --git a/sdk/python/v1beta1/setup.py b/sdk/python/v1beta1/setup.py
index 32538052baf..fbf6b865156 100644
--- a/sdk/python/v1beta1/setup.py
+++ b/sdk/python/v1beta1/setup.py
@@ -38,6 +38,8 @@
         "kubeflow/katib/katib_api_pb2.py",
     )
 
+# TODO(Electronic-Waste): Remove the import rewrite when protobuf supports `python_package` option.
+# REF: https://github.com/protocolbuffers/protobuf/issues/7061
 if os.path.exists(katib_grpc_svc_file):
     shutil.copy(
         katib_grpc_svc_file,

From ff5031a5ea133146c08332039e6de71b0fa7bb22 Mon Sep 17 00:00:00 2001
From: Electronic-Waste <2690692950@qq.com>
Date: Thu, 22 Aug 2024 17:41:32 +0000
Subject: [PATCH 13/15] fix(sdk): fix lint error with black.

Signed-off-by: Electronic-Waste <2690692950@qq.com>
---
 .../kubeflow/katib/api/katib_client_test.py   | 43 +++++-------
 .../kubeflow/katib/api/report_metrics.py      |  4 +-
 .../kubeflow/katib/api/report_metrics_test.py | 68 ++++++-------------
 sdk/python/v1beta1/setup.py                   |  2 +-
 4 files changed, 41 insertions(+), 76 deletions(-)

diff --git a/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py b/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py
index e7a8663af18..9f853f62a9c 100644
--- a/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py
+++ b/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py
@@ -50,7 +50,7 @@ def get_observation_log_response(*args, **kwargs):
                 metric_logs=[
                     katib_api_pb2.MetricLog(
                         time_stamp="2024-07-29T15:09:08Z",
-                        metric=katib_api_pb2.Metric(name="result",value="0.99")
+                        metric=katib_api_pb2.Metric(name="result", value="0.99"),
                     )
                 ]
             )
@@ -245,36 +245,28 @@ def create_experiment(
 test_get_trial_metrics_data = [
     (
         "valid trial name",
-        {
-            "name": "example",
-            "namespace": "valid",
-            "timeout": constants.DEFAULT_TIMEOUT
-        },
+        {"name": "example", "namespace": "valid", "timeout": constants.DEFAULT_TIMEOUT},
         [
             katib_api_pb2.MetricLog(
                 time_stamp="2024-07-29T15:09:08Z",
-                metric=katib_api_pb2.Metric(name="result",value="0.99")
+                metric=katib_api_pb2.Metric(name="result", value="0.99"),
             )
-        ]
+        ],
     ),
     (
         "invalid trial name",
         {
             "name": "invalid",
             "namespace": "invalid",
-            "timeout": constants.DEFAULT_TIMEOUT
+            "timeout": constants.DEFAULT_TIMEOUT,
         },
-        RuntimeError
+        RuntimeError,
     ),
     (
         "GetObservationLog timeout error",
-        {
-            "name": "example",
-            "namespace": "valid",
-            "timeout": 0
-        },
-        RuntimeError
-    )
+        {"name": "example", "namespace": "valid", "timeout": 0},
+        RuntimeError,
+    ),
 ]
 
 
@@ -287,16 +279,11 @@ def katib_client():
                 side_effect=create_namespaced_custom_object_response
             )
         ),
-    ), patch(
-        "kubernetes.config.load_kube_config",
-        return_value=Mock()
-    ), patch(
-        "kubeflow.katib.katib_api_pb2_grpc.DBManagerStub", 
+    ), 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
-            )
-        )
+            GetObservationLog=Mock(side_effect=get_observation_log_response)
+        ),
     ):
         client = KatibClient()
         yield client
@@ -318,7 +305,9 @@ def test_create_experiment(katib_client, test_name, kwargs, expected_output):
     print("test execution complete")
 
 
-@pytest.mark.parametrize("test_name,kwargs,expected_output", test_get_trial_metrics_data)
+@pytest.mark.parametrize(
+    "test_name,kwargs,expected_output", test_get_trial_metrics_data
+)
 def test_get_trial_metrics(katib_client, test_name, kwargs, expected_output):
     """
     test get_trial_metrics function of katib client
diff --git a/sdk/python/v1beta1/kubeflow/katib/api/report_metrics.py b/sdk/python/v1beta1/kubeflow/katib/api/report_metrics.py
index b8c513e9b78..50dd02c9c20 100644
--- a/sdk/python/v1beta1/kubeflow/katib/api/report_metrics.py
+++ b/sdk/python/v1beta1/kubeflow/katib/api/report_metrics.py
@@ -68,11 +68,11 @@ def report_metrics(
                     metric_logs=[
                         katib_api_pb2.MetricLog(
                             time_stamp=timestamp,
-                            metric=katib_api_pb2.Metric(name=name,value=str(value))   
+                            metric=katib_api_pb2.Metric(name=name, value=str(value)),
                         )
                         for name, value in metrics.items()
                     ]
-                )
+                ),
             ),
             timeout=timeout,
         )
diff --git a/sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py b/sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py
index 129afcf4577..69f13698b13 100644
--- a/sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py
+++ b/sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py
@@ -17,71 +17,40 @@ def report_observation_log_response(*args, **kwargs):
 test_report_metrics_data = [
     (
         "valid metrics with float type",
-        {
-            "metrics": {
-                "result": 0.99
-            },
-            "timeout": constants.DEFAULT_TIMEOUT
-
-        },
+        {"metrics": {"result": 0.99}, "timeout": constants.DEFAULT_TIMEOUT},
         TEST_RESULT_SUCCESS,
-        ENV_VARIABLE_NOT_EMPTY
+        ENV_VARIABLE_NOT_EMPTY,
     ),
     (
         "valid metrics with string type",
-        {
-            "metrics": {
-                "result": "0.99"
-            },
-            "timeout": constants.DEFAULT_TIMEOUT
-        },
+        {"metrics": {"result": "0.99"}, "timeout": constants.DEFAULT_TIMEOUT},
         TEST_RESULT_SUCCESS,
-        ENV_VARIABLE_NOT_EMPTY
+        ENV_VARIABLE_NOT_EMPTY,
     ),
     (
         "valid metrics with int type",
-        {
-            "metrics": {
-                "result": 1
-            },
-            "timeout": constants.DEFAULT_TIMEOUT
-        },
+        {"metrics": {"result": 1}, "timeout": constants.DEFAULT_TIMEOUT},
         TEST_RESULT_SUCCESS,
-        ENV_VARIABLE_NOT_EMPTY
+        ENV_VARIABLE_NOT_EMPTY,
     ),
     (
         "ReportObservationLog timeout error",
-        {
-            "metrics": {
-                "result": 0.99
-            },
-            "timeout": 0
-        },
+        {"metrics": {"result": 0.99}, "timeout": 0},
         RuntimeError,
-        ENV_VARIABLE_NOT_EMPTY
+        ENV_VARIABLE_NOT_EMPTY,
     ),
     (
         "invalid metrics with type string",
-        {
-            "metrics": {
-                "result": "abc"
-            },
-            "timeout": constants.DEFAULT_TIMEOUT
-        },
+        {"metrics": {"result": "abc"}, "timeout": constants.DEFAULT_TIMEOUT},
         ValueError,
-        ENV_VARIABLE_NOT_EMPTY
+        ENV_VARIABLE_NOT_EMPTY,
     ),
     (
         "Trial name is not passed to env variables",
-        {
-            "metrics": {
-                "result": 0.99
-            },
-            "timeout": constants.DEFAULT_TIMEOUT
-        },
+        {"metrics": {"result": 0.99}, "timeout": constants.DEFAULT_TIMEOUT},
         ValueError,
-        ENV_VARIABLE_EMPTY
-    )
+        ENV_VARIABLE_EMPTY,
+    ),
 ]
 
 
@@ -113,9 +82,16 @@ def mock_report_observation_log():
 @pytest.mark.parametrize(
     "test_name,kwargs,expected_output,mock_getenv",
     test_report_metrics_data,
-    indirect=["mock_getenv"]
+    indirect=["mock_getenv"],
 )
-def test_report_metrics(test_name, kwargs, expected_output, mock_getenv, mock_get_current_k8s_namespace, mock_report_observation_log):
+def test_report_metrics(
+    test_name,
+    kwargs,
+    expected_output,
+    mock_getenv,
+    mock_get_current_k8s_namespace,
+    mock_report_observation_log,
+):
     """
     test report_metrics function
     """
diff --git a/sdk/python/v1beta1/setup.py b/sdk/python/v1beta1/setup.py
index fbf6b865156..49c689a235c 100644
--- a/sdk/python/v1beta1/setup.py
+++ b/sdk/python/v1beta1/setup.py
@@ -46,7 +46,7 @@
         "kubeflow/katib/katib_api_pb2_grpc.py",
     )
 
-    with open("kubeflow/katib/katib_api_pb2_grpc.py", 'r+') as file:
+    with open("kubeflow/katib/katib_api_pb2_grpc.py", "r+") as file:
         content = file.read()
         new_content = content.replace("api_pb2", "kubeflow.katib.katib_api_pb2")
         file.seek(0)

From 74eabaa3bdf2baef79d5f684dd4432de3c68c055 Mon Sep 17 00:00:00 2001
From: Electronic-Waste <2690692950@qq.com>
Date: Thu, 22 Aug 2024 17:43:53 +0000
Subject: [PATCH 14/15] fix(sdk): fix lint error with isort.

Signed-off-by: Electronic-Waste <2690692950@qq.com>
---
 .../v1beta1/kubeflow/katib/api/katib_client.py |  3 +--
 .../kubeflow/katib/api/katib_client_test.py    | 18 ++++++------------
 .../kubeflow/katib/api/report_metrics.py       |  2 +-
 .../kubeflow/katib/api/report_metrics_test.py  |  2 +-
 4 files changed, 9 insertions(+), 16 deletions(-)

diff --git a/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py b/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py
index 5cb535cd40b..78808d17f05 100644
--- a/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py
+++ b/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py
@@ -21,11 +21,10 @@
 
 import grpc
 import kubeflow.katib.katib_api_pb2 as katib_api_pb2
+import kubeflow.katib.katib_api_pb2_grpc as katib_api_pb2_grpc
 from kubeflow.katib import models
 from kubeflow.katib.api_client import ApiClient
 from kubeflow.katib.constants import constants
-import kubeflow.katib.katib_api_pb2 as katib_api_pb2
-import kubeflow.katib.katib_api_pb2_grpc as katib_api_pb2_grpc
 from kubeflow.katib.utils import utils
 from kubernetes import client, config
 
diff --git a/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py b/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py
index 9f853f62a9c..bf94d45d076 100644
--- a/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py
+++ b/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py
@@ -2,20 +2,14 @@
 from typing import List, Optional
 from unittest.mock import Mock, patch
 
+import kubeflow.katib.katib_api_pb2 as katib_api_pb2
 import pytest
-from kubeflow.katib import (
-    KatibClient,
-    V1beta1AlgorithmSpec,
-    V1beta1Experiment,
-    V1beta1ExperimentSpec,
-    V1beta1FeasibleSpace,
-    V1beta1ObjectiveSpec,
-    V1beta1ParameterSpec,
-    V1beta1TrialParameterSpec,
-    V1beta1TrialTemplate,
-)
+from kubeflow.katib import (KatibClient, V1beta1AlgorithmSpec,
+                            V1beta1Experiment, V1beta1ExperimentSpec,
+                            V1beta1FeasibleSpace, V1beta1ObjectiveSpec,
+                            V1beta1ParameterSpec, V1beta1TrialParameterSpec,
+                            V1beta1TrialTemplate)
 from kubeflow.katib.constants import constants
-import kubeflow.katib.katib_api_pb2 as katib_api_pb2
 from kubernetes.client import V1ObjectMeta
 
 TEST_RESULT_SUCCESS = "success"
diff --git a/sdk/python/v1beta1/kubeflow/katib/api/report_metrics.py b/sdk/python/v1beta1/kubeflow/katib/api/report_metrics.py
index 50dd02c9c20..aec62d3f6a7 100644
--- a/sdk/python/v1beta1/kubeflow/katib/api/report_metrics.py
+++ b/sdk/python/v1beta1/kubeflow/katib/api/report_metrics.py
@@ -17,9 +17,9 @@
 from typing import Any, Dict
 
 import grpc
-from kubeflow.katib.constants import constants
 import kubeflow.katib.katib_api_pb2 as katib_api_pb2
 import kubeflow.katib.katib_api_pb2_grpc as katib_api_pb2_grpc
+from kubeflow.katib.constants import constants
 from kubeflow.katib.utils import utils
 
 
diff --git a/sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py b/sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py
index 69f13698b13..4ceba92ef2a 100644
--- a/sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py
+++ b/sdk/python/v1beta1/kubeflow/katib/api/report_metrics_test.py
@@ -1,8 +1,8 @@
 from unittest.mock import patch
 
+import pytest
 from kubeflow.katib import report_metrics
 from kubeflow.katib.constants import constants
-import pytest
 
 TEST_RESULT_SUCCESS = "success"
 ENV_VARIABLE_EMPTY = True

From 7b6fd85666734b8b8e1c30f57cfd29ef1cc5096e Mon Sep 17 00:00:00 2001
From: Electronic-Waste <2690692950@qq.com>
Date: Thu, 22 Aug 2024 17:47:35 +0000
Subject: [PATCH 15/15] fix(sdk): reformat import in katib_client_test.py.

Signed-off-by: Electronic-Waste <2690692950@qq.com>
---
 .../kubeflow/katib/api/katib_client_test.py      | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py b/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py
index bf94d45d076..fef18adfa0f 100644
--- a/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py
+++ b/sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py
@@ -4,11 +4,17 @@
 
 import kubeflow.katib.katib_api_pb2 as katib_api_pb2
 import pytest
-from kubeflow.katib import (KatibClient, V1beta1AlgorithmSpec,
-                            V1beta1Experiment, V1beta1ExperimentSpec,
-                            V1beta1FeasibleSpace, V1beta1ObjectiveSpec,
-                            V1beta1ParameterSpec, V1beta1TrialParameterSpec,
-                            V1beta1TrialTemplate)
+from kubeflow.katib import (
+    KatibClient,
+    V1beta1AlgorithmSpec,
+    V1beta1Experiment,
+    V1beta1ExperimentSpec,
+    V1beta1FeasibleSpace,
+    V1beta1ObjectiveSpec,
+    V1beta1ParameterSpec,
+    V1beta1TrialParameterSpec,
+    V1beta1TrialTemplate,
+)
 from kubeflow.katib.constants import constants
 from kubernetes.client import V1ObjectMeta