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

sky/server/plugins.py

@@ -0,0 +1,222 @@
+"""Load plugins for the SkyPilot API server."""
+import abc
+import dataclasses
+import importlib
+import os
+from typing import Dict, List, Optional, Tuple
+
+from fastapi import FastAPI
+
+from sky import sky_logging
+from sky.skylet import constants as skylet_constants
+from sky.utils import common_utils
+from sky.utils import config_utils
+from sky.utils import yaml_utils
+
+logger = sky_logging.init_logger(__name__)
+
+_DEFAULT_PLUGINS_CONFIG_PATH = '~/.sky/plugins.yaml'
+_PLUGINS_CONFIG_ENV_VAR = (
+    f'{skylet_constants.SKYPILOT_SERVER_ENV_VAR_PREFIX}PLUGINS_CONFIG')
+
+
+class ExtensionContext:
+    """Context provided to plugins during installation.
+
+    Attributes:
+        app: The FastAPI application instance.
+        rbac_rules: List of RBAC rules registered by the plugin.
+        Example:
+        [
+            ('user', RBACRule(path='/plugins/api/xx/*', method='POST')),
+            ('user', RBACRule(path='/plugins/api/xx/*', method='DELETE'))
+        ]
+    """
+
+    def __init__(self, app: Optional[FastAPI] = None):
+        self.app = app
+        self.rbac_rules: List[Tuple[str, RBACRule]] = []
+
+    def register_rbac_rule(self,
+                           path: str,
+                           method: str,
+                           description: Optional[str] = None,
+                           role: str = 'user') -> None:
+        """Register an RBAC rule for this plugin.
+
+        This method allows plugins to declare which endpoints should be
+        restricted to admin users during the install phase.
+
+        Args:
+            path: The path pattern to restrict (supports wildcards with
+                  keyMatch2).
+                  Example: '/plugins/api/credentials/*'
+            method: The HTTP method to restrict. Example: 'POST', 'DELETE'
+            description: Optional description of what this rule protects.
+            role: The role to add this rule to (default: 'user').
+                  Rules added to 'user' role block regular users but allow
+                  admins.
+
+        Example:
+            def install(self, ctx: ExtensionContext):
+                # Only admin can upload credentials
+                ctx.register_rbac_rule(
+                    path='/plugins/api/credentials/*',
+                    method='POST',
+                    description='Only admin can upload credentials'
+                )
+        """
+        rule = RBACRule(path=path, method=method, description=description)
+        self.rbac_rules.append((role, rule))
+        logger.debug(f'Registered RBAC rule for {role}: {method} {path}'
+                     f'{f" - {description}" if description else ""}')
+
+
+@dataclasses.dataclass
+class RBACRule:
+    """RBAC rule for a plugin endpoint.
+
+    Attributes:
+        path: The path pattern to match (supports wildcards with keyMatch2).
+              Example: '/plugins/api/credentials/*'
+        method: The HTTP method to restrict. Example: 'POST', 'DELETE'
+        description: Optional description of what this rule protects.
+    """
+    path: str
+    method: str
+    description: Optional[str] = None
+
+
+class BasePlugin(abc.ABC):
+    """Base class for all SkyPilot server plugins."""
+
+    @property
+    def js_extension_path(self) -> Optional[str]:
+        """Optional API route to the JavaScript extension to load."""
+        return None
+
+    @abc.abstractmethod
+    def install(self, extension_context: ExtensionContext):
+        """Hook called by API server to let the plugin install itself."""
+        raise NotImplementedError
+
+    def shutdown(self):
+        """Hook called by API server to let the plugin shutdown."""
+        pass
+
+
+def _config_schema():
+    plugin_schema = {
+        'type': 'object',
+        'required': ['class'],
+        'additionalProperties': False,
+        'properties': {
+            'class': {
+                'type': 'string',
+            },
+            'parameters': {
+                'type': 'object',
+                'required': [],
+                'additionalProperties': True,
+            },
+        },
+    }
+    return {
+        'type': 'object',
+        'required': [],
+        'additionalProperties': False,
+        'properties': {
+            'plugins': {
+                'type': 'array',
+                'items': plugin_schema,
+                'default': [],
+            },
+        },
+    }
+
+
+def _load_plugin_config() -> Optional[config_utils.Config]:
+    """Load plugin config."""
+    config_path = os.getenv(_PLUGINS_CONFIG_ENV_VAR,
+                            _DEFAULT_PLUGINS_CONFIG_PATH)
+    config_path = os.path.expanduser(config_path)
+    if not os.path.exists(config_path):
+        return None
+    config = yaml_utils.read_yaml(config_path) or {}
+    common_utils.validate_schema(config,
+                                 _config_schema(),
+                                 err_msg_prefix='Invalid plugins config: ')
+    return config_utils.Config.from_dict(config)
+
+
+_PLUGINS: Dict[str, BasePlugin] = {}
+_EXTENSION_CONTEXT: Optional[ExtensionContext] = None
+
+
+def load_plugins(extension_context: ExtensionContext):
+    """Load and initialize plugins from the config."""
+    global _EXTENSION_CONTEXT
+    _EXTENSION_CONTEXT = extension_context
+
+    config = _load_plugin_config()
+    if not config:
+        return
+
+    for plugin_config in config.get('plugins', []):
+        class_path = plugin_config['class']
+        module_path, class_name = class_path.rsplit('.', 1)
+        try:
+            module = importlib.import_module(module_path)
+        except ImportError as e:
+            raise ImportError(
+                f'Failed to import plugin module: {module_path}. '
+                'Please check if the module is installed in your Python '
+                'environment.') from e
+        try:
+            plugin_cls = getattr(module, class_name)
+        except AttributeError as e:
+            raise AttributeError(
+                f'Could not find plugin {class_name} class in module '
+                f'{module_path}. ') from e
+        if not issubclass(plugin_cls, BasePlugin):
+            raise TypeError(
+                f'Plugin {class_path} must inherit from BasePlugin.')
+        parameters = plugin_config.get('parameters') or {}
+        plugin = plugin_cls(**parameters)
+        plugin.install(extension_context)
+        _PLUGINS[class_path] = plugin
+
+
+def get_plugins() -> List[BasePlugin]:
+    """Return shallow copies of the registered plugins."""
+    return list(_PLUGINS.values())
+
+
+def get_plugin_rbac_rules() -> Dict[str, List[Dict[str, str]]]:
+    """Collect RBAC rules from all loaded plugins.
+
+    Collects rules from the ExtensionContext.
+
+    Returns:
+        Dictionary mapping role names to lists of blocklist rules.
+        Example:
+        {
+            'user': [
+                {'path': '/plugins/api/credentials/*', 'method': 'POST'},
+                {'path': '/plugins/api/credentials/*', 'method': 'DELETE'}
+            ]
+        }
+    """
+    rules_by_role: Dict[str, List[Dict[str, str]]] = {}
+
+    # Collect rules registered via ExtensionContext
+    if _EXTENSION_CONTEXT:
+        for role, rule in _EXTENSION_CONTEXT.rbac_rules:
+            if role not in rules_by_role:
+                rules_by_role[role] = []
+            rules_by_role[role].append({
+                'path': rule.path,
+                'method': rule.method,
+            })
+
+    return rules_by_role

sky/server/requests/executor.py

@@ -44,6 +44,7 @@ from sky.server import common as server_common
 from sky.server import config as server_config
 from sky.server import constants as server_constants
 from sky.server import metrics as metrics_lib
+from sky.server import plugins
 from sky.server.requests import payloads
 from sky.server.requests import preconditions
 from sky.server.requests import process
@@ -159,6 +160,8 @@ queue_backend = server_config.QueueBackend.MULTIPROCESSING
 def executor_initializer(proc_group: str):
     setproctitle.setproctitle(f'SkyPilot:executor:{proc_group}:'
                               f'{multiprocessing.current_process().pid}')
+    # Load plugins for executor process.
+    plugins.load_plugins(plugins.ExtensionContext())
     # Executor never stops, unless the whole process is killed.
     threading.Thread(target=metrics_lib.process_monitor,
                      args=(f'worker:{proc_group}', threading.Event()),
@@ -533,8 +536,8 @@ def _request_execution_wrapper(request_id: str,
         # so that the "Request xxxx failed due to ..." log message will be
         # written to the original stdout and stderr file descriptors.
         _restore_output()
-        logger.info(f'Request {request_id} failed due to '
-                    f'{common_utils.format_exception(e)}')
+        logger.error(f'Request {request_id} failed due to '
+                     f'{common_utils.format_exception(e)}')
         return
     else:
         api_requests.set_request_succeeded(

sky/server/requests/payloads.py

@@ -82,7 +82,7 @@ def request_body_env_vars() -> dict:
         if common.is_api_server_local() and env_var in EXTERNAL_LOCAL_ENV_VARS:
             env_vars[env_var] = os.environ[env_var]
     env_vars[constants.USER_ID_ENV_VAR] = common_utils.get_user_hash()
-    env_vars[constants.USER_ENV_VAR] = common_utils.get_current_user_name()
+    env_vars[constants.USER_ENV_VAR] = common_utils.get_local_user_name()
     env_vars[
         usage_constants.USAGE_RUN_ID_ENV_VAR] = usage_lib.messages.usage.run_id
     if not common.is_api_server_local():
@@ -670,6 +670,11 @@ class KubernetesNodeInfoRequestBody(RequestBody):
     context: Optional[str] = None
 
 
+class SlurmNodeInfoRequestBody(RequestBody):
+    """The request body for the slurm node info endpoint."""
+    slurm_cluster_name: Optional[str] = None
+
+
 class ListAcceleratorsBody(RequestBody):
     """The request body for the list accelerators endpoint."""
     gpus_only: bool = True
@@ -854,3 +859,9 @@ class RequestPayload(BasePayload):
     status_msg: Optional[str] = None
     should_retry: bool = False
     finished_at: Optional[float] = None
+
+
+class SlurmGpuAvailabilityRequestBody(RequestBody):
+    """Request body for getting Slurm real-time GPU availability."""
+    name_filter: Optional[str] = None
+    quantity_filter: Optional[int] = None

sky/server/requests/request_names.py

@@ -10,6 +10,8 @@ class RequestName(str, enum.Enum):
     REALTIME_KUBERNETES_GPU_AVAILABILITY = (
         'realtime_kubernetes_gpu_availability')
     KUBERNETES_NODE_INFO = 'kubernetes_node_info'
+    REALTIME_SLURM_GPU_AVAILABILITY = 'realtime_slurm_gpu_availability'
+    SLURM_NODE_INFO = 'slurm_node_info'
     STATUS_KUBERNETES = 'status_kubernetes'
     LIST_ACCELERATORS = 'list_accelerators'
     LIST_ACCELERATOR_COUNTS = 'list_accelerator_counts'

sky/server/requests/requests.py

@@ -33,6 +33,7 @@ from sky.server import daemons
 from sky.server.requests import payloads
 from sky.server.requests.serializers import decoders
 from sky.server.requests.serializers import encoders
+from sky.server.requests.serializers import return_value_serializers
 from sky.utils import asyncio_utils
 from sky.utils import common_utils
 from sky.utils import ux_utils
@@ -231,13 +232,16 @@ class Request:
         assert isinstance(self.request_body,
                           payloads.RequestBody), (self.name, self.request_body)
         try:
+            # Use version-aware serializer to handle backward compatibility
+            # for old clients that don't recognize new fields.
+            serializer = return_value_serializers.get_serializer(self.name)
             return payloads.RequestPayload(
                 request_id=self.request_id,
                 name=self.name,
                 entrypoint=encoders.pickle_and_encode(self.entrypoint),
                 request_body=encoders.pickle_and_encode(self.request_body),
                 status=self.status.value,
-                return_value=orjson.dumps(self.return_value).decode('utf-8'),
+                return_value=serializer(self.return_value),
                 error=orjson.dumps(self.error).decode('utf-8'),
                 pid=self.pid,
                 created_at=self.created_at,

sky/server/requests/serializers/encoders.py

@@ -266,6 +266,23 @@ def encode_realtime_gpu_availability(
     return encoded
 
 
+@register_encoder('realtime_slurm_gpu_availability')
+def encode_realtime_slurm_gpu_availability(
+    return_value: List[Tuple[str,
+                             List[Any]]]) -> List[Tuple[str, List[List[Any]]]]:
+    # Convert RealtimeGpuAvailability namedtuples to lists
+    # for JSON serialization.
+    encoded = []
+    for context, gpu_list in return_value:
+        converted_gpu_list = []
+        for gpu in gpu_list:
+            assert isinstance(gpu, models.RealtimeGpuAvailability), (
+                f'Expected RealtimeGpuAvailability, got {type(gpu)}')
+            converted_gpu_list.append(list(gpu))
+        encoded.append((context, converted_gpu_list))
+    return encoded
+
+
 @register_encoder('list_accelerators')
 def encode_list_accelerators(
         return_value: Dict[str, List[Any]]) -> Dict[str, Any]:

sky/server/requests/serializers/return_value_serializers.py

@@ -0,0 +1,60 @@
+"""Version-aware serializers for request return values.
+
+These serializers run at encode() time when remote_api_version is available,
+to handle backward compatibility for old clients.
+
+The existing encoders.py handles object -> dict conversion at set_return_value()
+time. This module handles dict -> JSON string serialization at encode() time,
+with version-aware field filtering for backward compatibility.
+"""
+from typing import Any, Callable, Dict
+
+import orjson
+
+from sky.server import constants as server_constants
+from sky.server import versions
+
+handlers: Dict[str, Callable[[Any], str]] = {}
+
+
+def register_serializer(*names: str):
+    """Decorator to register a version-aware serializer."""
+
+    def decorator(func):
+        for name in names:
+            if name != server_constants.DEFAULT_HANDLER_NAME:
+                name = server_constants.REQUEST_NAME_PREFIX + name
+            if name in handlers:
+                raise ValueError(f'Serializer {name} already registered: '
+                                 f'{handlers[name]}')
+            handlers[name] = func
+        return func
+
+    return decorator
+
+
+def get_serializer(name: str) -> Callable[[Any], str]:
+    """Get the serializer for a request name."""
+    return handlers.get(name, handlers[server_constants.DEFAULT_HANDLER_NAME])
+
+
+@register_serializer(server_constants.DEFAULT_HANDLER_NAME)
+def default_serializer(return_value: Any) -> str:
+    """The default serializer."""
+    return orjson.dumps(return_value).decode('utf-8')
+
+
+@register_serializer('kubernetes_node_info')
+def serialize_kubernetes_node_info(return_value: Dict[str, Any]) -> str:
+    """Serialize kubernetes node info with version compatibility.
+
+    The is_ready field was added in API version 25. Remove it for old clients
+    that don't recognize it.
+    """
+    remote_api_version = versions.get_remote_api_version()
+    if (return_value and remote_api_version is not None and
+            remote_api_version < 25):
+        # Remove is_ready field for old clients that don't recognize it
+        for node_info in return_value.get('node_info_dict', {}).values():
+            node_info.pop('is_ready', None)
+    return orjson.dumps(return_value).decode('utf-8')

sky/server/server.py

@@ -20,7 +20,7 @@ import struct
 import sys
 import threading
 import traceback
-from typing import Dict, List, Literal, Optional, Set, Tuple
+from typing import Any, Dict, List, Literal, Optional, Set, Tuple
 import uuid
 import zipfile
 
@@ -48,6 +48,7 @@ from sky.jobs.server import server as jobs_rest
 from sky.metrics import utils as metrics_utils
 from sky.provision import metadata_utils
 from sky.provision.kubernetes import utils as kubernetes_utils
+from sky.provision.slurm import utils as slurm_utils
 from sky.schemas.api import responses
 from sky.serve.server import server as serve_rest
 from sky.server import common
@@ -56,6 +57,8 @@ from sky.server import constants as server_constants
 from sky.server import daemons
 from sky.server import metrics
 from sky.server import middleware_utils
+from sky.server import plugins
+from sky.server import server_utils
 from sky.server import state
 from sky.server import stream_utils
 from sky.server import versions
@@ -470,7 +473,8 @@ async def schedule_on_boot_check_async():
         await executor.schedule_request_async(
             request_id='skypilot-server-on-boot-check',
             request_name=request_names.RequestName.CHECK,
-            request_body=payloads.CheckBody(),
+            request_body=server_utils.build_body_at_server(
+                request=None, body_type=payloads.CheckBody),
             func=sky_check.check,
             schedule_type=requests_lib.ScheduleType.SHORT,
             is_skypilot_system=True,
@@ -493,7 +497,8 @@ async def lifespan(app: fastapi.FastAPI):  # pylint: disable=redefined-outer-nam
             await executor.schedule_request_async(
                 request_id=event.id,
                 request_name=event.name,
-                request_body=payloads.RequestBody(),
+                request_body=server_utils.build_body_at_server(
+                    request=None, body_type=payloads.RequestBody),
                 func=event.run_event,
                 schedule_type=requests_lib.ScheduleType.SHORT,
                 is_skypilot_system=True,
@@ -652,6 +657,17 @@ app.add_middleware(BearerTokenMiddleware)
 # middleware above.
 app.add_middleware(InitializeRequestAuthUserMiddleware)
 app.add_middleware(RequestIDMiddleware)
+
+# Load plugins after all the middlewares are added, to keep the core
+# middleware stack intact if a plugin adds new middlewares.
+# Note: server.py will be imported twice in server process, once as
+# the top-level entrypoint module and once imported by uvicorn, we only
+# load the plugin when imported by uvicorn for server process.
+# TODO(aylei): move uvicorn app out of the top-level module to avoid
+# duplicate app initialization.
+if __name__ == 'sky.server.server':
+    plugins.load_plugins(plugins.ExtensionContext(app=app))
+
 app.include_router(jobs_rest.router, prefix='/jobs', tags=['jobs'])
 app.include_router(serve_rest.router, prefix='/serve', tags=['serve'])
 app.include_router(users_rest.router, prefix='/users', tags=['users'])
@@ -746,8 +762,11 @@ async def enabled_clouds(request: fastapi.Request,
     await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name=request_names.RequestName.ENABLED_CLOUDS,
-        request_body=payloads.EnabledCloudsBody(workspace=workspace,
-                                                expand=expand),
+        request_body=server_utils.build_body_at_server(
+            request=request,
+            body_type=payloads.EnabledCloudsBody,
+            workspace=workspace,
+            expand=expand),
         func=core.enabled_clouds,
         schedule_type=requests_lib.ScheduleType.SHORT,
     )
@@ -784,6 +803,35 @@ async def kubernetes_node_info(
     )
 
 
+@app.post('/slurm_gpu_availability')
+async def slurm_gpu_availability(
+    request: fastapi.Request,
+    slurm_gpu_availability_body: payloads.SlurmGpuAvailabilityRequestBody
+) -> None:
+    """Gets real-time Slurm GPU availability."""
+    await executor.schedule_request_async(
+        request_id=request.state.request_id,
+        request_name=request_names.RequestName.REALTIME_SLURM_GPU_AVAILABILITY,
+        request_body=slurm_gpu_availability_body,
+        func=core.realtime_slurm_gpu_availability,
+        schedule_type=requests_lib.ScheduleType.SHORT,
+    )
+
+
+@app.get('/slurm_node_info')
+async def slurm_node_info(
+        request: fastapi.Request,
+        slurm_node_info_body: payloads.SlurmNodeInfoRequestBody) -> None:
+    """Gets detailed information for each node in the Slurm cluster."""
+    await executor.schedule_request_async(
+        request_id=request.state.request_id,
+        request_name=request_names.RequestName.SLURM_NODE_INFO,
+        request_body=slurm_node_info_body,
+        func=slurm_utils.slurm_node_info,
+        schedule_type=requests_lib.ScheduleType.SHORT,
+    )
+
+
 @app.get('/status_kubernetes')
 async def status_kubernetes(request: fastapi.Request) -> None:
     """[Experimental] Get all SkyPilot resources (including from other '
@@ -791,7 +839,8 @@ async def status_kubernetes(request: fastapi.Request) -> None:
     await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name=request_names.RequestName.STATUS_KUBERNETES,
-        request_body=payloads.RequestBody(),
+        request_body=server_utils.build_body_at_server(
+            request=request, body_type=payloads.RequestBody),
         func=core.status_kubernetes,
         schedule_type=requests_lib.ScheduleType.SHORT,
     )
@@ -1460,7 +1509,8 @@ async def storage_ls(request: fastapi.Request) -> None:
     await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name=request_names.RequestName.STORAGE_LS,
-        request_body=payloads.RequestBody(),
+        request_body=server_utils.build_body_at_server(
+            request=request, body_type=payloads.RequestBody),
         func=core.storage_ls,
         schedule_type=requests_lib.ScheduleType.SHORT,
     )
@@ -1752,6 +1802,15 @@ async def api_status(
         return encoded_request_tasks
 
 
+@app.get('/api/plugins', response_class=fastapi_responses.ORJSONResponse)
+async def list_plugins() -> Dict[str, List[Dict[str, Any]]]:
+    """Return metadata about loaded backend plugins."""
+    plugin_info = [{
+        'js_extension_path': plugin.js_extension_path,
+    } for plugin in plugins.get_plugins()]
+    return {'plugins': plugin_info}
+
+
 @app.get(
     '/api/health',
     # response_model_exclude_unset omits unset fields
@@ -2007,7 +2066,8 @@ async def all_contexts(request: fastapi.Request) -> None:
     await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name=request_names.RequestName.ALL_CONTEXTS,
-        request_body=payloads.RequestBody(),
+        request_body=server_utils.build_body_at_server(
+            request=request, body_type=payloads.RequestBody),
         func=core.get_all_contexts,
         schedule_type=requests_lib.ScheduleType.SHORT,
     )
@@ -2057,6 +2117,14 @@ async def serve_dashboard(full_path: str):
     if os.path.isfile(file_path):
         return fastapi.responses.FileResponse(file_path)
 
+    # Serve plugin catch-all page for any /plugins/* paths so client-side
+    # routing can bootstrap correctly.
+    if full_path == 'plugins' or full_path.startswith('plugins/'):
+        plugin_catchall = os.path.join(server_constants.DASHBOARD_DIR,
+                                       'plugins', '[...slug].html')
+        if os.path.isfile(plugin_catchall):
+            return fastapi.responses.FileResponse(plugin_catchall)
+
     # Serve index.html for client-side routing
     # e.g. /clusters, /jobs
     index_path = os.path.join(server_constants.DASHBOARD_DIR, 'index.html')
@@ -2220,6 +2288,8 @@ if __name__ == '__main__':
 
         for gt in global_tasks:
             gt.cancel()
+        for plugin in plugins.get_plugins():
+            plugin.shutdown()
         subprocess_utils.run_in_parallel(lambda worker: worker.cancel(),
                                          workers,
                                          num_threads=len(workers))

sky/server/server_utils.py

@@ -0,0 +1,30 @@
+"""Utilities for the API server."""
+
+from typing import Optional, Type, TypeVar
+
+import fastapi
+
+from sky.server.requests import payloads
+from sky.skylet import constants
+
+_BodyT = TypeVar('_BodyT', bound=payloads.RequestBody)
+
+
+# TODO(aylei): remove this and disable request body construction at server-side
+def build_body_at_server(request: Optional[fastapi.Request],
+                         body_type: Type[_BodyT], **data) -> _BodyT:
+    """Builds the request body at the server.
+
+    For historical reasons, some handlers mimic a client request body
+    at server-side in order to coordinate with the interface of executor.
+    This will cause issues where the client info like user identity is not
+    respected in these handlers. This function is a helper to build the request
+    body at server-side with the auth user overridden.
+    """
+    request_body = body_type(**data)
+    if request is not None:
+        auth_user = getattr(request.state, 'auth_user', None)
+        if auth_user:
+            request_body.env_vars[constants.USER_ID_ENV_VAR] = auth_user.id
+            request_body.env_vars[constants.USER_ENV_VAR] = auth_user.name
+    return request_body

sky/setup_files/dependencies.py

@@ -84,6 +84,7 @@ install_requires = [
     'bcrypt==4.0.1',
     'pyjwt',
     'gitpython',
+    'paramiko',
     'types-paramiko',
     'alembic',
     'aiohttp',
@@ -234,6 +235,7 @@ cloud_dependencies: Dict[str, List[str]] = {
     'hyperbolic': [],  # No dependencies needed for hyperbolic
     'seeweb': ['ecsapi==0.4.0'],
     'shadeform': [],  # No dependencies needed for shadeform
+    'slurm': [],  # No dependencies needed for slurm
 }
 
 # Calculate which clouds should be included in the [all] installation.

sky/skylet/attempt_skylet.py

@@ -9,6 +9,7 @@ import psutil
 
 from sky.skylet import constants
 from sky.skylet import runtime_utils
+from sky.utils import common_utils
 
 VERSION_FILE = runtime_utils.get_runtime_dir_path(constants.SKYLET_VERSION_FILE)
 SKYLET_LOG_FILE = runtime_utils.get_runtime_dir_path(constants.SKYLET_LOG_FILE)
@@ -97,8 +98,13 @@ def restart_skylet():
     for pid in _find_running_skylet_pids():
         try:
             os.kill(pid, signal.SIGKILL)
-        except (OSError, ProcessLookupError):
-            # Process died between detection and kill
+            # Wait until process fully terminates so its socket gets released.
+            # Without this, find_free_port may race with the kernel closing the
+            # socket and fail to bind to the port that's supposed to be free.
+            psutil.Process(pid).wait(timeout=5)
+        except (OSError, ProcessLookupError, psutil.NoSuchProcess,
+                psutil.TimeoutExpired):
+            # Process died between detection and kill, or timeout waiting
             pass
     # Clean up the PID file
     try:
@@ -106,7 +112,11 @@ def restart_skylet():
     except OSError:
         pass  # Best effort cleanup
 
-    port = constants.SKYLET_GRPC_PORT
+    # TODO(kevin): Handle race conditions here. Race conditions can only
+    # happen on Slurm, where there could be multiple clusters running in
+    # one network namespace. For other clouds, the behaviour will be that
+    # it always gets port 46590 (default port).
+    port = common_utils.find_free_port(constants.SKYLET_GRPC_PORT)
     subprocess.run(
         # We have made sure that `attempt_skylet.py` is executed with the
         # skypilot runtime env activated, so that skylet can access the cloud