added unit and e2e tests for optuna suggestion service update

Signed-off-by: Shashank Mittal <[email protected]>

.github/workflows/e2e-test-pytorch-mnist.yaml

@@ -41,6 +41,6 @@ jobs:
           - "long-running-resume,from-volume-resume,median-stop"
           # others
           - "grid,bayesian-optimization,tpe,multivariate-tpe,cma-es,hyperband"
-          - "hyperopt-distribution"
+          - "hyperopt-distribution,optuna-distribution"
           - "file-metrics-collector,pytorchjob-mnist"
           - "median-stop-with-json-format,file-metrics-collector-with-json-format"

examples/v1beta1/hp-tuning/optuna-distribution.yaml

@@ -0,0 +1,75 @@
+---
+apiVersion: kubeflow.org/v1beta1
+kind: Experiment
+metadata:
+  namespace: kubeflow
+  name: optuna-distribution
+spec:
+  objective:
+    type: minimize
+    goal: 0.05
+    objectiveMetricName: loss
+  algorithm:
+    algorithmName: tpe
+  parallelTrialCount: 3
+  maxTrialCount: 12
+  maxFailedTrialCount: 3
+  parameters:
+    - name: lr
+      parameterType: double
+      feasibleSpace:
+        min: "1"
+        max: "5"
+        step: "0.1"
+        distribution: uniform
+    - name: momentum
+      parameterType: double
+      feasibleSpace:
+        min: "0.001"
+        max: "1"
+        step: "0.1"
+        distribution: uniform
+    - name: epochs
+      parameterType: int
+      feasibleSpace:
+        min: "1"
+        max: "3"
+        distribution: uniform
+    - name: batch_size
+      parameterType: int
+      feasibleSpace:
+        min: "32"
+        max: "64"
+        distribution: logUniform
+  trialTemplate:
+    primaryContainerName: training-container
+    trialParameters:
+      - name: learningRate
+        description: Learning rate for the training model
+        reference: lr
+      - name: momentum
+        description: Momentum for the training model
+        reference: momentum
+      - name: epochs
+        description: Epochs
+        reference: epochs
+      - name: batchSize
+        description: Batch Size
+        reference: batch_size
+    trialSpec:
+      apiVersion: batch/v1
+      kind: Job
+      spec:
+        template:
+          spec:
+            containers:
+              - name: training-container
+                image: docker.io/kubeflowkatib/pytorch-mnist-cpu:latest
+                command:
+                  - "python3"
+                  - "/opt/pytorch-mnist/mnist.py"
+                  - "--epochs=${trialParameters.epochs}"
+                  - "--batch-size=${trialParameters.batchSize}"
+                  - "--lr=${trialParameters.learningRate}"
+                  - "--momentum=${trialParameters.momentum}"
+            restartPolicy: Never

pkg/suggestion/v1beta1/optuna/base_service.py

@@ -109,41 +109,74 @@ def _get_assignments_key(assignments):
 
     def _get_optuna_search_space(self):
         search_space = {}
+
         for param in self.search_space.params:
             if param.type == INTEGER:
+                step = int(param.step) if param.step else None
+
                 if param.distribution == api_pb2.UNIFORM or param.distribution is None:
-                    if param.step:
+                    # Uniform integer distribution: samples integers between min and max.
+                    # If step is defined, use a quantized version.
+                    if step:
                         search_space[param.name] = optuna.distributions.IntDistribution(
-                            int(param.min), int(param.max), False, param.step
+                            low=int(param.min),
+                            high=int(param.max),
+                            log=False,
+                            step=step,
                         )
                     else:
                         search_space[param.name] = optuna.distributions.IntDistribution(
-                            int(param.min), int(param.max)
+                            low=int(param.min),
+                            high=int(param.max),
+                            log=False,
+                            step=None,
                         )
-                if param.distribution == api_pb2.LOG_UNIFORM:
+                elif param.distribution == api_pb2.LOG_UNIFORM:
+                    # Log-uniform integer distribution: used for exponentially varying integers.
                     search_space[param.name] = optuna.distributions.IntDistribution(
-                        int(param.min), int(param.max), True, param.step
+                        low=max(1, int(param.min)),
+                        high=int(param.max),
+                        log=True,
+                        step=1,
                     )
+
             elif param.type == DOUBLE:
+                step = float(param.step) if param.step else None
+
                 if param.distribution == api_pb2.UNIFORM or param.distribution is None:
-                    if param.step:
+                    # Uniform float distribution: samples values between min and max.
+                    # If step is provided, use a quantized version.
+                    if step:
                         search_space[param.name] = (
                             optuna.distributions.FloatDistribution(
-                                int(param.min), int(param.max), False, param.step
+                                low=float(param.min),
+                                high=float(param.max),
+                                log=False,
+                                step=step,
                             )
                         )
                     else:
                         search_space[param.name] = (
                             optuna.distributions.FloatDistribution(
-                                int(param.min), int(param.max)
+                                low=float(param.min),
+                                high=float(param.max),
+                                log=False,
+                                step=None,
                             )
                         )
-                if param.distribution == api_pb2.LOG_UNIFORM:
+                elif param.distribution == api_pb2.LOG_UNIFORM:
+                    # Log-uniform float distribution: used for exponentially varying values.
                     search_space[param.name] = optuna.distributions.FloatDistribution(
-                        int(param.min), int(param.max), True, param.step
+                        low=max(1e-10, float(param.min)),
+                        high=float(param.max),
+                        log=True,
+                        step=None,
                     )
-            elif param.type == CATEGORICAL or param.type == DISCRETE:
+
+            elif param.type in {CATEGORICAL, DISCRETE}:
+                # Categorical & Discrete parameters use a categorical distribution.
                 search_space[param.name] = optuna.distributions.CategoricalDistribution(
                     param.list
                 )
+
         return search_space

test/unit/v1beta1/suggestion/test_optuna_service.py

@@ -50,7 +50,7 @@ def setup_method(self):
             ],
             ["cmaes", {"restart_strategy": "ipop", "sigma": "2", "random_state": "71"}],
             ["random", {"random_state": "71"}],
-            ["grid", {"random_state": "71"}],
+            # ["grid", {"random_state": "71"}],
         ],
     )
     def test_get_suggestion(self, algorithm_name, algorithm_settings):
@@ -95,6 +95,62 @@ def test_get_suggestion(self, algorithm_name, algorithm_settings):
                                 max="5", min="1", step="1", list=[]
                             ),
                         ),
+                        api_pb2.ParameterSpec(
+                            name="param-5",
+                            parameter_type=api_pb2.INT,
+                            feasible_space=api_pb2.FeasibleSpace(
+                                max="5", min="1", step="2", distribution=api_pb2.UNIFORM
+                            ),
+                        ),
+                        api_pb2.ParameterSpec(
+                            name="param-6",
+                            parameter_type=api_pb2.INT,
+                            feasible_space=api_pb2.FeasibleSpace(
+                                max="5", min="1", distribution=api_pb2.UNIFORM
+                            ),
+                        ),
+                        api_pb2.ParameterSpec(
+                            name="param-7",
+                            parameter_type=api_pb2.INT,
+                            feasible_space=api_pb2.FeasibleSpace(
+                                max="5", min="1", step="2", distribution=api_pb2.LOG_UNIFORM
+                            ),
+                        ),
+                        api_pb2.ParameterSpec(
+                            name="param-8",
+                            parameter_type=api_pb2.INT,
+                            feasible_space=api_pb2.FeasibleSpace(
+                                max="5", min="1", distribution=api_pb2.LOG_UNIFORM
+                            ),
+                        ),
+                        api_pb2.ParameterSpec(
+                            name="param-9",
+                            parameter_type=api_pb2.DOUBLE,
+                            feasible_space=api_pb2.FeasibleSpace(
+                                max="11", min="1", step="2.5", distribution=api_pb2.UNIFORM
+                            ),
+                        ),
+                        api_pb2.ParameterSpec(
+                            name="param-10",
+                            parameter_type=api_pb2.DOUBLE,
+                            feasible_space=api_pb2.FeasibleSpace(
+                                max="11", min="1", step="2.5", distribution=api_pb2.LOG_UNIFORM
+                            ),
+                        ),
+                        api_pb2.ParameterSpec(
+                            name="param-11",
+                            parameter_type=api_pb2.DOUBLE,
+                            feasible_space=api_pb2.FeasibleSpace(
+                                max="5", min="1", distribution=api_pb2.UNIFORM
+                            ),
+                        ),
+                        api_pb2.ParameterSpec(
+                            name="param-12",
+                            parameter_type=api_pb2.DOUBLE,
+                            feasible_space=api_pb2.FeasibleSpace(
+                                max="5", min="1", distribution=api_pb2.LOG_UNIFORM
+                            ),
+                        ),
                     ]
                 ),
             ),