import kubeflow.katib as katib

# Step 1. Create an objective function with push-based metrics collection.
def objective(parameters):
    # Import required packages.
    import time
    import kubeflow.katib as katib
    time.sleep(5)
    # Calculate objective function.
    result = 4 * int(parameters["a"]) - float(parameters["b"]) ** 2
    # Push metrics to Katib DB.
    katib.report_metrics({"result": result})

# Step 2. Create HyperParameter search space.
parameters = {
    "a": katib.search.int(min=10, max=20),
    "b": katib.search.double(min=0.1, max=0.2)
}

# Step 3. Create Katib Experiment with 4 Trials and 2 CPUs per Trial.
# We choose to install the latest changes of Python SDK because `report_metrics` has not been
# supported yet. Thus, the base image must have `git` command to download the package.
katib_client = katib.KatibClient(namespace="kubeflow")
name = "tune-experiment"
katib_client.tune(
    name=name,
    objective=objective,
    parameters=parameters,
    base_image="electronicwaste/push-metrics-collector:v0.0.9", # python:3.11-slim + git
    objective_metric_name="result",
    max_trial_count=4,
    resources_per_trial={"cpu": "2"},
    packages_to_install=["git+https://github.com/kubeflow/katib.git@master#subdirectory=sdk/python/v1beta1"],
    # packages_to_install=["kubeflow-katib==0.18.0"],
    metrics_collector_config={"kind": "Push"},
)

# Step 4. Wait until Katib Experiment is complete
katib_client.wait_for_experiment_condition(name=name)

# Step 5. Get the best HyperParameters.
print(katib_client.get_optimal_hyperparameters(name))