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.

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

[SkyServe] Add CLI args for resources override in sky serve {up, update} #2979

Merged
merged 3 commits into from
Feb 9, 2024
Merged
Changes from all commits
Commits
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
158 changes: 139 additions & 19 deletions sky/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -367,12 +367,6 @@ def _merge_env_vars(env_dict: Optional[Dict[str, str]],


_TASK_OPTIONS = [
click.option('--name',
'-n',
required=False,
type=str,
help=('Task name. Overrides the "name" '
'config in the YAML if both are supplied.')),
click.option(
'--workdir',
required=False,
Expand Down Expand Up @@ -475,6 +469,14 @@ def _merge_env_vars(env_dict: Optional[Dict[str, str]],
same value of ``$MY_ENV3`` in the local environment.""",
)
]
_TASK_OPTIONS_WITH_NAME = [
click.option('--name',
'-n',
required=False,
type=str,
help=('Task name. Overrides the "name" '
'config in the YAML if both are supplied.')),
] + _TASK_OPTIONS
_EXTRA_RESOURCES_OPTIONS = [
click.option(
'--gpus',
Expand Down Expand Up @@ -1335,7 +1337,7 @@ def cli():
flag_value=backends.LocalDockerBackend.NAME,
default=False,
help='If used, runs locally inside a docker container.')
@_add_click_options(_TASK_OPTIONS + _EXTRA_RESOURCES_OPTIONS)
@_add_click_options(_TASK_OPTIONS_WITH_NAME + _EXTRA_RESOURCES_OPTIONS)
@click.option(
'--idle-minutes-to-autostop',
'-i',
Expand Down Expand Up @@ -1515,7 +1517,7 @@ def launch(
is_flag=True,
help=('If True, as soon as a job is submitted, return from this call '
'and do not stream execution logs.'))
@_add_click_options(_TASK_OPTIONS + _EXTRA_RESOURCES_OPTIONS)
@_add_click_options(_TASK_OPTIONS_WITH_NAME + _EXTRA_RESOURCES_OPTIONS)
@usage_lib.entrypoint
# pylint: disable=redefined-builtin
def exec(
Expand Down Expand Up @@ -3946,7 +3948,7 @@ def spot():
nargs=-1,
**_get_shell_complete_args(_complete_file_name))
# TODO(zhwu): Add --dryrun option to test the launch command.
@_add_click_options(_TASK_OPTIONS + _EXTRA_RESOURCES_OPTIONS)
@_add_click_options(_TASK_OPTIONS_WITH_NAME + _EXTRA_RESOURCES_OPTIONS)
@click.option('--spot-recovery',
default=None,
type=str,
Expand Down Expand Up @@ -4349,15 +4351,54 @@ def serve():
pass


def _generate_task_with_service(service_yaml_args: List[str],
not_supported_cmd: str) -> sky.Task:
def _generate_task_with_service(
service_name: str,
service_yaml_args: List[str],
workdir: Optional[str],
cloud: Optional[str],
region: Optional[str],
zone: Optional[str],
num_nodes: Optional[int],
use_spot: Optional[bool],
image_id: Optional[str],
env_file: Optional[Dict[str, str]],
env: List[Tuple[str, str]],
gpus: Optional[str],
instance_type: Optional[str],
ports: Tuple[str],
cpus: Optional[str],
memory: Optional[str],
disk_size: Optional[int],
disk_tier: Optional[str],
not_supported_cmd: str,
) -> sky.Task:
"""Generate a task with service section from a service YAML file."""
is_yaml, _ = _check_yaml(''.join(service_yaml_args))
if not is_yaml:
raise click.UsageError('SERVICE_YAML must be a valid YAML file.')
env = _merge_env_vars(env_file, env)
# We keep nargs=-1 in service_yaml argument to reuse this function.
task = _make_task_or_dag_from_entrypoint_with_overrides(
service_yaml_args, entrypoint_name='Service')
service_yaml_args,
# For Service YAML, we override the `name` field with service name.
name=service_name,
workdir=workdir,
cloud=cloud,
region=region,
zone=zone,
gpus=gpus,
cpus=cpus,
memory=memory,
instance_type=instance_type,
num_nodes=num_nodes,
use_spot=use_spot,
image_id=image_id,
env=env,
disk_size=disk_size,
disk_tier=disk_tier,
ports=ports,
entrypoint_name='Service',
)
if isinstance(task, sky.Dag):
raise click.UsageError(
_DAG_NOT_SUPPORTED_MESSAGE.format(command=not_supported_cmd))
Expand Down Expand Up @@ -4407,16 +4448,34 @@ def _generate_task_with_service(service_yaml_args: List[str],
type=str,
help='A service name. Unique for each service. If not provided, '
'a unique name is autogenerated.')
@_add_click_options(_TASK_OPTIONS + _EXTRA_RESOURCES_OPTIONS)
@click.option('--yes',
'-y',
is_flag=True,
default=False,
required=False,
help='Skip confirmation prompt.')
# TODO(tian): Support the task_option overrides for the service.
@timeline.event
@usage_lib.entrypoint
def serve_up(
service_yaml: List[str],
service_name: Optional[str],
workdir: Optional[str],
cloud: Optional[str],
region: Optional[str],
zone: Optional[str],
num_nodes: Optional[int],
use_spot: Optional[bool],
image_id: Optional[str],
env_file: Optional[Dict[str, str]],
env: List[Tuple[str, str]],
gpus: Optional[str],
instance_type: Optional[str],
ports: Tuple[str],
cpus: Optional[str],
memory: Optional[str],
disk_size: Optional[int],
disk_tier: Optional[str],
yes: bool,
):
"""Launch a SkyServe service.
Expand Down Expand Up @@ -4450,8 +4509,27 @@ def serve_up(
if service_name is None:
service_name = serve_lib.generate_service_name()

task = _generate_task_with_service(service_yaml_args=service_yaml,
not_supported_cmd='sky serve up')
task = _generate_task_with_service(
service_name=service_name,
service_yaml_args=service_yaml,
workdir=workdir,
cloud=cloud,
region=region,
zone=zone,
gpus=gpus,
cpus=cpus,
memory=memory,
instance_type=instance_type,
num_nodes=num_nodes,
use_spot=use_spot,
image_id=image_id,
env_file=env_file,
env=env,
disk_size=disk_size,
disk_tier=disk_tier,
ports=ports,
not_supported_cmd='sky serve up',
)
click.secho('Service Spec:', fg='cyan')
click.echo(task.service)

Expand Down Expand Up @@ -4479,13 +4557,36 @@ def serve_up(
type=str,
nargs=-1,
**_get_shell_complete_args(_complete_file_name))
@_add_click_options(_TASK_OPTIONS + _EXTRA_RESOURCES_OPTIONS)
@click.option('--yes',
'-y',
is_flag=True,
default=False,
required=False,
help='Skip confirmation prompt.')
def serve_update(service_name: str, service_yaml: List[str], yes: bool):
@timeline.event
@usage_lib.entrypoint
def serve_update(
service_name: str,
service_yaml: List[str],
workdir: Optional[str],
cloud: Optional[str],
region: Optional[str],
zone: Optional[str],
num_nodes: Optional[int],
use_spot: Optional[bool],
image_id: Optional[str],
env_file: Optional[Dict[str, str]],
env: List[Tuple[str, str]],
gpus: Optional[str],
instance_type: Optional[str],
ports: Tuple[str],
cpus: Optional[str],
memory: Optional[str],
disk_size: Optional[int],
disk_tier: Optional[str],
yes: bool,
):
"""Update a SkyServe service.

service_yaml must point to a valid YAML file.
Expand All @@ -4496,8 +4597,27 @@ def serve_update(service_name: str, service_yaml: List[str], yes: bool):

sky serve update sky-service-16aa new_service.yaml
"""
task = _generate_task_with_service(service_yaml_args=service_yaml,
not_supported_cmd='sky serve update')
task = _generate_task_with_service(
service_name=service_name,
service_yaml_args=service_yaml,
workdir=workdir,
cloud=cloud,
region=region,
zone=zone,
gpus=gpus,
cpus=cpus,
memory=memory,
instance_type=instance_type,
num_nodes=num_nodes,
use_spot=use_spot,
image_id=image_id,
env_file=env_file,
env=env,
disk_size=disk_size,
disk_tier=disk_tier,
ports=ports,
not_supported_cmd='sky serve update',
)
click.secho('Service Spec:', fg='cyan')
click.echo(task.service)

Expand Down Expand Up @@ -4835,7 +4955,7 @@ def _get_candidate_configs(yaml_path: str) -> Optional[List[Dict[str, str]]]:
required=True,
type=str,
help='Benchmark name.')
@_add_click_options(_TASK_OPTIONS)
@_add_click_options(_TASK_OPTIONS_WITH_NAME)
@click.option('--gpus',
required=False,
type=str,
Expand Down
Loading