skypilot-nightly 1.0.0.dev20250812__py3-none-any.whl → 1.0.0.dev20250815__py3-none-any.whl

sky/setup_files/dependencies.py

@@ -88,17 +88,18 @@ local_ray = [
     'ray[default] >= 2.2.0, != 2.6.0',
 ]
 
+# See requirements-dev.txt for the version of grpc and protobuf
+# used to generate the code during development.
 remote = [
-    # Adopted from ray's setup.py:
-    # https://github.com/ray-project/ray/blob/ray-2.9.3/python/setup.py#L251-L252
-    # SkyPilot: != 1.48.0 is required to avoid the error where ray dashboard
-    # fails to start when ray start is called (#2054).
-    # Tracking issue: https://github.com/ray-project/ray/issues/30984
-    'grpcio >= 1.32.0, != 1.48.0; python_version < \'3.10\'',
-    'grpcio >= 1.42.0, != 1.48.0; python_version >= \'3.10\'',
-    # Adopted from ray's setup.py:
-    # https://github.com/ray-project/ray/blob/ray-2.9.3/python/setup.py#L343
-    'protobuf >= 3.15.3, != 3.19.5',
+    # The grpc version at runtime has to be newer than the version
+    # used to generate the code.
+    'grpcio>=1.63.0',
+    # >= 5.26.1 because the runtime version can't be older than the version
+    # used to generate the code.
+    # < 7.0.0 because code generated for a major version V will be supported by
+    # protobuf runtimes of version V and V+1.
+    # https://protobuf.dev/support/cross-version-runtime-guarantee
+    'protobuf >= 5.26.1, < 7.0.0',
 ]
 
 # NOTE: Change the templates/jobs-controller.yaml.j2 file if any of the

sky/skylet/autostop_lib.py

@@ -16,8 +16,13 @@ from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
     import psutil
+
+    from sky.schemas.generated import autostopv1_pb2
 else:
     psutil = adaptors_common.LazyImport('psutil')
+    # To avoid requiring protobuf to be installed on the client side.
+    autostopv1_pb2 = adaptors_common.LazyImport(
+        'sky.schemas.generated.autostopv1_pb2')
 
 logger = sky_logging.init_logger(__name__)
 
@@ -55,11 +60,9 @@ Determines the condition for resetting the idleness timer.
 This option works in conjunction with ``--{pair}``. Options:
 
 \b
-1. ``jobs_and_ssh`` (default): Wait for all jobs to complete AND all SSH
-sessions to disconnect.
-2. ``jobs``: Wait for all jobs to complete.
-3. ``none``: Stop immediately after idle time expires, regardless of running
-jobs or SSH connections."""
+1. ``jobs_and_ssh`` (default): Wait for in-progress jobs and SSH connections to finish.
+2. ``jobs``: Only wait for in-progress jobs.
+3. ``none``: Wait for nothing; autostop right after ``{pair}``."""
 
     @classmethod
     def from_str(cls, mode: str) -> 'AutostopWaitFor':
@@ -78,6 +81,36 @@ jobs or SSH connections."""
                                  f'\'{cls.JOBS.value}\', or '
                                  f'\'{cls.NONE.value}\'. ')
 
+    @classmethod
+    def from_protobuf(
+        cls, protobuf_value: 'autostopv1_pb2.AutostopWaitFor'
+    ) -> Optional['AutostopWaitFor']:
+        """Convert protobuf AutostopWaitFor enum to Python enum value."""
+        protobuf_to_enum = {
+            autostopv1_pb2.AUTOSTOP_WAIT_FOR_JOBS_AND_SSH: cls.JOBS_AND_SSH,
+            autostopv1_pb2.AUTOSTOP_WAIT_FOR_JOBS: cls.JOBS,
+            autostopv1_pb2.AUTOSTOP_WAIT_FOR_NONE: cls.NONE,
+            autostopv1_pb2.AUTOSTOP_WAIT_FOR_UNSPECIFIED: None,
+        }
+        if protobuf_value not in protobuf_to_enum:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    f'Unknown protobuf AutostopWaitFor value: {protobuf_value}')
+        return protobuf_to_enum[protobuf_value]
+
+    def to_protobuf(self) -> 'autostopv1_pb2.AutostopWaitFor':
+        """Convert this Python enum value to protobuf enum value."""
+        enum_to_protobuf = {
+            AutostopWaitFor.JOBS_AND_SSH:
+                autostopv1_pb2.AUTOSTOP_WAIT_FOR_JOBS_AND_SSH,
+            AutostopWaitFor.JOBS: autostopv1_pb2.AUTOSTOP_WAIT_FOR_JOBS,
+            AutostopWaitFor.NONE: autostopv1_pb2.AUTOSTOP_WAIT_FOR_NONE,
+        }
+        if self not in enum_to_protobuf:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(f'Unknown AutostopWaitFor value: {self}')
+        return enum_to_protobuf[self]
+
 
 DEFAULT_AUTOSTOP_WAIT_FOR: AutostopWaitFor = AutostopWaitFor.JOBS_AND_SSH
 

sky/skylet/constants.py

@@ -90,12 +90,14 @@ TASK_ID_LIST_ENV_VAR = f'{SKYPILOT_ENV_VAR_PREFIX}TASK_IDS'
 # cluster yaml is updated.
 #
 # TODO(zongheng,zhanghao): make the upgrading of skylet automatic?
-SKYLET_VERSION = '16'
+SKYLET_VERSION = '17'
 # The version of the lib files that skylet/jobs use. Whenever there is an API
 # change for the job_lib or log_lib, we need to bump this version, so that the
 # user can be notified to update their SkyPilot version on the remote cluster.
 SKYLET_LIB_VERSION = 4
 SKYLET_VERSION_FILE = '~/.sky/skylet_version'
+SKYLET_GRPC_PORT = 46590
+SKYLET_GRPC_TIMEOUT_SECONDS = 5
 
 # Docker default options
 DEFAULT_DOCKER_CONTAINER_NAME = 'sky_container'

sky/skylet/services.py

@@ -0,0 +1,44 @@
+"""gRPC service implementations for skylet."""
+
+import grpc
+
+from sky import sky_logging
+from sky.schemas.generated import autostopv1_pb2
+from sky.schemas.generated import autostopv1_pb2_grpc
+from sky.skylet import autostop_lib
+
+logger = sky_logging.init_logger(__name__)
+
+
+class AutostopServiceImpl(autostopv1_pb2_grpc.AutostopServiceServicer):
+    """Implementation of the AutostopService gRPC service."""
+
+    def SetAutostop(  # type: ignore[return]
+            self, request: autostopv1_pb2.SetAutostopRequest,
+            context: grpc.ServicerContext
+    ) -> autostopv1_pb2.SetAutostopResponse:
+        """Sets autostop configuration for the cluster."""
+        try:
+            wait_for = autostop_lib.AutostopWaitFor.from_protobuf(
+                request.wait_for)
+            autostop_lib.set_autostop(
+                idle_minutes=request.idle_minutes,
+                backend=request.backend,
+                wait_for=wait_for if wait_for is not None else
+                autostop_lib.DEFAULT_AUTOSTOP_WAIT_FOR,
+                down=request.down)
+            return autostopv1_pb2.SetAutostopResponse()
+        except Exception as e:  # pylint: disable=broad-except
+            context.abort(grpc.StatusCode.INTERNAL, str(e))
+
+    def IsAutostopping(  # type: ignore[return]
+            self, request: autostopv1_pb2.IsAutostoppingRequest,
+            context: grpc.ServicerContext
+    ) -> autostopv1_pb2.IsAutostoppingResponse:
+        """Checks if the cluster is currently autostopping."""
+        try:
+            is_autostopping = autostop_lib.get_is_autostopping()
+            return autostopv1_pb2.IsAutostoppingResponse(
+                is_autostopping=is_autostopping)
+        except Exception as e:  # pylint: disable=broad-except
+            context.abort(grpc.StatusCode.INTERNAL, str(e))

sky/skylet/skylet.py

@@ -1,11 +1,17 @@
 """skylet: a daemon running on the head node of a cluster."""
 
+import concurrent.futures
+import os
 import time
 
+import grpc
+
 import sky
 from sky import sky_logging
+from sky.schemas.generated import autostopv1_pb2_grpc
 from sky.skylet import constants
 from sky.skylet import events
+from sky.skylet import services
 
 # Use the explicit logger name so that the logger is under the
 # `sky.skylet.skylet` namespace when executed directly, so as
@@ -31,7 +37,46 @@ EVENTS = [
     events.UsageHeartbeatReportEvent(),
 ]
 
-while True:
-    time.sleep(events.EVENT_CHECKING_INTERVAL_SECONDS)
-    for event in EVENTS:
-        event.run()
+
+def start_grpc_server(port: int = constants.SKYLET_GRPC_PORT) -> grpc.Server:
+    """Start the gRPC server."""
+    # This is the default value in Python 3.8 - 3.12,
+    # putting it here for visibility.
+    # TODO(kevin): Determine the optimal max number of threads.
+    max_workers = min(32, (os.cpu_count() or 1) + 4)
+    server = grpc.server(
+        concurrent.futures.ThreadPoolExecutor(max_workers=max_workers))
+
+    autostopv1_pb2_grpc.add_AutostopServiceServicer_to_server(
+        services.AutostopServiceImpl(), server)
+
+    listen_addr = f'127.0.0.1:{port}'
+    server.add_insecure_port(listen_addr)
+
+    server.start()
+    logger.info(f'gRPC server started on {listen_addr}')
+
+    return server
+
+
+def run_event_loop():
+    """Run the existing event loop."""
+
+    while True:
+        time.sleep(events.EVENT_CHECKING_INTERVAL_SECONDS)
+        for event in EVENTS:
+            event.run()
+
+
+def main():
+    grpc_server = start_grpc_server()
+    try:
+        run_event_loop()
+    except KeyboardInterrupt:
+        logger.info('Shutting down skylet...')
+    finally:
+        grpc_server.stop(grace=5)
+
+
+if __name__ == '__main__':
+    main()

sky/task.py

@@ -10,26 +10,25 @@ from typing import (Any, Callable, Dict, Iterable, List, Optional, Set, Tuple,
 
 import colorama
 
-import sky
 from sky import clouds
+from sky import dag as dag_lib
 from sky import exceptions
+from sky import resources as resources_lib
 from sky import sky_logging
 from sky.adaptors import common as adaptors_common
-import sky.dag
 from sky.data import data_utils
 from sky.data import storage as storage_lib
 from sky.provision import docker_utils
 from sky.serve import service_spec
 from sky.skylet import constants
 from sky.utils import common_utils
+from sky.utils import registry
 from sky.utils import schemas
 from sky.utils import ux_utils
 from sky.utils import volume as volume_lib
 
 if typing.TYPE_CHECKING:
     import yaml
-
-    from sky import resources as resources_lib
 else:
     yaml = adaptors_common.LazyImport('yaml')
 
@@ -382,26 +381,28 @@ class Task:
         self.estimated_inputs_size_gigabytes: Optional[float] = None
         self.estimated_outputs_size_gigabytes: Optional[float] = None
         # Default to CPU VM
-        self.resources: Union[List[sky.Resources],
-                              Set[sky.Resources]] = {sky.Resources()}
+        self.resources: Union[List['resources_lib.Resources'],
+                              Set['resources_lib.Resources']] = {
+                                  resources_lib.Resources()
+                              }
         self._service: Optional[service_spec.SkyServiceSpec] = None
 
         # Resources that this task cannot run on.
         self.blocked_resources = blocked_resources
 
-        self.time_estimator_func: Optional[Callable[['sky.Resources'],
+        self.time_estimator_func: Optional[Callable[['resources_lib.Resources'],
                                                     int]] = None
         self.file_mounts: Optional[Dict[str, str]] = None
 
         # Only set when 'self' is a jobs controller task: 'self.managed_job_dag'
         # is the underlying managed job dag (sky.Dag object).
-        self.managed_job_dag: Optional['sky.Dag'] = None
+        self.managed_job_dag: Optional['dag_lib.Dag'] = None
 
         # Only set when 'self' is a sky serve controller task.
         self.service_name: Optional[str] = None
 
         # Filled in by the optimizer.  If None, this Task is not planned.
-        self.best_resources: Optional[sky.Resources] = None
+        self.best_resources: Optional['resources_lib.Resources'] = None
 
         # For internal use only.
         self.file_mounts_mapping: Optional[Dict[str,
@@ -418,7 +419,7 @@ class Task:
         if file_mounts is not None:
             self.set_file_mounts(file_mounts)
 
-        dag = sky.dag.get_current_dag()
+        dag = dag_lib.get_current_dag()
         if dag is not None:
             dag.add(self)
 
@@ -783,7 +784,8 @@ class Task:
                 '_cluster_config_overrides'] = cluster_config_override
         if volumes:
             resources_config['volumes'] = volumes
-        task.set_resources(sky.Resources.from_yaml_config(resources_config))
+        task.set_resources(
+            resources_lib.Resources.from_yaml_config(resources_config))
 
         service = config.pop('service', None)
         pool = config.pop('pool', None)
@@ -931,7 +933,8 @@ class Task:
         for key, (vol_name, vol_req) in topology.items():
             if vol_req is not None:
                 if key == 'cloud':
-                    override_params[key] = sky.CLOUD_REGISTRY.from_str(vol_req)
+                    override_params[key] = registry.CLOUD_REGISTRY.from_str(
+                        vol_req)
                 else:
                     override_params[key] = vol_req
         self.set_resources_override(override_params)
@@ -1142,7 +1145,7 @@ class Task:
         Returns:
           self: The current task, with resources set.
         """
-        if isinstance(resources, sky.Resources):
+        if isinstance(resources, resources_lib.Resources):
             resources = {resources}
         # TODO(woosuk): Check if the resources are None.
         self.resources = _with_docker_login_config(resources, self.envs,
@@ -1187,8 +1190,8 @@ class Task:
         self._service = service
         return self
 
-    def set_time_estimator(self, func: Callable[['sky.Resources'],
-                                                int]) -> 'Task':
+    def set_time_estimator(
+            self, func: Callable[['resources_lib.Resources'], int]) -> 'Task':
         """Sets a func mapping resources to estimated time (secs).
 
         This is EXPERIMENTAL.
@@ -1712,7 +1715,7 @@ class Task:
         return required_features
 
     def __rshift__(self, b):
-        sky.dag.get_current_dag().add_edge(self, b)
+        dag_lib.get_current_dag().add_edge(self, b)
 
     def __repr__(self):
         if isinstance(self.run, str):

sky/templates/aws-ray.yml.j2

@@ -50,7 +50,7 @@ provider:
   disable_launch_config_check: true
 
 auth:
-  ssh_user: ubuntu
+  ssh_user: {{ssh_user}}
   ssh_private_key: {{ssh_private_key}}
 {% if ssh_proxy_command is not none %}
   ssh_proxy_command: {{ssh_proxy_command}}
@@ -68,7 +68,7 @@ available_node_types:
       ImageId: {{image_id}}  # Deep Learning AMI (Ubuntu 18.04); see aws.py.
       # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ec2.html#EC2.ServiceResource.create_instances
       BlockDeviceMappings:
-        - DeviceName: /dev/sda1
+        - DeviceName: {{root_device_name}}
           Ebs:
             VolumeSize: {{disk_size}}
             VolumeType: {{disk_tier}}

sky/templates/jobs-controller.yaml.j2

@@ -15,6 +15,12 @@ file_mounts:
   {{controller_file_mount_path}}: {{local_file_mount_path}}
   {%- endfor %}
 
+# NOTE(dev): This needs to be a subset of sky/templates/sky-serve-controller.yaml.j2.
+# It is because we use the --fast flag to submit jobs and no --fast flag to launch pools.
+# So when we launch a new pool, it will install the required dependencies.
+# TODO(tian): Add --fast to launch pools as well, and figure out the dependency installation.
+# Maybe in the --fast implementation, we can store the hash of setup commands that used to be
+# run and don't skip setup phase if the hash is different.
 setup: |
   {{ sky_activate_python_env }}
   # Disable the pip version check to avoid the warning message, which makes the

sky/templates/kubernetes-ray.yml.j2

@@ -378,6 +378,7 @@ available_node_types:
         {% if volume_mounts %}
         securityContext:
           fsGroup: 1000
+          fsGroupChangePolicy: OnRootMismatch
         {% endif %}
 
         # Add node selector if GPU/TPUs are requested:

sky/utils/command_runner.py

@@ -674,7 +674,7 @@ class SSHCommandRunner(CommandRunner):
             ssh += ['-tt']
         if port_forward is not None:
             for local, remote in port_forward:
-                logger.info(
+                logger.debug(
                     f'Forwarding local port {local} to remote port {remote}.')
                 ssh += ['-NL', f'{local}:localhost:{remote}']
         if self._docker_ssh_proxy_command is not None:

sky/utils/common_utils.py

@@ -1,6 +1,7 @@
 """Utils shared between all of sky"""
 
 import difflib
+import enum
 import functools
 import getpass
 import hashlib
@@ -55,6 +56,25 @@ _VALID_ENV_VAR_REGEX = '[a-zA-Z_][a-zA-Z0-9_]*'
 logger = sky_logging.init_logger(__name__)
 
 
+class ProcessStatus(enum.Enum):
+    """Process status."""
+
+    # The process is scheduled to run, but not started yet.
+    SCHEDULED = 'SCHEDULED'
+
+    # The process is running
+    RUNNING = 'RUNNING'
+
+    # The process is finished and succeeded
+    SUCCEEDED = 'SUCCEEDED'
+
+    # The process is interrupted
+    INTERRUPTED = 'INTERRUPTED'
+
+    # The process failed
+    FAILED = 'FAILED'
+
+
 @annotations.lru_cache(scope='request')
 def get_usage_run_id() -> str:
     """Returns a unique run id for each 'run'.

sky/utils/config_utils.py

@@ -8,6 +8,26 @@ logger = sky_logging.init_logger(__name__)
 
 _REGION_CONFIG_CLOUDS = ['nebius', 'oci']
 
+# Kubernetes API use list to represent dictionary fields with patch strategy
+# merge and each item is indexed by the patch merge key. The following map
+# maps the field name to the patch merge key.
+# pylint: disable=line-too-long
+# Ref: https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.27/#podspec-v1-core
+# NOTE: field containers and imagePullSecrets are not included deliberately for
+# backward compatibility (we only support one container per pod now).
+_PATCH_MERGE_KEYS = {
+    'initContainers': 'name',
+    'ephemeralContainers': 'name',
+    'volumes': 'name',
+    'volumeMounts': 'name',
+    'resourceClaims': 'name',
+    'env': 'name',
+    'hostAliases': 'ip',
+    'topologySpreadConstraints': 'topologyKey',
+    'ports': 'containerPort',
+    'volumeDevices': 'devicePath',
+}
+
 
 class Config(Dict[str, Any]):
     """SkyPilot config that supports setting/getting values with nested keys."""
@@ -211,19 +231,23 @@ def merge_k8s_configs(
                 merge_k8s_configs(base_config[key][0], value[0],
                                   next_allowed_override_keys,
                                   next_disallowed_override_keys)
-            elif key in ['volumes', 'volumeMounts', 'initContainers']:
-                # If the key is 'volumes', 'volumeMounts', or 'initContainers',
-                # we search for item with the same name and merge it.
+            # For list fields with patch strategy "merge", we merge the list
+            # by the patch merge key.
+            elif key in _PATCH_MERGE_KEYS:
+                patch_merge_key = _PATCH_MERGE_KEYS[key]
                 for override_item in value:
-                    override_item_name = override_item.get('name')
+                    override_item_name = override_item.get(patch_merge_key)
                     if override_item_name is not None:
                         existing_base_item = next(
                             (v for v in base_config[key]
-                             if v.get('name') == override_item_name), None)
+                             if v.get(patch_merge_key) == override_item_name),
+                            None)
                         if existing_base_item is not None:
                             merge_k8s_configs(existing_base_item, override_item)
                         else:
                             base_config[key].append(override_item)
