skypilot-nightly 1.0.0.dev20250808__py3-none-any.whl → 1.0.0.dev20250814__py3-none-any.whl

sky/serve/server/impl.py

@@ -11,7 +11,6 @@ import uuid
 import colorama
 import filelock
 
-import sky
 from sky import backends
 from sky import exceptions
 from sky import execution
@@ -25,6 +24,7 @@ from sky.serve import constants as serve_constants
 from sky.serve import serve_state
 from sky.serve import serve_utils
 from sky.skylet import constants
+from sky.skylet import job_lib
 from sky.utils import admin_policy_utils
 from sky.utils import command_runner
 from sky.utils import common
@@ -39,7 +39,7 @@ logger = sky_logging.init_logger(__name__)
 
 
 def _rewrite_tls_credential_paths_and_get_tls_env_vars(
-        service_name: str, task: 'sky.Task') -> Dict[str, Any]:
+        service_name: str, task: 'task_lib.Task') -> Dict[str, Any]:
     """Rewrite the paths of TLS credentials in the task.
 
     Args:
@@ -103,15 +103,11 @@ def _get_service_record(
 
 
 def up(
-    task: 'sky.Task',
+    task: 'task_lib.Task',
     service_name: Optional[str] = None,
     pool: bool = False,
 ) -> Tuple[str, str]:
     """Spins up a service or a pool."""
-    if pool and not serve_utils.is_consolidation_mode(pool):
-        raise ValueError(
-            'Pool is only supported in consolidation mode. To fix, set '
-            '`jobs.controller.consolidation_mode: true` in SkyPilot config.')
     task.validate()
     serve_utils.validate_service_task(task, pool=pool)
     assert task.service is not None
@@ -191,8 +187,7 @@ def up(
         controller_log_file = (
             serve_utils.generate_remote_controller_log_file_name(service_name))
         controller_resources = controller_utils.get_controller_resources(
-            controller=controller_utils.Controllers.SKY_SERVE_CONTROLLER,
-            task_resources=task.resources)
+            controller=controller, task_resources=task.resources)
         controller_job_id = None
         if serve_utils.is_consolidation_mode(pool):
             # We need a unique integer per sky.serve.up call to avoid name
@@ -228,10 +223,11 @@ def up(
         # balancer port from the controller? So we don't need to open so many
         # ports here. Or, we should have a nginx traffic control to refuse
         # any connection to the unregistered ports.
-        controller_resources = {
-            r.copy(ports=[serve_constants.LOAD_BALANCER_PORT_RANGE])
-            for r in controller_resources
-        }
+        if not pool:
+            controller_resources = {
+                r.copy(ports=[serve_constants.LOAD_BALANCER_PORT_RANGE])
+                for r in controller_resources
+            }
         controller_task.set_resources(controller_resources)
 
         # # Set service_name so the backend will know to modify default ray
@@ -325,7 +321,7 @@ def up(
                                               [controller_job_id],
                                               stream_logs=False)
             controller_job_status = list(statuses.values())[0]
-            if controller_job_status == sky.JobStatus.PENDING:
+            if controller_job_status == job_lib.JobStatus.PENDING:
                 # Max number of services reached due to vCPU constraint.
                 # The controller job is pending due to ray job scheduling.
                 # We manually cancel the job here.
@@ -350,7 +346,7 @@ def up(
         else:
             lb_port = serve_utils.load_service_initialization_result(
                 lb_port_payload)
-            if not serve_utils.is_consolidation_mode(pool):
+            if not serve_utils.is_consolidation_mode(pool) and not pool:
                 socket_endpoint = backend_utils.get_endpoints(
                     controller_handle.cluster_name,
                     lb_port,
@@ -374,10 +370,10 @@ def up(
                 f'\n📋 Useful Commands'
                 f'\n{ux_utils.INDENT_SYMBOL}To submit jobs to the pool:\t'
                 f'{ux_utils.BOLD}sky jobs launch --pool {service_name} '
-                f'<run-command>{ux_utils.RESET_BOLD}'
+                f'<yaml_file>{ux_utils.RESET_BOLD}'
                 f'\n{ux_utils.INDENT_SYMBOL}To submit multiple jobs:\t'
                 f'{ux_utils.BOLD}sky jobs launch --pool {service_name} '
-                f'--num-jobs 10 <run-command>{ux_utils.RESET_BOLD}'
+                f'--num-jobs 10 <yaml_file>{ux_utils.RESET_BOLD}'
                 f'\n{ux_utils.INDENT_SYMBOL}To check the pool status:\t'
                 f'{ux_utils.BOLD}sky jobs pool status {service_name}'
                 f'{ux_utils.RESET_BOLD}'
@@ -421,7 +417,7 @@ def up(
 
 
 def update(
-    task: 'sky.Task',
+    task: 'task_lib.Task',
     service_name: str,
     mode: serve_utils.UpdateMode = serve_utils.DEFAULT_UPDATE_MODE,
     pool: bool = False,
@@ -576,7 +572,7 @@ def update(
 
 
 def apply(
-    task: 'sky.Task',
+    task: 'task_lib.Task',
     service_name: str,
     mode: serve_utils.UpdateMode = serve_utils.DEFAULT_UPDATE_MODE,
     pool: bool = False,

sky/serve/service.py

@@ -15,11 +15,13 @@ import filelock
 
 from sky import authentication
 from sky import exceptions
+from sky import global_user_state
 from sky import sky_logging
 from sky import task as task_lib
 from sky.backends import backend_utils
 from sky.backends import cloud_vm_ray_backend
 from sky.data import data_utils
+from sky.jobs import scheduler as jobs_scheduler
 from sky.serve import constants
 from sky.serve import controller
 from sky.serve import load_balancer
@@ -28,6 +30,7 @@ from sky.serve import serve_state
 from sky.serve import serve_utils
 from sky.skylet import constants as skylet_constants
 from sky.utils import common_utils
+from sky.utils import controller_utils
 from sky.utils import subprocess_utils
 from sky.utils import ux_utils
 
@@ -120,7 +123,16 @@ def _cleanup(service_name: str) -> bool:
     replica_infos = serve_state.get_replica_infos(service_name)
     info2proc: Dict[replica_managers.ReplicaInfo,
                     multiprocessing.Process] = dict()
+    # NOTE(dev): This relies on `sky/serve/serve_utils.py::
+    # generate_replica_cluster_name`. Change it if you change the function.
+    existing_cluster_names = global_user_state.get_cluster_names_start_with(
+        service_name)
     for info in replica_infos:
+        if info.cluster_name not in existing_cluster_names:
+            logger.info(f'Cluster {info.cluster_name} for replica '
+                        f'{info.replica_id} not found. Might be a failed '
+                        'cluster. Skipping.')
+            continue
         p = multiprocessing.Process(target=replica_managers.terminate_cluster,
                                     args=(info.cluster_name,))
         p.start()
@@ -214,22 +226,25 @@ def _start(service_name: str, tmp_task_yaml: str, job_id: int, entrypoint: str):
         service_name, version)
 
     if not is_recovery:
-        if (len(serve_state.get_services()) >=
-                serve_utils.get_num_service_threshold()):
-            cleanup_storage(tmp_task_yaml)
-            with ux_utils.print_exception_no_traceback():
-                raise RuntimeError('Max number of services reached.')
-        success = serve_state.add_service(
-            service_name,
-            controller_job_id=job_id,
-            policy=service_spec.autoscaling_policy_str(),
-            requested_resources_str=backend_utils.get_task_resources_str(task),
-            load_balancing_policy=service_spec.load_balancing_policy,
-            status=serve_state.ServiceStatus.CONTROLLER_INIT,
-            tls_encrypted=service_spec.tls_credential is not None,
-            pool=service_spec.pool,
-            controller_pid=os.getpid(),
-            entrypoint=entrypoint)
+        with filelock.FileLock(controller_utils.get_resources_lock_path()):
+            if not controller_utils.can_start_new_process():
+                cleanup_storage(tmp_task_yaml)
+                with ux_utils.print_exception_no_traceback():
+                    raise RuntimeError(
+                        constants.MAX_NUMBER_OF_SERVICES_REACHED_ERROR)
+            success = serve_state.add_service(
+                service_name,
+                controller_job_id=job_id,
+                policy=service_spec.autoscaling_policy_str(),
+                requested_resources_str=backend_utils.get_task_resources_str(
+                    task),
+                load_balancing_policy=service_spec.load_balancing_policy,
+                status=serve_state.ServiceStatus.CONTROLLER_INIT,
+                tls_encrypted=service_spec.tls_credential is not None,
+                pool=service_spec.pool,
+                controller_pid=os.getpid(),
+                entrypoint=entrypoint)
+        jobs_scheduler.maybe_schedule_next_jobs()
         # Directly throw an error here. See sky/serve/api.py::up
         # for more details.
         if not success:

sky/server/server.py

@@ -17,7 +17,7 @@ import resource
 import shutil
 import sys
 import threading
-from typing import Any, Dict, List, Literal, Optional, Set, Tuple
+from typing import Dict, List, Literal, Optional, Set, Tuple
 import uuid
 import zipfile
 
@@ -42,6 +42,7 @@ from sky.data import storage_utils
 from sky.jobs.server import server as jobs_rest
 from sky.metrics import utils as metrics_utils
 from sky.provision.kubernetes import utils as kubernetes_utils
+from sky.schemas.api import responses
 from sky.serve.server import server as serve_rest
 from sky.server import common
 from sky.server import config as server_config
@@ -1531,8 +1532,12 @@ async def api_status(
         return encoded_request_tasks
 
 
-@app.get('/api/health')
-async def health(request: fastapi.Request) -> Dict[str, Any]:
+@app.get(
+    '/api/health',
+    # response_model_exclude_unset omits unset fields
+    # in the response JSON.
+    response_model_exclude_unset=True)
+async def health(request: fastapi.Request) -> responses.APIHealthResponse:
     """Checks the health of the API server.
 
     Returns:
