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

sky/provision/primeintellect/utils.py

@@ -0,0 +1,398 @@
+"""Prime Intellect library wrapper for SkyPilot."""
+
+import json
+import os
+import shlex
+import time
+from typing import Any, Dict, List, Optional, Tuple, Union
+import uuid
+
+import requests
+
+from sky.catalog import common as catalog_common
+from sky.utils import common_utils
+
+_df = None
+_lookup_dict = None
+
+# Base URL for Prime Intellect API (used as default if not configured).
+DEFAULT_BASE_URL = 'https://api.primeintellect.ai'
+CREDENTIALS_PATH = '~/.prime/config.json'
+INITIAL_BACKOFF_SECONDS = 10
+MAX_BACKOFF_FACTOR = 10
+MAX_ATTEMPTS = 6
+
+
+class PrimeintellectAPIError(Exception):
+    """Base exception for Prime Intellect API errors."""
+
+    def __init__(self,
+                 message: str,
+                 status_code: Optional[int] = None,
+                 response_data: Optional[Dict[str, Any]] = None):
+        super().__init__(message)
+        self.status_code = status_code
+        self.response_data = response_data
+
+
+class PrimeintellectResourcesUnavailableError(PrimeintellectAPIError):
+    """Exception for when resources are unavailable on Prime Intellect."""
+    pass
+
+
+def _parse_api_error(response: Any) -> Tuple[str, bool]:
+    """Parse API error response to extract meaningful error messages.
+
+    Returns:
+        Tuple[str, bool]:
+        - str: A human-readable error message parsed from the API response.
+        - bool: True if the error indicates resource unavailability (e.g.,
+          capacity issues or quota/limit exceeded), otherwise False.
+    """
+    try:
+        if hasattr(response, 'json'):
+            error_data = response.json()
+        else:
+            error_data = response
+
+        if isinstance(error_data, dict):
+            # Try to extract error message from common error response fields
+            error_message = error_data.get('message', '')
+            if not error_message:
+                error_message = error_data.get('error', '')
+            if not error_message:
+                error_message = error_data.get('detail', '')
+
+            # Check if it's a resource unavailability error
+            if any(keyword in error_message.lower() for keyword in [
+                    'no capacity', 'capacity', 'unavailable', 'out of stock',
+                    'insufficient', 'not available', 'quota exceeded',
+                    'limit exceeded'
+            ]):
+                return error_message, True
+
+            return error_message, False
+
+        return str(error_data), False
+    except Exception:  # pylint: disable=broad-except
+        return f'HTTP {response.status_code} {response.reason}', False
+
+
+def _try_request_with_backoff(
+        method: str,
+        url: str,
+        headers: Dict[str, str],
+        data: Optional[Union[str, Dict[str, Any]]] = None) -> Dict[str, Any]:
+    backoff = common_utils.Backoff(initial_backoff=INITIAL_BACKOFF_SECONDS,
+                                   max_backoff_factor=MAX_BACKOFF_FACTOR)
+    for i in range(MAX_ATTEMPTS):
+        timeout = 30
+        if method == 'get':
+            response = requests.get(url,
+                                    headers=headers,
+                                    params=data,
+                                    timeout=timeout)
+        elif method == 'post':
+            response = requests.post(url,
+                                     headers=headers,
+                                     json=data,
+                                     timeout=timeout)
+        elif method == 'put':
+            response = requests.put(url,
+                                    headers=headers,
+                                    json=data,
+                                    timeout=timeout)
+        elif method == 'patch':
+            response = requests.patch(url,
+                                      headers=headers,
+                                      json=data,
+                                      timeout=timeout)
+        elif method == 'delete':
+            response = requests.delete(url, headers=headers, timeout=timeout)
+        else:
+            raise ValueError(f'Unsupported requests method: {method}')
+        # If rate limited, wait and try again
+        if response.status_code == 429 and i != MAX_ATTEMPTS - 1:
+            time.sleep(backoff.current_backoff())
+            continue
+        if response.ok:
+            return response.json()
+        else:
+            # Parse the error response for meaningful messages
+            err, is_resource_unavailable = _parse_api_error(response)
+
+            # Create a more informative error message
+            if not err:
+                err = (f'API request failed: {method} {url}: '
+                       f'{response.status_code} {response.reason}')
+            else:
+                err = f'API request failed: {err}'
+
+            # Raise appropriate exception based on error type
+            if is_resource_unavailable:
+                raise PrimeintellectResourcesUnavailableError(
+                    err,
+                    status_code=response.status_code,
+                    response_data=response.json()
+                    if hasattr(response, 'json') else None)
+            else:
+                raise PrimeintellectAPIError(
+                    err,
+                    status_code=response.status_code,
+                    response_data=response.json()
+                    if hasattr(response, 'json') else None)
+    return {}
+
+
+def get_upstream_cloud_id(instance_type: str) -> Optional[str]:
+    global _df, _lookup_dict
+    if _df is None:
+        _df = catalog_common.read_catalog('primeintellect/vms.csv')
+        _lookup_dict = (
+            _df.set_index('InstanceType')['UpstreamCloudId'].to_dict())
+    return _lookup_dict.get(instance_type)
+
+
+class PrimeIntellectAPIClient:
+    """Client for interacting with Prime Intellect API."""
+
+    def __init__(self) -> None:
+        self.credentials = os.path.expanduser(CREDENTIALS_PATH)
+        assert os.path.exists(self.credentials), 'Credentials not found'
+        with open(self.credentials, 'r', encoding='utf-8') as f:
+            self._credentials = json.load(f)
+        self.api_key = self._credentials.get('api_key')
+        self.team_id = self._credentials.get('team_id')
+        self.base_url = self._credentials.get('base_url', DEFAULT_BASE_URL)
+        self.headers = {
+            'Authorization': f'Bearer {self.api_key}',
+            'Content-Type': 'application/json'
+        }
+
+    def list_instances(self, **search_kwargs) -> List[Dict[str, Any]]:
+        response = _try_request_with_backoff('get',
+                                             f'{self.base_url}/api/v1/pods',
+                                             headers=self.headers,
+                                             data=search_kwargs)
+        return response['data']
+
+    def get_instance_details(self, instance_id: str) -> Dict[str, Any]:
+        return _try_request_with_backoff(
+            'get',
+            f'{self.base_url}/api/v1/pods/{instance_id}',
+            headers=self.headers)
+
+    def launch(self,
+               name: str,
+               instance_type: str,
+               region: str,
+               availability_zone: str,
+               disk_size: int,
+               vcpus: int = 0,
+               memory: int = 0) -> Dict[str, Any]:
+        """Create a pod/instance via Prime Intellect API.
+
+        Args:
+            name: User-visible name of the pod.
+            instance_type: A catalog instance type string. The expected format
+                is:
+                "<provider>__<accelerator>__<vcpus>__<memory>[_SPOT]".
+
+                - <provider>: Upstream provider tag (e.g., "primecompute").
+                - <accelerator>:
+                  * GPU nodes: "<N>x<GPU_MODEL>", e.g., "8xH100_80GB".
+                  * CPU-only nodes: the literal "CPU_NODE".
+                - <vcpus>: Integer string for vCPU count (e.g., "104").
+                - <memory>: Integer string for memory in GB (e.g., "752").
+                - Optional suffix "_SPOT" may be present in the full string
+                  (ignored here; pricing/spot behavior is not controlled by
+                  this method).
+
+                Notes:
+                - Parsing: only the first two components (provider,
+                  accelerator) are needed to build the payload. The vCPU
+                  and memory values are provided via the ``vcpus`` and
+                  ``memory`` arguments.
+                - Catalog lookup: the full instance_type string is used to
+                  map to the catalog's UpstreamCloudId.
+                - CPU-only: accelerator "CPU_NODE" is a sentinel for
+                  "no GPU". We set gpuType='CPU_NODE' and gpuCount=1 to
+                  represent CPU-only pods.
+                - Spot: the optional "__SPOT" suffix (if present) is ignored
+                  here; pricing/spot behavior is handled elsewhere.
+
+            region: Country/region code used by Prime Intellect.
+            availability_zone: Data center ID (zone) within the region.
+            disk_size: Boot disk size in GB.
+            vcpus: Optional explicit vCPU override; if >0 it will be sent.
+            memory: Optional explicit memory override in GB; if >0 it will be
+                sent.
+
+        Returns:
+            The API response JSON as a dict.
+        """
+        cloud_id = get_upstream_cloud_id(instance_type)
+        assert cloud_id, 'cloudId cannot be None'
+        assert availability_zone, 'availability_zone cannot be None'
+
+        # Parse the instance_type. We only need the first two components:
+        # provider and accelerator info (see docstring above).
+        provider, gpu_parts, _, _ = instance_type.split('__', 3)
+        if 'CPU_NODE' in gpu_parts:
+            # Prime Intellect API uses the same schema for CPU-only and GPU
+            # pods. For CPU-only instances, we set gpuType='CPU_NODE' and
+            # gpuCount=1 as a sentinel to indicate "no GPU". This is how CPU
+            # instances are represented internally on our platform; the
+            # backend does not interpret this as having a physical GPU.
+            gpu_type = 'CPU_NODE'
+            gpu_count = 1
+        else:
+            parts = gpu_parts.split('x', 1)
+            gpu_count = int(parts[0])
+            gpu_type = parts[1]
+
+        payload: Dict[str, Any] = {
+            'pod': {
+                'name': name,
+                'cloudId': cloud_id,
+                'socket': 'PCIe',
+                'gpuType': gpu_type,
+                'gpuCount': int(gpu_count),
+                'diskSize': disk_size,
+                # Prime Intellect API historically required maxPrice.
+                # Set to 0 to indicate on-demand/non-spot pricing.
+                'maxPrice': 0,
+            },
+            'provider': {
+                'type': provider,
+            }
+        }
+
+        if vcpus > 0:
+            payload['pod']['vcpus'] = vcpus
+        if memory > 0:
+            payload['pod']['memory'] = memory
+
+        if region != 'UNSPECIFIED':
+            payload['pod']['country'] = region
+        if availability_zone != 'UNSPECIFIED':
+            payload['pod']['dataCenterId'] = availability_zone
+
+        if self.team_id is not None and self.team_id != '':
+            payload['team'] = {'teamId': self.team_id}
+
+        response = _try_request_with_backoff(
+            'post',
+            f'{self.base_url}/api/v1/pods',
+            headers=self.headers,
+            data=payload,
+        )
+        return response
+
+    def remove(self, instance_id: str) -> Dict[str, Any]:
+        return _try_request_with_backoff(
+            'delete',
+            f'{self.base_url}/api/v1/pods/{instance_id}',
+            headers=self.headers,
+        )
+
+    def list_ssh_keys(self) -> List[Dict[str, Any]]:
+        response = _try_request_with_backoff('get',
+                                             f'{self.base_url}/api/v1/ssh_keys',
+                                             headers=self.headers)
+        return response['data']
+
+    def get_or_add_ssh_key(self, ssh_pub_key: str = '') -> Dict[str, str]:
+        """Add ssh key if not already added."""
+        # Check if the public key is already added
+        ssh_keys = self.list_ssh_keys()
+        for key in ssh_keys:
+            if key['publicKey'].strip().split()[:2] == ssh_pub_key.strip(
+            ).split()[:2]:
+                return {'name': key['name'], 'ssh_key': ssh_pub_key}
+
+        # Add the public key to Prime Intellect account if not already added
+        ssh_key_name = 'skypilot-' + str(uuid.uuid4()).replace('-', '')[:8]
+        _try_request_with_backoff(
+            'post',
+            f'{self.base_url}/api/v1/ssh_keys',
+            headers=self.headers,
+            data={
+                'name': ssh_key_name,
+                'publicKey': ssh_pub_key
+            },
+        )
+        return {'name': ssh_key_name, 'ssh_key': ssh_pub_key}
+
+
+def parse_ssh_connection(ssh_connection: Any) -> Tuple[Optional[str], int]:
+    """Parse and extract SSH username and port from a connection field.
+
+    The provider may return the SSH connection in multiple shapes. This helper
+    robustly extracts the SSH username and port while tolerating extra flags or
+    various tokenizations.
+
+    Accepted formats (examples):
+    - String with port flag:
+      "ubuntu@1.2.3.4 -p 2222 [-o <flag> ...]"
+    - String without explicit port (defaults to 22):
+      "ubuntu@1.2.3.4"
+    - String with host:port:
+      "ubuntu@1.2.3.4:2222"
+    - List with a single target:
+      ["ubuntu@1.2.3.4"]
+    - List of tokens (e.g., split form):
+      ["ubuntu@1.2.3.4", "-p", "2222"]
+
+    Args:
+        ssh_connection: The raw field from the API; can be a string or a list
+            of strings.
+
+    Returns:
+        (ssh_user, ssh_port): username if found, else None; port if found,
+        else 22.
+    """
+    ssh_user: Optional[str] = None
+    ssh_port: int = 22
+
+    # Normalize into a list of tokens for easier processing.
+    tokens: List[str] = []
+    if isinstance(ssh_connection, str):
+        try:
+            tokens = shlex.split(ssh_connection)
+        except Exception:  # pylint: disable=broad-except
+            tokens = [ssh_connection]
+    elif isinstance(ssh_connection, list):
+        for elem in ssh_connection:
+            if isinstance(elem, str):
+                try:
+                    tokens.extend(shlex.split(elem))
+                except Exception:  # pylint: disable=broad-except
+                    tokens.append(elem)
+    else:
+        # Unknown type; return defaults.
+        return ssh_user, ssh_port
+
+    # Find the first token containing '@' as the user@host candidate.
+    user_host: Optional[str] = next((t for t in tokens if '@' in t), None)
+    if user_host:
+        ssh_user = user_host.split('@', 1)[0].strip()
+        # Try host:port format (after '@').
+        host_part = user_host.split('@', 1)[1]
+        if ':' in host_part:
+            _, maybe_port = host_part.rsplit(':', 1)
+            try:
+                ssh_port = int(maybe_port)
+            except ValueError:
+                pass
+
+    # Check for '-p <port>' pair anywhere in the tokens. This takes priority.
+    if '-p' in tokens:
+        idx = tokens.index('-p')
+        if idx + 1 < len(tokens):
+            try:
+                ssh_port = int(tokens[idx + 1])
+            except ValueError:
+                pass
+
+    return ssh_user, ssh_port

