From e0ea84e4d1e73cac3b930110caa887246defc452 Mon Sep 17 00:00:00 2001 From: akhilsaivenkata Date: Wed, 29 May 2024 23:49:13 +0000 Subject: [PATCH] Support Docker image as objective in the tune API Signed-off-by: akhilsaivenkata --- .../kubeflow/katib/api/katib_client.py | 133 +++++++++--------- 1 file changed, 70 insertions(+), 63 deletions(-) diff --git a/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py b/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py index 8be9e52f6da..9a3c90ea84d 100644 --- a/sdk/python/v1beta1/kubeflow/katib/api/katib_client.py +++ b/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