skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/execution.py

@@ -3,12 +3,12 @@
 See `Stage` for a Task's life cycle.
 """
 import enum
-import os
+import typing
 from typing import List, Optional, Tuple, Union
 
 import colorama
 
-import sky
+from sky import admin_policy
 from sky import backends
 from sky import clouds
 from sky import global_user_state
@@ -16,14 +16,19 @@ from sky import optimizer
 from sky import sky_logging
 from sky.backends import backend_utils
 from sky.usage import usage_lib
+from sky.utils import admin_policy_utils
+from sky.utils import common
 from sky.utils import controller_utils
 from sky.utils import dag_utils
-from sky.utils import env_options
+from sky.utils import resources_utils
 from sky.utils import rich_utils
-from sky.utils import subprocess_utils
+from sky.utils import status_lib
 from sky.utils import timeline
 from sky.utils import ux_utils
 
+if typing.TYPE_CHECKING:
+    import sky
+
 logger = sky_logging.init_logger(__name__)
 
 
@@ -55,8 +60,9 @@ def _maybe_clone_disk_from_cluster(clone_disk_from: Optional[str],
     with rich_utils.safe_status('Creating image from source cluster '
                                 f'{clone_disk_from!r}'):
         image_id = original_cloud.create_image_from_cluster(
-            clone_disk_from,
-            handle.cluster_name_on_cloud,
+            cluster_name=resources_utils.ClusterName(
+                display_name=clone_disk_from,
+                name_on_cloud=handle.cluster_name_on_cloud),
             region=handle.launched_resources.region,
             zone=handle.launched_resources.zone,
         )
@@ -98,7 +104,7 @@ def _execute(
     handle: Optional[backends.ResourceHandle] = None,
     backend: Optional[backends.Backend] = None,
     retry_until_up: bool = False,
-    optimize_target: optimizer.OptimizeTarget = optimizer.OptimizeTarget.COST,
+    optimize_target: common.OptimizeTarget = common.OptimizeTarget.COST,
     stages: Optional[List[Stage]] = None,
     cluster_name: Optional[str] = None,
     detach_setup: bool = False,
@@ -106,8 +112,10 @@ def _execute(
     idle_minutes_to_autostop: Optional[int] = None,
     no_setup: bool = False,
     clone_disk_from: Optional[str] = None,
+    skip_unnecessary_provisioning: bool = False,
     # Internal only:
     # pylint: disable=invalid-name
+    _quiet_optimizer: bool = False,
     _is_launched_by_jobs_controller: bool = False,
     _is_launched_by_sky_serve_controller: bool = False,
 ) -> Tuple[Optional[int], Optional[backends.ResourceHandle]]:
@@ -126,8 +134,9 @@ def _execute(
         Note that if errors occur during provisioning/data syncing/setting up,
         the cluster will not be torn down for debugging purposes.
       stream_logs: bool; whether to stream all tasks' outputs to the client.
-      handle: Optional[backends.ResourceHandle]; if provided, execution will use
-        an existing backend cluster handle instead of provisioning a new one.
+      handle: Optional[backends.ResourceHandle]; if provided, execution will
+        attempt to use an existing backend cluster handle instead of
+        provisioning a new one.
       backend: Backend; backend to use for executing the tasks. Defaults to
         CloudVmRayBackend()
       retry_until_up: bool; whether to retry the provisioning until the cluster
@@ -148,6 +157,11 @@ def _execute(
       idle_minutes_to_autostop: int; if provided, the cluster will be set to
         autostop after this many minutes of idleness.
       no_setup: bool; whether to skip setup commands or not when (re-)launching.
+      clone_disk_from: Optional[str]; if set, clone the disk from the specified
+        cluster.
+      skip_unecessary_provisioning: bool; if True, compare the calculated
+        cluster config to the current cluster's config. If they match, shortcut
+        provisioning even if we have Stage.PROVISION.
 
     Returns:
       job_id: Optional[int]; the job ID of the submitted job. None if the
@@ -156,21 +170,35 @@ def _execute(
       handle: Optional[backends.ResourceHandle]; the handle to the cluster. None
         if dryrun.
     """