sky/resources.py

@@ -1331,10 +1331,18 @@ class Resources:
                     clouds.CloudImplementationFeatures.IMAGE_ID
                 })
         except exceptions.NotSupportedError as e:
+            # Provide a more helpful error message for Lambda cloud
+            if self.cloud.is_same_cloud(clouds.Lambda()):
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(
+                        'Lambda cloud only supports Docker images. '
+                        'Please prefix your image with "docker:" '
+                        '(e.g., image_id: docker:your-image-name).') from e
             with ux_utils.print_exception_no_traceback():
                 raise ValueError(
                     'image_id is only supported for AWS/GCP/Azure/IBM/OCI/'
-                    'Kubernetes, please explicitly specify the cloud.') from e
+                    'Kubernetes. For Lambda cloud, use "docker:" prefix for '
+                    'Docker images.') from e
 
         if self._region is not None:
             # If the image_id has None as key (region-agnostic),

sky/schemas/generated/servev1_pb2.py

@@ -0,0 +1,58 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: sky/schemas/generated/servev1.proto
+# Protobuf Python Version: 5.26.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n#sky/schemas/generated/servev1.proto\x12\x08serve.v1\"\x1d\n\x0cServiceNames\x12\r\n\x05names\x18\x01 \x03(\t\"s\n\rServiceStatus\x12\x33\n\x06status\x18\x01 \x03(\x0b\x32#.serve.v1.ServiceStatus.StatusEntry\x1a-\n\x0bStatusEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"m\n\x17GetServiceStatusRequest\x12\x32\n\rservice_names\x18\x01 \x01(\x0b\x32\x16.serve.v1.ServiceNamesH\x00\x88\x01\x01\x12\x0c\n\x04pool\x18\x02 \x01(\x08\x42\x10\n\x0e_service_names\"E\n\x18GetServiceStatusResponse\x12)\n\x08statuses\x18\x01 \x03(\x0b\x32\x17.serve.v1.ServiceStatus\")\n\x11\x41\x64\x64VersionRequest\x12\x14\n\x0cservice_name\x18\x01 \x01(\t\"%\n\x12\x41\x64\x64VersionResponse\x12\x0f\n\x07version\x18\x01 \x01(\x03\"}\n\x18TerminateServicesRequest\x12\x32\n\rservice_names\x18\x01 \x01(\x0b\x32\x16.serve.v1.ServiceNamesH\x00\x88\x01\x01\x12\r\n\x05purge\x18\x02 \x01(\x08\x12\x0c\n\x04pool\x18\x03 \x01(\x08\x42\x10\n\x0e_service_names\",\n\x19TerminateServicesResponse\x12\x0f\n\x07message\x18\x01 \x01(\t\"R\n\x17TerminateReplicaRequest\x12\x14\n\x0cservice_name\x18\x01 \x01(\t\x12\x12\n\nreplica_id\x18\x02 \x01(\x03\x12\r\n\x05purge\x18\x03 \x01(\x08\"+\n\x18TerminateReplicaResponse\x12\x0f\n\x07message\x18\x01 \x01(\t\"T\n\x1eWaitServiceRegistrationRequest\x12\x14\n\x0cservice_name\x18\x01 \x01(\t\x12\x0e\n\x06job_id\x18\x02 \x01(\x03\x12\x0c\n\x04pool\x18\x03 \x01(\x08\"2\n\x1fWaitServiceRegistrationResponse\x12\x0f\n\x07lb_port\x18\x01 \x01(\x03\"Y\n\x14UpdateServiceRequest\x12\x14\n\x0cservice_name\x18\x01 \x01(\t\x12\x0f\n\x07version\x18\x02 \x01(\x03\x12\x0c\n\x04mode\x18\x03 \x01(\t\x12\x0c\n\x04pool\x18\x04 \x01(\x08\"\x17\n\x15UpdateServiceResponse2\xad\x04\n\x0cServeService\x12Y\n\x10GetServiceStatus\x12!.serve.v1.GetServiceStatusRequest\x1a\".serve.v1.GetServiceStatusResponse\x12G\n\nAddVersion\x12\x1b.serve.v1.AddVersionRequest\x1a\x1c.serve.v1.AddVersionResponse\x12\\\n\x11TerminateServices\x12\".serve.v1.TerminateServicesRequest\x1a#.serve.v1.TerminateServicesResponse\x12Y\n\x10TerminateReplica\x12!.serve.v1.TerminateReplicaRequest\x1a\".serve.v1.TerminateReplicaResponse\x12n\n\x17WaitServiceRegistration\x12(.serve.v1.WaitServiceRegistrationRequest\x1a).serve.v1.WaitServiceRegistrationResponse\x12P\n\rUpdateService\x12\x1e.serve.v1.UpdateServiceRequest\x1a\x1f.serve.v1.UpdateServiceResponseb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'sky.schemas.generated.servev1_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  DESCRIPTOR._loaded_options = None
+  _globals['_SERVICESTATUS_STATUSENTRY']._loaded_options = None
+  _globals['_SERVICESTATUS_STATUSENTRY']._serialized_options = b'8\001'
+  _globals['_SERVICENAMES']._serialized_start=49
+  _globals['_SERVICENAMES']._serialized_end=78
+  _globals['_SERVICESTATUS']._serialized_start=80
+  _globals['_SERVICESTATUS']._serialized_end=195
+  _globals['_SERVICESTATUS_STATUSENTRY']._serialized_start=150
+  _globals['_SERVICESTATUS_STATUSENTRY']._serialized_end=195
+  _globals['_GETSERVICESTATUSREQUEST']._serialized_start=197
+  _globals['_GETSERVICESTATUSREQUEST']._serialized_end=306
+  _globals['_GETSERVICESTATUSRESPONSE']._serialized_start=308
+  _globals['_GETSERVICESTATUSRESPONSE']._serialized_end=377
+  _globals['_ADDVERSIONREQUEST']._serialized_start=379
+  _globals['_ADDVERSIONREQUEST']._serialized_end=420
+  _globals['_ADDVERSIONRESPONSE']._serialized_start=422
+  _globals['_ADDVERSIONRESPONSE']._serialized_end=459
+  _globals['_TERMINATESERVICESREQUEST']._serialized_start=461
+  _globals['_TERMINATESERVICESREQUEST']._serialized_end=586
+  _globals['_TERMINATESERVICESRESPONSE']._serialized_start=588
+  _globals['_TERMINATESERVICESRESPONSE']._serialized_end=632
+  _globals['_TERMINATEREPLICAREQUEST']._serialized_start=634
+  _globals['_TERMINATEREPLICAREQUEST']._serialized_end=716
+  _globals['_TERMINATEREPLICARESPONSE']._serialized_start=718
+  _globals['_TERMINATEREPLICARESPONSE']._serialized_end=761
+  _globals['_WAITSERVICEREGISTRATIONREQUEST']._serialized_start=763
+  _globals['_WAITSERVICEREGISTRATIONREQUEST']._serialized_end=847
+  _globals['_WAITSERVICEREGISTRATIONRESPONSE']._serialized_start=849
+  _globals['_WAITSERVICEREGISTRATIONRESPONSE']._serialized_end=899
+  _globals['_UPDATESERVICEREQUEST']._serialized_start=901
+  _globals['_UPDATESERVICEREQUEST']._serialized_end=990
+  _globals['_UPDATESERVICERESPONSE']._serialized_start=992
+  _globals['_UPDATESERVICERESPONSE']._serialized_end=1015
+  _globals['_SERVESERVICE']._serialized_start=1018
+  _globals['_SERVESERVICE']._serialized_end=1575
+# @@protoc_insertion_point(module_scope)

sky/schemas/generated/servev1_pb2.pyi

@@ -0,0 +1,115 @@
+from google.protobuf.internal import containers as _containers
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from typing import ClassVar as _ClassVar, Iterable as _Iterable, Mapping as _Mapping, Optional as _Optional, Union as _Union
+
+DESCRIPTOR: _descriptor.FileDescriptor
+
+class ServiceNames(_message.Message):
+    __slots__ = ("names",)
+    NAMES_FIELD_NUMBER: _ClassVar[int]
+    names: _containers.RepeatedScalarFieldContainer[str]
+    def __init__(self, names: _Optional[_Iterable[str]] = ...) -> None: ...
+
+class ServiceStatus(_message.Message):
+    __slots__ = ("status",)
+    class StatusEntry(_message.Message):
+        __slots__ = ("key", "value")
+        KEY_FIELD_NUMBER: _ClassVar[int]
+        VALUE_FIELD_NUMBER: _ClassVar[int]
+        key: str
+        value: str
+        def __init__(self, key: _Optional[str] = ..., value: _Optional[str] = ...) -> None: ...
+    STATUS_FIELD_NUMBER: _ClassVar[int]
+    status: _containers.ScalarMap[str, str]
+    def __init__(self, status: _Optional[_Mapping[str, str]] = ...) -> None: ...
+
+class GetServiceStatusRequest(_message.Message):
+    __slots__ = ("service_names", "pool")
+    SERVICE_NAMES_FIELD_NUMBER: _ClassVar[int]
+    POOL_FIELD_NUMBER: _ClassVar[int]
+    service_names: ServiceNames
+    pool: bool
+    def __init__(self, service_names: _Optional[_Union[ServiceNames, _Mapping]] = ..., pool: bool = ...) -> None: ...
+
+class GetServiceStatusResponse(_message.Message):
+    __slots__ = ("statuses",)
+    STATUSES_FIELD_NUMBER: _ClassVar[int]
+    statuses: _containers.RepeatedCompositeFieldContainer[ServiceStatus]
+    def __init__(self, statuses: _Optional[_Iterable[_Union[ServiceStatus, _Mapping]]] = ...) -> None: ...
+
+class AddVersionRequest(_message.Message):
+    __slots__ = ("service_name",)
+    SERVICE_NAME_FIELD_NUMBER: _ClassVar[int]
+    service_name: str
+    def __init__(self, service_name: _Optional[str] = ...) -> None: ...
+
+class AddVersionResponse(_message.Message):
+    __slots__ = ("version",)
+    VERSION_FIELD_NUMBER: _ClassVar[int]
+    version: int
+    def __init__(self, version: _Optional[int] = ...) -> None: ...
+
+class TerminateServicesRequest(_message.Message):
+    __slots__ = ("service_names", "purge", "pool")
+    SERVICE_NAMES_FIELD_NUMBER: _ClassVar[int]
+    PURGE_FIELD_NUMBER: _ClassVar[int]
+    POOL_FIELD_NUMBER: _ClassVar[int]
+    service_names: ServiceNames
+    purge: bool
+    pool: bool
+    def __init__(self, service_names: _Optional[_Union[ServiceNames, _Mapping]] = ..., purge: bool = ..., pool: bool = ...) -> None: ...
+
+class TerminateServicesResponse(_message.Message):
+    __slots__ = ("message",)
+    MESSAGE_FIELD_NUMBER: _ClassVar[int]
+    message: str
+    def __init__(self, message: _Optional[str] = ...) -> None: ...
+
+class TerminateReplicaRequest(_message.Message):
+    __slots__ = ("service_name", "replica_id", "purge")
+    SERVICE_NAME_FIELD_NUMBER: _ClassVar[int]
+    REPLICA_ID_FIELD_NUMBER: _ClassVar[int]
+    PURGE_FIELD_NUMBER: _ClassVar[int]
+    service_name: str
+    replica_id: int
+    purge: bool
+    def __init__(self, service_name: _Optional[str] = ..., replica_id: _Optional[int] = ..., purge: bool = ...) -> None: ...
+
+class TerminateReplicaResponse(_message.Message):
+    __slots__ = ("message",)
+    MESSAGE_FIELD_NUMBER: _ClassVar[int]
+    message: str
+    def __init__(self, message: _Optional[str] = ...) -> None: ...
+
+class WaitServiceRegistrationRequest(_message.Message):
+    __slots__ = ("service_name", "job_id", "pool")
+    SERVICE_NAME_FIELD_NUMBER: _ClassVar[int]
+    JOB_ID_FIELD_NUMBER: _ClassVar[int]
+    POOL_FIELD_NUMBER: _ClassVar[int]
+    service_name: str
+    job_id: int
+    pool: bool
+    def __init__(self, service_name: _Optional[str] = ..., job_id: _Optional[int] = ..., pool: bool = ...) -> None: ...
+
+class WaitServiceRegistrationResponse(_message.Message):
+    __slots__ = ("lb_port",)
+    LB_PORT_FIELD_NUMBER: _ClassVar[int]
+    lb_port: int
+    def __init__(self, lb_port: _Optional[int] = ...) -> None: ...
+
+class UpdateServiceRequest(_message.Message):
+    __slots__ = ("service_name", "version", "mode", "pool")
+    SERVICE_NAME_FIELD_NUMBER: _ClassVar[int]
+    VERSION_FIELD_NUMBER: _ClassVar[int]
+    MODE_FIELD_NUMBER: _ClassVar[int]
+    POOL_FIELD_NUMBER: _ClassVar[int]
+    service_name: str
+    version: int
+    mode: str
+    pool: bool
+    def __init__(self, service_name: _Optional[str] = ..., version: _Optional[int] = ..., mode: _Optional[str] = ..., pool: bool = ...) -> None: ...
+
+class UpdateServiceResponse(_message.Message):
+    __slots__ = ()
+    def __init__(self) -> None: ...