feat(optimizer): Add get_best_results API to OptimizerClient (kubeflow#152)

* feat(optimizer): Add get_best_results API to OptimizerClient

Signed-off-by: kramaranya <kramaranya15@gmail.com>

* Update kubeflow/optimizer/api/optimizer_client.py

Co-authored-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
Signed-off-by: Anya Kramar <akramar@redhat.com>

---------

Signed-off-by: kramaranya <kramaranya15@gmail.com>
Signed-off-by: Anya Kramar <akramar@redhat.com>
Co-authored-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>

Author: kramaranya

kubeflow/optimizer/__init__.py

@@ -20,7 +20,12 @@
 
 # Import the Kubeflow Optimizer types.
 from kubeflow.optimizer.types.algorithm_types import GridSearch, RandomSearch
-from kubeflow.optimizer.types.optimization_types import Objective, OptimizationJob, TrialConfig
+from kubeflow.optimizer.types.optimization_types import (
+    Objective,
+    OptimizationJob,
+    Result,
+    TrialConfig,
+)
 from kubeflow.optimizer.types.search_types import Search
 
 # Import the Kubeflow Trainer types.
@@ -33,6 +38,7 @@
     "OptimizationJob",
     "OptimizerClient",
     "RandomSearch",
+    "Result",
     "Search",
     "TrainJobTemplate",
     "TrialConfig",

kubeflow/optimizer/api/optimizer_client.py

@@ -23,7 +23,7 @@
 from kubeflow.optimizer.types.optimization_types import (
     Objective,
     OptimizationJob,
-    Trial,
+    Result,
     TrialConfig,
 )
 from kubeflow.trainer.types.types import TrainJobTemplate
@@ -160,6 +160,25 @@ def get_job_logs(
         """
         return self.backend.get_job_logs(name=name, trial_name=trial_name, follow=follow)
 
+    def get_best_results(self, name: str) -> Optional[Result]:
+        """Get the best hyperparameters and metrics from an OptimizationJob.
+
+        This method retrieves the optimal hyperparameters and their corresponding metrics
+        from the best trial found during the optimization process.
+
+        Args:
+            name: Name of the OptimizationJob.
+
+        Returns:
+            A Result object containing the best hyperparameters and metrics, or None if
+            no best trial is available yet.
+
+        Raises:
+            TimeoutError: Timeout to get an OptimizationJob.
+            RuntimeError: Failed to get an OptimizationJob.
+        """
+        return self.backend.get_best_results(name=name)
+
     def wait_for_job_status(
         self,
         name: str,
@@ -193,22 +212,6 @@ def wait_for_job_status(
             polling_interval=polling_interval,
         )
 
-    def get_best_trial(self, name: str) -> Optional[Trial]:
-        """Get the current best Trial for an OptimizationJob.
-
-        Args:
-            name: Name of the OptimizationJob.
-
-        Returns:
-            The current best Trial with parameters, metrics, and associated TrainJob.
-            Returns None if the best trial is not available yet.
-
-        Raises:
-            TimeoutError: Timeout to get OptimizationJob.
-            RuntimeError: Failed to get OptimizationJob.
-        """
-        return self.backend.get_best_trial(name=name)
-
     def delete_job(self, name: str):
         """Delete the OptimizationJob.
 

kubeflow/optimizer/backends/base.py

@@ -21,7 +21,7 @@
 from kubeflow.optimizer.types.optimization_types import (
     Objective,
     OptimizationJob,
-    Trial,
+    Result,
     TrialConfig,
 )
 from kubeflow.trainer.types.types import TrainJobTemplate
@@ -57,6 +57,10 @@ def get_job_logs(
     ) -> Iterator[str]:
         raise NotImplementedError()
 
+    @abc.abstractmethod
+    def get_best_results(self, name: str) -> Optional[Result]:
+        raise NotImplementedError()
+
     @abc.abstractmethod
     def wait_for_job_status(
         self,
@@ -67,10 +71,6 @@ def wait_for_job_status(
     ) -> OptimizationJob:
         raise NotImplementedError()
 
-    @abc.abstractmethod
-    def get_best_trial(self, name: str) -> Optional[Trial]:
-        raise NotImplementedError()
-
     @abc.abstractmethod
     def delete_job(self, name: str):
         raise NotImplementedError()

kubeflow/optimizer/backends/kubernetes/backend.py

@@ -35,6 +35,7 @@
     Metric,
     Objective,
     OptimizationJob,
+    Result,
     Trial,
     TrialConfig,
 )
@@ -223,7 +224,7 @@ def get_job_logs(
         # Determine what trial to get logs from.
         if trial_name is None:
             # Get logs from the best current trial.
-            best_trial = self.get_best_trial(name)
+            best_trial = self._get_best_trial(name)
             if best_trial is None:
                 # Get first trial if available.
                 optimization_job = self.get_job(name)
@@ -249,6 +250,18 @@ def get_job_logs(
             pod_name=pod_name, container_name=container_name, follow=follow
         )
 
+    def get_best_results(self, name: str) -> Optional[Result]:
+        """Get the best hyperparameters and metrics from an OptimizationJob"""
+        best_trial = self._get_best_trial(name)
+
+        if best_trial is None:
+            return None
+
+        return Result(
+            parameters=best_trial.parameters,
+            metrics=best_trial.metrics,
+        )
+
     def wait_for_job_status(
         self,
         name: str,
@@ -293,7 +306,29 @@ def wait_for_job_status(
             f"{status}"
         )
 
-    def get_best_trial(self, name: str) -> Optional[Trial]:
+    def delete_job(self, name: str):
+        """Delete the OptimizationJob"""
+
+        try:
+            self.custom_api.delete_namespaced_custom_object(
+                constants.GROUP,
+                constants.VERSION,
+                self.namespace,
+                constants.EXPERIMENT_PLURAL,
+                name=name,
+            )
+        except multiprocessing.TimeoutError as e:
+            raise TimeoutError(
+                f"Timeout to delete {constants.OPTIMIZATION_JOB_KIND}: {self.namespace}/{name}"
+            ) from e
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to delete {constants.OPTIMIZATION_JOB_KIND}: {self.namespace}/{name}"
+            ) from e
+
+        logger.debug(f"{constants.OPTIMIZATION_JOB_KIND} {self.namespace}/{name} has been deleted")
+
+    def _get_best_trial(self, name: str) -> Optional[Trial]:
         """Get the best current Trial for the OptimizationJob"""
         optimization_job = self.__get_experiment_cr(name)
 
@@ -338,28 +373,6 @@ def get_best_trial(self, name: str) -> Optional[Trial]:
 
         return None
 
-    def delete_job(self, name: str):
-        """Delete the OptimizationJob"""
-
-        try:
-            self.custom_api.delete_namespaced_custom_object(
-                constants.GROUP,
-                constants.VERSION,
-                self.namespace,
-                constants.EXPERIMENT_PLURAL,
-                name=name,
-            )
-        except multiprocessing.TimeoutError as e:
-            raise TimeoutError(
-                f"Timeout to delete {constants.OPTIMIZATION_JOB_KIND}: {self.namespace}/{name}"
-            ) from e
-        except Exception as e:
-            raise RuntimeError(
-                f"Failed to delete {constants.OPTIMIZATION_JOB_KIND}: {self.namespace}/{name}"
-            ) from e
-
-        logger.debug(f"{constants.OPTIMIZATION_JOB_KIND} {self.namespace}/{name} has been deleted")
-
     def __get_experiment_cr(self, name: str) -> models.V1beta1Experiment:
         """Get the Experiment CR from Kubernetes API"""
         try:

kubeflow/optimizer/types/optimization_types.py

@@ -73,6 +73,19 @@ class Metric:
     latest: str
 
 
+@dataclass
+class Result:
+    """Result containing the best hyperparameters and metrics.
+
+    Args:
+        parameters (`dict[str, str]`): The best hyperparameters found during optimization.
+        metrics (`list[Metric]`): The metrics achieved with these hyperparameters.
+    """
+
+    parameters: dict[str, str]
+    metrics: list[Metric]
+
+
 # Representation of the single trial
 @dataclass
 class Trial: