skypilot-nightly 1.0.0.dev20250916__py3-none-any.whl → 1.0.0.dev20250919__py3-none-any.whl

sky/server/metrics.py

@@ -1,11 +1,11 @@
 """Instrumentation for the API server."""
 
-import contextlib
-import functools
+import asyncio
 import multiprocessing
 import os
 import threading
 import time
+from typing import List
 
 import fastapi
 from prometheus_client import generate_latest
@@ -15,112 +15,12 @@ import psutil
 import starlette.middleware.base
 import uvicorn
 
+from sky import core
 from sky import sky_logging
-from sky.skylet import constants
-
-# Whether the metrics are enabled, cannot be changed at runtime.
-METRICS_ENABLED = os.environ.get(constants.ENV_VAR_SERVER_METRICS_ENABLED,
-                                 'false').lower() == 'true'
-
-_KB = 2**10
-_MB = 2**20
-_MEM_BUCKETS = [
-    _KB,
-    256 * _KB,
-    512 * _KB,
-    _MB,
-    2 * _MB,
-    4 * _MB,
-    8 * _MB,
-    16 * _MB,
-    32 * _MB,
-    64 * _MB,
-    128 * _MB,
-    256 * _MB,
-    float('inf'),
-]
+from sky.metrics import utils as metrics_utils
 
 logger = sky_logging.init_logger(__name__)
 
-# Total number of API server requests, grouped by path, method, and status.
-SKY_APISERVER_REQUESTS_TOTAL = prom.Counter(
-    'sky_apiserver_requests_total',
-    'Total number of API server requests',
-    ['path', 'method', 'status'],
-)
-
-# Time spent processing API server requests, grouped by path, method, and
-# status.
-SKY_APISERVER_REQUEST_DURATION_SECONDS = prom.Histogram(
-    'sky_apiserver_request_duration_seconds',
-    'Time spent processing API server requests',
-    ['path', 'method', 'status'],
-    buckets=(0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 20.0, 30.0,
-             60.0, 120.0, float('inf')),
-)
-
-# Time spent processing a piece of code, refer to time_it().
-SKY_APISERVER_CODE_DURATION_SECONDS = prom.Histogram(
-    'sky_apiserver_code_duration_seconds',
-    'Time spent processing code',
-    ['name', 'group'],
-    buckets=(0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 20.0, 30.0,
-             60.0, 120.0, float('inf')),
-)
-
-SKY_APISERVER_EVENT_LOOP_LAG_SECONDS = prom.Histogram(
-    'sky_apiserver_event_loop_lag_seconds',
-    'Scheduling delay of the server event loop',
-    ['pid'],
-    buckets=(0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2, 5, 20.0,
-             60.0, float('inf')),
-)
-
-SKY_APISERVER_WEBSOCKET_CONNECTIONS = prom.Gauge(
-    'sky_apiserver_websocket_connections',
-    'Number of websocket connections',
-    ['pid'],
-    multiprocess_mode='livesum',
-)
-
-SKY_APISERVER_WEBSOCKET_CLOSED_TOTAL = prom.Counter(
-    'sky_apiserver_websocket_closed_total',
-    'Number of websocket closed',
-    ['pid', 'reason'],
-)
-
-# The number of execution starts in each worker process, we do not record
-# histogram here as the duration has been measured in
-# SKY_APISERVER_CODE_DURATION_SECONDS without the worker label (process id).
-# Recording histogram WITH worker label will cause high cardinality.
-SKY_APISERVER_PROCESS_EXECUTION_START_TOTAL = prom.Counter(
-    'sky_apiserver_process_execution_start_total',
-    'Total number of execution starts in each worker process',
-    ['request', 'pid'],
-)
-
-SKY_APISERVER_PROCESS_PEAK_RSS = prom.Gauge(
-    'sky_apiserver_process_peak_rss',
-    'Peak RSS we saw in each process in last 30 seconds',
-    ['pid', 'type'],
-)
-
-SKY_APISERVER_PROCESS_CPU_TOTAL = prom.Gauge(
-    'sky_apiserver_process_cpu_total',
-    'Total CPU times a worker process has been running',
-    ['pid', 'type', 'mode'],
-)
-
-SKY_APISERVER_REQUEST_MEMORY_USAGE_BYTES = prom.Histogram(
-    'sky_apiserver_request_memory_usage_bytes',
-    'Peak memory usage of requests', ['name'],
-    buckets=_MEM_BUCKETS)
-
-SKY_APISERVER_REQUEST_RSS_INCR_BYTES = prom.Histogram(
-    'sky_apiserver_request_rss_incr_bytes',
-    'RSS increment after requests', ['name'],
-    buckets=_MEM_BUCKETS)
-
 metrics_app = fastapi.FastAPI()
 
 
@@ -139,6 +39,42 @@ async def metrics() -> fastapi.Response:
                             headers={'Cache-Control': 'no-cache'})
 
 
+@metrics_app.get('/gpu-metrics')
+async def gpu_metrics() -> fastapi.Response:
+    """Gets the GPU metrics from multiple external k8s clusters"""
+    contexts = core.get_all_contexts()
+    all_metrics: List[str] = []
+    successful_contexts = 0
+
+    tasks = [
+        asyncio.create_task(metrics_utils.get_metrics_for_context(context))
+        for context in contexts
+        if context != 'in-cluster'
+    ]
+
+    results = await asyncio.gather(*tasks, return_exceptions=True)
+
+    for i, result in enumerate(results):
+        if isinstance(result, Exception):
+            logger.error(
+                f'Failed to get metrics for context {contexts[i]}: {result}')
+        elif isinstance(result, BaseException):
+            # Avoid changing behavior for non-Exception BaseExceptions
+            # like KeyboardInterrupt/SystemExit: re-raise them.
+            raise result
+        else:
+            metrics_text = result
+            all_metrics.append(metrics_text)
+            successful_contexts += 1
+
+    combined_metrics = '\n\n'.join(all_metrics)
+
+    # Return as plain text for Prometheus compatibility
+    return fastapi.Response(
+        content=combined_metrics,
+        media_type='text/plain; version=0.0.4; charset=utf-8')
+
+
 def build_metrics_server(host: str, port: int) -> uvicorn.Server:
     metrics_config = uvicorn.Config(
         'sky.server.metrics:metrics_app',
@@ -182,61 +118,17 @@ class PrometheusMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
             status_code_group = '5xx'
             raise
         finally:
-            SKY_APISERVER_REQUESTS_TOTAL.labels(path=path,
-                                                method=method,
-                                                status=status_code_group).inc()
+            metrics_utils.SKY_APISERVER_REQUESTS_TOTAL.labels(
+                path=path, method=method, status=status_code_group).inc()
             if not streaming:
                 duration = time.time() - start_time
-                SKY_APISERVER_REQUEST_DURATION_SECONDS.labels(
+                metrics_utils.SKY_APISERVER_REQUEST_DURATION_SECONDS.labels(
                     path=path, method=method,
                     status=status_code_group).observe(duration)
 
         return response
 
 
-@contextlib.contextmanager
-def time_it(name: str, group: str = 'default'):
-    """Context manager to measure and record code execution duration."""
-    if not METRICS_ENABLED:
-        yield
-    else:
-        start_time = time.time()
-        try:
-            yield
-        finally:
-            duration = time.time() - start_time
-            SKY_APISERVER_CODE_DURATION_SECONDS.labels(
-                name=name, group=group).observe(duration)
-
-
-def time_me(func):
-    """Measure the duration of decorated function."""
-
-    @functools.wraps(func)
-    def wrapper(*args, **kwargs):
-        if not METRICS_ENABLED:
-            return func(*args, **kwargs)
-        name = f'{func.__module__}/{func.__name__}'
-        with time_it(name, group='function'):
-            return func(*args, **kwargs)
-
-    return wrapper
-
-
-def time_me_async(func):
-    """Measure the duration of decorated async function."""
-
-    @functools.wraps(func)
-    async def async_wrapper(*args, **kwargs):
-        if not METRICS_ENABLED:
-            return await func(*args, **kwargs)
-        name = f'{func.__module__}/{func.__name__}'
-        with time_it(name, group='function'):
-            return await func(*args, **kwargs)
-
-    return async_wrapper
-
-
 peak_rss_bytes = 0
 
 
@@ -252,13 +144,15 @@ def process_monitor(process_type: str, stop: threading.Event):
             last_bucket_end = time.time()
             bucket_peak = 0
         peak_rss_bytes = max(bucket_peak, proc.memory_info().rss)
-        SKY_APISERVER_PROCESS_PEAK_RSS.labels(
+        metrics_utils.SKY_APISERVER_PROCESS_PEAK_RSS.labels(
             pid=pid, type=process_type).set(peak_rss_bytes)
         ctimes = proc.cpu_times()
-        SKY_APISERVER_PROCESS_CPU_TOTAL.labels(pid=pid,
-                                               type=process_type,
-                                               mode='user').set(ctimes.user)
-        SKY_APISERVER_PROCESS_CPU_TOTAL.labels(pid=pid,
-                                               type=process_type,
-                                               mode='system').set(ctimes.system)
+        metrics_utils.SKY_APISERVER_PROCESS_CPU_TOTAL.labels(pid=pid,
+                                                             type=process_type,
+                                                             mode='user').set(
+                                                                 ctimes.user)
+        metrics_utils.SKY_APISERVER_PROCESS_CPU_TOTAL.labels(pid=pid,
+                                                             type=process_type,
+                                                             mode='system').set(
+                                                                 ctimes.system)
         time.sleep(1)

sky/server/requests/executor.py

@@ -39,6 +39,7 @@ from sky import global_user_state
 from sky import models
 from sky import sky_logging
 from sky import skypilot_config
+from sky.metrics import utils as metrics_utils
 from sky.server import common as server_common
 from sky.server import config as server_config
 from sky.server import constants as server_constants
@@ -422,10 +423,10 @@ def _request_execution_wrapper(request_id: str,
                     config = skypilot_config.to_dict()
                     logger.debug(f'request config: \n'
                                  f'{yaml_utils.dump_yaml_str(dict(config))}')
-                metrics_lib.SKY_APISERVER_PROCESS_EXECUTION_START_TOTAL.labels(
-                    request=request_name, pid=pid).inc()
-                with metrics_lib.time_it(name=request_name,
-                                         group='request_execution'):
+                (metrics_utils.SKY_APISERVER_PROCESS_EXECUTION_START_TOTAL.
+                 labels(request=request_name, pid=pid).inc())
+                with metrics_utils.time_it(name=request_name,
+                                           group='request_execution'):
                     return_value = func(**request_body.to_kwargs())
                 f.flush()
         except KeyboardInterrupt:
@@ -465,8 +466,11 @@ 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)
-                with metrics_lib.time_it(name='release_memory',
-                                         group='internal'):
+                # Clear request level cache to release all memory used by
+                # the request.
+                annotations.clear_request_level_cache()
+                with metrics_utils.time_it(name='release_memory',
+                                           group='internal'):
                     common_utils.release_memory()
                 _record_memory_metrics(request_name, proc, rss_begin, peak_rss)
             except Exception as e:  # pylint: disable=broad-except
@@ -490,11 +494,11 @@ def _record_memory_metrics(request_name: str, proc: psutil.Process,
     rss_end = proc.memory_info().rss
 
     # Answer "how much RSS this request contributed?"
-    metrics_lib.SKY_APISERVER_REQUEST_RSS_INCR_BYTES.labels(
+    metrics_utils.SKY_APISERVER_REQUEST_RSS_INCR_BYTES.labels(
         name=request_name).observe(max(rss_end - rss_begin, 0))
     # Estimate the memory usage by the request by capturing the
     # peak memory delta during the request execution.
-    metrics_lib.SKY_APISERVER_REQUEST_MEMORY_USAGE_BYTES.labels(
+    metrics_utils.SKY_APISERVER_REQUEST_MEMORY_USAGE_BYTES.labels(
         name=request_name).observe(max(peak_rss - rss_begin, 0))
 
 

sky/server/requests/payloads.py

@@ -792,6 +792,12 @@ class GetConfigBody(RequestBody):
 class CostReportBody(RequestBody):
     """The request body for the cost report endpoint."""
     days: Optional[int] = 30
+    # we use hashes instead of names to avoid the case where
+    # the name is not unique
+    cluster_hashes: Optional[List[str]] = None
+    # Only return fields that are needed for the dashboard
+    # summary page
+    dashboard_summary_response: bool = False
 
 
 class RequestPayload(BasePayload):

sky/server/requests/requests.py

@@ -25,10 +25,10 @@ from sky import exceptions
 from sky import global_user_state
 from sky import sky_logging
 from sky import skypilot_config
+from sky.metrics import utils as metrics_lib
 from sky.server import common as server_common
 from sky.server import constants as server_constants
 from sky.server import daemons
-from sky.server import metrics as metrics_lib
 from sky.server.requests import payloads
 from sky.server.requests.serializers import decoders
 from sky.server.requests.serializers import encoders

sky/server/requests/serializers/encoders.py

@@ -185,8 +185,9 @@ def encode_cost_report(
     for cluster_report in cost_report:
         if cluster_report['status'] is not None:
             cluster_report['status'] = cluster_report['status'].value
-        cluster_report['resources'] = pickle_and_encode(
-            cluster_report['resources'])
+        if 'resources' in cluster_report:
+            cluster_report['resources'] = pickle_and_encode(
+                cluster_report['resources'])
     return cost_report
 
 

sky/server/server.py

@@ -437,7 +437,7 @@ async def loop_lag_monitor(loop: asyncio.AbstractEventLoop,
         if lag_threshold is not None and lag > lag_threshold:
             logger.warning(f'Event loop lag {lag} seconds exceeds threshold '
                            f'{lag_threshold} seconds.')
-        metrics.SKY_APISERVER_EVENT_LOOP_LAG_SECONDS.labels(
+        metrics_utils.SKY_APISERVER_EVENT_LOOP_LAG_SECONDS.labels(
             pid=pid).observe(lag)
         target = now + interval
         loop.call_at(target, tick)
@@ -470,7 +470,7 @@ async def lifespan(app: fastapi.FastAPI):  # pylint: disable=redefined-outer-nam
             # can safely ignore the error if the task is already scheduled.
             logger.debug(f'Request {event.id} already exists.')
     asyncio.create_task(cleanup_upload_ids())
-    if metrics.METRICS_ENABLED:
+    if metrics_utils.METRICS_ENABLED:
         # Start monitoring the event loop lag in each server worker
         # event loop (process).
         asyncio.create_task(loop_lag_monitor(asyncio.get_event_loop()))
@@ -1743,7 +1743,7 @@ async def kubernetes_pod_ssh_proxy(websocket: fastapi.WebSocket,
             return
 
     logger.info(f'Starting port-forward to local port: {local_port}')
-    conn_gauge = metrics.SKY_APISERVER_WEBSOCKET_CONNECTIONS.labels(
+    conn_gauge = metrics_utils.SKY_APISERVER_WEBSOCKET_CONNECTIONS.labels(
         pid=os.getpid())
     ssh_failed = False
     websocket_closed = False
@@ -1807,14 +1807,14 @@ async def kubernetes_pod_ssh_proxy(websocket: fastapi.WebSocket,
                          'ssh websocket connection was closed. Remaining '
                          f'output: {str(stdout)}')
             reason = 'KubectlPortForwardExit'
-            metrics.SKY_APISERVER_WEBSOCKET_CLOSED_TOTAL.labels(
+            metrics_utils.SKY_APISERVER_WEBSOCKET_CLOSED_TOTAL.labels(
                 pid=os.getpid(), reason='KubectlPortForwardExit').inc()
         else:
             if ssh_failed:
                 reason = 'SSHToPodDisconnected'
             else:
                 reason = 'ClientClosed'
-        metrics.SKY_APISERVER_WEBSOCKET_CLOSED_TOTAL.labels(
+        metrics_utils.SKY_APISERVER_WEBSOCKET_CLOSED_TOTAL.labels(
             pid=os.getpid(), reason=reason).inc()
 
 
@@ -1831,42 +1831,6 @@ async def all_contexts(request: fastapi.Request) -> None:
     )
 
 
-@app.get('/gpu-metrics')
-async def gpu_metrics() -> fastapi.Response:
-    """Gets the GPU metrics from multiple external k8s clusters"""
-    contexts = core.get_all_contexts()
-    all_metrics: List[str] = []
-    successful_contexts = 0
-
-    tasks = [
-        asyncio.create_task(metrics_utils.get_metrics_for_context(context))
-        for context in contexts
-        if context != 'in-cluster'
-    ]
-
-    results = await asyncio.gather(*tasks, return_exceptions=True)
-
-    for i, result in enumerate(results):
-        if isinstance(result, Exception):
-            logger.error(
-                f'Failed to get metrics for context {contexts[i]}: {result}')
-        elif isinstance(result, BaseException):
-            # Avoid changing behavior for non-Exception BaseExceptions
-            # like KeyboardInterrupt/SystemExit: re-raise them.
-            raise result
-        else:
-            metrics_text = result
-            all_metrics.append(metrics_text)
-            successful_contexts += 1
-
-    combined_metrics = '\n\n'.join(all_metrics)
-
-    # Return as plain text for Prometheus compatibility
-    return fastapi.Response(
-        content=combined_metrics,
-        media_type='text/plain; version=0.0.4; charset=utf-8')
-
-
 # === Internal APIs ===
 @app.get('/api/completion/cluster_name')
 async def complete_cluster_name(incomplete: str,) -> List[str]:

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

@@ -29,6 +29,7 @@ SKY_REMOTE_RAY_PORT_FILE = '~/.sky/ray_port.json'
 SKY_REMOTE_RAY_TEMPDIR = '/tmp/ray_skypilot'
 SKY_REMOTE_RAY_VERSION = '2.9.3'
 
+SKY_UNSET_PYTHONPATH = 'env -u PYTHONPATH'
 # We store the absolute path of the python executable (/opt/conda/bin/python3)
 # in this file, so that any future internal commands that need to use python
 # can use this path. This is useful for the case where the user has a custom
@@ -40,7 +41,7 @@ SKY_GET_PYTHON_PATH_CMD = (f'[ -s {SKY_PYTHON_PATH_FILE} ] && '
                            f'cat {SKY_PYTHON_PATH_FILE} 2> /dev/null || '
                            'which python3')
 # Python executable, e.g., /opt/conda/bin/python3
-SKY_PYTHON_CMD = f'$({SKY_GET_PYTHON_PATH_CMD})'
+SKY_PYTHON_CMD = f'{SKY_UNSET_PYTHONPATH} $({SKY_GET_PYTHON_PATH_CMD})'
 # Prefer SKY_UV_PIP_CMD, which is faster.
 # TODO(cooperc): remove remaining usage (GCP TPU setup).
 SKY_PIP_CMD = f'{SKY_PYTHON_CMD} -m pip'
@@ -56,13 +57,16 @@ SKY_REMOTE_PYTHON_ENV: str = f'~/{SKY_REMOTE_PYTHON_ENV_NAME}'
 ACTIVATE_SKY_REMOTE_PYTHON_ENV = f'source {SKY_REMOTE_PYTHON_ENV}/bin/activate'
 # uv is used for venv and pip, much faster than python implementations.
 SKY_UV_INSTALL_DIR = '"$HOME/.local/bin"'
-SKY_UV_CMD = f'UV_SYSTEM_PYTHON=false {SKY_UV_INSTALL_DIR}/uv'
+SKY_UV_CMD = ('UV_SYSTEM_PYTHON=false '
+              f'{SKY_UNSET_PYTHONPATH} {SKY_UV_INSTALL_DIR}/uv')
 # This won't reinstall uv if it's already installed, so it's safe to re-run.
 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 '
+    '--no-project --no-config')
 # 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 +157,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 +460,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.

sky/skylet/job_lib.py

@@ -559,21 +559,20 @@ def get_jobs_info(user_hash: Optional[str] = None,
     jobs_info = []
     for job in jobs:
         jobs_info.append(
-            jobsv1_pb2.JobInfo(
-                job_id=job['job_id'],
-                job_name=job['job_name'],
-                username=job['username'],
-                submitted_at=job['submitted_at'],
-                status=job['status'].to_protobuf(),
-                run_timestamp=job['run_timestamp'],
-                start_at=job['start_at']
-                if job['start_at'] is not None else -1.0,
-                end_at=job['end_at'] if job['end_at'] is not None else 0.0,
-                resources=job['resources'] or '',
-                pid=job['pid'],
-                log_path=os.path.join(constants.SKY_LOGS_DIRECTORY,
-                                      job['run_timestamp']),
-                metadata=json.dumps(job['metadata'])))
+            jobsv1_pb2.JobInfo(job_id=job['job_id'],
+                               job_name=job['job_name'],
+                               username=job['username'],
+                               submitted_at=job['submitted_at'],
+                               status=job['status'].to_protobuf(),
+                               run_timestamp=job['run_timestamp'],
+                               start_at=job['start_at'],
+                               end_at=job['end_at'],
+                               resources=job['resources'],
+                               pid=job['pid'],
+                               log_path=os.path.join(
+                                   constants.SKY_LOGS_DIRECTORY,
+                                   job['run_timestamp']),
+                               metadata=json.dumps(job['metadata'])))
     return jobs_info
 
 

sky/skylet/services.py

@@ -10,7 +10,11 @@ from sky.schemas.generated import autostopv1_pb2
 from sky.schemas.generated import autostopv1_pb2_grpc
 from sky.schemas.generated import jobsv1_pb2
 from sky.schemas.generated import jobsv1_pb2_grpc
+from sky.schemas.generated import servev1_pb2
+from sky.schemas.generated import servev1_pb2_grpc
+from sky.serve import serve_rpc_utils
 from sky.serve import serve_state
+from sky.serve import serve_utils
 from sky.skylet import autostop_lib
 from sky.skylet import constants
 from sky.skylet import job_lib
@@ -52,6 +56,100 @@ class AutostopServiceImpl(autostopv1_pb2_grpc.AutostopServiceServicer):
             context.abort(grpc.StatusCode.INTERNAL, str(e))
 
 
+class ServeServiceImpl(servev1_pb2_grpc.ServeServiceServicer):
+    """Implementation of the ServeService gRPC service."""
+
+    # NOTE (kyuds): this grpc service will run cluster-side,
+    # thus guaranteeing that SERVE_VERSION is above 5.
+    # Therefore, we removed some SERVE_VERSION checks
+    # present in the original codegen.
+
+    def GetServiceStatus(  # type: ignore[return]
+            self, request: servev1_pb2.GetServiceStatusRequest,
+            context: grpc.ServicerContext
+    ) -> servev1_pb2.GetServiceStatusResponse:
+        """Gets serve status."""
+        try:
+            service_names, pool = (
+                serve_rpc_utils.GetServiceStatusRequestConverter.from_proto(request))  # pylint: disable=line-too-long
+            statuses = serve_utils.get_service_status_pickled(
+                service_names, pool)
+            return serve_rpc_utils.GetServiceStatusResponseConverter.to_proto(
+                statuses)
+        except Exception as e:  # pylint: disable=broad-except
+            context.abort(grpc.StatusCode.INTERNAL, str(e))
+
+    def AddVersion(  # type: ignore[return]
+            self, request: servev1_pb2.AddVersionRequest,
+            context: grpc.ServicerContext) -> servev1_pb2.AddVersionResponse:
+        """Adds serve version"""
+        try:
+            service_name = request.service_name
+            version = serve_state.add_version(service_name)
+            return servev1_pb2.AddVersionResponse(version=version)
+        except Exception as e:  # pylint: disable=broad-except
+            context.abort(grpc.StatusCode.INTERNAL, str(e))
+
+    def TerminateServices(  # type: ignore[return]
+            self, request: servev1_pb2.TerminateServicesRequest,
+            context: grpc.ServicerContext
+    ) -> servev1_pb2.TerminateServicesResponse:
+        """Terminates serve"""
+        try:
+            service_names, purge, pool = (
+                serve_rpc_utils.TerminateServicesRequestConverter.from_proto(request))  # pylint: disable=line-too-long
+            message = serve_utils.terminate_services(service_names, purge, pool)
+            return servev1_pb2.TerminateServicesResponse(message=message)
+        except Exception as e:  # pylint: disable=broad-except
+            context.abort(grpc.StatusCode.INTERNAL, str(e))
+
+    def TerminateReplica(  # type: ignore[return]
+            self, request: servev1_pb2.TerminateReplicaRequest,
+            context: grpc.ServicerContext
+    ) -> servev1_pb2.TerminateReplicaResponse:
+        """Terminate replica"""
+        try:
+            service_name = request.service_name
+            replica_id = request.replica_id
+            purge = request.purge
+            message = serve_utils.terminate_replica(service_name, replica_id,
+                                                    purge)
+            return servev1_pb2.TerminateReplicaResponse(message=message)
+        except Exception as e:  # pylint: disable=broad-except
+            context.abort(grpc.StatusCode.INTERNAL, str(e))
+
+    def WaitServiceRegistration(  # type: ignore[return]
+        self, request: servev1_pb2.WaitServiceRegistrationRequest,
+        context: grpc.ServicerContext
+    ) -> servev1_pb2.WaitServiceRegistrationResponse:
+        """Wait for service to be registered"""
+        try:
+            service_name = request.service_name
+            job_id = request.job_id
+            pool = request.pool
+            encoded = serve_utils.wait_service_registration(
+                service_name, job_id, pool)
+            lb_port = serve_utils.load_service_initialization_result(encoded)
+            return servev1_pb2.WaitServiceRegistrationResponse(lb_port=lb_port)
+        except Exception as e:  # pylint: disable=broad-except
+            context.abort(grpc.StatusCode.INTERNAL, str(e))
+
+    def UpdateService(  # type: ignore[return]
+            self, request: servev1_pb2.UpdateServiceRequest,
+            context: grpc.ServicerContext) -> servev1_pb2.UpdateServiceResponse:
+        """Update service"""
+        try:
+            service_name = request.service_name
+            version = request.version
+            mode = request.mode
+            pool = request.pool
+            serve_utils.update_service_encoded(service_name, version, mode,
+                                               pool)
+            return servev1_pb2.UpdateServiceResponse()
+        except Exception as e:  # pylint: disable=broad-except
+            context.abort(grpc.StatusCode.INTERNAL, str(e))
+
+
 class JobsServiceImpl(jobsv1_pb2_grpc.JobsServiceServicer):
     """Implementation of the JobsService gRPC service."""
 

sky/skylet/skylet.py

@@ -10,6 +10,7 @@ import sky
 from sky import sky_logging
 from sky.schemas.generated import autostopv1_pb2_grpc
 from sky.schemas.generated import jobsv1_pb2_grpc
+from sky.schemas.generated import servev1_pb2_grpc
 from sky.skylet import constants
 from sky.skylet import events
 from sky.skylet import services
@@ -50,9 +51,10 @@ def start_grpc_server(port: int = constants.SKYLET_GRPC_PORT) -> grpc.Server:
 
     autostopv1_pb2_grpc.add_AutostopServiceServicer_to_server(
         services.AutostopServiceImpl(), server)
-
     jobsv1_pb2_grpc.add_JobsServiceServicer_to_server(
         services.JobsServiceImpl(), server)
+    servev1_pb2_grpc.add_ServeServiceServicer_to_server(
+        services.ServeServiceImpl(), server)
 
     listen_addr = f'127.0.0.1:{port}'
     server.add_insecure_port(listen_addr)