skypilot-nightly 1.0.0.dev20250502__py3-none-any.whl → 1.0.0.dev20251203__py3-none-any.whl

sky/utils/controller_utils.py

@@ -2,11 +2,10 @@
 import copy
 import dataclasses
 import enum
-import getpass
 import os
 import tempfile
 import typing
-from typing import Any, Dict, Iterable, List, Optional, Set, Tuple
+from typing import Any, Callable, Dict, Iterable, List, Optional, Set
 import uuid
 
 import colorama
@@ -24,10 +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.provision.kubernetes import constants as kubernetes_constants
 from sky.serve import constants as serve_constants
+from sky.serve import serve_state
+from sky.server import config as server_config
 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
@@ -35,10 +38,16 @@ from sky.utils import env_options
 from sky.utils import registry
 from sky.utils import rich_utils
 from sky.utils import ux_utils
+from sky.utils import yaml_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__)
 
@@ -63,8 +72,9 @@ class _ControllerSpec:
     """Spec for skypilot controllers."""
     controller_type: str
     name: str
-    cluster_name: str
-    in_progress_hint: str
+    _cluster_name_func: Callable[[], str]
+    _cluster_name_from_server: Optional[str]  # For client-side only
+    in_progress_hint: Callable[[bool], str]
     decline_cancel_hint: str
     _decline_down_when_failed_to_fetch_status_hint: str
     decline_down_for_dirty_controller_hint: str
@@ -84,6 +94,24 @@ class _ControllerSpec:
         return self._check_cluster_name_hint.format(
             cluster_name=self.cluster_name)
 
+    @property
+    def cluster_name(self) -> str:
+        """The cluster name of the controller.
+
+        On the server-side, the cluster name is the actual cluster name,
+        which is read from common.(JOB|SKY_SERVE)_CONTROLLER_NAME.
+
+        On the client-side, the cluster name may not be accurate,
+        as we may not know the exact name, because we are missing
+        the server-side common.SERVER_ID. We have to wait until
+        we get the actual cluster name from the server.
+        """
+        return (self._cluster_name_from_server if self._cluster_name_from_server
+                is not None else self._cluster_name_func())
+
+    def set_cluster_name_from_server(self, cluster_name: str) -> None:
+        self._cluster_name_from_server = cluster_name
+
 
 # TODO: refactor controller class to not be an enum.
 class Controllers(enum.Enum):