@@ -1570,7 +1575,8 @@ async def health(request: fastapi.Request) -> Dict[str, Any]:
             # - There is no harm when an malicious client calls /api/health
             #   without authentication since no sensitive information is
             #   returned.
-            return {'status': common.ApiServerStatus.HEALTHY}
+            return responses.APIHealthResponse(
+                status=common.ApiServerStatus.HEALTHY,)
         # TODO(aylei): remove this after min_compatible_api_version >= 14.
         if client_version < 14:
             # For Client with API version < 14, the NEEDS_AUTH status is not
@@ -1579,19 +1585,19 @@ async def health(request: fastapi.Request) -> Dict[str, Any]:
                                         detail='Authentication required')
 
     logger.debug(f'Health endpoint: request.state.auth_user = {user}')
-    return {
-        'status': server_status,
+    return responses.APIHealthResponse(
+        status=server_status,
         # Kept for backward compatibility, clients before 0.11.0 will read this
         # field to check compatibility and hint the user to upgrade the CLI.
         # TODO(aylei): remove this field after 0.13.0
-        'api_version': str(server_constants.API_VERSION),
-        'version': sky.__version__,
-        'version_on_disk': common.get_skypilot_version_on_disk(),
-        'commit': sky.__commit__,
-        'user': user.to_dict() if user is not None else None,
-        'basic_auth_enabled': os.environ.get(
-            constants.ENV_VAR_ENABLE_BASIC_AUTH, 'false').lower() == 'true',
-    }
+        api_version=str(server_constants.API_VERSION),
+        version=sky.__version__,
+        version_on_disk=common.get_skypilot_version_on_disk(),
+        commit=sky.__commit__,
+        basic_auth_enabled=os.environ.get(constants.ENV_VAR_ENABLE_BASIC_AUTH,
+                                          'false').lower() == 'true',
+        user=user if user is not None else None,
+    )
 
 
 @app.websocket('/kubernetes-pod-ssh-proxy')

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/skypilot_config.py

@@ -575,8 +575,8 @@ def _reload_config_as_server() -> None:
         with _DB_USE_LOCK:
             sqlalchemy_engine = sqlalchemy.create_engine(db_url,
                                                          poolclass=NullPool)
-            db_utils.add_tables_to_db_sqlalchemy(Base.metadata,
-                                                 sqlalchemy_engine)
+            db_utils.add_all_tables_to_db_sqlalchemy(Base.metadata,
+                                                     sqlalchemy_engine)
 
             def _get_config_yaml_from_db(
                     key: str) -> Optional[config_utils.Config]:
@@ -867,8 +867,8 @@ def update_api_server_config_no_lock(config: config_utils.Config) -> None:
             with _DB_USE_LOCK:
                 sqlalchemy_engine = sqlalchemy.create_engine(existing_db_url,
                                                              poolclass=NullPool)
-                db_utils.add_tables_to_db_sqlalchemy(Base.metadata,
-                                                     sqlalchemy_engine)
+                db_utils.add_all_tables_to_db_sqlalchemy(
+                    Base.metadata, sqlalchemy_engine)
 
                 def _set_config_yaml_to_db(key: str,
                                            config: config_utils.Config):

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}}