cli.py

"""The 'sky' command line tool.

Example usage:

 # See available commands.
 >> sky

 # Run a task, described in a yaml file.
 # Provisioning, setup, file syncing are handled.
 >> sky launch task.yaml
 >> sky launch [-c cluster_name] task.yaml

 # Show the list of running clusters.
 >> sky status

 # Tear down a specific cluster.
 >> sky down cluster_name

 # Tear down all existing clusters.
 >> sky down -a

NOTE: the order of command definitions in this file corresponds to how they are
listed in "sky --help". Take care to put logically connected commands close to
each other.
"""
importcollections
importcopy
importdatetime
importfunctools
importgetpass
importos
importpathlib
importshlex
importshutil
importsubprocess
importsys
importtextwrap
importtraceback
importtyping
fromtypingimport (Any, Callable, Dict, Generator, List, Optional, Set, Tuple,
Union)

importclick
importcolorama
importdotenv
importrequestsasrequests_lib
fromrichimportprogressasrich_progress
importyaml

importsky
fromskyimportbackends
fromskyimportclouds
fromskyimportexceptions
fromskyimportglobal_user_state
fromskyimportjobsasmanaged_jobs
fromskyimportmodels
fromskyimportserveasserve_lib
fromskyimportsky_logging
fromskyimportskypilot_config
fromsky.adaptorsimportcommonasadaptors_common
fromsky.benchmarkimportbenchmark_state
fromsky.benchmarkimportbenchmark_utils
fromsky.clientimportsdk
fromsky.cloudsimportservice_catalog
fromsky.dataimportstorage_utils
fromsky.provision.kubernetesimportconstantsaskubernetes_constants
fromsky.provision.kubernetesimportutilsaskubernetes_utils
fromsky.serverimportcommonasserver_common
fromsky.serverimportconstantsasserver_constants
fromsky.server.requestsimportrequests
fromsky.skyletimportconstants
fromsky.skyletimportjob_lib
fromsky.usageimportusage_lib
fromsky.utilsimportannotations
fromsky.utilsimportcluster_utils
fromsky.utilsimportcommon
fromsky.utilsimportcommon_utils
fromsky.utilsimportcontroller_utils
fromsky.utilsimportdag_utils
fromsky.utilsimportenv_options
fromsky.utilsimportlog_utils
fromsky.utilsimportregistry
fromsky.utilsimportresources_utils
fromsky.utilsimportrich_utils
fromsky.utilsimportstatus_lib
fromsky.utilsimportsubprocess_utils
fromsky.utilsimporttimeline
fromsky.utilsimportux_utils
fromsky.utils.cli_utilsimportstatus_utils

iftyping.TYPE_CHECKING:
importtypes

pd=adaptors_common.LazyImport('pandas')
logger=sky_logging.init_logger(__name__)

_CONTEXT_SETTINGS=dict(help_option_names=['-h', '--help'])

_CLUSTER_FLAG_HELP="""\
A cluster name. If provided, either reuse an existing cluster with that name or
provision a new cluster with that name. Otherwise provision a new cluster with
an autogenerated name."""

# The maximum number of in-progress managed jobs to show in the status
# command.
_NUM_MANAGED_JOBS_TO_SHOW_IN_STATUS=5
_NUM_MANAGED_JOBS_TO_SHOW=50

_STATUS_PROPERTY_CLUSTER_NUM_ERROR_MESSAGE= (
'{cluster_num} cluster{plural} {verb}. Please specify {cause} '
'cluster to show its {property}.\nUsage: `sky status --{flag} <cluster>`')

_DAG_NOT_SUPPORTED_MESSAGE= ('YAML specifies a DAG which is only supported by '
'`sky jobs launch`. `{command}` supports a '
'single task only.')