@@ -93,10 +121,11 @@ class Controllers(enum.Enum):
     JOBS_CONTROLLER = _ControllerSpec(
         controller_type='jobs',
         name='managed jobs controller',
-        cluster_name=common.JOB_CONTROLLER_NAME,
-        in_progress_hint=(
-            '* {job_info}To see all managed jobs: '
-            f'{colorama.Style.BRIGHT}sky jobs queue{colorama.Style.RESET_ALL}'),
+        _cluster_name_func=lambda: common.JOB_CONTROLLER_NAME,
+        _cluster_name_from_server=None,
+        in_progress_hint=lambda _:
+        ('* {job_info}To see all managed jobs: '
+         f'{colorama.Style.BRIGHT}sky jobs queue{colorama.Style.RESET_ALL}'),
         decline_cancel_hint=(
             'Cancelling the jobs controller\'s jobs is not allowed.\nTo cancel '
             f'managed jobs, use: {colorama.Style.BRIGHT}sky jobs cancel '
@@ -124,10 +153,14 @@ class Controllers(enum.Enum):
     SKY_SERVE_CONTROLLER = _ControllerSpec(
         controller_type='serve',
         name='serve controller',
-        cluster_name=common.SKY_SERVE_CONTROLLER_NAME,
+        _cluster_name_func=lambda: common.SKY_SERVE_CONTROLLER_NAME,
+        _cluster_name_from_server=None,
         in_progress_hint=(
-            f'* To see detailed service status: {colorama.Style.BRIGHT}'
-            f'sky serve status -v{colorama.Style.RESET_ALL}'),
+            lambda pool:
+            (f'* To see detailed pool status: {colorama.Style.BRIGHT}'
+             f'sky jobs pool status -v{colorama.Style.RESET_ALL}') if pool else
+            (f'* To see detailed service status: {colorama.Style.BRIGHT}'
+             f'sky serve status -v{colorama.Style.RESET_ALL}')),
         decline_cancel_hint=(
             'Cancelling the sky serve controller\'s jobs is not allowed.'),
         _decline_down_when_failed_to_fetch_status_hint=(
@@ -154,7 +187,9 @@ class Controllers(enum.Enum):
         default_autostop_config=serve_constants.CONTROLLER_AUTOSTOP)
 
     @classmethod
-    def from_name(cls, name: Optional[str]) -> Optional['Controllers']:
+    def from_name(cls,
+                  name: Optional[str],
+                  expect_exact_match: bool = True) -> Optional['Controllers']:
         """Check if the cluster name is a controller name.
 
         Returns:
@@ -168,15 +203,32 @@ class Controllers(enum.Enum):
         # we may not know the exact name, because we are missing the server-side
         # common.SERVER_ID. So, we will assume anything that matches the prefix
         # is a controller.
+        prefix = None
         if name.startswith(common.SKY_SERVE_CONTROLLER_PREFIX):
             controller = cls.SKY_SERVE_CONTROLLER
+            prefix = common.SKY_SERVE_CONTROLLER_PREFIX
         elif name.startswith(common.JOB_CONTROLLER_PREFIX):
             controller = cls.JOBS_CONTROLLER
-        if controller is not None and name != controller.value.cluster_name:
+            prefix = common.JOB_CONTROLLER_PREFIX
+
+        if controller is not None and expect_exact_match:
+            assert name == controller.value.cluster_name, (
+                name, controller.value.cluster_name)
+        elif controller is not None and name != controller.value.cluster_name:
             # The client-side cluster_name is not accurate. Assume that `name`
             # is the actual cluster name, so need to set the controller's
             # cluster name to the input name.
-            controller.value.cluster_name = name
+
+            # Assert that the cluster name is well-formed. It should be
+            # {prefix}{hash}, where prefix is set above, and hash is a valid
+            # user hash.
+            assert prefix is not None, prefix
+            assert name.startswith(prefix), name
+            assert common_utils.is_valid_user_hash(name[len(prefix):]), (name,
+                                                                         prefix)
+
+            # Update the cluster name.
+            controller.value.set_cluster_name_from_server(name)
         return controller
 
     @classmethod
@@ -193,27 +245,35 @@ class Controllers(enum.Enum):
         return None
 
 
-def high_availability_specified(cluster_name: Optional[str],
-                                skip_warning: bool = True) -> bool:
+def get_controller_for_pool(pool: bool) -> Controllers:
+    """Get the controller type."""
+    if pool:
+        return Controllers.JOBS_CONTROLLER
+    return Controllers.SKY_SERVE_CONTROLLER
+
+
+def high_availability_specified(cluster_name: Optional[str]) -> bool:
     """Check if the controller high availability is specified in user config.
     """
-    controller = Controllers.from_name(cluster_name)
+    controller = Controllers.from_name(cluster_name, expect_exact_match=False)
     if controller is None:
         return False
 
+    if controller.value.controller_type == 'jobs':
+        # pylint: disable-next=import-outside-toplevel
+        from sky.jobs import utils as managed_job_utils
+        if managed_job_utils.is_consolidation_mode():
+            return True
+    elif controller.value.controller_type == 'serve':
+        # pylint: disable-next=import-outside-toplevel
+        from sky.serve import serve_utils
+        if serve_utils.is_consolidation_mode():
+            return True
+
     if skypilot_config.loaded():
-        high_availability = skypilot_config.get_nested(
-            (controller.value.controller_type, 'controller',
-             'high_availability'), False)
-        if high_availability:
-            if controller.value.controller_type != 'serve':
-                if not skip_warning:
-                    print(f'{colorama.Fore.RED}High availability controller is'
-                          'only supported for SkyServe controller. It cannot'
-                          f'be enabled for {controller.value.name}.'
-                          f'Skipping this flag.{colorama.Style.RESET_ALL}')
-            else:
-                return True
+        return skypilot_config.get_nested((controller.value.controller_type,
+                                           'controller', 'high_availability'),
+                                          False)
     return False
 
 
@@ -250,6 +310,13 @@ def _get_cloud_dependencies_installation_commands(
         sky_check.get_cached_enabled_clouds_or_refresh(
             sky_cloud.CloudCapability.STORAGE))
     enabled_clouds = enabled_compute_clouds.union(enabled_storage_clouds)
+    enabled_k8s_and_ssh = [
+        repr(cloud)
+        for cloud in enabled_clouds
+        if isinstance(cloud, clouds.Kubernetes)
+    ]
+    k8s_and_ssh_label = ' and '.join(sorted(enabled_k8s_and_ssh))
+    k8s_dependencies_installed = False
 
     for cloud in enabled_clouds:
         cloud_python_dependencies: List[str] = copy.deepcopy(
@@ -269,10 +336,33 @@ def _get_cloud_dependencies_installation_commands(
             step_prefix = prefix_str.replace('<step>', str(len(commands) + 1))
             commands.append(f'echo -en "\\r{step_prefix}GCP SDK{empty_str}" &&'
                             f'{gcp.GOOGLE_SDK_INSTALLATION_COMMAND}')
-        elif isinstance(cloud, clouds.Kubernetes):
+            if clouds.cloud_in_iterable(clouds.Kubernetes(), enabled_clouds):
+                # Install gke-gcloud-auth-plugin used for exec-auth with GKE.
+                # We install the plugin here instead of the next elif branch
+                # because gcloud is required to install the plugin, so the order
+                # of command execution is critical.
+
+                # We install plugin here regardless of whether exec-auth is
+                # actually used as exec-auth may be used in the future.
+                # TODO (kyuds): how to implement conservative installation?
+                commands.append(
+                    '(command -v gke-gcloud-auth-plugin &>/dev/null || '
+                    '(gcloud components install gke-gcloud-auth-plugin --quiet &>/dev/null))')  # pylint: disable=line-too-long
+        elif isinstance(cloud, clouds.Nebius):
             step_prefix = prefix_str.replace('<step>', str(len(commands) + 1))
             commands.append(
-                f'echo -en "\\r{step_prefix}Kubernetes{empty_str}" && '
+                f'echo -en "\\r{step_prefix}Nebius{empty_str}" && '
+                'curl -sSL https://storage.eu-north1.nebius.cloud/cli/install.sh '  # pylint: disable=line-too-long
+                '| sudo NEBIUS_INSTALL_FOLDER=/usr/local/bin bash &> /dev/null && '
+                'nebius profile create --profile sky '
+                '--endpoint api.nebius.cloud '
+                '--service-account-file $HOME/.nebius/credentials.json '
+                '&> /dev/null || echo "Unable to create Nebius profile."')
+        elif (isinstance(cloud, clouds.Kubernetes) and
+              not k8s_dependencies_installed):
+            step_prefix = prefix_str.replace('<step>', str(len(commands) + 1))
+            commands.append(
+                f'echo -en "\\r{step_prefix}{k8s_and_ssh_label}{empty_str}" && '
                 # Install k8s + skypilot dependencies
                 'sudo bash -c "if '
                 '! command -v curl &> /dev/null || '
@@ -292,7 +382,10 @@ def _get_cloud_dependencies_installation_commands(
                 '(curl -s -LO "https://dl.k8s.io/release/v1.31.6'
                 '/bin/linux/$ARCH/kubectl" && '
                 'sudo install -o root -g root -m 0755 '
-                'kubectl /usr/local/bin/kubectl))')
+                'kubectl /usr/local/bin/kubectl)) && '
+                f'echo -e \'#!/bin/bash\\nexport PATH="{kubernetes_constants.SKY_K8S_EXEC_AUTH_PATH}"\\nexec "$@"\' | sudo tee /usr/local/bin/{kubernetes_constants.SKY_K8S_EXEC_AUTH_WRAPPER} > /dev/null && '  # pylint: disable=line-too-long
+                f'sudo chmod +x /usr/local/bin/{kubernetes_constants.SKY_K8S_EXEC_AUTH_WRAPPER}')  # pylint: disable=line-too-long
+            k8s_dependencies_installed = True
         elif isinstance(cloud, clouds.Cudo):
             step_prefix = prefix_str.replace('<step>', str(len(commands) + 1))
             commands.append(
@@ -345,7 +438,7 @@ def check_cluster_name_not_controller(
     Returns:
       None, if the cluster name is not a controller name.
     """
-    controller = Controllers.from_name(cluster_name)
+    controller = Controllers.from_name(cluster_name, expect_exact_match=False)
     if controller is not None:
         msg = controller.value.check_cluster_name_hint
         if operation_str is not None:
@@ -355,10 +448,11 @@ def check_cluster_name_not_controller(
 
 
 # Internal only:
-def download_and_stream_latest_job_log(
+def download_and_stream_job_log(
         backend: 'cloud_vm_ray_backend.CloudVmRayBackend',
         handle: 'cloud_vm_ray_backend.CloudVmRayResourceHandle',
-        local_dir: str) -> Optional[str]:
+        local_dir: str,
+        job_ids: Optional[List[str]] = None) -> Optional[str]:
     """Downloads and streams the latest job log.
 
     This function is only used by jobs controller and sky serve controller.
@@ -376,7 +470,7 @@ def download_and_stream_latest_job_log(
             # multi-node cluster is preempted, and we recover the managed job
             # on the existing cluster, which leads to a larger job_id. Those
             # job_ids all represent the same logical managed job.
-            job_ids=None,
+            job_ids=job_ids,
             local_dir=local_dir)
     except Exception as e:  # pylint: disable=broad-except
         # We want to avoid crashing the controller. sync_down_logs() is pretty
@@ -394,7 +488,7 @@ def download_and_stream_latest_job_log(
         return None
 
     log_dir = list(log_dirs.values())[0]
-    log_file = os.path.join(log_dir, 'run.log')
+    log_file = os.path.expanduser(os.path.join(log_dir, 'run.log'))
 
     # Print the logs to the console.
     # TODO(zhwu): refactor this into log_utils, along with the refactoring for
@@ -439,10 +533,13 @@ def shared_controller_vars_to_fill(
         # before popping allowed_contexts. If it is not on Kubernetes,
         # we may be able to use allowed_contexts.
         local_user_config.pop('allowed_contexts', None)
+        # Remove api_server config so that the controller does not try to use
+        # a remote API server.
+        local_user_config.pop('api_server', None)
         with tempfile.NamedTemporaryFile(
                 delete=False,
                 suffix=_LOCAL_SKYPILOT_CONFIG_PATH_SUFFIX) as temp_file:
-            common_utils.dump_yaml(temp_file.name, dict(**local_user_config))
+            yaml_utils.dump_yaml(temp_file.name, dict(**local_user_config))
         local_user_config_path = temp_file.name
 
     vars_to_fill: Dict[str, Any] = {
@@ -461,7 +558,7 @@ def shared_controller_vars_to_fill(
     env_vars.update({
         # Should not use $USER here, as that env var can be empty when
         # running in a container.
-        constants.USER_ENV_VAR: getpass.getuser(),
+        constants.USER_ENV_VAR: common_utils.get_current_user_name(),
         constants.USER_ID_ENV_VAR: common_utils.get_user_hash(),
         # Skip cloud identity check to avoid the overhead.
         env_options.Options.SKIP_CLOUD_IDENTITY_CHECK.env_key: '1',
@@ -472,7 +569,15 @@ def shared_controller_vars_to_fill(
         # with a remote API server.
         constants.USING_REMOTE_API_SERVER_ENV_VAR: str(
             common_utils.get_using_remote_api_server()),
+        constants.IS_SKYPILOT_SERVE_CONTROLLER:
+            ('true'
+             if controller == Controllers.SKY_SERVE_CONTROLLER else 'false'),
     })
+    override_concurrent_launches = os.environ.get(
+        constants.SERVE_OVERRIDE_CONCURRENT_LAUNCHES, None)
+    if override_concurrent_launches is not None:
+        env_vars[constants.SERVE_OVERRIDE_CONCURRENT_LAUNCHES] = str(
+            int(override_concurrent_launches))
     if skypilot_config.loaded():
         # Only set the SKYPILOT_CONFIG env var if the user has a config file.
         env_vars[
@@ -504,6 +609,30 @@ def get_controller_resources(
         if custom_controller_resources_config is not None:
             controller_resources_config_copied.update(
                 custom_controller_resources_config)
+        # Compatibility with the old way of specifying the controller autostop
+        # config. TODO(cooperc): Remove this before 0.12.0.
+        custom_controller_autostop_config = skypilot_config.get_nested(
+            (controller.value.controller_type, 'controller', 'autostop'), None)
+        if custom_controller_autostop_config is not None:
+            logger.warning(
+                f'{colorama.Fore.YELLOW}Warning: Config value '
+                f'`{controller.value.controller_type}.controller.autostop` '
+                'is deprecated. Please use '
+                f'`{controller.value.controller_type}.controller.resources.'
+                f'autostop` instead.{colorama.Style.RESET_ALL}')
+            # Only set the autostop config if it is not already specified.
+            if controller_resources_config_copied.get('autostop') is None:
+                controller_resources_config_copied['autostop'] = (
+                    custom_controller_autostop_config)
+            else:
+                logger.warning(f'{colorama.Fore.YELLOW}Ignoring the old '
+                               'config, since it is already specified in '
+                               f'resources.{colorama.Style.RESET_ALL}')
+    # Set the default autostop config for the controller, if not already
+    # specified.
+    if controller_resources_config_copied.get('autostop') is None:
+        controller_resources_config_copied['autostop'] = (
+            controller.value.default_autostop_config)
 
     try:
         controller_resources = resources.Resources.from_yaml_config(
@@ -529,12 +658,16 @@ def get_controller_resources(
     controller_resources_to_use: resources.Resources = list(
         controller_resources)[0]
 
-    controller_record = global_user_state.get_cluster_from_name(
+    controller_handle = global_user_state.get_handle_from_cluster_name(
         controller.value.cluster_name)
-    if controller_record is not None:
-        handle = controller_record.get('handle', None)
-        if handle is not None:
-            controller_resources_to_use = handle.launched_resources
+    if controller_handle is not None:
+        if controller_handle is not None:
+            # Use the existing resources, but override the autostop config with
+            # the one currently specified in the config.
+            controller_resources_to_use = (
+                controller_handle.launched_resources.copy(
+                    autostop=controller_resources_config_copied.get('autostop'))
+            )
 
     # If the controller and replicas are from the same cloud (and region/zone),
     # it should provide better connectivity. We will let the controller choose
@@ -595,8 +728,9 @@ def get_controller_resources(
     controller_zone = controller_resources_to_use.zone
 
     # Filter clouds if controller_resources_to_use.cloud is specified.
-    filtered_clouds = ({controller_cloud} if controller_cloud is not None else
-                       requested_clouds_with_region_zone.keys())
+    filtered_clouds: Set[str] = {controller_cloud
+                                } if controller_cloud is not None else set(
+                                    requested_clouds_with_region_zone.keys())
 
     # Filter regions and zones and construct the result.
     result: Set[resources.Resources] = set()
@@ -605,15 +739,17 @@ def get_controller_resources(
                                                         {None: {None}})
 
         # Filter regions if controller_resources_to_use.region is specified.
-        filtered_regions = ({controller_region} if controller_region is not None
-                            else regions.keys())
+        filtered_regions: Set[Optional[str]] = ({
+            controller_region
+        } if controller_region is not None else set(regions.keys()))
 
         for region in filtered_regions:
             zones = regions.get(region, {None})
 
             # Filter zones if controller_resources_to_use.zone is specified.
-            filtered_zones = ({controller_zone}
-                              if controller_zone is not None else zones)
+            filtered_zones: Set[Optional[str]] = ({
+                controller_zone
+            } if controller_zone is not None else set(zones))
 
             # Create combinations of cloud, region, and zone.
             for zone in filtered_zones:
@@ -628,38 +764,15 @@ def get_controller_resources(
     return result
 
 
-def get_controller_autostop_config(
-        controller: Controllers) -> Tuple[Optional[int], bool]:
-    """Get the autostop config for the controller.
-
-    Returns:
-      A tuple of (idle_minutes_to_autostop, down), which correspond to the
-      values passed to execution.launch().
-    """
-    controller_autostop_config_copied: Dict[str, Any] = copy.copy(
-        controller.value.default_autostop_config)
-    if skypilot_config.loaded():
-        custom_controller_autostop_config = skypilot_config.get_nested(
-            (controller.value.controller_type, 'controller', 'autostop'), None)
-        if custom_controller_autostop_config is False:
-            # Disabled with `autostop: false` in config.
-            # To indicate autostop is disabled, we return None for
-            # idle_minutes_to_autostop.
-            return None, False
-        elif custom_controller_autostop_config is True:
-            # Enabled with default values. There is no change in behavior, but
-            # this is included by for completeness, since `False` is valid.
-            pass
-        elif custom_controller_autostop_config is not None:
-            # We have specific config values.
-            # Override the controller autostop config with the ones specified in
-            # the config.
-            assert isinstance(custom_controller_autostop_config, dict)
-            controller_autostop_config_copied.update(
-                custom_controller_autostop_config)
-
-    return (controller_autostop_config_copied['idle_minutes'],
-            controller_autostop_config_copied['down'])
+def get_controller_mem_size_gb() -> float:
+    try:
+        with open(os.path.expanduser(constants.CONTROLLER_K8S_MEMORY_FILE),
+                  'r',
+                  encoding='utf-8') as f:
+            return float(f.read())
+    except FileNotFoundError:
+        pass
+    return common_utils.get_mem_size_gb()
 
 
 def _setup_proxy_command_on_controller(
@@ -690,7 +803,7 @@ def _setup_proxy_command_on_controller(
     # NOTE: suppose that we have a controller in old VPC, then user
     # changes 'vpc_name' in the config and does a 'job launch' /
     # 'serve up'. In general, the old controller may not successfully
-    # launch the job in the new VPC. This happens if the two VPCs don’t
+    # launch the job in the new VPC. This happens if the two VPCs don't
     # have peering set up. Like other places in the code, we assume
     # properly setting up networking is user's responsibilities.
     # TODO(zongheng): consider adding a basic check that checks
@@ -701,7 +814,11 @@ def _setup_proxy_command_on_controller(
     config = config_utils.Config.from_dict(user_config)
     proxy_command_key = (str(controller_launched_cloud).lower(),
                          'ssh_proxy_command')
-    ssh_proxy_command = config.get_nested(proxy_command_key, None)
+    ssh_proxy_command = skypilot_config.get_effective_region_config(
+        cloud=str(controller_launched_cloud).lower(),
+        region=None,
+        keys=('ssh_proxy_command',),
+        default_value=None)
     if isinstance(ssh_proxy_command, str):
         config.set_nested(proxy_command_key, None)
     elif isinstance(ssh_proxy_command, dict):
@@ -731,9 +848,9 @@ def replace_skypilot_config_path_in_file_mounts(
             continue
         if local_path.endswith(_LOCAL_SKYPILOT_CONFIG_PATH_SUFFIX):
             with tempfile.NamedTemporaryFile('w', delete=False) as f:
-                user_config = common_utils.read_yaml(local_path)
+                user_config = yaml_utils.read_yaml(local_path)
                 config = _setup_proxy_command_on_controller(cloud, user_config)
-                common_utils.dump_yaml(f.name, dict(**config))
+                yaml_utils.dump_yaml(f.name, dict(**config))
                 file_mounts[remote_path] = f.name
                 replaced = True
     if replaced:
@@ -776,7 +893,7 @@ def translate_local_file_mounts_to_two_hop(
     file_mount_id = 0
 
     file_mounts_to_translate = task.file_mounts or {}
-    if task.workdir is not None:
+    if task.workdir is not None and isinstance(task.workdir, str):
         file_mounts_to_translate[constants.SKY_REMOTE_WORKDIR] = task.workdir
         task.workdir = None
 
@@ -844,7 +961,8 @@ def maybe_translate_local_file_mounts_and_sync_up(task: 'task_lib.Task',
         copy_mounts = {}
 
     has_local_source_paths_file_mounts = bool(copy_mounts)
-    has_local_source_paths_workdir = task.workdir is not None
+    has_local_source_paths_workdir = (task.workdir is not None and
+                                      isinstance(task.workdir, str))
 
     msg = None
     if has_local_source_paths_workdir and has_local_source_paths_file_mounts:
@@ -892,7 +1010,7 @@ def maybe_translate_local_file_mounts_and_sync_up(task: 'task_lib.Task',
 
     # Step 1: Translate the workdir to SkyPilot storage.
     new_storage_mounts = {}
-    if task.workdir is not None:
+    if task.workdir is not None and isinstance(task.workdir, str):
         workdir = task.workdir
         task.workdir = None
         if (constants.SKY_REMOTE_WORKDIR in original_file_mounts or
@@ -1113,3 +1231,179 @@ 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 =======================
+
+# Monitoring process for service is 512MB. This is based on an old
+# estimation but we keep it here for now.
+# TODO(tian): Remeasure this.
+SERVE_MONITORING_MEMORY_MB = 512
+# The resource consumption ratio of service launch to serve down.
+SERVE_LAUNCH_RATIO = 2.0
+
+# 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'
+
+# keep 2GB reserved after the controllers
+MAXIMUM_CONTROLLER_RESERVED_MEMORY_MB = 2048
+
+# NOTE: In the current implementation, we only consider the memory
+# The ratio of resources consumption for managed jobs and pool/serve.
+# This measures pool_resources / jobs_resources. If 2 GB memory is allocated to
+# jobs, then 2 * POOL_JOBS_RESOURCES_RATIO GB memory is allocated to pool/serve.
+POOL_JOBS_RESOURCES_RATIO = 1
+# Number of ongoing launches launches allowed per worker. Can probably be
+# increased a bit to around 16 but keeping it lower to just to be safe
+LAUNCHES_PER_WORKER = 8
+# Number of ongoing launches allowed per service. Can probably be increased
+# a bit as well.
+LAUNCHES_PER_SERVICE = 4
+
+# Based on testing, each worker takes around 200-300MB memory. Keeping it
+# higher to be safe.
+JOB_WORKER_MEMORY_MB = 400
+# this can probably be increased to around 300-400 but keeping it lower to just
+# to be safe
+MAX_JOBS_PER_WORKER = 200
+# Maximum number of controllers that can be running. Hard to handle more than
+# 512 launches at once.
+MAX_CONTROLLERS = 512 // LAUNCHES_PER_WORKER
+# Limit the number of jobs that can be running at once on the entire jobs
+# controller cluster. It's hard to handle cancellation of more than 2000 jobs at
+# once.
+# TODO(cooperc): Once we eliminate static bottlenecks (e.g. sqlite), remove this
+# hardcoded max limit.
+MAX_TOTAL_RUNNING_JOBS = 2000
+
+
+def compute_memory_reserved_for_controllers(
+        reserve_for_controllers: bool, reserve_extra_for_pool: bool) -> float:
+    reserved_memory_mb = 0.0
+    if reserve_for_controllers:
+        reserved_memory_mb = float(MAXIMUM_CONTROLLER_RESERVED_MEMORY_MB)
+        if reserve_extra_for_pool:
+            reserved_memory_mb *= (1. + POOL_JOBS_RESOURCES_RATIO)
+    return reserved_memory_mb
+
+
+def _get_total_usable_memory_mb(pool: bool, consolidation_mode: bool) -> float:
+    controller_reserved = compute_memory_reserved_for_controllers(
+        reserve_for_controllers=True, reserve_extra_for_pool=pool)
+    total_memory_mb = (common_utils.get_mem_size_gb() * 1024 -
+                       controller_reserved)
+    if not consolidation_mode:
+        return total_memory_mb
+    config = server_config.compute_server_config(
+        deploy=True, quiet=True, reserved_memory_mb=controller_reserved)
+    used = 0.0
+    used += ((config.long_worker_config.garanteed_parallelism +
+              config.long_worker_config.burstable_parallelism) *
+             server_config.LONG_WORKER_MEM_GB * 1024)
+    used += ((config.short_worker_config.garanteed_parallelism +
+              config.short_worker_config.burstable_parallelism) *
+             server_config.SHORT_WORKER_MEM_GB * 1024)
+    return total_memory_mb - used
+
+
+def _is_consolidation_mode(pool: bool) -> bool:
+    return skypilot_config.get_nested(
+        ('jobs' if pool else 'serve', 'controller', 'consolidation_mode'),
+        default_value=False)
+
+
+@annotations.lru_cache(scope='request')
+def _get_parallelism(pool: bool, raw_resource_per_unit: float) -> int:
+    """Returns the number of jobs controllers / services that should be running.
+
+    This is the number of controllers / services that should be running
+    to maximize resource utilization.
+
+    In consolidation mode, we use the existing API server so our resource
+    requirements are just for the job controllers / services. We try taking
+    up as much memory as possible left over from the API server.
+
+    In non-consolidation mode, we have to take into account the memory of the
+    API server workers. We limit to only 8 launches per worker, so our logic is
+    each controller will take CONTROLLER_MEMORY_MB + 8 * WORKER_MEMORY_MB. We
+    leave some leftover room for ssh codegen and ray status overhead.
+    """
+    consolidation_mode = _is_consolidation_mode(pool)
+
+    total_memory_mb = _get_total_usable_memory_mb(pool, consolidation_mode)
+
+    # In consolidation mode, we assume the API server is running in deployment
+    # mode, hence resource management (i.e. how many requests are allowed) is
+    # done by the API server.
+    resource_per_unit_worker = 0.
+    # Otherwise, it runs a local API server on the jobs/serve controller.
+    # We need to do the resource management ourselves.
+    if not consolidation_mode:
+        launches_per_worker = (LAUNCHES_PER_WORKER
+                               if pool else LAUNCHES_PER_SERVICE)
+        resource_per_unit_worker = (launches_per_worker *
+                                    server_config.LONG_WORKER_MEM_GB * 1024)
+
+    # If running pool on jobs controller, we need to account for the resources
+    # consumed by the jobs.
+    ratio = (1. + POOL_JOBS_RESOURCES_RATIO) if pool else 1.
+    resource_per_unit = ratio * (raw_resource_per_unit +
+                                 resource_per_unit_worker)
+
+    return max(int(total_memory_mb / resource_per_unit), 1)
+
+
+def get_number_of_jobs_controllers() -> int:
+    return min(
+        MAX_CONTROLLERS,
+        _get_parallelism(pool=True, raw_resource_per_unit=JOB_WORKER_MEMORY_MB))
+
+
+@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
+
+
+def _get_number_of_services(pool: bool) -> int:
+    return _get_parallelism(pool=pool,
+                            raw_resource_per_unit=SERVE_MONITORING_MEMORY_MB *
+                            POOL_JOBS_RESOURCES_RATIO)
+
+
+@annotations.lru_cache(scope='request')
+def _get_request_parallelism(pool: bool) -> int:
+    # NOTE(dev): One smoke test depends on this value.
+    # tests/smoke_tests/test_sky_serve.py::test_skyserve_new_autoscaler_update
+    # assumes 4 concurrent launches.
+    override_concurrent_launches = os.environ.get(
+        constants.SERVE_OVERRIDE_CONCURRENT_LAUNCHES, None)
+    if override_concurrent_launches is not None and not pool:
+        return int(override_concurrent_launches)
+    # Limitation per service x number of services
+    launches_per_worker = (LAUNCHES_PER_WORKER
+                           if pool else LAUNCHES_PER_SERVICE)
+    return (launches_per_worker * POOL_JOBS_RESOURCES_RATIO *
+            _get_number_of_services(pool))
+
+
+def can_provision(pool: bool) -> bool:
+    # TODO(tian): probe API server to see if there is any pending provision
+    # requests.
+    return can_terminate(pool)
+
+
+def can_start_new_process(pool: bool) -> bool:
+    return serve_state.get_num_services() < _get_number_of_services(pool)
+
+
+def can_terminate(pool: bool) -> bool:
+    # TODO(tian): probe API server to see if there is any pending terminate
+    # requests.
+    num_terminating = (
+        serve_state.total_number_provisioning_replicas() +
+        serve_state.total_number_terminating_replicas() / SERVE_LAUNCH_RATIO)
+    return num_terminating < _get_request_parallelism(pool)