# Install required packages (Kubeflow Pipelines and Katib SDK).
!pip install kfp==1.8.12
!pip install kubeflow-katib==0.13.0

import kfp
import kfp.dsl as dsl
from kfp import components

from kubeflow.katib import ApiClient
from kubeflow.katib import V1beta1ExperimentSpec
from kubeflow.katib import V1beta1AlgorithmSpec
from kubeflow.katib import V1beta1EarlyStoppingSpec
from kubeflow.katib import V1beta1EarlyStoppingSetting
from kubeflow.katib import V1beta1ObjectiveSpec
from kubeflow.katib import V1beta1ParameterSpec
from kubeflow.katib import V1beta1FeasibleSpace
from kubeflow.katib import V1beta1TrialTemplate
from kubeflow.katib import V1beta1TrialParameterSpec

"""
Explanation: Kubeflow Pipelines with Katib component
In this notebook you will:
- Create Katib Experiment using random algorithm.
- Use median stopping rule as an early stopping algorithm.
- Use Kubernetes Job with mxnet mnist training container as a Trial template.
- Create Pipeline to get the optimal hyperparameters.
Reference documentation:
- https://kubeflow.org/docs/components/katib/experiment/#random-search
- https://kubeflow.org/docs/components/katib/early-stopping/
- https://kubeflow.org/docs/pipelines/overview/concepts/component/
Note: This Pipeline runs in the multi-user mode. Follow this guide to give your Notebook access to Kubeflow Pipelines.
End of explanation
"""


# Experiment name and namespace.
experiment_name = "median-stop"
experiment_namespace = "kubeflow-user-example-com"

# Trial count specification.
max_trial_count = 18
max_failed_trial_count = 3
parallel_trial_count = 2

# Objective specification.
objective=V1beta1ObjectiveSpec(
    type="maximize",
    goal= 0.99,
    objective_metric_name="Validation-accuracy",
    additional_metric_names=[
        "Train-accuracy"
    ]
)

# Algorithm specification.
algorithm=V1beta1AlgorithmSpec(
    algorithm_name="random",
)

# Early Stopping specification.
early_stopping=V1beta1EarlyStoppingSpec(
    algorithm_name="medianstop",
    algorithm_settings=[
        V1beta1EarlyStoppingSetting(
            name="min_trials_required",
            value="2"
        )
    ]
)


# Experiment search space.
# In this example we tune learning rate, number of layer and optimizer.
# Learning rate has bad feasible space to show more early stopped Trials.
parameters=[
    V1beta1ParameterSpec(
        name="lr",
        parameter_type="double",
        feasible_space=V1beta1FeasibleSpace(
            min="0.01",
            max="0.3"
        ),
    ),
    V1beta1ParameterSpec(
        name="num-layers",
        parameter_type="int",
        feasible_space=V1beta1FeasibleSpace(
            min="2",
            max="5"
        ),
    ),
    V1beta1ParameterSpec(
        name="optimizer",
        parameter_type="categorical",
        feasible_space=V1beta1FeasibleSpace(
            list=[
                "sgd", 
                "adam",
                "ftrl"
            ]
        ),
    ),
]


"""
Explanation: Define an Experiment
You have to create an Experiment object before deploying it. This Experiment is similar to this YAML.
End of explanation
"""


# JSON template specification for the Trial's Worker Kubernetes Job.
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/mxnet-mnist:v0.13.0",
                        "command": [
                            "python3",
                            "/opt/mxnet-mnist/mnist.py",
                            "--batch-size=64",
                            "--lr=${trialParameters.learningRate}",
                            "--num-layers=${trialParameters.numberLayers}",
                            "--optimizer=${trialParameters.optimizer}"
                        ]
                    }
                ],
                "restartPolicy": "Never"
            }
        }
    }
}

# Configure parameters for the Trial template.
# We set the retain parameter to "True" to not clean-up the Trial Job's Kubernetes Pods.
trial_template=V1beta1TrialTemplate(
    retain=True,
    primary_container_name="training-container",
    trial_parameters=[
        V1beta1TrialParameterSpec(
            name="learningRate",
            description="Learning rate for the training model",
            reference="lr"
        ),
        V1beta1TrialParameterSpec(
            name="numberLayers",
            description="Number of training model layers",
            reference="num-layers"
        ),
        V1beta1TrialParameterSpec(
            name="optimizer",
            description="Training model optimizer (sdg, adam or ftrl)",
            reference="optimizer"
        ),
    ],
    trial_spec=trial_spec
)

"""
Explanation: Define a Trial template
In this example, the Trial's Worker is the Kubernetes Job.
End of explanation
"""


experiment_spec=V1beta1ExperimentSpec(
    max_trial_count=max_trial_count,
    max_failed_trial_count=max_failed_trial_count,
    parallel_trial_count=parallel_trial_count,
    objective=objective,
    algorithm=algorithm,
    early_stopping=early_stopping,
    parameters=parameters,
    trial_template=trial_template
)

"""
Explanation: Define an Experiment specification
Create an Experiment specification from the above parameters.
End of explanation
"""


# Get the Katib launcher.
katib_experiment_launcher_op = components.load_component_from_url(
    "https://raw.githubusercontent.com/kubeflow/pipelines/master/components/kubeflow/katib-launcher/component.yaml")

@dsl.pipeline(
    name="Launch Katib early stopping Experiment",
    description="An example to launch Katib Experiment with early stopping"
)

def median_stop():

    # Katib launcher component.
    # Experiment Spec should be serialized to a valid Kubernetes object.
    op = katib_experiment_launcher_op(
        experiment_name=experiment_name,
        experiment_namespace=experiment_namespace,
        experiment_spec=ApiClient().sanitize_for_serialization(experiment_spec),
        experiment_timeout_minutes=60,
        delete_finished_experiment=False)

    # Output container to print the results.
    op_out = dsl.ContainerOp(
        name="best-hp",
        image="library/bash:4.4.23",
        command=["sh", "-c"],
        arguments=["echo Best HyperParameters: %s" % op.output],
    )

"""
Explanation: Create a Pipeline using Katib component
The best hyperparameters are printed after Experiment is finished.
The Experiment is not deleted after the Pipeline is finished.
End of explanation
"""


# Run the Kubeflow Pipeline in the user's namespace.
kfp.Client().create_run_from_pipeline_func(median_stop,  namespace=experiment_namespace, arguments={})

"""
Explanation: Run the Kubeflow Pipeline
You can check the Katib Experiment info in the Katib UI.
End of explanation
"""


