Support Docker image as objective in the tune API

Signed-off-by: akhilsaivenkata <[email protected]>

sdk/python/v1beta1/kubeflow/katib/api/katib_client.py

@@ -153,9 +153,9 @@ def tune(
         self,
         # TODO (andreyvelich): How to be consistent with other APIs (name) ?
         name: str,
-        objective: Callable,
+        objective: Union[Callable, str],
         parameters: Dict[str, Any],
-        base_image: str = constants.BASE_IMAGE_TENSORFLOW,
+        #base_image: str = constants.BASE_IMAGE_TENSORFLOW,
         namespace: Optional[str] = None,
         env_per_trial: Optional[
             Union[Dict[str, str], List[Union[client.V1EnvVar, client.V1EnvFromSource]]]
@@ -283,65 +283,72 @@ def tune(
         if max_failed_trial_count is not None:
             experiment.spec.max_failed_trial_count = max_failed_trial_count
 
-        # Validate objective function.
-        utils.validate_objective_function(objective)
-
-        # Extract objective function implementation.
-        objective_code = inspect.getsource(objective)
-
-        # Objective function might be defined in some indented scope
-        # (e.g. in another function). We need to dedent the function code.
-        objective_code = textwrap.dedent(objective_code)
-
-        # Iterate over input parameters.
-        input_params = {}
-        experiment_params = []
-        trial_params = []
-        for p_name, p_value in parameters.items():
-            # If input parameter value is Katib Experiment parameter sample.
-            if isinstance(p_value, models.V1beta1ParameterSpec):
-                # Wrap value for the function input.
-                input_params[p_name] = f"${{trialParameters.{p_name}}}"
-
-                # Add value to the Katib Experiment parameters.
-                p_value.name = p_name
-                experiment_params.append(p_value)
-
-                # Add value to the Katib Experiment's Trial parameters.
-                trial_params.append(
-                    models.V1beta1TrialParameterSpec(name=p_name, reference=p_name)
-                )
-            else:
-                # Otherwise, add value to the function input.
-                input_params[p_name] = p_value
-
-        # Wrap objective function to execute it from the file. For example
-        # def objective(parameters):
-        #     print(f'Parameters are {parameters}')
-        # objective({'lr': '${trialParameters.lr}', 'epochs': '${trialParameters.epochs}', 'is_dist': False})
-        objective_code = f"{objective_code}\n{objective.__name__}({input_params})\n"
-
-        # Prepare execute script template.
-        exec_script = textwrap.dedent(
-            """
-            program_path=$(mktemp -d)
-            read -r -d '' SCRIPT << EOM\n
-            {objective_code}
-            EOM
-            printf "%s" "$SCRIPT" > $program_path/ephemeral_objective.py
-            python3 -u $program_path/ephemeral_objective.py"""
-        )
-
-        # Add objective code to the execute script.
-        exec_script = exec_script.format(objective_code=objective_code)
-
-        # Install Python packages if that is required.
-        if packages_to_install is not None:
-            exec_script = (
-                utils.get_script_for_python_packages(packages_to_install, pip_index_url)
-                + exec_script
+        # Handle different types of objective input
+        if callable(objective):
+            # Validate objective function.
+            utils.validate_objective_function(objective)
+
+            # Extract objective function implementation.
+            objective_code = inspect.getsource(objective)
+
+            # Objective function might be defined in some indented scope
+            # (e.g. in another function). We need to dedent the function code.
+            objective_code = textwrap.dedent(objective_code)
+
+            # Iterate over input parameters.
+            input_params = {}
+            experiment_params = []
+            trial_params = []
+            base_image = constants.BASE_IMAGE_TENSORFLOW,
+            for p_name, p_value in parameters.items():
+                # If input parameter value is Katib Experiment parameter sample.
+                if isinstance(p_value, models.V1beta1ParameterSpec):
+                    # Wrap value for the function input.
+                    input_params[p_name] = f"${{trialParameters.{p_name}}}"
+
+                    # Add value to the Katib Experiment parameters.
+                    p_value.name = p_name
+                    experiment_params.append(p_value)
+
+                    # Add value to the Katib Experiment's Trial parameters.
+                    trial_params.append(
+                        models.V1beta1TrialParameterSpec(name=p_name, reference=p_name)
+                    )
+                else:
+                    # Otherwise, add value to the function input.
+                    input_params[p_name] = p_value
+
+            # Wrap objective function to execute it from the file. For example
+            # def objective(parameters):
+            #     print(f'Parameters are {parameters}')
+            # objective({'lr': '${trialParameters.lr}', 'epochs': '${trialParameters.epochs}', 'is_dist': False})
+            objective_code = f"{objective_code}\n{objective.__name__}({input_params})\n"
+
+            # Prepare execute script template.
+            exec_script = textwrap.dedent(
+                """
+                program_path=$(mktemp -d)
+                read -r -d '' SCRIPT << EOM\n
+                {objective_code}
+                EOM
+                printf "%s" "$SCRIPT" > $program_path/ephemeral_objective.py
+                python3 -u $program_path/ephemeral_objective.py"""
             )
 
+            # Add objective code to the execute script.
+            exec_script = exec_script.format(objective_code=objective_code)
+
+            # Install Python packages if that is required.
+            if packages_to_install is not None:
+                exec_script = (
+                    utils.get_script_for_python_packages(packages_to_install, pip_index_url)
+                    + exec_script
+                )
+        elif isinstance(objective, str):
+            base_image=objective
+        else:
+            raise ValueError("The objective must be a callable function or a docker image.")
+
         if isinstance(resources_per_trial, dict):
             if "gpu" in resources_per_trial:
                 resources_per_trial["nvidia.com/gpu"] = resources_per_trial.pop("gpu")
@@ -384,8 +391,8 @@ def tune(
                             client.V1Container(
                                 name=constants.DEFAULT_PRIMARY_CONTAINER_NAME,
                                 image=base_image,
-                                command=["bash", "-c"],
-                                args=[exec_script],
+                                command=["bash", "-c"] if callable(objective) else None,
+                                args=[exec_script] if callable(objective) else None,
                                 env=env,
                                 env_from=env_from,
                                 resources=resources_per_trial,
@@ -400,12 +407,12 @@ def tune(
         trial_template = models.V1beta1TrialTemplate(
             primary_container_name=constants.DEFAULT_PRIMARY_CONTAINER_NAME,
             retain=retain_trials,
-            trial_parameters=trial_params,
+            trial_parameters=trial_params if callable(objective) else [],
             trial_spec=trial_spec,
         )
 
         # Add parameters to the Katib Experiment.
-        experiment.spec.parameters = experiment_params
+        experiment.spec.parameters = experiment_params if callable(objective) else []
 
         # Add Trial template to the Katib Experiment.
         experiment.spec.trial_template = trial_template