+
     dag = dag_utils.convert_entrypoint_to_dag(entrypoint)
+    for task in dag.tasks:
+        if task.storage_mounts is not None:
+            for storage in task.storage_mounts.values():
+                # Ensure the storage is constructed.
+                storage.construct()
+    dag, _ = admin_policy_utils.apply(
+        dag,
+        request_options=admin_policy.RequestOptions(
+            cluster_name=cluster_name,
+            idle_minutes_to_autostop=idle_minutes_to_autostop,
+            down=down,
+            dryrun=dryrun,
+        ))
     assert len(dag) == 1, f'We support 1 task for now. {dag}'
     task = dag.tasks[0]
 
     if any(r.job_recovery is not None for r in task.resources):
-        with ux_utils.print_exception_no_traceback():
-            raise ValueError(
-                'Job recovery is specified in the task. To launch a '
-                'managed job, please use: sky jobs launch')
+        logger.warning(
+            f'{colorama.Style.DIM}The task has `job_recovery` specified, '
+            'but is launched as an unmanaged job. It will be ignored.'
+            'To enable job recovery, use managed jobs: sky jobs launch.'
+            f'{colorama.Style.RESET_ALL}')
 
     cluster_exists = False
     if cluster_name is not None:
-        existing_handle = global_user_state.get_handle_from_cluster_name(
-            cluster_name)
-        cluster_exists = existing_handle is not None
+        cluster_record = global_user_state.get_cluster_from_name(cluster_name)
+        cluster_exists = cluster_record is not None
         # TODO(woosuk): If the cluster exists, print a warning that
         # `cpus` and `memory` are not used as a job scheduling constraint,
         # unlike `gpus`.
@@ -206,7 +234,8 @@ def _execute(
                             '(after all jobs finish).'
                             f'{colorama.Style.RESET_ALL}')
                 idle_minutes_to_autostop = 1
-            stages.remove(Stage.DOWN)
+            if Stage.DOWN in stages:
+                stages.remove(Stage.DOWN)
             if idle_minutes_to_autostop >= 0:
                 requested_features.add(
                     clouds.CloudImplementationFeatures.AUTO_TERMINATE)
@@ -238,8 +267,8 @@ def _execute(
             bold = colorama.Style.BRIGHT
             reset = colorama.Style.RESET_ALL
             logger.info(
-                f'{yellow}Launching an unmanaged spot task, which does not '
-                f'automatically recover from preemptions.{reset}\n{yellow}To '
+                f'{yellow}Launching a spot job that does not '
+                f'automatically recover from preemptions. To '
                 'get automatic recovery, use managed job instead: '
                 f'{reset}{bold}sky jobs launch{reset} {yellow}or{reset} '
                 f'{bold}sky.jobs.launch(){reset}.')
@@ -253,7 +282,15 @@ def _execute(
                     # no-credential machine should not enter optimize(), which
                     # would directly error out ('No cloud is enabled...').  Fix
                     # by moving `sky check` checks out of optimize()?
-                    dag = sky.optimize(dag, minimize=optimize_target)
+                    controller = controller_utils.Controllers.from_name(
+                        cluster_name)
+                    if controller is not None:
+                        logger.info(
+                            f'Choosing resources for {controller.value.name}...'
+                        )
+                    dag = optimizer.Optimizer.optimize(dag,
+                                                       minimize=optimize_target,
+                                                       quiet=_quiet_optimizer)
                     task = dag.tasks[0]  # Keep: dag may have been deep-copied.
                     assert task.best_resources is not None, task
 
@@ -267,13 +304,18 @@ def _execute(
 
     try:
         if Stage.PROVISION in stages:
-            if handle is None:
-                handle = backend.provision(task,
-                                           task.best_resources,
-                                           dryrun=dryrun,
-                                           stream_logs=stream_logs,
-                                           cluster_name=cluster_name,
-                                           retry_until_up=retry_until_up)
+            assert handle is None or skip_unnecessary_provisioning, (
+                'Provisioning requested, but handle is already set. PROVISION '
+                'should be excluded from stages or '
+                'skip_unecessary_provisioning should be set. ')
+            handle = backend.provision(
+                task,
+                task.best_resources,
+                dryrun=dryrun,
+                stream_logs=stream_logs,
+                cluster_name=cluster_name,
+                retry_until_up=retry_until_up,
+                skip_unnecessary_provisioning=skip_unnecessary_provisioning)
 
         if handle is None:
             assert dryrun, ('If not dryrun, handle must be set or '
@@ -281,11 +323,18 @@ def _execute(
             logger.info('Dryrun finished.')
             return None, None
 
-        if Stage.SYNC_WORKDIR in stages and not dryrun:
-            if task.workdir is not None:
-                backend.sync_workdir(handle, task.workdir)
+        do_workdir = (Stage.SYNC_WORKDIR in stages and not dryrun and
+                      task.workdir is not None)
+        do_file_mounts = (Stage.SYNC_FILE_MOUNTS in stages and not dryrun and
+                          (task.file_mounts is not None or
+                           task.storage_mounts is not None))
+        if do_workdir or do_file_mounts:
+            logger.info(ux_utils.starting_message('Syncing files.'))
+
+        if do_workdir:
+            backend.sync_workdir(handle, task.workdir)
 
-        if Stage.SYNC_FILE_MOUNTS in stages and not dryrun:
+        if do_file_mounts:
             backend.sync_file_mounts(handle, task.file_mounts,
                                      task.storage_mounts)
 
@@ -318,23 +367,6 @@ def _execute(
                 backend.teardown_ephemeral_storage(task)
                 backend.teardown(handle, terminate=True)
     finally:
-        controller = controller_utils.Controllers.from_name(cluster_name)
-        if controller is None and not _is_launched_by_sky_serve_controller:
-            # UX: print live clusters to make users aware (to save costs).
-            #
-            # Don't print if this job is launched by the jobs controller,
-            # because managed jobs are serverless, there can be many of them,
-            # and users tend to continuously monitor managed jobs using `sky
-            # job queue`. Also don't print if this job is a skyserve controller
-            # job or launched by a skyserve controller job, because the
-            # redirect for this subprocess.run won't success and it will
-            # pollute the controller logs.
-            #
-            # Disable the usage collection for this status command.
-            env = dict(os.environ,
-                       **{env_options.Options.DISABLE_LOGGING.value: '1'})
-            subprocess_utils.run(
-                'sky status --no-show-managed-jobs --no-show-services', env=env)
         print()
         print('\x1b[?25h', end='')  # Show cursor.
     return job_id, handle
@@ -351,19 +383,19 @@ def launch(
     down: bool = False,
     stream_logs: bool = True,
     backend: Optional[backends.Backend] = None,
-    optimize_target: optimizer.OptimizeTarget = optimizer.OptimizeTarget.COST,
-    detach_setup: bool = False,
-    detach_run: bool = False,
+    optimize_target: common.OptimizeTarget = common.OptimizeTarget.COST,
     no_setup: bool = False,
     clone_disk_from: Optional[str] = None,
+    fast: bool = False,
     # Internal only:
     # pylint: disable=invalid-name
+    _quiet_optimizer: bool = False,
     _is_launched_by_jobs_controller: bool = False,
     _is_launched_by_sky_serve_controller: bool = False,
     _disable_controller_check: bool = False,
 ) -> Tuple[Optional[int], Optional[backends.ResourceHandle]]:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
-    """Launch a cluster or task.
+    """Launches a cluster or task.
 
     The task's setup and run commands are executed under the task's workdir
     (when specified, it is synced to remote cluster).  The task undergoes job
@@ -373,6 +405,16 @@ def launch(
     usage) a sky.Dag. In the latter case, currently it must contain a single
     task; support for pipelines/general DAGs are in experimental branches.
 
+    Example:
+        .. code-block:: python
+
+            import sky
+            task = sky.Task(run='echo hello SkyPilot')
+            task.set_resources(
+                sky.Resources(cloud=sky.AWS(), accelerators='V100:4'))
+            sky.launch(task, cluster_name='my-cluster')
+
+
     Args:
         task: sky.Task, or sky.Dag (experimental; 1-task only) to launch.
         cluster_name: name of the cluster to create/reuse.  If None,
@@ -384,7 +426,7 @@ def launch(
             cluster's job queue. Idleness gets reset whenever setting-up/
             running/pending jobs are found in the job queue. Setting this
             flag is equivalent to running
-            ``sky.launch(..., detach_run=True, ...)`` and then
+            ``sky.launch(...)`` and then
             ``sky.autostop(idle_minutes=<minutes>)``. If not set, the cluster
             will not be autostopped.
         down: Tear down the cluster after all jobs finish (successfully or
@@ -398,27 +440,12 @@ def launch(
             (CloudVMRayBackend).
         optimize_target: target to optimize for. Choices: OptimizeTarget.COST,
             OptimizeTarget.TIME.
-        detach_setup: If True, run setup in non-interactive mode as part of the
-            job itself. You can safely ctrl-c to detach from logging, and it
-            will not interrupt the setup process. To see the logs again after
-            detaching, use `sky logs`. To cancel setup, cancel the job via
-            `sky cancel`. Useful for long-running setup
-            commands.
-        detach_run: If True, as soon as a job is submitted, return from this
-            function and do not stream execution logs.
         no_setup: if True, do not re-run setup commands.
         clone_disk_from: [Experimental] if set, clone the disk from the
             specified cluster. This is useful to migrate the cluster to a
             different availability zone or region.
-
-    Example:
-        .. code-block:: python
-
-            import sky
-            task = sky.Task(run='echo hello SkyPilot')
-            task.set_resources(
-                sky.Resources(cloud=sky.AWS(), accelerators='V100:4'))
-            sky.launch(task, cluster_name='my-cluster')
+        fast: [Experimental] If the cluster is already up and available,
+            skip provisioning and setup steps.
 
     Raises:
         exceptions.ClusterOwnerIdentityMismatchError: if the cluster is
@@ -448,26 +475,78 @@ def launch(
       handle: Optional[backends.ResourceHandle]; the handle to the cluster. None
         if dryrun.
     """
+
     entrypoint = task
+    entrypoint.validate()
     if not _disable_controller_check:
         controller_utils.check_cluster_name_not_controller(
             cluster_name, operation_str='sky.launch')
 
+    handle = None
+    stages = None
+    skip_unnecessary_provisioning = False
+    # Check if cluster exists and we are doing fast provisioning
+    if fast and cluster_name is not None:
+        cluster_status, maybe_handle = (
+            backend_utils.refresh_cluster_status_handle(cluster_name))
+        if cluster_status == status_lib.ClusterStatus.INIT:
+            # If the cluster is INIT, it may be provisioning. We want to prevent
+            # concurrent calls from queueing up many sequential reprovision
+            # attempts. Since provisioning will hold the cluster status lock, we
+            # wait to hold that lock by force refreshing the status. This will
+            # block until the cluster finishes provisioning, then correctly see
+            # that it is UP.
+            # TODO(cooperc): If multiple processes launched in parallel see that
+            # the cluster is STOPPED or does not exist, they will still all try
+            # to provision it, since we do not hold the lock continuously from
+            # the status check until the provision call. Fixing this requires a
+            # bigger refactor.
+            cluster_status, maybe_handle = (
+                backend_utils.refresh_cluster_status_handle(
+                    cluster_name,
+                    force_refresh_statuses=[
+                        # If the cluster is INIT, we want to try to grab the
+                        # status lock, which should block until provisioning is
+                        # finished.
+                        status_lib.ClusterStatus.INIT,
+                    ],
+                    # Wait indefinitely to obtain the lock, so that we don't
+                    # have multiple processes launching the same cluster at
+                    # once.
+                    cluster_status_lock_timeout=-1,
+                ))
+        if cluster_status == status_lib.ClusterStatus.UP:
+            handle = maybe_handle
+            stages = [
+                # Provisioning will be short-circuited if the existing
+                # cluster config hash matches the calculated one.
+                Stage.PROVISION,
+                Stage.SYNC_WORKDIR,
+                Stage.SYNC_FILE_MOUNTS,
+                Stage.PRE_EXEC,
+                Stage.EXEC,
+                Stage.DOWN,
+            ]
+            skip_unnecessary_provisioning = True
+
     return _execute(
         entrypoint=entrypoint,
         dryrun=dryrun,
         down=down,
         stream_logs=stream_logs,
-        handle=None,
+        handle=handle,
         backend=backend,
         retry_until_up=retry_until_up,
         optimize_target=optimize_target,
+        stages=stages,
         cluster_name=cluster_name,
-        detach_setup=detach_setup,
-        detach_run=detach_run,
+        detach_setup=True,
+        detach_run=True,
         idle_minutes_to_autostop=idle_minutes_to_autostop,
         no_setup=no_setup,
         clone_disk_from=clone_disk_from,
+        skip_unnecessary_provisioning=skip_unnecessary_provisioning,
+        _quiet_optimizer=_quiet_optimizer,
         _is_launched_by_jobs_controller=_is_launched_by_jobs_controller,
         _is_launched_by_sky_serve_controller=
         _is_launched_by_sky_serve_controller,
@@ -482,10 +561,9 @@ def exec(  # pylint: disable=redefined-builtin
     down: bool = False,
     stream_logs: bool = True,
     backend: Optional[backends.Backend] = None,
-    detach_run: bool = False,
 ) -> Tuple[Optional[int], Optional[backends.ResourceHandle]]:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
-    """Execute a task on an existing cluster.
+    """Executes a task on an existing cluster.
 
     This function performs two actions:
 
@@ -520,12 +598,11 @@ def exec(  # pylint: disable=redefined-builtin
         stream_logs: if True, show the logs in the terminal.
         backend: backend to use.  If None, use the default backend
             (CloudVMRayBackend).
-        detach_run: if True, detach from logging once the task has been
-            submitted.
 
     Raises:
-        ValueError: if the specified cluster does not exist or is not in UP
-            status.
+        ValueError: if the specified cluster is not in UP status.
+        sky.exceptions.ClusterDoesNotExist: if the specified cluster does not
+            exist.
         sky.exceptions.NotSupportedError: if the specified cluster is a
             controller that does not support this operation.
 
@@ -537,11 +614,7 @@ def exec(  # pylint: disable=redefined-builtin
         if dryrun.
     """
     entrypoint = task
-    if isinstance(entrypoint, sky.Dag):
-        logger.warning(
-            f'{colorama.Fore.YELLOW}Passing a sky.Dag to sky.exec() is '
-            'deprecated. Pass sky.Task instead.'
-            f'{colorama.Style.RESET_ALL}')
+    entrypoint.validate(workdir_only=True)
     controller_utils.check_cluster_name_not_controller(cluster_name,
                                                        operation_str='sky.exec')
 
@@ -562,5 +635,5 @@ def exec(  # pylint: disable=redefined-builtin
             Stage.EXEC,
         ],
         cluster_name=cluster_name,
-        detach_run=detach_run,
+        detach_run=True,
     )