From 4e8c6fab7537b716763cbd2f2f291a4badc869ef Mon Sep 17 00:00:00 2001 From: Prakhar Singhal Date: Sat, 10 Aug 2024 05:13:47 +0530 Subject: [PATCH] * refactor(sdk): added option for custom metric collector for tune in katlib_client.py Signed-off-by: Prakhar Singhal <153047595+prakhar479@users.noreply.github.com> --- sdk/python/v1beta1/kubeflow/katib/api/katib_client.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 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..659d808373a 100644 --- a/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py +++ b/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py @@ -251,8 +251,8 @@ def tune( pip_index_url: The PyPI url 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. - + for using custom metric collectors use "custom_collector" key, + for example, `metrics_collector_config = {"custom_collector": "prometheus "}`. Raises: ValueError: Function arguments have incorrect type or value. TimeoutError: Timeout to create Katib Experiment. @@ -387,7 +387,9 @@ 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_spec = models.V1beta1MetricsCollectorSpec( - collector=models.V1beta1CollectorSpec(kind=metrics_collector_config["kind"]) + collector=models.V1beta1CollectorSpec( + kind=metrics_collector_config["kind"], + custom_collector=metrics_collector_config["custom_collector"]) ) # Create Trial specification.