-
Notifications
You must be signed in to change notification settings - Fork 1k
Expand file tree
/
Copy pathcli.py
More file actions
4066 lines (3623 loc) · 150 KB
/
cli.py
File metadata and controls
4066 lines (3623 loc) · 150 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
"""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
TODO:
- Add support for local Docker backend. Currently this module is very coupled
with CloudVmRayBackend, as seen by the many use of ray commands.
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.
"""
import copy
import datetime
import functools
import multiprocessing
import os
import shlex
import subprocess
import sys
import textwrap
import typing
from typing import Any, Dict, List, Optional, Sequence, Tuple
import click
import colorama
from rich import progress as rich_progress
import yaml
import sky
from sky import backends
from sky import check as sky_check
from sky import clouds
from sky import exceptions
from sky import global_user_state
from sky import sky_logging
from sky import spot as spot_lib
from sky import core
from sky.backends import backend_utils
from sky.backends import onprem_utils
from sky.benchmark import benchmark_state
from sky.benchmark import benchmark_utils
from sky.clouds import service_catalog
from sky.data import storage_utils
from sky.skylet import constants
from sky.skylet import job_lib
from sky.utils import log_utils
from sky.utils import common_utils
from sky.utils import command_runner
from sky.utils import schemas
from sky.utils import subprocess_utils
from sky.utils import timeline
from sky.utils import ux_utils
from sky.utils.cli_utils import status_utils
from sky.usage import usage_lib
if typing.TYPE_CHECKING:
from sky.backends import backend as backend_lib
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."""
_INTERACTIVE_NODE_TYPES = ('cpunode', 'gpunode', 'tpunode')
_INTERACTIVE_NODE_DEFAULT_RESOURCES = {
'cpunode': sky.Resources(cloud=None,
instance_type=None,
accelerators=None,
use_spot=False),
'gpunode': sky.Resources(cloud=None,
instance_type=None,
accelerators={'K80': 1},
use_spot=False),
'tpunode': sky.Resources(cloud=sky.GCP(),
instance_type=None,
accelerators={'tpu-v2-8': 1},
accelerator_args={'runtime_version': '2.5.0'},
use_spot=False),
}
# The maximum number of in-progress spot jobs to show in the status
# command.
_NUM_SPOT_JOBS_TO_SHOW_IN_STATUS = 5
def _get_glob_clusters(clusters: Sequence[str]) -> List[str]:
"""Returns a list of clusters that match the glob pattern."""
glob_clusters = []
for cluster in clusters:
glob_cluster = global_user_state.get_glob_cluster_names(cluster)
if len(glob_cluster) == 0:
if onprem_utils.check_if_local_cloud(cluster):
click.echo(
constants.UNINITIALIZED_ONPREM_CLUSTER_MESSAGE.format(
cluster=cluster))
else:
click.echo(f'Cluster {cluster} not found.')
glob_clusters.extend(glob_cluster)
return list(set(glob_clusters))
def _get_glob_storages(storages: Sequence[str]) -> List[str]:
"""Returns a list of storages that match the glob pattern."""
glob_storages = []
for storage_object in storages:
glob_storage = global_user_state.get_glob_storage_name(storage_object)
if len(glob_storage) == 0:
click.echo(f'Storage {storage_object} not found.')
else:
plural = 's' if len(glob_storage) > 1 else ''
click.echo(f'Deleting {len(glob_storage)} storage object{plural}.')
glob_storages.extend(glob_storage)
return list(set(glob_storages))
def _warn_if_local_cluster(cluster: str, local_clusters: List[str],
message: str) -> bool:
"""Raises warning if the cluster name is a local cluster."""
if cluster in local_clusters:
click.echo(message)
return False
return True
def _interactive_node_cli_command(cli_func):
"""Click command decorator for interactive node commands."""
assert cli_func.__name__ in _INTERACTIVE_NODE_TYPES, cli_func.__name__
cluster_option = click.option('--cluster',
'-c',
default=None,
type=str,
required=False,
help=_CLUSTER_FLAG_HELP)
port_forward_option = click.option(
'--port-forward',
'-p',
multiple=True,
default=[],
type=int,
required=False,
help=('Port to be forwarded. To forward multiple ports, '
'use this option multiple times.'))
screen_option = click.option('--screen',
default=False,
is_flag=True,
help='If true, attach using screen.')
tmux_option = click.option('--tmux',
default=False,
is_flag=True,
help='If true, attach using tmux.')
cloud_option = click.option('--cloud',
default=None,
type=str,
help='Cloud provider to use.')
instance_type_option = click.option('--instance-type',
'-t',
default=None,
type=str,
help='Instance type to use.')
cpus = click.option(
'--cpus',
default=None,
type=str,
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.'))
memory = 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))'))
gpus = click.option('--gpus',
default=None,
type=str,
help=('Type and number of GPUs to use '
'(e.g., ``--gpus=V100:8`` or ``--gpus=V100``).'))
tpus = click.option(
'--tpus',
default=None,
type=str,
help=('Type and number of TPUs to use (e.g., ``--tpus=tpu-v3-8:4`` or '
'``--tpus=tpu-v3-8``).'))
spot_option = click.option('--use-spot',
default=None,
is_flag=True,
help='If true, use spot instances.')
tpuvm_option = click.option('--tpu-vm',
default=False,
is_flag=True,
help='If true, use TPU VMs.')
disk_size = click.option('--disk-size',
default=None,
type=int,
required=False,
help=('OS disk size in GBs.'))
no_confirm = click.option('--yes',
'-y',
is_flag=True,
default=False,
required=False,
help='Skip confirmation prompt.')
idle_autostop = click.option('--idle-minutes-to-autostop',
'-i',
default=None,
type=int,
required=False,
help=('Automatically stop the cluster after '
'this many minutes of idleness, i.e. '
'no running or pending jobs in the '
'cluster\'s job queue. Idleness gets '
'reset whenever setting-up/running/'
'pending jobs are found in the job '
'queue. If not set, the cluster '
'will not be auto-stopped.'))
autodown = click.option('--down',
default=False,
is_flag=True,
required=False,
help=('Autodown the cluster: tear down the '
'cluster after all jobs finish '
'(successfully or abnormally). If '
'--idle-minutes-to-autostop is also set, '
'the cluster will be torn down after the '
'specified idle time. Note that if errors '
'occur during provisioning/data syncing/'
'setting up, the cluster will not be torn '
'down for debugging purposes.'))
retry_until_up = click.option('--retry-until-up',
'-r',
is_flag=True,
default=False,
required=False,
help=('Whether to retry provisioning '
'infinitely until the cluster is up '
'if we fail to launch the cluster on '
'any possible region/cloud due to '
'unavailability errors.'))
region_option = click.option('--region',
default=None,
type=str,
required=False,
help='The region to use.')
zone_option = click.option('--zone',
default=None,
type=str,
required=False,
help='The zone to use.')
click_decorators = [
cli.command(cls=_DocumentedCodeCommand),
cluster_option,
no_confirm,
port_forward_option,
idle_autostop,
autodown,
retry_until_up,
# Resource options
*([cloud_option] if cli_func.__name__ != 'tpunode' else []),
region_option,
zone_option,
instance_type_option,
cpus,
memory,
*([gpus] if cli_func.__name__ == 'gpunode' else []),
*([tpus] if cli_func.__name__ == 'tpunode' else []),
spot_option,
*([tpuvm_option] if cli_func.__name__ == 'tpunode' else []),
# Attach options
screen_option,
tmux_option,
disk_size,
]
decorator = functools.reduce(lambda res, f: f(res),
reversed(click_decorators), cli_func)
return decorator
def _parse_env_var(env_var: str) -> Tuple[str, str]:
"""Parse env vars into a (KEY, VAL) pair."""
if '=' not in env_var:
value = os.environ.get(env_var)
if value is None:
raise click.UsageError(
f'{env_var} is not set in local environment.')
return (env_var, value)
ret = tuple(env_var.split('=', 1))
if len(ret) != 2:
raise click.UsageError(
f'Invalid env var: {env_var}. Must be in the form of KEY=VAL '
'or KEY.')
return ret[0], ret[1]
_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,
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(
'--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',
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.""",
)
]
_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.'),
),
]
def _complete_cluster_name(ctx: click.Context, param: click.Parameter,
incomplete: str) -> List[str]:
"""Handle shell completion for cluster names."""
del ctx, param # Unused.
return global_user_state.get_cluster_names_start_with(incomplete)
def _complete_storage_name(ctx: click.Context, param: click.Parameter,
incomplete: str) -> List[str]:
"""Handle shell completion for storage names."""
del ctx, param # Unused.
return global_user_state.get_storage_names_start_with(incomplete)
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.
"""
del ctx, param # Unused.
return [click.shell_completion.CompletionItem(incomplete, type='file')]
def _get_click_major_version():
return int(click.__version__.split('.')[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:
return dict(shell_complete=complete_fn)
return {}
_RELOAD_ZSH_CMD = 'source ~/.zshrc'
_RELOAD_FISH_CMD = 'source ~/.config/fish/config.fish'
_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."""
del param # Unused.
if not value or ctx.resilient_parsing:
return
if value == 'auto':
if 'SHELL' not in os.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 = '# For SkyPilot shell completion\n. ~/.sky/.sky-complete.zsh'
bashrc_diff = '# For SkyPilot shell completion\n. ~/.sky/.sky-complete.bash'
if value == 'bash':
install_cmd = f'_SKY_COMPLETE=bash_source sky > \
~/.sky/.sky-complete.bash && \
echo "{bashrc_diff}" >> ~/.bashrc'
cmd = f'(grep -q "SkyPilot" ~/.bashrc) || ({install_cmd})'
reload_cmd = _RELOAD_BASH_CMD
elif value == 'fish':
cmd = '_SKY_COMPLETE=fish_source sky > \
~/.config/fish/completions/sky.fish'
reload_cmd = _RELOAD_FISH_CMD
elif value == '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)
click.secho(f'Shell completion installed for {value}', fg='green')
click.echo(
'Completion will take effect once you restart the terminal: ' +
click.style(f'{reload_cmd}', bold=True))
except subprocess.CalledProcessError as e:
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."""
del param # Unused.
if not value or ctx.resilient_parsing:
return
if value == 'auto':
if 'SHELL' not in os.environ:
click.secho(
'Cannot auto-detect shell. Please specify shell explicitly.',
fg='red')
ctx.exit()
else:
value = os.path.basename(os.environ['SHELL'])
if value == '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
elif value == 'fish':
cmd = 'rm -f ~/.config/fish/completions/sky.fish'
reload_cmd = _RELOAD_FISH_CMD
elif value == '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')
click.echo('Changes will take effect once you restart the terminal: ' +
click.style(f'{reload_cmd}', bold=True))
except subprocess.CalledProcessError as e:
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):
for option in reversed(options):
func = option(func)
return func
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) -> Dict[str, Any]:
"""Parses the override parameters into a dictionary."""
override_params: Dict[str, Any] = {}
if cloud is not None:
if cloud.lower() == 'none':
override_params['cloud'] = None
else:
override_params['cloud'] = clouds.CLOUD_REGISTRY.from_str(cloud)
if region is not None:
if region.lower() == 'none':
override_params['region'] = None
else:
override_params['region'] = region
if zone is not None:
if zone.lower() == 'none':
override_params['zone'] = None
else:
override_params['zone'] = zone
if gpus is not None:
if gpus.lower() == 'none':
override_params['accelerators'] = None
else:
override_params['accelerators'] = gpus
if cpus is not None:
if cpus.lower() == 'none':
override_params['cpus'] = None
else:
override_params['cpus'] = cpus
if memory is not None:
if memory.lower() == 'none':
override_params['memory'] = None
else:
override_params['memory'] = memory
if instance_type is not None:
if instance_type.lower() == 'none':
override_params['instance_type'] = None
else:
override_params['instance_type'] = instance_type
if use_spot is not None:
override_params['use_spot'] = use_spot
if image_id is not None:
if image_id.lower() == 'none':
override_params['image_id'] = None
else:
override_params['image_id'] = image_id
if disk_size is not None:
override_params['disk_size'] = disk_size
return override_params
def _default_interactive_node_name(node_type: str):
"""Returns a deterministic name to refer to the same node."""
# FIXME: this technically can collide in Azure/GCP with another
# same-username user. E.g., sky-gpunode-ubuntu. Not a problem on AWS
# which is the current cloud for interactive nodes.
assert node_type in _INTERACTIVE_NODE_TYPES, node_type
return f'sky-{node_type}-{backend_utils.get_cleaned_username()}'
def _infer_interactive_node_type(resources: sky.Resources):
"""Determine interactive node type from resources."""
accelerators = resources.accelerators
cloud = resources.cloud
if accelerators:
# We only support homogenous accelerators for now.
assert len(accelerators) == 1, resources
acc, _ = list(accelerators.items())[0]
is_gcp = cloud is not None and cloud.is_same_cloud(sky.GCP())
if is_gcp and 'tpu' in acc:
return 'tpunode'
return 'gpunode'
return 'cpunode'
def _check_resources_match(backend: backends.Backend,
cluster_name: str,
task: 'sky.Task',
node_type: Optional[str] = None) -> None:
"""Check matching resources when reusing an existing cluster.
The only exception is when [cpu|tpu|gpu]node -c cluster_name is used with no
additional arguments, then login succeeds.
Args:
cluster_name: The name of the cluster.
task: The task requested to be run on the cluster.
node_type: Only used for interactive node. Node type to attach to VM.
"""
handle = global_user_state.get_handle_from_cluster_name(cluster_name)
if handle is None:
return
if node_type is not None:
inferred_node_type = _infer_interactive_node_type(
handle.launched_resources)
if node_type != inferred_node_type:
name_arg = ''
if cluster_name != _default_interactive_node_name(
inferred_node_type):
name_arg = f' -c {cluster_name}'
raise click.UsageError(
f'Failed to attach to interactive node {cluster_name}. '
f'Please use: {colorama.Style.BRIGHT}'
f'sky {inferred_node_type}{name_arg}{colorama.Style.RESET_ALL}')
return
backend.check_resources_fit_cluster(handle, task)
def _launch_with_confirm(
task: sky.Task,
backend: backends.Backend,
cluster: Optional[str],
*,
dryrun: bool,
detach_run: bool,
detach_setup: bool = False,
no_confirm: bool = False,
idle_minutes_to_autostop: Optional[int] = None,
down: bool = False, # pylint: disable=redefined-outer-name
retry_until_up: bool = False,
no_setup: bool = False,
node_type: Optional[str] = None,
):
"""Launch a cluster with a Task."""
with sky.Dag() as dag:
dag.add(task)
if cluster is None:
cluster = backend_utils.generate_cluster_name()
maybe_status, _ = backend_utils.refresh_cluster_status_handle(cluster)
if maybe_status is None:
# Show the optimize log before the prompt if the cluster does not exist.
try:
backend_utils.check_public_cloud_enabled()
except RuntimeError as e:
# Catch the exception where the public cloud is not enabled, and
# only print the error message without the error type.
click.secho(e, fg='yellow')
sys.exit(1)
dag = sky.optimize(dag)
task = dag.tasks[0]
_check_resources_match(backend, cluster, task, node_type=node_type)
confirm_shown = False
if not no_confirm:
# Prompt if (1) --cluster is None, or (2) cluster doesn't exist, or (3)
# it exists but is STOPPED.
prompt = None
if maybe_status is None:
cluster_str = '' if cluster is None else f' {cluster!r}'
if onprem_utils.check_if_local_cloud(cluster):
prompt = f'Initializing local cluster{cluster_str}. Proceed?'
else:
prompt = f'Launching a new cluster{cluster_str}. Proceed?'
elif maybe_status == global_user_state.ClusterStatus.STOPPED:
prompt = f'Restarting the stopped cluster {cluster!r}. Proceed?'
if prompt is not None:
confirm_shown = True
click.confirm(prompt, default=True, abort=True, show_default=True)
if node_type is not None:
if maybe_status != global_user_state.ClusterStatus.UP:
click.secho(f'Setting up interactive node {cluster}...',
fg='yellow')
# We do not sky.launch if interactive node is already up, so we need
# to update idle timeout and autodown here.
elif idle_minutes_to_autostop is not None:
core.autostop(cluster, idle_minutes_to_autostop, down)
elif down:
core.autostop(cluster, 1, down)
elif not confirm_shown:
click.secho(f'Running task on cluster {cluster}...', fg='yellow')
if node_type is None or maybe_status != global_user_state.ClusterStatus.UP:
# No need to sky.launch again when interactive node is already up.
sky.launch(
dag,
dryrun=dryrun,
stream_logs=True,
cluster_name=cluster,
detach_setup=detach_setup,
detach_run=detach_run,
backend=backend,
idle_minutes_to_autostop=idle_minutes_to_autostop,
down=down,
retry_until_up=retry_until_up,
no_setup=no_setup,
)
# TODO: skip installing ray to speed up provisioning.
def _create_and_ssh_into_node(
node_type: str,
resources: sky.Resources,
cluster_name: str,
backend: Optional['backend_lib.Backend'] = None,
port_forward: Optional[List[int]] = None,
session_manager: Optional[str] = None,
user_requested_resources: Optional[bool] = False,
no_confirm: bool = False,
idle_minutes_to_autostop: Optional[int] = None,
down: bool = False, # pylint: disable=redefined-outer-name
retry_until_up: bool = False,
):
"""Creates and attaches to an interactive node.
Args:
node_type: Type of the interactive node: { 'cpunode', 'gpunode' }.
resources: Resources to attach to VM.
cluster_name: a cluster name to identify the interactive node.
backend: the Backend to use (currently only CloudVmRayBackend).
port_forward: List of ports to forward.
session_manager: Attach session manager: { 'screen', 'tmux' }.
user_requested_resources: If true, user requested resources explicitly.
no_confirm: If true, skips confirmation prompt presented to user.
idle_minutes_to_autostop: Automatically stop the cluster after
specified minutes of idleness. Idleness gets
reset whenever setting-up/running/pending
jobs are found in the job queue.
down: If true, autodown the cluster after all jobs finish. If
idle_minutes_to_autostop is also set, the cluster will be torn
down after the specified idle time.
retry_until_up: Whether to retry provisioning infinitely until the
cluster is up if we fail to launch due to
unavailability errors.
"""
assert node_type in _INTERACTIVE_NODE_TYPES, node_type
assert session_manager in (None, 'screen', 'tmux'), session_manager
if onprem_utils.check_if_local_cloud(cluster_name):
raise click.BadParameter(
f'Name {cluster_name!r} taken by a local cluster and cannot '
f'be used for a {node_type}.')
# TODO: Add conda environment replication
# should be setup =
# 'conda env export | grep -v "^prefix: " > environment.yml'
# && conda env create -f environment.yml
task = sky.Task(
node_type,
workdir=None,
setup=None,
)
task.set_resources(resources)
backend = backend if backend is not None else backends.CloudVmRayBackend()
if not isinstance(backend, backends.CloudVmRayBackend):
raise click.UsageError('Interactive nodes are only supported for '
f'{backends.CloudVmRayBackend.__name__} '
f'backend. Got {type(backend).__name__}.')
maybe_status, _ = backend_utils.refresh_cluster_status_handle(cluster_name)
if maybe_status is not None and user_requested_resources:
name_arg = ''
if cluster_name != _default_interactive_node_name(node_type):
name_arg = f' -c {cluster_name}'
raise click.UsageError(
'Resources cannot be specified for an existing interactive node '
f'{cluster_name!r}. To login to the cluster, use: '
f'{colorama.Style.BRIGHT}'
f'sky {node_type}{name_arg}{colorama.Style.RESET_ALL}')
_launch_with_confirm(
task,
backend,
cluster_name,
dryrun=False,
detach_run=True,
no_confirm=no_confirm,
idle_minutes_to_autostop=idle_minutes_to_autostop,
down=down,
retry_until_up=retry_until_up,
node_type=node_type,
)
handle = global_user_state.get_handle_from_cluster_name(cluster_name)
# Use ssh rather than 'ray attach' to suppress ray messages, speed up
# connection, and for allowing adding 'cd workdir' in the future.
# Disable check, since the returncode could be non-zero if the user Ctrl-D.
commands = []
if session_manager == 'screen':
commands += ['screen', '-D', '-R']
elif session_manager == 'tmux':
commands += ['tmux', 'attach', '||', 'tmux', 'new']
backend.run_on_head(handle,
commands,
port_forward=port_forward,
ssh_mode=command_runner.SshMode.LOGIN)
cluster_name = handle.cluster_name
click.echo('To attach to it again: ', nl=False)
if cluster_name == _default_interactive_node_name(node_type):
option = ''
else:
option = f' -c {cluster_name}'
click.secho(f'sky {node_type}{option}', bold=True)
click.echo('To stop the node:\t', nl=False)
click.secho(f'sky stop {cluster_name}', bold=True)
click.echo('To tear down the node:\t', nl=False)
click.secho(f'sky down {cluster_name}', bold=True)
click.echo('To upload a folder:\t', nl=False)
click.secho(f'rsync -rP /local/path {cluster_name}:/remote/path', bold=True)
click.echo('To download a folder:\t', nl=False)
click.secho(f'rsync -rP {cluster_name}:/remote/path /local/path', bold=True)
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 = None
shell_splits = shlex.split(entrypoint)
yaml_file_provided = len(shell_splits) == 1 and \
(shell_splits[0].endswith('yaml') or shell_splits[0].endswith('.yml'))
try:
with open(entrypoint, 'r') as f:
try:
config = yaml.safe_load(f)
if isinstance(config, str):
# 'sky exec cluster ./my_script.sh'
is_yaml = False
except yaml.YAMLError as e:
if yaml_file_provided:
logger.debug(e)
invalid_reason = ('contains an invalid configuration. '
' Please check syntax.')
is_yaml = False
except OSError:
if yaml_file_provided:
entry_point_path = os.path.expanduser(entrypoint)
if not os.path.exists(entry_point_path):
invalid_reason = ('does not exist. Please check if the path'
' is correct.')
elif not os.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
if not is_yaml:
if yaml_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)
return is_yaml, config
def _make_task_from_entrypoint_with_overrides(
entrypoint: List[str],
*,
name: Optional[str] = None,
cluster: 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,
env: Optional[List[Tuple[str, str]]] = None,
# spot launch specific
spot_recovery: Optional[str] = None,
) -> sky.Task:
entrypoint = ' '.join(entrypoint)
is_yaml, yaml_config = _check_yaml(entrypoint)
entrypoint: Optional[str]
if is_yaml:
# Treat entrypoint as a yaml.
click.secho('Task from YAML spec: ', fg='yellow', nl=False)
click.secho(entrypoint, bold=True)
else:
if not entrypoint:
entrypoint = None
else:
# Treat entrypoint as a bash command.
click.secho('Task from command: ', fg='yellow', nl=False)
click.secho(entrypoint, bold=True)
if onprem_utils.check_local_cloud_args(cloud, cluster, yaml_config):
cloud = 'local'
if is_yaml:
usage_lib.messages.usage.update_user_task_yaml(entrypoint)
task = sky.Task.from_yaml(entrypoint)
else:
task = sky.Task(name='sky-cmd', run=entrypoint)
task.set_resources({sky.Resources()})
# Override.
if workdir is not None:
task.workdir = workdir
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,