Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[SDK]Support Docker image as objective in the tune API #2338

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open

[SDK]Support Docker image as objective in the tune API #2338

Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e0ea84e
Support Docker image as objective in the tune API
akhilsaivenkata May 29, 2024
2ef15e1
resolving review comments
akhilsaivenkata Jul 16, 2024
b2290f9
Merge branch 'master' of github.com:kubeflow/katib
akhilsaivenkata Jul 16, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
133 changes: 70 additions & 63 deletions sdk/python/v1beta1/kubeflow/katib/api/katib_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think, we should keep the base_image, since we use it when user set objective as train function.

namespace: Optional[str] = None,
env_per_trial: Optional[
Union[Dict[str, str], List[Union[client.V1EnvVar, client.V1EnvFromSource]]]
Expand Down Expand Up @@ -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,
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes @andreyvelich , i did keep the base_image here at this line, I have added it in the if block so this code change got mixed up with all the other lines

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")
Expand Down Expand Up @@ -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,
Comment on lines +406 to +407
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also I'm not sure if we can assign None to command and args here when we use Docker image as objective.

As @andreyvelich shows an example for us, we sometimes need to pass command and args to the training container to execute python scripts with some parameters.

Could you explain your idea in details so that I can understand more? WDYT👀 @akhilsaivenkata @andreyvelich

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think, initially we can just allow user to set image as objective without command and args.
Similar to how we allow create training job using base_image parameter: https://github.com/kubeflow/training-operator/blob/master/sdk/python/kubeflow/training/api/training_client.py#L327C35-L327C45.

env=env,
env_from=env_from,
resources=resources_per_trial,
Expand All @@ -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 [],
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

trial_parameters still be required even when user sets Docker image.
You can check example here: https://github.com/kubeflow/katib/blob/master/examples/v1beta1/hp-tuning/random.yaml#L31-L36

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure @andreyvelich , I will revert this change and keep the trial_parameters.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, I would like to know if we need to write new unit test cases or change existing ones ?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@akhilsaivenkata Yes, since we merge this PR: #2325, please add unit test for tune function.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cc @tariq-hasan

trial_spec=trial_spec,
)

# Add parameters to the Katib Experiment.
experiment.spec.parameters = experiment_params
experiment.spec.parameters = experiment_params if callable(objective) else []
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think parameters field is also needed since trial_paramaters is still required.

WDTY👀 @akhilsaivenkata @andreyvelich

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are right @Electronic-Waste , I have just reverted these two code changes and pushed it now.


# Add Trial template to the Katib Experiment.
experiment.spec.trial_template = trial_template
Expand Down
Loading