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

sky/serve/serve_rpc_utils.py

@@ -0,0 +1,179 @@
+"""Rpc Utilities for SkyServe"""
+
+import typing
+from typing import Any, Dict, List, Optional, Tuple
+
+from sky import backends
+from sky.adaptors import common as adaptors_common
+from sky.backends import backend_utils
+from sky.serve import serve_utils
+
+if typing.TYPE_CHECKING:
+    from sky.schemas.generated import servev1_pb2
+else:
+    servev1_pb2 = adaptors_common.LazyImport(
+        'sky.schemas.generated.servev1_pb2')
+
+# ======================= gRPC Converters for Sky Serve =======================
+
+
+class GetServiceStatusRequestConverter:
+    """Converter for GetServiceStatusRequest"""
+
+    @classmethod
+    def to_proto(cls, service_names: Optional[List[str]],
+                 pool: bool) -> 'servev1_pb2.GetServiceStatusRequest':
+        request = servev1_pb2.GetServiceStatusRequest()
+        request.pool = pool
+        if service_names is not None:
+            request.service_names.names.extend(service_names)
+        return request
+
+    @classmethod
+    def from_proto(
+        cls, proto: 'servev1_pb2.GetServiceStatusRequest'
+    ) -> Tuple[Optional[List[str]], bool]:
+        pool = proto.pool
+        if proto.HasField('service_names'):
+            service_names = list(proto.service_names.names)
+        else:
+            service_names = None
+        return service_names, pool
+
+
+class GetServiceStatusResponseConverter:
+    """Converter for GetServiceStatusResponse"""
+
+    @classmethod
+    def to_proto(
+        cls,
+        statuses: List[Dict[str,
+                            str]]) -> 'servev1_pb2.GetServiceStatusResponse':
+        response = servev1_pb2.GetServiceStatusResponse()
+        for status in statuses:
+            added = response.statuses.add()
+            added.status.update(status)
+        return response
+
+    @classmethod
+    def from_proto(
+            cls, proto: 'servev1_pb2.GetServiceStatusResponse'
+    ) -> List[Dict[str, str]]:
+        pickled = [dict(status.status) for status in proto.statuses]
+        return pickled
+
+
+class TerminateServicesRequestConverter:
+    """Converter for TerminateServicesRequest"""
+
+    @classmethod
+    def to_proto(cls, service_names: Optional[List[str]], purge: bool,
+                 pool: bool) -> 'servev1_pb2.TerminateServicesRequest':
+        request = servev1_pb2.TerminateServicesRequest()
+        request.purge = purge
+        request.pool = pool
+        if service_names is not None:
+            request.service_names.names.extend(service_names)
+        return request
+
+    @classmethod
+    def from_proto(
+        cls, proto: 'servev1_pb2.TerminateServicesRequest'
+    ) -> Tuple[Optional[List[str]], bool, bool]:
+        purge = proto.purge
+        pool = proto.pool
+        if proto.HasField('service_names'):
+            service_names = list(proto.service_names.names)
+        else:
+            service_names = None
+        return service_names, purge, pool
+
+
+# ========================= gRPC Runner for Sky Serve =========================
+
+
+class RpcRunner:
+    """gRPC Runner for Sky Serve
+
+    The RPC runner does not catch errors, and assumes that backend handle has
+    grpc enabled.
+
+    Common exceptions raised:
+        exceptions.FetchClusterInfoError
+        exceptions.SkyletInternalError
+        grpc.RpcError
+        grpc.FutureTimeoutError
+        AssertionError
+    """
+
+    @classmethod
+    def get_service_status(cls, handle: backends.CloudVmRayResourceHandle,
+                           service_names: Optional[List[str]],
+                           pool: bool) -> List[Dict[str, Any]]:
+        assert handle.is_grpc_enabled_with_flag
+        request = GetServiceStatusRequestConverter.to_proto(service_names, pool)
+        response = backend_utils.invoke_skylet_with_retries(
+            lambda: backends.SkyletClient(handle.get_grpc_channel()
+                                         ).get_service_status(request))
+        pickled = GetServiceStatusResponseConverter.from_proto(response)
+        return serve_utils.unpickle_service_status(pickled)
+
+    @classmethod
+    def add_version(cls, handle: backends.CloudVmRayResourceHandle,
+                    service_name: str) -> int:
+        assert handle.is_grpc_enabled_with_flag
+        request = servev1_pb2.AddVersionRequest(service_name=service_name)
+        response = backend_utils.invoke_skylet_with_retries(
+            lambda: backends.SkyletClient(handle.get_grpc_channel()
+                                         ).add_serve_version(request))
+        return response.version
+
+    @classmethod
+    def terminate_services(cls, handle: backends.CloudVmRayResourceHandle,
+                           service_names: Optional[List[str]], purge: bool,
+                           pool: bool) -> str:
+        assert handle.is_grpc_enabled_with_flag
+        request = TerminateServicesRequestConverter.to_proto(
+            service_names, purge, pool)
+        response = backend_utils.invoke_skylet_with_retries(
+            lambda: backends.SkyletClient(handle.get_grpc_channel()
+                                         ).terminate_services(request))
+        return response.message
+
+    @classmethod
+    def terminate_replica(cls, handle: backends.CloudVmRayResourceHandle,
+                          service_name: str, replica_id: int,
+                          purge: bool) -> str:
+        assert handle.is_grpc_enabled_with_flag
+        request = servev1_pb2.TerminateReplicaRequest(service_name=service_name,
+                                                      replica_id=replica_id,
+                                                      purge=purge)
+        response = backend_utils.invoke_skylet_with_retries(
+            lambda: backends.SkyletClient(handle.get_grpc_channel()
+                                         ).terminate_replica(request))
+        return response.message
+
+    @classmethod
+    def wait_service_registration(cls,
+                                  handle: backends.CloudVmRayResourceHandle,
+                                  service_name: str, job_id: int,
+                                  pool: bool) -> int:
+        assert handle.is_grpc_enabled_with_flag
+        request = servev1_pb2.WaitServiceRegistrationRequest(
+            service_name=service_name, job_id=job_id, pool=pool)
+        response = backend_utils.invoke_skylet_with_retries(
+            lambda: backends.SkyletClient(handle.get_grpc_channel()
+                                         ).wait_service_registration(request))
+        return response.lb_port
+
+    @classmethod
+    def update_service(cls, handle: backends.CloudVmRayResourceHandle,
+                       service_name: str, version: int,
+                       mode: serve_utils.UpdateMode, pool: bool) -> None:
+        assert handle.is_grpc_enabled_with_flag
+        request = servev1_pb2.UpdateServiceRequest(service_name=service_name,
+                                                   version=version,
+                                                   mode=mode.value,
+                                                   pool=pool)
+        backend_utils.invoke_skylet_with_retries(lambda: backends.SkyletClient(
+            handle.get_grpc_channel()).update_service(request))

sky/serve/serve_state.py

@@ -86,6 +86,7 @@ version_specs_table = sqlalchemy.Table(
     sqlalchemy.Column('service_name', sqlalchemy.Text, primary_key=True),
     sqlalchemy.Column('version', sqlalchemy.Integer, primary_key=True),
     sqlalchemy.Column('spec', sqlalchemy.LargeBinary),
+    sqlalchemy.Column('yaml_content', sqlalchemy.Text, server_default=None),
 )
 
 serve_ha_recovery_script_table = sqlalchemy.Table(
@@ -472,6 +473,7 @@ def _get_service_from_row(r: 'row.RowMapping') -> Dict[str, Any]:
         'controller_pid': r['controller_pid'],
         'hash': r['hash'],
         'entrypoint': r['entrypoint'],
+        'yaml_content': r.get('yaml_content'),
     }
     latest_spec = get_spec(r['name'], current_version)
     if latest_spec is not None:
@@ -480,21 +482,48 @@ def _get_service_from_row(r: 'row.RowMapping') -> Dict[str, Any]:
     return record
 
 
+def _build_services_with_latest_version_query(
+        service_name: Optional[str] = None) -> sqlalchemy.sql.Select:
+    """Builds a query joining services with their latest version and yaml.
+
+    Args:
+        service_name: If provided, filter to this service only.
+
+    Returns:
+        A SQLAlchemy selectable for fetching rows, including columns:
+        - max_version (latest version per service)
+        - services_table.*
+        - yaml_content (from version_specs_table for latest version)
+    """
+    subquery = sqlalchemy.select(
+        version_specs_table.c.service_name,
+        sqlalchemy.func.max(version_specs_table.c.version).label('max_version'),
+    ).group_by(version_specs_table.c.service_name).alias('v')
+
+    query = sqlalchemy.select(
+        subquery.c.max_version,
+        services_table,
+        version_specs_table.c.yaml_content,
+    ).select_from(
+        services_table.join(
+            subquery, services_table.c.name == subquery.c.service_name).join(
+                version_specs_table,
+                sqlalchemy.and_(
+                    version_specs_table.c.service_name == services_table.c.name,
+                    version_specs_table.c.version == subquery.c.max_version,
+                ),
+            ))
+    if service_name is not None:
+        query = query.where(services_table.c.name == service_name)
+    return query
+
+
 @init_db
 def get_services() -> List[Dict[str, Any]]:
     """Get all existing service records."""
     assert _SQLALCHEMY_ENGINE is not None
     with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        subquery = sqlalchemy.select(
-            version_specs_table.c.service_name,
-            sqlalchemy.func.max(
-                version_specs_table.c.version).label('max_version')).group_by(
-                    version_specs_table.c.service_name).alias('v')
-
-        query = sqlalchemy.select(
-            subquery.c.max_version, services_table).select_from(
-                services_table.join(
-                    subquery, services_table.c.name == subquery.c.service_name))
+        query = _build_services_with_latest_version_query()
         rows = session.execute(query).fetchall()
     records = []
     for row in rows:
@@ -517,20 +546,7 @@ def get_service_from_name(service_name: str) -> Optional[Dict[str, Any]]:
     """Get all existing service records."""
     assert _SQLALCHEMY_ENGINE is not None
     with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        subquery = sqlalchemy.select(
-            version_specs_table.c.service_name,
-            sqlalchemy.func.max(
-                version_specs_table.c.version).label('max_version')
-        ).where(version_specs_table.c.service_name == service_name).group_by(
-            version_specs_table.c.service_name).alias('v')
-
-        query = sqlalchemy.select(
-            subquery.c.max_version, services_table).select_from(
-                services_table.join(
-                    subquery,
-                    services_table.c.name == subquery.c.service_name)).where(
-                        services_table.c.name == service_name)
-
+        query = _build_services_with_latest_version_query(service_name)
         rows = session.execute(query).fetchall()
     for row in rows:
         return _get_service_from_row(row._mapping)  # pylint: disable=protected-access
@@ -686,22 +702,6 @@ def total_number_terminating_replicas() -> int:
     return terminating_count
 
 
-@init_db
-def total_number_scheduled_to_terminate_replicas() -> int:
-    """Returns the total number of terminating replicas."""
-    assert _SQLALCHEMY_ENGINE is not None
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        rows = session.execute(sqlalchemy.select(
-            replicas_table.c.replica_info)).fetchall()
-    terminating_count = 0
-    for row in rows:
-        replica_info: 'replica_managers.ReplicaInfo' = pickle.loads(row[0])
-        if (replica_info.status_property.sky_down_status ==
-                common_utils.ProcessStatus.SCHEDULED):
-            terminating_count += 1
-    return terminating_count
-
-
 def get_replicas_at_status(
     service_name: str,
     status: ReplicaStatus,
@@ -737,7 +737,8 @@ def add_version(service_name: str) -> int:
 
 @init_db
 def add_or_update_version(service_name: str, version: int,
-                          spec: 'service_spec.SkyServiceSpec') -> None:
+                          spec: 'service_spec.SkyServiceSpec',
+                          yaml_content: str) -> None:
     assert _SQLALCHEMY_ENGINE is not None
     with orm.Session(_SQLALCHEMY_ENGINE) as session:
         if (_SQLALCHEMY_ENGINE.dialect.name ==
@@ -750,27 +751,22 @@ def add_or_update_version(service_name: str, version: int,
             raise ValueError('Unsupported database dialect')
 
         insert_stmt = insert_func(version_specs_table).values(
-            service_name=service_name, version=version, spec=pickle.dumps(spec))
+            service_name=service_name,
+            version=version,
+            spec=pickle.dumps(spec),
+            yaml_content=yaml_content)
 
         insert_stmt = insert_stmt.on_conflict_do_update(
             index_elements=['service_name', 'version'],
-            set_={'spec': insert_stmt.excluded.spec})
+            set_={
+                'spec': insert_stmt.excluded.spec,
+                'yaml_content': insert_stmt.excluded.yaml_content
+            })
 
         session.execute(insert_stmt)
         session.commit()
 
 
-@init_db
-def remove_service_versions(service_name: str) -> None:
-    """Removes a replica from the database."""
-    assert _SQLALCHEMY_ENGINE is not None
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        session.execute(
-            sqlalchemy.delete(version_specs_table).where(
-                version_specs_table.c.service_name == service_name))
-        session.commit()
-
-
 @init_db
 def get_spec(service_name: str,
              version: int) -> Optional['service_spec.SkyServiceSpec']:
@@ -785,6 +781,19 @@ def get_spec(service_name: str,
     return pickle.loads(result[0]) if result else None
 
 
+@init_db
+def get_yaml_content(service_name: str, version: int) -> Optional[str]:
+    """Gets the yaml content of a version."""
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        result = session.execute(
+            sqlalchemy.select(version_specs_table.c.yaml_content).where(
+                sqlalchemy.and_(
+                    version_specs_table.c.service_name == service_name,
+                    version_specs_table.c.version == version))).fetchone()
+    return result[0] if result else None
+
+
 @init_db
 def delete_version(service_name: str, version: int) -> None:
     """Deletes a version from the database."""