def_get_cluster_records_and_set_ssh_config(
clusters: Optional[List[str]],
refresh: common.StatusRefreshMode=common.StatusRefreshMode.NONE,
all_users: bool=False,
) ->List[dict]:
"""Returns a list of clusters that match the glob pattern.

 Args:
 clusters: A list of cluster names to query. If None, query all clusters.
 refresh: The refresh mode for the status command.
 all_users: Whether to query clusters from all users.
 If clusters is not None, this field is ignored because cluster list
 can include other users' clusters.
 """
# TODO(zhwu): we should move this function into SDK.
# TODO(zhwu): this additional RTT makes CLIs slow. We should optimize this.
ifclustersisnotNone:
all_users=True
request_id=sdk.status(clusters, refresh=refresh, all_users=all_users)
cluster_records=sdk.stream_and_get(request_id)
# Update the SSH config for all clusters
forrecordincluster_records:
handle=record['handle']

ifnot (handleisnotNoneandhandle.cached_external_ipsisnotNone
and'credentials'inrecord):
# If the cluster is not UP or does not have credentials available,
# we need to remove the cluster from the SSH config.
cluster_utils.SSHConfigHelper.remove_cluster(record['name'])
continue

# During the failover, even though a cluster does not exist, the handle
# can still exist in the record, and we check for credentials to avoid
# updating the SSH config for non-existent clusters.
credentials=record['credentials']
ifisinstance(handle.launched_resources.cloud, clouds.Kubernetes):
# Replace the proxy command to proxy through the SkyPilot API
# server with websocket.
key_path= (cluster_utils.SSHConfigHelper.generate_local_key_file(
handle.cluster_name, credentials))
# Instead of directly use websocket_proxy.py, we add an
# additional proxy, so that ssh can use the head pod in the
# cluster to jump to worker pods.
proxy_command= (
f'ssh -tt -i {key_path} '
'-o StrictHostKeyChecking=no '
'-o UserKnownHostsFile=/dev/null '
'-o IdentitiesOnly=yes '
'-W \'[%h]:%p\' '
f'{handle.ssh_user}@127.0.0.1 '
'-o ProxyCommand='
# TODO(zhwu): write the template to a temp file, don't use
# the one in skypilot repo, to avoid changing the file when
# updating skypilot.
f'\'{sys.executable}{sky.__root_dir__}/templates/'
f'websocket_proxy.py '
f'{server_common.get_server_url()} '
f'{handle.cluster_name}\'')
credentials['ssh_proxy_command'] =proxy_command

cluster_utils.SSHConfigHelper.add_cluster(
handle.cluster_name,
handle.cached_external_ips,
credentials,
handle.cached_external_ssh_ports,
handle.docker_user,
handle.ssh_user,
 )

# Clean up SSH configs for clusters that do not exist.
#
# We do this in a conservative way: only when a query is made for all users
# or specific clusters. Without those, the table returned only contains the
# current user's clusters, and the information is not enough for
# removing clusters, because SkyPilot has no idea whether to remove
# ssh config of a cluster from another user.
clusters_exists=set(record['name'] forrecordincluster_records)
clusters_to_remove: Set[str] =set()
ifclustersisnotNone:
clusters_to_remove=set(clusters) -clusters_exists
elifall_users:
clusters_to_remove=set(cluster_utils.SSHConfigHelper.
list_cluster_names()) -clusters_exists

forcluster_nameinclusters_to_remove:
cluster_utils.SSHConfigHelper.remove_cluster(cluster_name)

returncluster_records


def_get_glob_storages(storages: List[str]) ->List[str]:
"""Returns a list of storages that match the glob pattern."""
glob_storages= []
forstorage_objectinstorages:
glob_storage=global_user_state.get_glob_storage_name(storage_object)
ifnotglob_storage:
click.echo(f'Storage {storage_object} not found.')
glob_storages.extend(glob_storage)
returnlist(set(glob_storages))


def_parse_env_var(env_var: str) ->Tuple[str, str]:
"""Parse env vars into a (KEY, VAL) pair."""
if'='notinenv_var:
value=os.environ.get(env_var)
ifvalueisNone:
raiseclick.UsageError(
f'{env_var} is not set in local environment.')
return (env_var, value)
ret=tuple(env_var.split('=', 1))
iflen(ret) !=2:
raiseclick.UsageError(
f'Invalid env var: {env_var}. Must be in the form of KEY=VAL '
'or KEY.')
returnret[0], ret[1]