+                    else:
+                        base_config[key].append(override_item)
             else:
                 base_config[key].extend(value)
         else:

sky/utils/controller_utils.py

@@ -23,11 +23,14 @@ from sky.clouds import gcp
 from sky.data import data_utils
 from sky.data import storage as storage_lib
 from sky.jobs import constants as managed_job_constants
+from sky.jobs import state as managed_job_state
 from sky.provision.kubernetes import constants as kubernetes_constants
 from sky.serve import constants as serve_constants
+from sky.serve import serve_state
 from sky.setup_files import dependencies
 from sky.skylet import constants
 from sky.skylet import log_lib
+from sky.utils import annotations
 from sky.utils import common
 from sky.utils import common_utils
 from sky.utils import config_utils
@@ -37,8 +40,13 @@ from sky.utils import rich_utils
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
+    import psutil
+
     from sky import task as task_lib
     from sky.backends import cloud_vm_ray_backend
+else:
+    from sky.adaptors import common as adaptors_common
+    psutil = adaptors_common.LazyImport('psutil')
 
 logger = sky_logging.init_logger(__name__)
 
@@ -1161,3 +1169,81 @@ def maybe_translate_local_file_mounts_and_sync_up(task: 'task_lib.Task',
     task.update_storage_mounts(updated_mount_storages)
     if msg:
         logger.info(ux_utils.finishing_message('Uploaded local files/folders.'))
+
+
+# ======================= Resources Management Functions =======================
+
+# Based on testing, assume a running job process uses 350MB memory. We use the
+# same estimation for service controller process.
+JOB_MEMORY_MB = 350
+# Monitoring process for service is 1GB. This is based on an old estimation but
+# we keep it here for now.
+# TODO(tian): Remeasure this.
+SERVE_MONITORING_MEMORY_MB = 1024
+# The ratio of service controller process to job process. We will treat each
+# service as SERVE_PROC_RATIO job processes.
+SERVE_PROC_RATIO = SERVE_MONITORING_MEMORY_MB / JOB_MEMORY_MB
+# Past 2000 simultaneous jobs, we become unstable.
+# See https://github.com/skypilot-org/skypilot/issues/4649.
+MAX_JOB_LIMIT = 2000
+# Number of ongoing launches launches allowed per CPU, for managed jobs.
+JOB_LAUNCHES_PER_CPU = 4
+# Number of ongoing launches launches allowed per CPU, for services. This is
+# also based on an old estimation, but SKyServe indeed spawn a new process
+# for each launch operation, so it should be slightly more resources demanding
+# than managed jobs.
+SERVE_LAUNCHES_PER_CPU = 2
+# The ratio of service launch to job launch. This is inverted as the parallelism
+# is determined by 1 / LAUNCHES_PER_CPU.
+SERVE_LAUNCH_RATIO = JOB_LAUNCHES_PER_CPU / SERVE_LAUNCHES_PER_CPU
+
+# The _RESOURCES_LOCK should be held whenever we are checking the parallelism
+# control or updating the schedule_state of any job or service. Any code that
+# takes this lock must conclude by calling maybe_schedule_next_jobs.
+_RESOURCES_LOCK = '~/.sky/locks/controller_resources.lock'
+
+
+@annotations.lru_cache(scope='global', maxsize=1)
+def get_resources_lock_path() -> str:
+    path = os.path.expanduser(_RESOURCES_LOCK)
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    return path
+
+
+@annotations.lru_cache(scope='request')
+def _get_job_parallelism() -> int:
+    job_memory = JOB_MEMORY_MB * 1024 * 1024
+    job_limit = min(psutil.virtual_memory().total // job_memory, MAX_JOB_LIMIT)
+    return max(job_limit, 1)
+
+
+@annotations.lru_cache(scope='request')
+def _get_launch_parallelism() -> int:
+    cpus = os.cpu_count()
+    return cpus * JOB_LAUNCHES_PER_CPU if cpus is not None else 1
+
+
+def can_provision() -> bool:
+    # We always prioritize terminating over provisioning, to save the cost on
+    # idle resources.
+    if serve_state.total_number_scheduled_to_terminate_replicas() > 0:
+        return False
+    return can_terminate()
+
+
+def can_start_new_process() -> bool:
+    num_procs = (serve_state.get_num_services() * SERVE_PROC_RATIO +
+                 managed_job_state.get_num_alive_jobs())
+    return num_procs < _get_job_parallelism()
+
+
+# We limit the number of terminating replicas to the number of CPUs. This is
+# just a temporary solution to avoid overwhelming the controller. After one job
+# controller PR, we should use API server to handle resources management.
+def can_terminate() -> bool:
+    num_terminating = (
+        serve_state.total_number_provisioning_replicas() * SERVE_LAUNCH_RATIO +
+        # Each terminate process will take roughly the same CPUs as job launch.
+        serve_state.total_number_terminating_replicas() +
+        managed_job_state.get_num_launching_jobs())
+    return num_terminating < _get_launch_parallelism()

sky/utils/db/db_utils.py

@@ -32,6 +32,23 @@ if typing.TYPE_CHECKING:
 _DB_TIMEOUT_S = 60
 
 
+class UniqueConstraintViolationError(Exception):
+    """Exception raised for unique constraint violation.
+    Attributes:
+        value -- the input value that caused the error
+        message -- explanation of the error
+    """
+
+    def __init__(self, value, message='Unique constraint violation'):
+        self.value = value
+        self.message = message
+        super().__init__(self.message)
+
+    def __str__(self):
+        return (f'UniqueConstraintViolationError: {self.message} '
+                f'(Value: {self.value})')
+
+
 class SQLAlchemyDialect(enum.Enum):
     SQLITE = 'sqlite'
     POSTGRESQL = 'postgresql'

sky/utils/db/migration_utils.py

@@ -19,7 +19,7 @@ logger = sky_logging.init_logger(__name__)
 DB_INIT_LOCK_TIMEOUT_SECONDS = 10
 
 GLOBAL_USER_STATE_DB_NAME = 'state_db'
-GLOBAL_USER_STATE_VERSION = '005'
+GLOBAL_USER_STATE_VERSION = '006'
 GLOBAL_USER_STATE_LOCK_PATH = '~/.sky/locks/.state_db.lock'
 
 SPOT_JOBS_DB_NAME = 'spot_jobs_db'