skypilot-nightly 1.0.0.dev20250915__py3-none-any.whl → 1.0.0.dev20250918__py3-none-any.whl

sky/serve/server/impl.py

@@ -5,6 +5,7 @@ import shlex
 import signal
 import tempfile
 import threading
+import typing
 from typing import Any, Dict, List, Optional, Set, Tuple, Union
 import uuid
 
@@ -17,10 +18,12 @@ from sky import execution
 from sky import sky_logging
 from sky import skypilot_config
 from sky import task as task_lib
+from sky.adaptors import common as adaptors_common
 from sky.backends import backend_utils
 from sky.catalog import common as service_catalog_common
 from sky.data import storage as storage_lib
 from sky.serve import constants as serve_constants
+from sky.serve import serve_rpc_utils
 from sky.serve import serve_state
 from sky.serve import serve_utils
 from sky.skylet import constants
@@ -36,6 +39,11 @@ from sky.utils import subprocess_utils
 from sky.utils import ux_utils
 from sky.utils import yaml_utils
 
+if typing.TYPE_CHECKING:
+    import grpc
+else:
+    grpc = adaptors_common.LazyImport('grpc')
+
 logger = sky_logging.init_logger(__name__)
 
 
@@ -78,24 +86,35 @@ def _get_service_record(
     """Get the service record."""
     noun = 'pool' if pool else 'service'
 
-    code = serve_utils.ServeCodeGen.get_service_status([service_name],
-                                                       pool=pool)
-    returncode, serve_status_payload, stderr = backend.run_on_head(
-        handle,
-        code,
-        require_outputs=True,
-        stream_logs=False,
-        separate_stderr=True)
-    try:
-        subprocess_utils.handle_returncode(returncode,
-                                           code,
-                                           f'Failed to get {noun} status',
-                                           stderr,
-                                           stream_logs=True)
-    except exceptions.CommandError as e:
-        raise RuntimeError(e.error_msg) from e
+    assert isinstance(handle, backends.CloudVmRayResourceHandle)
+    use_legacy = not handle.is_grpc_enabled_with_flag
 
-    service_statuses = serve_utils.load_service_status(serve_status_payload)
+    if handle.is_grpc_enabled_with_flag:
+        try:
+            service_statuses = serve_rpc_utils.RpcRunner.get_service_status(
+                handle, [service_name], pool)
+        except exceptions.SkyletMethodNotImplementedError:
+            use_legacy = True
+
+    if use_legacy:
+        code = serve_utils.ServeCodeGen.get_service_status([service_name],
+                                                           pool=pool)
+        returncode, serve_status_payload, stderr = backend.run_on_head(
+            handle,
+            code,
+            require_outputs=True,
+            stream_logs=False,
+            separate_stderr=True)
+        try:
+            subprocess_utils.handle_returncode(returncode,
+                                               code,
+                                               f'Failed to get {noun} status',
+                                               stderr,
+                                               stream_logs=True)
+        except exceptions.CommandError as e:
+            raise RuntimeError(e.error_msg) from e
+
+        service_statuses = serve_utils.load_service_status(serve_status_payload)
 
     assert len(service_statuses) <= 1, service_statuses
     if not service_statuses:
@@ -287,30 +306,44 @@ def up(
         fore = colorama.Fore
 
         assert controller_job_id is not None and controller_handle is not None
+        assert isinstance(controller_handle, backends.CloudVmRayResourceHandle)
+        backend = backend_utils.get_backend_from_handle(controller_handle)
+        assert isinstance(backend, backends.CloudVmRayBackend)
         # TODO(tian): Cache endpoint locally to speedup. Endpoint won't
         # change after the first time, so there is no consistency issue.
-        with rich_utils.safe_status(
-                ux_utils.spinner_message(
-                    f'Waiting for the {noun} to register')):
-            # This function will check the controller job id in the database
-            # and return the endpoint if the job id matches. Otherwise it will
-            # return None.
-            code = serve_utils.ServeCodeGen.wait_service_registration(
-                service_name, controller_job_id, pool)
-            backend = backend_utils.get_backend_from_handle(controller_handle)
-            assert isinstance(backend, backends.CloudVmRayBackend)
-            assert isinstance(controller_handle,
-                              backends.CloudVmRayResourceHandle)
-            returncode, lb_port_payload, _ = backend.run_on_head(
-                controller_handle,
-                code,
-                require_outputs=True,
-                stream_logs=False)
         try:
-            subprocess_utils.handle_returncode(
-                returncode, code, f'Failed to wait for {noun} initialization',
-                lb_port_payload)
-        except exceptions.CommandError:
+            with rich_utils.safe_status(
+                    ux_utils.spinner_message(
+                        f'Waiting for the {noun} to register')):
+                # This function will check the controller job id in the database
+                # and return the endpoint if the job id matches. Otherwise it
+                # will return None.
+                use_legacy = not controller_handle.is_grpc_enabled_with_flag
+
+                if controller_handle.is_grpc_enabled_with_flag:
+                    try:
+                        lb_port = serve_rpc_utils.RpcRunner.wait_service_registration(  # pylint: disable=line-too-long
+                            controller_handle, service_name, controller_job_id,
+                            pool)
+                    except exceptions.SkyletMethodNotImplementedError:
+                        use_legacy = True
+
+                if use_legacy:
+                    code = serve_utils.ServeCodeGen.wait_service_registration(
+                        service_name, controller_job_id, pool)
+                    returncode, lb_port_payload, _ = backend.run_on_head(
+                        controller_handle,
+                        code,
+                        require_outputs=True,
+                        stream_logs=False)
+                    subprocess_utils.handle_returncode(
+                        returncode, code,
+                        f'Failed to wait for {noun} initialization',
+                        lb_port_payload)
+                    lb_port = serve_utils.load_service_initialization_result(
+                        lb_port_payload)
+        except (exceptions.CommandError, grpc.FutureTimeoutError,
+                grpc.RpcError):
             if serve_utils.is_consolidation_mode(pool):
                 with ux_utils.print_exception_no_traceback():
                     raise RuntimeError(
@@ -344,8 +377,6 @@ def up(
                         'Failed to spin up the service. Please '
                         'check the logs above for more details.') from None
         else:
-            lb_port = serve_utils.load_service_initialization_result(
-                lb_port_payload)
             if not serve_utils.is_consolidation_mode(pool) and not pool:
                 socket_endpoint = backend_utils.get_endpoints(
                     controller_handle.cluster_name,
@@ -461,6 +492,7 @@ def update(
         f'use {ux_utils.BOLD}sky serve up{ux_utils.RESET_BOLD}',
     )
 
+    assert isinstance(handle, backends.CloudVmRayResourceHandle)
     backend = backend_utils.get_backend_from_handle(handle)
     assert isinstance(backend, backends.CloudVmRayBackend)
 
@@ -503,29 +535,39 @@ def update(
         controller_utils.maybe_translate_local_file_mounts_and_sync_up(
             task, task_type='serve')
 
-    code = serve_utils.ServeCodeGen.add_version(service_name)
-    returncode, version_string_payload, stderr = backend.run_on_head(
-        handle,
-        code,
-        require_outputs=True,
-        stream_logs=False,
-        separate_stderr=True)
-    try:
-        subprocess_utils.handle_returncode(returncode,
-                                           code,
-                                           'Failed to add version',
-                                           stderr,
-                                           stream_logs=True)
-    except exceptions.CommandError as e:
-        raise RuntimeError(e.error_msg) from e
+    use_legacy = not handle.is_grpc_enabled_with_flag
 
-    version_string = serve_utils.load_version_string(version_string_payload)
-    try:
-        current_version = int(version_string)
-    except ValueError as e:
-        with ux_utils.print_exception_no_traceback():
-            raise ValueError(f'Failed to parse version: {version_string}; '
-                             f'Returncode: {returncode}') from e
+    if handle.is_grpc_enabled_with_flag:
+        try:
+            current_version = serve_rpc_utils.RpcRunner.add_version(
+                handle, service_name)
+        except exceptions.SkyletMethodNotImplementedError:
+            use_legacy = True
+
+    if use_legacy:
+        code = serve_utils.ServeCodeGen.add_version(service_name)
+        returncode, version_string_payload, stderr = backend.run_on_head(
+            handle,
+            code,
+            require_outputs=True,
+            stream_logs=False,
+            separate_stderr=True)
+        try:
+            subprocess_utils.handle_returncode(returncode,
+                                               code,
+                                               'Failed to add version',
+                                               stderr,
+                                               stream_logs=True)
+        except exceptions.CommandError as e:
+            raise RuntimeError(e.error_msg) from e
+
+        version_string = serve_utils.load_version_string(version_string_payload)
+        try:
+            current_version = int(version_string)
+        except ValueError as e:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(f'Failed to parse version: {version_string}; '
+                                 f'Returncode: {returncode}') from e
 
     with tempfile.NamedTemporaryFile(
             prefix=f'{service_name}-v{current_version}',
@@ -540,23 +582,33 @@ def update(
                                      {remote_task_yaml_path: service_file.name},
                                      storage_mounts=None)
 
-        code = serve_utils.ServeCodeGen.update_service(service_name,
-                                                       current_version,
-                                                       mode=mode.value,
-                                                       pool=pool)
-        returncode, _, stderr = backend.run_on_head(handle,
-                                                    code,
-                                                    require_outputs=True,
-                                                    stream_logs=False,
-                                                    separate_stderr=True)
-        try:
-            subprocess_utils.handle_returncode(returncode,
-                                               code,
-                                               f'Failed to update {noun}s',
-                                               stderr,
-                                               stream_logs=True)
-        except exceptions.CommandError as e:
-            raise RuntimeError(e.error_msg) from e
+        use_legacy = not handle.is_grpc_enabled_with_flag
+
+        if handle.is_grpc_enabled_with_flag:
+            try:
+                serve_rpc_utils.RpcRunner.update_service(
+                    handle, service_name, current_version, mode, pool)
+            except exceptions.SkyletMethodNotImplementedError:
+                use_legacy = True
+
+        if use_legacy:
+            code = serve_utils.ServeCodeGen.update_service(service_name,
+                                                           current_version,
+                                                           mode=mode.value,
+                                                           pool=pool)
+            returncode, _, stderr = backend.run_on_head(handle,
+                                                        code,
+                                                        require_outputs=True,
+                                                        stream_logs=False,
+                                                        separate_stderr=True)
+            try:
+                subprocess_utils.handle_returncode(returncode,
+                                                   code,
+                                                   f'Failed to update {noun}s',
+                                                   stderr,
+                                                   stream_logs=True)
+            except exceptions.CommandError as e:
+                raise RuntimeError(e.error_msg) from e
 
     cmd = 'sky jobs pool status' if pool else 'sky serve status'
     logger.info(
@@ -619,29 +671,44 @@ def down(
         raise ValueError(f'Can only specify one of {noun}_names or all. '
                          f'Provided {argument_str!r}.')
 
-    backend = backend_utils.get_backend_from_handle(handle)
-    assert isinstance(backend, backends.CloudVmRayBackend)
     service_names = None if all else service_names
-    code = serve_utils.ServeCodeGen.terminate_services(service_names, purge,
-                                                       pool)
 
     try:
-        returncode, stdout, _ = backend.run_on_head(handle,
-                                                    code,
-                                                    require_outputs=True,
-                                                    stream_logs=False)
+        assert isinstance(handle, backends.CloudVmRayResourceHandle)
+        use_legacy = not handle.is_grpc_enabled_with_flag
+
+        if handle.is_grpc_enabled_with_flag:
+            try:
+                stdout = serve_rpc_utils.RpcRunner.terminate_services(
+                    handle, service_names, purge, pool)
+            except exceptions.SkyletMethodNotImplementedError:
+                use_legacy = True
+
+        if use_legacy:
+            backend = backend_utils.get_backend_from_handle(handle)
+            assert isinstance(backend, backends.CloudVmRayBackend)
+            code = serve_utils.ServeCodeGen.terminate_services(
+                service_names, purge, pool)
+
+            returncode, stdout, _ = backend.run_on_head(handle,
+                                                        code,
+                                                        require_outputs=True,
+                                                        stream_logs=False)
+
+            subprocess_utils.handle_returncode(returncode, code,
+                                               f'Failed to terminate {noun}',
+                                               stdout)
     except exceptions.FetchClusterInfoError as e:
         raise RuntimeError(
             'Failed to fetch controller IP. Please refresh controller status '
-            f'by `sky status -r {controller_type.value.cluster_name}` '
-            'and try again.') from e
-
-    try:
-        subprocess_utils.handle_returncode(returncode, code,
-                                           f'Failed to terminate {noun}',
-                                           stdout)
+            f'by `sky status -r {controller_type.value.cluster_name}` and try '
+            'again.') from e
     except exceptions.CommandError as e:
         raise RuntimeError(e.error_msg) from e
+    except grpc.RpcError as e:
+        raise RuntimeError(f'{e.details()} ({e.code()})') from e
+    except grpc.FutureTimeoutError as e:
+        raise RuntimeError('gRPC timed out') from e
 
     logger.info(stdout)
 
@@ -669,27 +736,40 @@ def status(
         stopped_message=controller_type.value.default_hint_if_non_existent.
         replace('service', noun))
 
-    backend = backend_utils.get_backend_from_handle(handle)
-    assert isinstance(backend, backends.CloudVmRayBackend)
+    assert isinstance(handle, backends.CloudVmRayResourceHandle)
+    use_legacy = not handle.is_grpc_enabled_with_flag
 
-    code = serve_utils.ServeCodeGen.get_service_status(service_names, pool=pool)
-    returncode, serve_status_payload, stderr = backend.run_on_head(
-        handle,
-        code,
-        require_outputs=True,
-        stream_logs=False,
-        separate_stderr=True)
+    if handle.is_grpc_enabled_with_flag:
+        try:
+            service_records = serve_rpc_utils.RpcRunner.get_service_status(
+                handle, service_names, pool)
+        except exceptions.SkyletMethodNotImplementedError:
+            use_legacy = True
+
+    if use_legacy:
+        backend = backend_utils.get_backend_from_handle(handle)
+        assert isinstance(backend, backends.CloudVmRayBackend)
+
+        code = serve_utils.ServeCodeGen.get_service_status(service_names,
+                                                           pool=pool)
+        returncode, serve_status_payload, stderr = backend.run_on_head(
+            handle,
+            code,
+            require_outputs=True,
+            stream_logs=False,
+            separate_stderr=True)
 
-    try:
-        subprocess_utils.handle_returncode(returncode,
-                                           code,
-                                           f'Failed to fetch {noun}s',
-                                           stderr,
-                                           stream_logs=True)
-    except exceptions.CommandError as e:
-        raise RuntimeError(e.error_msg) from e
+        try:
+            subprocess_utils.handle_returncode(returncode,
+                                               code,
+                                               f'Failed to fetch {noun}s',
+                                               stderr,
+                                               stream_logs=True)
+        except exceptions.CommandError as e:
+            raise RuntimeError(e.error_msg) from e
+
+        service_records = serve_utils.load_service_status(serve_status_payload)
 
-    service_records = serve_utils.load_service_status(serve_status_payload)
     # Get the endpoint for each service
     for service_record in service_records:
         service_record['endpoint'] = None
@@ -792,25 +872,37 @@ def _get_all_replica_targets(
         handle: backends.CloudVmRayResourceHandle,
         pool: bool) -> Set[serve_utils.ServiceComponentTarget]:
     """Helper function to get targets for all live replicas."""
-    code = serve_utils.ServeCodeGen.get_service_status([service_name],
-                                                       pool=pool)
-    returncode, serve_status_payload, stderr = backend.run_on_head(
-        handle,
-        code,
-        require_outputs=True,
-        stream_logs=False,
-        separate_stderr=True)
+    assert isinstance(handle, backends.CloudVmRayResourceHandle)
+    use_legacy = not handle.is_grpc_enabled_with_flag
 
-    try:
-        subprocess_utils.handle_returncode(returncode,
-                                           code,
-                                           'Failed to fetch services',
-                                           stderr,
-                                           stream_logs=True)
-    except exceptions.CommandError as e:
-        raise RuntimeError(e.error_msg) from e
+    if handle.is_grpc_enabled_with_flag:
+        try:
+            service_records = serve_rpc_utils.RpcRunner.get_service_status(
+                handle, [service_name], pool)
+        except exceptions.SkyletMethodNotImplementedError:
+            use_legacy = True
+
+    if use_legacy:
+        code = serve_utils.ServeCodeGen.get_service_status([service_name],
+                                                           pool=pool)
+        returncode, serve_status_payload, stderr = backend.run_on_head(
+            handle,
+            code,
+            require_outputs=True,
+            stream_logs=False,
+            separate_stderr=True)
+
+        try:
+            subprocess_utils.handle_returncode(returncode,
+                                               code,
+                                               'Failed to fetch services',
+                                               stderr,
+                                               stream_logs=True)
+        except exceptions.CommandError as e:
+            raise RuntimeError(e.error_msg) from e
+
+        service_records = serve_utils.load_service_status(serve_status_payload)
 
-    service_records = serve_utils.load_service_status(serve_status_payload)
     if not service_records:
         raise ValueError(f'Service {service_name!r} not found.')
     assert len(service_records) == 1

sky/server/common.py

@@ -515,6 +515,19 @@ def get_request_id(response: 'requests.Response') -> RequestId[T]:
     return RequestId[T](request_id)
 
 
+def get_stream_request_id(
+        response: 'requests.Response') -> Optional[RequestId[T]]:
+    """This is same as the above function, but just for `sdk.stream_and_get.
+    We do this because `/api/stream` may choose the latest request id, and
+    we need to keep track of that information. Request id in this case can
+    be None."""
+    handle_request_error(response)
+    request_id = response.headers.get(server_constants.STREAM_REQUEST_HEADER)
+    if request_id is not None:
+        return RequestId[T](request_id)
+    return None
+
+
 def _start_api_server(deploy: bool = False,
                       host: str = '127.0.0.1',
                       foreground: bool = False,

sky/server/constants.py

@@ -61,3 +61,6 @@ DASHBOARD_DIR = os.path.join(os.path.dirname(__file__), '..', 'dashboard',
 
 # The interval (seconds) for the event to be restarted in the background.
 DAEMON_RESTART_INTERVAL_SECONDS = 20
+
+# Cookie header for stream request id.
+STREAM_REQUEST_HEADER = 'X-SkyPilot-Stream-Request-ID'

sky/server/requests/executor.py

@@ -282,8 +282,8 @@ def _get_queue(schedule_type: api_requests.ScheduleType) -> RequestQueue:
 
 @contextlib.contextmanager
 def override_request_env_and_config(
-        request_body: payloads.RequestBody,
-        request_id: str) -> Generator[None, None, None]:
+        request_body: payloads.RequestBody, request_id: str,
+        request_name: str) -> Generator[None, None, None]:
     """Override the environment and SkyPilot config for a request."""
     original_env = os.environ.copy()
     try:
@@ -319,9 +319,22 @@ def override_request_env_and_config(
         with skypilot_config.override_skypilot_config(
                 request_body.override_skypilot_config,
                 request_body.override_skypilot_config_path):
-            # Rejecting requests to workspaces that the user does not have
-            # permission to access.
-            workspaces_core.reject_request_for_unauthorized_workspace(user)
+            # Skip permission check for sky.workspaces.get request
+            # as it is used to determine which workspaces the user
+            # has access to.
+            if request_name != 'sky.workspaces.get':
+                try:
+                    # Reject requests that the user does not have permission
+                    # to access.
+                    workspaces_core.reject_request_for_unauthorized_workspace(
+                        user)
+                except exceptions.PermissionDeniedError as e:
+                    logger.debug(
+                        f'{request_id} permission denied to workspace: '
+                        f'{skypilot_config.get_active_workspace()}: {e}')
+                    raise e
+            logger.debug(
+                f'{request_id} permission granted to {request_name} request')
             yield
     finally:
         # We need to call the save_timeline() since atexit will not be
@@ -402,7 +415,8 @@ def _request_execution_wrapper(request_id: str,
         # captured in the log file.
         try:
             with sky_logging.add_debug_log_handler(request_id), \
-                override_request_env_and_config(request_body, request_id), \
+                override_request_env_and_config(
+                    request_body, request_id, request_name), \
                 tempstore.tempdir():
                 if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
                     config = skypilot_config.to_dict()
@@ -451,6 +465,9 @@ def _request_execution_wrapper(request_id: str,
                 # Capture the peak RSS before GC.
                 peak_rss = max(proc.memory_info().rss,
                                metrics_lib.peak_rss_bytes)
+                # Clear request level cache to release all memory used by
+                # the request.
+                annotations.clear_request_level_cache()
                 with metrics_lib.time_it(name='release_memory',
                                          group='internal'):
                     common_utils.release_memory()

sky/server/server.py

@@ -1571,6 +1571,15 @@ async def stream(
                     detail=f'Log path {log_path!r} does not exist')
 
         log_path_to_stream = resolved_log_path
+
+    headers = {
+        'Cache-Control': 'no-cache, no-transform',
+        'X-Accel-Buffering': 'no',
+        'Transfer-Encoding': 'chunked'
+    }
+    if request_id is not None:
+        headers[server_constants.STREAM_REQUEST_HEADER] = request_id
+
     return fastapi.responses.StreamingResponse(
         content=stream_utils.log_streamer(request_id,
                                           log_path_to_stream,
@@ -1578,11 +1587,7 @@ async def stream(
                                           tail=tail,
                                           follow=follow),
         media_type='text/plain',
-        headers={
-            'Cache-Control': 'no-cache, no-transform',
-            'X-Accel-Buffering': 'no',
-            'Transfer-Encoding': 'chunked'
-        },
+        headers=headers,
     )
 
 

sky/setup_files/dependencies.py

@@ -189,6 +189,7 @@ extras_require: Dict[str, List[str]] = {
     'fluidstack': [],  # No dependencies needed for fluidstack
     'cudo': ['cudo-compute>=0.1.10'],
     'paperspace': [],  # No dependencies needed for paperspace
+    'primeintellect': [],  # No dependencies needed for primeintellect
     'do': ['pydo>=0.3.0', 'azure-core>=1.24.0', 'azure-common'],
     'vast': ['vastai-sdk>=0.1.12'],
     'vsphere': [

sky/skylet/constants.py

@@ -62,7 +62,8 @@ SKY_UV_INSTALL_CMD = (f'{SKY_UV_CMD} -V >/dev/null 2>&1 || '
                       'curl -LsSf https://astral.sh/uv/install.sh '
                       f'| UV_INSTALL_DIR={SKY_UV_INSTALL_DIR} sh')
 SKY_UV_PIP_CMD: str = (f'VIRTUAL_ENV={SKY_REMOTE_PYTHON_ENV} {SKY_UV_CMD} pip')
-SKY_UV_RUN_CMD: str = (f'VIRTUAL_ENV={SKY_REMOTE_PYTHON_ENV} {SKY_UV_CMD} run')
+SKY_UV_RUN_CMD: str = (
+    f'VIRTUAL_ENV={SKY_REMOTE_PYTHON_ENV} {SKY_UV_CMD} run --active')
 # Deleting the SKY_REMOTE_PYTHON_ENV_NAME from the PATH and unsetting relevant
 # VIRTUAL_ENV envvars to deactivate the environment. `deactivate` command does
 # not work when conda is used.
@@ -153,7 +154,7 @@ CONDA_INSTALLATION_COMMANDS = (
     # because for some images, conda is already installed, but not initialized.
     # In this case, we need to initialize conda and set auto_activate_base to
     # true.
-    '{ bash Miniconda3-Linux.sh -b; '
+    '{ bash Miniconda3-Linux.sh -b || true; '
     'eval "$(~/miniconda3/bin/conda shell.bash hook)" && conda init && '
     # Caller should replace {conda_auto_activate} with either true or false.
     'conda config --set auto_activate_base {conda_auto_activate} && '
@@ -456,7 +457,8 @@ CATALOG_SCHEMA_VERSION = 'v8'
 CATALOG_DIR = '~/.sky/catalogs'
 ALL_CLOUDS = ('aws', 'azure', 'gcp', 'ibm', 'lambda', 'scp', 'oci',
               'kubernetes', 'runpod', 'vast', 'vsphere', 'cudo', 'fluidstack',
-              'paperspace', 'do', 'nebius', 'ssh', 'hyperbolic', 'seeweb')
+              'paperspace', 'primeintellect', 'do', 'nebius', 'ssh',
+              'hyperbolic', 'seeweb')
 # END constants used for service catalog.
 
 # The user ID of the SkyPilot system.