def_async_call_or_wait(request_id: str, async_call: bool,
request_name: str) ->Any:
short_request_id=request_id[:8]
ifnotasync_call:
try:
returnsdk.stream_and_get(request_id)
exceptKeyboardInterrupt:
logger.info(
ux_utils.starting_message('Request will continue running '
'asynchronously.') +
f'\n{ux_utils.INDENT_SYMBOL}{colorama.Style.DIM}View logs: '
f'{ux_utils.BOLD}sky api logs {short_request_id}'
f'{colorama.Style.RESET_ALL}'
f'\n{ux_utils.INDENT_SYMBOL}{colorama.Style.DIM}Or, '
'visit: '
f'{server_common.get_server_url()}/api/stream?'
f'request_id={short_request_id}'
f'\n{ux_utils.INDENT_LAST_SYMBOL}{colorama.Style.DIM}To cancel '
'the request, run: '
f'{ux_utils.BOLD}sky api cancel {short_request_id}'
f'{colorama.Style.RESET_ALL}'
f'\n{colorama.Style.RESET_ALL}')
raise
else:
click.secho(f'Submitted {request_name} request: {request_id}',
fg='green')
click.echo(
f'{ux_utils.INDENT_SYMBOL}{colorama.Style.DIM}Check logs with: '
f'{ux_utils.BOLD}sky api logs {short_request_id}'
f'{colorama.Style.RESET_ALL}\n'
f'{ux_utils.INDENT_SYMBOL}{colorama.Style.DIM}Or, visit: '
f'{server_common.get_server_url()}/api/stream?'
f'request_id={short_request_id}'
f'\n{ux_utils.INDENT_LAST_SYMBOL}{colorama.Style.DIM}To cancel '
'the request, run: '
f'{ux_utils.BOLD}sky api cancel {short_request_id}'
f'{colorama.Style.RESET_ALL}\n')


def_merge_env_vars(env_dict: Optional[Dict[str, str]],
env_list: List[Tuple[str, str]]) ->List[Tuple[str, str]]:
"""Merges all values from env_list into env_dict."""
ifnotenv_dict:
returnenv_list
for (key, value) inenv_list:
env_dict[key] =value
returnlist(env_dict.items())


defconfig_option(expose_value: bool):
"""A decorator for the --config option.

 This decorator is used to parse the --config option.

 Any overrides specified in the command line will be applied to the skypilot
 config before the decorated function is called.

 If expose_value is True, the decorated function will receive the parsed
 config overrides as 'config_override' parameter.

 Args:
 expose_value: Whether to expose the value of the option to the decorated
 function.
 """

defpreprocess_config_options(ctx, param, value):
delctx# Unused.
param.name='config_override'
try:
iflen(value) ==0:
returnNone
else:
# Apply the config overrides to the skypilot config.
returnskypilot_config.apply_cli_config(value)
exceptValueErrorase:
raiseclick.BadParameter(f'{str(e)}') frome

defreturn_option_decorator(func):
returnclick.option(
'--config',
required=False,
type=str,
multiple=True,
expose_value=expose_value,
callback=preprocess_config_options,
help=('Path to a config file or a comma-separated '
'list of key-value pairs '
'(e.g. "nested.key1=val1,another.key2=val2").'),
 )(func)

returnreturn_option_decorator


_COMMON_OPTIONS= [
click.option('--async/--no-async',
'async_call',
required=False,
is_flag=True,
default=False,
help=('Run the command asynchronously.'))
]

_TASK_OPTIONS= [
click.option(
'--workdir',
required=False,
type=click.Path(exists=True, file_okay=False),
help=('If specified, sync this dir to the remote working directory, '
'where the task will be invoked. '
'Overrides the "workdir" config in the YAML if both are supplied.'
 )),
click.option(
'--cloud',
required=False,
type=str,
help=('The cloud to use. If specified, overrides the "resources.cloud" '
'config. Passing "none" resets the config.')),
click.option(
'--region',
required=False,
type=str,
help=('The region to use. If specified, overrides the '
'"resources.region" config. Passing "none" resets the config.')),
click.option(
'--zone',
required=False,
type=str,
help=('The zone to use. If specified, overrides the '
'"resources.zone" config. Passing "none" resets the config.')),
click.option(
'--num-nodes',
required=False,
type=int,
help=('Number of nodes to execute the task on. '
'Overrides the "num_nodes" config in the YAML if both are '
'supplied.')),
click.option(
'--cpus',
default=None,
type=str,
required=False,
help=('Number of vCPUs each instance must have (e.g., '
'``--cpus=4`` (exactly 4) or ``--cpus=4+`` (at least 4)). '
'This is used to automatically select the instance type.')),
click.option(
'--memory',
default=None,
type=str,
required=False,
help=(
'Amount of memory each instance must have in GB (e.g., '
'``--memory=16`` (exactly 16GB), ``--memory=16+`` (at least 16GB))'
 )),
click.option('--disk-size',
default=None,
type=int,
required=False,
help=('OS disk size in GBs.')),
click.option('--disk-tier',
default=None,
type=click.Choice(resources_utils.DiskTier.supported_tiers(),
case_sensitive=False),
required=False,
help=resources_utils.DiskTier.cli_help_message()),
click.option(
'--use-spot/--no-use-spot',
required=False,
default=None,
help=('Whether to request spot instances. If specified, overrides the '
'"resources.use_spot" config.')),
click.option('--image-id',
required=False,
default=None,
help=('Custom image id for launching the instances. '
'Passing "none" resets the config.')),
click.option('--env-file',
required=False,
type=dotenv.dotenv_values,
help="""\
 Path to a dotenv file with environment variables to set on the remote
 node.

 If any values from ``--env-file`` conflict with values set by
 ``--env``, the ``--env`` value will be preferred."""),
click.option(
'--env',
required=False,
type=_parse_env_var,
multiple=True,
help="""\
 Environment variable to set on the remote node.
 It can be specified multiple times.
 Examples:

\b
 1. ``--env MY_ENV=1``: set ``$MY_ENV`` on the cluster to be 1.

 2. ``--env MY_ENV2=$HOME``: set ``$MY_ENV2`` on the cluster to be the
 same value of ``$HOME`` in the local environment where the CLI command
 is run.

 3. ``--env MY_ENV3``: set ``$MY_ENV3`` on the cluster to be the
 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',
required=False,
type=str,
help=
 ('Type and number of GPUs to use. Example values: '
'"V100:8", "V100" (short for a count of 1), or "V100:0.5" '
'(fractional counts are supported by the scheduling framework). '
'If a new cluster is being launched by this command, this is the '
'resources to provision. If an existing cluster is being reused, this'
' is seen as the task demand, which must fit the cluster\'s total '
'resources and is used for scheduling the task. '
'Overrides the "accelerators" '
'config in the YAML if both are supplied. '
'Passing "none" resets the config.')),
click.option(
'--instance-type',
'-t',
required=False,
type=str,
help=('The instance type to use. If specified, overrides the '
'"resources.instance_type" config. Passing "none" resets the '
'config.'),
 ),
click.option(
'--ports',
required=False,
type=str,
multiple=True,
help=('Ports to open on the cluster. '
'If specified, overrides the "ports" config in the YAML. '),
 ),
]


def_complete_cluster_name(ctx: click.Context, param: click.Parameter,
incomplete: str) ->List[str]:
"""Handle shell completion for cluster names."""
delctx, param# Unused.
# TODO(zhwu): we send requests to API server for completion, which can cause
# large latency. We should investigate caching mechanism if needed.
response=requests_lib.get(
f'{server_common.get_server_url()}'
f'/api/completion/cluster_name?incomplete={incomplete}',
timeout=2.0,
 )
response.raise_for_status()
returnresponse.json()


def_complete_storage_name(ctx: click.Context, param: click.Parameter,
incomplete: str) ->List[str]:
"""Handle shell completion for storage names."""
delctx, param# Unused.
response=requests_lib.get(
f'{server_common.get_server_url()}'
f'/api/completion/storage_name?incomplete={incomplete}',
timeout=2.0,
 )
response.raise_for_status()
returnresponse.json()


def_complete_file_name(ctx: click.Context, param: click.Parameter,
incomplete: str) ->List[str]:
"""Handle shell completion for file names.

 Returns a special completion marker that tells click to use
 the shell's default file completion.
 """
delctx, param# Unused.
return [click.shell_completion.CompletionItem(incomplete, type='file')]


def_get_click_major_version():
returnint(click.__version__.split('.', maxsplit=1)[0])


def_get_shell_complete_args(complete_fn):
# The shell_complete argument is only valid on click >= 8.0.
if_get_click_major_version() >=8:
returndict(shell_complete=complete_fn)
return {}


_RELOAD_ZSH_CMD='source ~/.zshrc'
_RELOAD_BASH_CMD='source ~/.bashrc'


def_install_shell_completion(ctx: click.Context, param: click.Parameter,
value: str):
"""A callback for installing shell completion for click."""
delparam# Unused.
ifnotvalueorctx.resilient_parsing:
return

ifvalue=='auto':
if'SHELL'notinos.environ:
click.secho(
'Cannot auto-detect shell. Please specify shell explicitly.',
fg='red')
ctx.exit()
else:
value=os.path.basename(os.environ['SHELL'])

zshrc_diff='\n# For SkyPilot shell completion\n. ~/.sky/.sky-complete.zsh'
bashrc_diff= ('\n# For SkyPilot shell completion'
'\n. ~/.sky/.sky-complete.bash')

ifvalue=='bash':
install_cmd=f'_SKY_COMPLETE=bash_source sky > \
 ~/.sky/.sky-complete.bash && \
 echo "{bashrc_diff}" >> ~/.bashrc'

cmd= (f'(grep -q "SkyPilot" ~/.bashrc) || '
f'([[ ${{BASH_VERSINFO[0]}} -ge 4 ]] && ({install_cmd}) || '
f'(echo "Bash must be version 4 or above." && exit 1))')

reload_cmd=_RELOAD_BASH_CMD

elifvalue=='fish':
cmd='_SKY_COMPLETE=fish_source sky > \
 ~/.config/fish/completions/sky.fish'

# Fish does not need to be reloaded and will automatically pick up
# completions.
reload_cmd=None

elifvalue=='zsh':
install_cmd=f'_SKY_COMPLETE=zsh_source sky > \
 ~/.sky/.sky-complete.zsh && \
 echo "{zshrc_diff}" >> ~/.zshrc'

cmd=f'(grep -q "SkyPilot" ~/.zshrc) || ({install_cmd})'
reload_cmd=_RELOAD_ZSH_CMD

else:
click.secho(f'Unsupported shell: {value}', fg='red')
ctx.exit()

try:
subprocess.run(cmd,
shell=True,
check=True,
executable=shutil.which('bash'))
click.secho(f'Shell completion installed for {value}', fg='green')
ifreload_cmdisnotNone:
click.echo(
'Completion will take effect once you restart the terminal: '+
click.style(f'{reload_cmd}', bold=True))
exceptsubprocess.CalledProcessErrorase:
click.secho(f'> Installation failed with code {e.returncode}', fg='red')
ctx.exit()


def_uninstall_shell_completion(ctx: click.Context, param: click.Parameter,
value: str):
"""A callback for uninstalling shell completion for click."""
delparam# Unused.
ifnotvalueorctx.resilient_parsing:
return

ifvalue=='auto':
if'SHELL'notinos.environ:
click.secho(
'Cannot auto-detect shell. Please specify shell explicitly.',
fg='red')
ctx.exit()
else:
value=os.path.basename(os.environ['SHELL'])

ifvalue=='bash':
cmd='sed -i"" -e "/# For SkyPilot shell completion/d" ~/.bashrc && \
 sed -i"" -e "/sky-complete.bash/d" ~/.bashrc && \
 rm -f ~/.sky/.sky-complete.bash'

reload_cmd=_RELOAD_BASH_CMD

elifvalue=='fish':
cmd='rm -f ~/.config/fish/completions/sky.fish'
# Fish does not need to be reloaded and will automatically pick up
# completions.
reload_cmd=None

elifvalue=='zsh':
cmd='sed -i"" -e "/# For SkyPilot shell completion/d" ~/.zshrc && \
 sed -i"" -e "/sky-complete.zsh/d" ~/.zshrc && \
 rm -f ~/.sky/.sky-complete.zsh'

reload_cmd=_RELOAD_ZSH_CMD

else:
click.secho(f'Unsupported shell: {value}', fg='red')
ctx.exit()

try:
subprocess.run(cmd, shell=True, check=True)
click.secho(f'Shell completion uninstalled for {value}', fg='green')
ifreload_cmdisnotNone:
click.echo(
'Changes will take effect once you restart the terminal: '+
click.style(f'{reload_cmd}', bold=True))
exceptsubprocess.CalledProcessErrorase:
click.secho(f'> Uninstallation failed with code {e.returncode}',
fg='red')
ctx.exit()


def_add_click_options(options: List[click.Option]):
"""A decorator for adding a list of click option decorators."""

def_add_options(func):
foroptioninreversed(options):
func=option(func)
returnfunc

return_add_options


def_parse_override_params(
cloud: Optional[str] =None,
region: Optional[str] =None,
zone: Optional[str] =None,
gpus: Optional[str] =None,
cpus: Optional[str] =None,
memory: Optional[str] =None,
instance_type: Optional[str] =None,
use_spot: Optional[bool] =None,
image_id: Optional[str] =None,
disk_size: Optional[int] =None,
disk_tier: Optional[str] =None,
ports: Optional[Tuple[str, ...]] =None,
config_override: Optional[Dict[str, Any]] =None) ->Dict[str, Any]:
"""Parses the override parameters into a dictionary."""
override_params: Dict[str, Any] = {}
ifcloudisnotNone:
ifcloud.lower() =='none':
override_params['cloud'] =None
else:
override_params['cloud'] =registry.CLOUD_REGISTRY.from_str(cloud)
ifregionisnotNone:
ifregion.lower() =='none':
override_params['region'] =None
else:
override_params['region'] =region
ifzoneisnotNone:
ifzone.lower() =='none':
override_params['zone'] =None
else:
override_params['zone'] =zone
ifgpusisnotNone:
ifgpus.lower() =='none':
override_params['accelerators'] =None
else:
override_params['accelerators'] =gpus
ifcpusisnotNone:
ifcpus.lower() =='none':
override_params['cpus'] =None
else:
override_params['cpus'] =cpus
ifmemoryisnotNone:
ifmemory.lower() =='none':
override_params['memory'] =None
else:
override_params['memory'] =memory
ifinstance_typeisnotNone:
ifinstance_type.lower() =='none':
override_params['instance_type'] =None
else:
override_params['instance_type'] =instance_type
ifuse_spotisnotNone:
override_params['use_spot'] =use_spot
ifimage_idisnotNone:
ifimage_id.lower() =='none':
override_params['image_id'] =None
else:
override_params['image_id'] =image_id
ifdisk_sizeisnotNone:
override_params['disk_size'] =disk_size
ifdisk_tierisnotNone:
ifdisk_tier.lower() =='none':
override_params['disk_tier'] =None
else:
override_params['disk_tier'] =disk_tier
ifports:
ifany(p.lower() =='none'forpinports):
iflen(ports) >1:
withux_utils.print_exception_no_traceback():
raiseValueError('Cannot specify both "none" and other '
'ports.')
override_params['ports'] =None
else:
override_params['ports'] =ports
ifconfig_override:
override_params['_cluster_config_overrides'] =config_override
returnoverride_params


def_check_yaml(entrypoint: str) ->Tuple[bool, Optional[Dict[str, Any]]]:
"""Checks if entrypoint is a readable YAML file.

 Args:
 entrypoint: Path to a YAML file.
 """
is_yaml=True
config: Optional[List[Dict[str, Any]]] =None
result=None
shell_splits=shlex.split(entrypoint)
yaml_file_provided= (len(shell_splits) ==1and
 (shell_splits[0].endswith('yaml') or
shell_splits[0].endswith('.yml')))
invalid_reason=''
try:
withopen(entrypoint, 'r', encoding='utf-8') asf:
try:
config=list(yaml.safe_load_all(f))
ifconfig:
# FIXME(zongheng): in a chain DAG YAML it only returns the
# first section. OK for downstream but is weird.
result=config[0]
else:
result= {}
ifisinstance(result, str):
# 'sky exec cluster ./my_script.sh'
is_yaml=False
exceptyaml.YAMLErrorase:
ifyaml_file_provided:
logger.debug(e)
detailed_error=f'\nYAML Error: {e}\n'
invalid_reason= ('contains an invalid configuration. '
'Please check syntax.\n'
f'{detailed_error}')
is_yaml=False

exceptOSError:
ifyaml_file_provided:
entry_point_path=os.path.expanduser(entrypoint)
ifnotos.path.exists(entry_point_path):
invalid_reason= ('does not exist. Please check if the path'
' is correct.')
elifnotos.path.isfile(entry_point_path):
invalid_reason= ('is not a file. Please check if the path'
' is correct.')
else:
invalid_reason= ('yaml.safe_load() failed. Please check if the'
' path is correct.')
is_yaml=False
ifnotis_yaml:
ifyaml_file_provided:
click.confirm(
f'{entrypoint!r} looks like a yaml path but {invalid_reason}\n'
'It will be treated as a command to be run remotely. Continue?',
abort=True)
returnis_yaml, result


def_pop_and_ignore_fields_in_override_params(
params: Dict[str, Any], field_to_ignore: List[str]) ->None:
"""Pops and ignores fields in override params.

 Args:
 params: Override params.
 field_to_ignore: Fields to ignore.

 Returns:
 Override params with fields ignored.
 """
iffield_to_ignoreisnotNone:
forfieldinfield_to_ignore:
field_value=params.pop(field, None)
iffield_valueisnotNone:
click.secho(f'Override param {field}={field_value} is ignored.',
fg='yellow')


def_make_task_or_dag_from_entrypoint_with_overrides(
entrypoint: Tuple[str, ...],
*,
name: Optional[str] =None,
workdir: Optional[str] =None,
cloud: Optional[str] =None,
region: Optional[str] =None,
zone: Optional[str] =None,
gpus: Optional[str] =None,
cpus: Optional[str] =None,
memory: Optional[str] =None,
instance_type: Optional[str] =None,
num_nodes: Optional[int] =None,
use_spot: Optional[bool] =None,
image_id: Optional[str] =None,
disk_size: Optional[int] =None,
disk_tier: Optional[str] =None,
ports: Optional[Tuple[str, ...]] =None,
env: Optional[List[Tuple[str, str]]] =None,
field_to_ignore: Optional[List[str]] =None,
# job launch specific
job_recovery: Optional[str] =None,
config_override: Optional[Dict[str, Any]] =None,
) ->Union[sky.Task, sky.Dag]:
"""Creates a task or a dag from an entrypoint with overrides.

 Returns:
 A dag iff the entrypoint is YAML and contains more than 1 task.
 Otherwise, a task.
 """
entrypoint=' '.join(entrypoint)
is_yaml, _=_check_yaml(entrypoint)
entrypoint: Optional[str]
ifis_yaml:
# Treat entrypoint as a yaml.
click.secho('YAML to run: ', fg='cyan', nl=False)
click.secho(entrypoint)
else:
ifnotentrypoint:
entrypoint=None
else:
# Treat entrypoint as a bash command.
click.secho('Command to run: ', fg='cyan', nl=False)
click.secho(entrypoint)

override_params=_parse_override_params(cloud=cloud,
region=region,
zone=zone,
gpus=gpus,
cpus=cpus,
memory=memory,
instance_type=instance_type,
use_spot=use_spot,
image_id=image_id,
disk_size=disk_size,
disk_tier=disk_tier,
ports=ports,
config_override=config_override)
iffield_to_ignoreisnotNone:
_pop_and_ignore_fields_in_override_params(override_params,
field_to_ignore)

ifis_yaml:
assertentrypointisnotNone
usage_lib.messages.usage.update_user_task_yaml(entrypoint)
dag=dag_utils.load_chain_dag_from_yaml(entrypoint, env_overrides=env)
iflen(dag.tasks) >1:
# When the dag has more than 1 task. It is unclear how to
# override the params for the dag. So we just ignore the
# override params.
ifoverride_params:
click.secho(
f'WARNING: override params {override_params} are ignored, '
'since the yaml file contains multiple tasks.',
fg='yellow')
returndag
assertlen(dag.tasks) ==1, (
f'If you see this, please file an issue; tasks: {dag.tasks}')
task=dag.tasks[0]
else:
task=sky.Task(name='sky-cmd', run=entrypoint)
task.set_resources({sky.Resources()})
# env update has been done for DAG in load_chain_dag_from_yaml for YAML.
task.update_envs(env)

# Override.
ifworkdirisnotNone:
task.workdir=workdir

# job launch specific.
ifjob_recoveryisnotNone:
override_params['job_recovery'] =job_recovery

task.set_resources_override(override_params)

ifnum_nodesisnotNone:
task.num_nodes=num_nodes
ifnameisnotNone:
task.name=name
returntask


class_NaturalOrderGroup(click.Group):
"""Lists commands in the order defined in this script.

 Reference: https://github.com/pallets/click/issues/513
 """

deflist_commands(self, ctx): # pylint: disable=unused-argument
returnself.commands.keys()

@usage_lib.entrypoint('sky.cli', fallback=True)
definvoke(self, ctx):
returnsuper().invoke(ctx)


class_DocumentedCodeCommand(click.Command):
"""Corrects help strings for documented commands such that --help displays
 properly and code blocks are rendered in the official web documentation.
 """

defget_help(self, ctx):
help_str=ctx.command.help
ctx.command.help=help_str.replace('.. code-block:: bash\n', '\b')
returnsuper().get_help(ctx)


def_with_deprecation_warning(
f,
original_name: str,
alias_name: str,
override_command_argument: Optional[Dict[str, Any]] =None):

@functools.wraps(f)
defwrapper(self, *args, **kwargs):
override_str=''
ifoverride_command_argumentisnotNone:
overrides= []
fork, vinoverride_command_argument.items():
ifisinstance(v, bool):
ifv:
overrides.append(f'--{k}')
else:
overrides.append(f'--no-{k}')
else:
overrides.append(f'--{k.replace("_", "-")}={v}')
override_str=' with additional arguments '+' '.join(overrides)
click.secho(
f'WARNING: `{alias_name}` has been renamed to `{original_name}` '
f'and will be removed in a future release. Please use the '
f'latter{override_str} instead.\n',
err=True,
fg='yellow')
returnf(self, *args, **kwargs)

returnwrapper


def_override_arguments(callback, override_command_argument: Dict[str, Any]):

defwrapper(*args, **kwargs):
logger.info(f'Overriding arguments: {override_command_argument}')
kwargs.update(override_command_argument)
returncallback(*args, **kwargs)

returnwrapper


def_add_command_alias(
group: click.Group,
command: click.Command,
hidden: bool=False,
new_group: Optional[click.Group] =None,
new_command_name: Optional[str] =None,
override_command_argument: Optional[Dict[str, Any]] =None,
with_warning: bool=True,
) ->None: