skypilot-nightly 1.0.0.dev20250703__py3-none-any.whl → 1.0.0.dev20250704__py3-none-any.whl

sky/provision/cudo/cudo_wrapper.py

@@ -28,12 +28,10 @@ def launch(name: str, data_center_id: str, ssh_key: str, machine_type: str,
                                  size_gib=disk_size),
         metadata=tags)
 
-    try:
-        api = cudo.cudo.cudo_api.virtual_machines()
-        vm = api.create_vm(cudo.cudo.cudo_api.project_id_throwable(), request)
-        return vm.to_dict()['id']
-    except cudo.cudo.rest.ApiException as e:
-        raise e
+    api = cudo.cudo.cudo_api.virtual_machines()
+    vm = api.create_vm(cudo.cudo.cudo_api.project_id_throwable(), request)
+
+    return vm.to_dict()['id']
 
 
 def remove(instance_id: str):
@@ -54,11 +52,8 @@ def remove(instance_id: str):
     state = 'unknown'
     project_id = cudo.cudo.cudo_api.project_id_throwable()
     while retry_count < max_retries:
-        try:
-            vm = api.get_vm(project_id, instance_id)
-            state = vm.to_dict()['vm']['short_state']
-        except cudo.cudo.rest.ApiException as e:
-            raise e
+        vm = api.get_vm(project_id, instance_id)
+        state = vm.to_dict()['vm']['short_state']
 
         if state in terminate_ok:
             break
@@ -69,76 +64,82 @@ def remove(instance_id: str):
             'Timeout error, could not terminate due to VM state: {}'.format(
                 state))
 
-    try:
-        api.terminate_vm(project_id, instance_id)
-    except cudo.cudo.rest.ApiException as e:
-        raise e
+    api.terminate_vm(project_id, instance_id)
 
 
 def set_tags(instance_id: str, tags: Dict):
     """Sets the tags for the given instance."""
-    try:
-        api = cudo.cudo.cudo_api.virtual_machines()
-        api.update_vm_metadata(
-            cudo.cudo.cudo_api.project_id(), instance_id,
-            cudo.cudo.UpdateVMMetadataBody(
-                metadata=tags,
-                merge=True))  # TODO (skypilot team) merge or overwrite?
-    except cudo.cudo.rest.ApiException as e:
-        raise e
+    api = cudo.cudo.cudo_api.virtual_machines()
+    api.update_vm_metadata(
+        cudo.cudo.cudo_api.project_id(), instance_id,
+        cudo.cudo.UpdateVMMetadataBody(
+            metadata=tags,
+            merge=True))  # TODO (skypilot team) merge or overwrite?
 
 
 def get_instance(vm_id):
-    try:
-        api = cudo.cudo.cudo_api.virtual_machines()
-        vm = api.get_vm(cudo.cudo.cudo_api.project_id_throwable(), vm_id)
-        vm_dict = vm.to_dict()
-        return vm_dict
-    except cudo.cudo.rest.ApiException as e:
-        raise e
+    api = cudo.cudo.cudo_api.virtual_machines()
+    vm = api.get_vm(cudo.cudo.cudo_api.project_id_throwable(), vm_id)
+    vm_dict = vm.to_dict()
+    return vm_dict
 
 
 def list_instances():
-    try:
-        api = cudo.cudo.cudo_api.virtual_machines()
-        vms = api.list_vms(cudo.cudo.cudo_api.project_id_throwable())
-        instances = {}
-        for vm in vms.to_dict()['vms']:
-            ex_ip = vm['external_ip_address']
-            in_ip = vm['internal_ip_address']
-            if not in_ip:
-                in_ip = ex_ip
-            instance = {
-                # active_state, init_state, lcm_state, short_state
-                'status': vm['short_state'],
-                'tags': vm['metadata'],
-                'name': vm['id'],
-                'ip': ex_ip,
-                'external_ip': ex_ip,
-                'internal_ip': in_ip
-            }
-            instances[vm['id']] = instance
-        return instances
-    except cudo.cudo.rest.ApiException as e:
-        raise e
+    api = cudo.cudo.cudo_api.virtual_machines()
+    vms = api.list_vms(cudo.cudo.cudo_api.project_id_throwable())
+    instances = {}
+    for vm in vms.to_dict()['vms']:
+        ex_ip = vm['external_ip_address']
+        in_ip = vm['internal_ip_address']
+        if not in_ip:
+            in_ip = ex_ip
+        instance = {
+            # active_state, init_state, lcm_state, short_state
+            'status': vm['short_state'],
+            'tags': vm['metadata'],
+            'name': vm['id'],
+            'ip': ex_ip,
+            'external_ip': ex_ip,
+            'internal_ip': in_ip
+        }
+        instances[vm['id']] = instance
+    return instances
 
 
 def vm_available(to_start_count, gpu_count, gpu_model, data_center_id, mem,
                  cpus):
-    try:
-        gpu_model = utils.skypilot_gpu_to_cudo_gpu(gpu_model)
-        api = cudo.cudo.cudo_api.virtual_machines()
-        types = api.list_vm_machine_types(mem,
-                                          cpus,
-                                          gpu=gpu_count,
-                                          gpu_model=gpu_model,
-                                          data_center_id=data_center_id)
-        types_dict = types.to_dict()
-        hc = types_dict['host_configs']
-        total_count = sum(item['count_vm_available'] for item in hc)
-        if total_count < to_start_count:
-            raise Exception(
-                'Too many VMs requested, try another gpu type or region')
-        return total_count
-    except cudo.cudo.rest.ApiException as e:
-        raise e
+    gpu_model = utils.skypilot_gpu_to_cudo_gpu(gpu_model)
+    api = cudo.cudo.cudo_api.virtual_machines()
+    types = api.list_vm_machine_types2()
+    types_dict = types.to_dict()
+    machine_types = types_dict['machine_types']
+
+    # Filter machine types based on requirements
+    matching_types = []
+    for machine_type in machine_types:
+        # Check if this machine type matches our requirements
+        if (machine_type['data_center_id'] == data_center_id and
+                machine_type['gpu_model'] == gpu_model and
+                machine_type['min_vcpu'] <= cpus <= machine_type.get(
+                    'max_vcpu_free', float('inf')) and
+                machine_type['min_memory_gib'] <= mem <= machine_type.get(
+                    'max_memory_gib_free', float('inf'))):
+
+            # Calculate available VMs based on resource constraints
+            max_vms_by_vcpu = machine_type[
+                'total_vcpu_free'] // cpus if cpus > 0 else float('inf')
+            max_vms_by_memory = machine_type[
+                'total_memory_gib_free'] // mem if mem > 0 else float('inf')
+            max_vms_by_gpu = machine_type[
+                'total_gpu_free'] // gpu_count if gpu_count > 0 else float(
+                    'inf')
+
+            available_vms = min(max_vms_by_vcpu, max_vms_by_memory,
+                                max_vms_by_gpu)
+            matching_types.append(available_vms)
+
+    total_count = sum(matching_types)
+    if total_count < to_start_count:
+        raise Exception(
+            'Too many VMs requested, try another gpu type or region')
+    return total_count

sky/server/common.py

@@ -22,7 +22,6 @@ import uuid
 import colorama
 import filelock
 
-import sky
 from sky import exceptions
 from sky import sky_logging
 from sky import skypilot_config
@@ -31,6 +30,7 @@ from sky.client import service_account_auth
 from sky.data import data_utils
 from sky.server import constants as server_constants
 from sky.server import rest
+from sky.server import versions
 from sky.skylet import constants
 from sky.usage import usage_lib
 from sky.utils import annotations
@@ -66,34 +66,11 @@ RETRY_COUNT_ON_TIMEOUT = 3
 # (e.g. in high contention env) and we will exit eagerly if server exit.
 WAIT_APISERVER_START_TIMEOUT_SEC = 60
 
-_VERSION_INFO = (
-    f'{colorama.Style.RESET_ALL}'
-    f'{colorama.Style.DIM}'
-    'client version: v{client_version} (API version: v{client_api_version})\n'
-    'server version: v{server_version} (API version: v{server_api_version})'
-    f'{colorama.Style.RESET_ALL}')
 _LOCAL_API_SERVER_RESTART_HINT = (
-    f'{colorama.Fore.YELLOW}Please restart the SkyPilot API server with:\n'
+    f'{colorama.Fore.YELLOW}The local SkyPilot API server is not compatible '
+    'with the client. Please restart the API server with:\n'
     f'{colorama.Style.BRIGHT}sky api stop; sky api start'
     f'{colorama.Style.RESET_ALL}')
-_LOCAL_SERVER_VERSION_MISMATCH_WARNING = (
-    f'{colorama.Fore.YELLOW}Client and local API server version mismatch:\n'
-    '{version_info}\n'
-    f'{_LOCAL_API_SERVER_RESTART_HINT}'
-    f'{colorama.Style.RESET_ALL}')
-_CLIENT_TOO_OLD_WARNING = (
-    f'{colorama.Fore.YELLOW}Your SkyPilot client is too old:\n'
-    '{version_info}\n'
-    f'{colorama.Fore.YELLOW}Upgrade your client with:\n'
-    '{command}'
-    f'{colorama.Style.RESET_ALL}')
-_REMOTE_SERVER_TOO_OLD_WARNING = (
-    f'{colorama.Fore.YELLOW}SkyPilot API server is too old:\n'
-    '{version_info}\n'
-    f'{colorama.Fore.YELLOW}Contact your administrator to upgrade the '
-    'remote API server or downgrade your local client with:\n'
-    '{command}\n'
-    f'{colorama.Style.RESET_ALL}')
 _SERVER_INSTALL_VERSION_MISMATCH_WARNING = (
     f'{colorama.Fore.YELLOW}SkyPilot API server version does not match the '
     'installation on disk:\n'
@@ -105,10 +82,6 @@ _SERVER_INSTALL_VERSION_MISMATCH_WARNING = (
     f'{colorama.Fore.YELLOW}This can happen if you upgraded SkyPilot without '
     'restarting the API server.'
     f'{colorama.Style.RESET_ALL}')
-# Parse local API version eargly to catch version format errors.
-_LOCAL_API_VERSION: int = int(server_constants.API_VERSION)
-# SkyPilot dev version.
-_DEV_VERSION = '1.0.0-dev0'
 
 RequestId = str
 ApiVersion = Optional[str]
@@ -134,6 +107,7 @@ class ApiServerInfo:
     commit: Optional[str] = None
     user: Optional[Dict[str, Any]] = None
     basic_auth_enabled: bool = False
+    error: Optional[str] = None
 
 
 def get_api_cookie_jar_path() -> pathlib.Path:
@@ -282,6 +256,23 @@ def is_api_server_local():
     return get_server_url() in AVAILABLE_LOCAL_API_SERVER_URLS
 
 
+def _handle_non_200_server_status(
+        response: 'requests.Response') -> ApiServerInfo:
+    if response.status_code == 401:
+        return ApiServerInfo(status=ApiServerStatus.NEEDS_AUTH)
+    if response.status_code == 400:
+        # Check if a version mismatch error is returned.
+        try:
+            body = response.json()
+            if (body.get('error',
+                         '') == ApiServerStatus.VERSION_MISMATCH.value):
+                return ApiServerInfo(status=ApiServerStatus.VERSION_MISMATCH,
+                                     error=body.get('message', ''))
+        except json.JSONDecodeError:
+            pass
+    return ApiServerInfo(status=ApiServerStatus.UNHEALTHY)
+
+
 def get_api_server_status(endpoint: Optional[str] = None) -> ApiServerInfo:
     """Retrieve the status of the API server.
 
@@ -315,10 +306,10 @@ def get_api_server_status(endpoint: Optional[str] = None) -> ApiServerInfo:
             return ApiServerInfo(status=ApiServerStatus.UNHEALTHY)
 
         logger.debug(f'Health check status: {response.status_code}')
-        if response.status_code == 401:
-            return ApiServerInfo(status=ApiServerStatus.NEEDS_AUTH)
-        elif response.status_code != 200:
-            return ApiServerInfo(status=ApiServerStatus.UNHEALTHY)
+
+        if response.status_code != 200:
+            return _handle_non_200_server_status(response)
+
         # The response is 200, so we can parse the response.
         try:
             result = response.json()
@@ -340,8 +331,24 @@ def get_api_server_status(endpoint: Optional[str] = None) -> ApiServerInfo:
                                f'version info. {server_url} may '
                                f'not be running SkyPilot API server.')
                 server_info.status = ApiServerStatus.UNHEALTHY
-            elif api_version != server_constants.API_VERSION:
-                server_info.status = ApiServerStatus.VERSION_MISMATCH
+            version_info = versions.check_compatibility_at_client(
+                response.headers)
+            if version_info is None:
+                # Backward compatibility for server prior to v0.11.0 which
+                # does not check compatibility at server side.
+                # TODO(aylei): remove this after v0.13.0 is released.
+                return ApiServerInfo(
+                    status=ApiServerStatus.VERSION_MISMATCH,
+                    error=versions.SERVER_TOO_OLD_ERROR.format(
+                        remote_version=version,
+                        local_version=versions.get_local_readable_version(),
+                        min_version=server_constants.MIN_COMPATIBLE_VERSION,
+                        command=versions.install_version_command(
+                            version, commit)))
+            if version_info.error is not None:
+                return ApiServerInfo(status=ApiServerStatus.VERSION_MISMATCH,
+                                     error=version_info.error)
+
             cookies = get_cookies_from_response(response)
             set_api_cookie_jar(cookies, create_if_not_exists=False)
             return server_info
@@ -490,7 +497,7 @@ def _start_api_server(deploy: bool = False,
         server_url = get_server_url(host)
         dashboard_msg = ''
         api_server_info = get_api_server_status(server_url)
-        if api_server_info.version == _DEV_VERSION:
+        if api_server_info.version == versions.DEV_VERSION:
             dashboard_msg += (
                 f'\n{colorama.Style.RESET_ALL}{ux_utils.INDENT_SYMBOL}'
                 f'{colorama.Fore.YELLOW}')
@@ -554,33 +561,11 @@ def check_server_healthy(
     api_server_info = get_api_server_status(endpoint)
     api_server_status = api_server_info.status
     if api_server_status == ApiServerStatus.VERSION_MISMATCH:
-        sv = api_server_info.api_version
-        assert sv is not None, 'Server API version is None'
-        try:
-            server_is_older = int(sv) < _LOCAL_API_VERSION
-        except ValueError:
-            # Raised when the server version using an unknown scheme.
-            # Version compatibility checking is expected to handle all legacy
-            # cases so we safely assume the server is newer when the version
-            # scheme is unknown.
-            logger.debug('API server version using unknown scheme: %s', sv)
-            server_is_older = False
-        version_info = _get_version_info_hint(api_server_info)
+        msg = api_server_info.error
         if is_api_server_local():
             # For local server, just hint user to restart the server to get
             # a consistent version.
-            msg = _LOCAL_SERVER_VERSION_MISMATCH_WARNING.format(
-                version_info=version_info)
-        else:
-            assert api_server_info.version is not None, 'Server version is None'
-            if server_is_older:
-                msg = _REMOTE_SERVER_TOO_OLD_WARNING.format(
-                    version_info=version_info,
-                    command=_install_server_version_command(api_server_info))
-            else:
-                msg = _CLIENT_TOO_OLD_WARNING.format(
-                    version_info=version_info,
-                    command=_install_server_version_command(api_server_info))
+            msg = _LOCAL_API_SERVER_RESTART_HINT
         with ux_utils.print_exception_no_traceback():
             raise exceptions.APIVersionMismatchError(msg)
     elif api_server_status == ApiServerStatus.UNHEALTHY:
@@ -614,37 +599,6 @@ def check_server_healthy(
     return api_server_status, api_server_info
 
 
-def _get_version_info_hint(server_info: ApiServerInfo) -> str:
-    assert server_info.version is not None, 'Server version is None'
-    # version_on_disk may be None if the server is older
-    assert server_info.commit is not None, 'Server commit is None'
-    sv = server_info.version
-    cv = sky.__version__
-    if server_info.version == _DEV_VERSION:
-        sv = f'{sv} with commit {server_info.commit}'
-    if cv == _DEV_VERSION:
-        cv = f'{cv} with commit {sky.__commit__}'
-    return _VERSION_INFO.format(client_version=cv,
-                                server_version=sv,
-                                client_api_version=server_constants.API_VERSION,
-                                server_api_version=server_info.api_version)
-
-
-def _install_server_version_command(server_info: ApiServerInfo) -> str:
-    assert server_info.version is not None, 'Server version is None'
-    assert server_info.commit is not None, 'Server commit is None'
-    if server_info.version == _DEV_VERSION:
-        # Dev build without valid version.
-        return ('pip install git+https://github.com/skypilot-org/skypilot@'
-                f'{server_info.commit}')
-    elif 'dev' in server_info.version:
-        # Nightly version.
-        return f'pip install -U "skypilot-nightly=={server_info.version}"'
-    else:
-        # Stable version.
-        return f'pip install -U "skypilot=={server_info.version}"'
-
-
 # Keep in sync with sky/setup_files/setup.py find_version()
 def get_skypilot_version_on_disk() -> str:
     """Get the version of the SkyPilot code on disk."""

sky/server/constants.py

@@ -4,10 +4,31 @@ import os
 
 from sky.skylet import constants
 
-# API server version, whenever there is a change in API server that requires a
-# restart of the local API server or error out when the client does not match
-# the server version.
-API_VERSION = '10'
+# pylint: disable=line-too-long
+# The SkyPilot API version that the code currently use.
+# Bump this version when the API is changed and special compatibility handling
+# based on version info is needed.
+# For more details and code guidelines, refer to:
+# https://docs.skypilot.co/en/latest/developers/CONTRIBUTING.html#backward-compatibility-guidelines
+API_VERSION = 11
+
+# The minimum peer API version that the code should still work with.
+# Notes (dev):
+# - This value is maintained by the CI pipeline, DO NOT EDIT this manually.
+# - Compatibility code for versions lower than this can be safely removed.
+# Refer to API_VERSION for more details.
+MIN_COMPATIBLE_API_VERSION = 11
+
+# The semantic version of the minimum compatible API version.
+# Refer to MIN_COMPATIBLE_API_VERSION for more details.
+# Note (dev): DO NOT EDIT this constant manually.
+MIN_COMPATIBLE_VERSION = '0.10.0'
+
+# The HTTP header name for the API version of the sender.
+API_VERSION_HEADER = 'X-SkyPilot-API-Version'
+
+# The HTTP header name for the SkyPilot version of the sender.
+VERSION_HEADER = 'X-SkyPilot-Version'
 
 # Prefix for API request names.
 REQUEST_NAME_PREFIX = 'sky.'

sky/server/requests/payloads.py

@@ -1,9 +1,27 @@
 """Payloads for the Sky API requests.
 
-TODO(zhwu): We can consider a better way to handle the default values of the
-kwargs for the payloads, otherwise, we have to keep the default values the sync
-with the backend functions. The benefit of having the default values in the
-payloads is that a user can find the default values in the Restful API docs.
+All the payloads that will be used between the client and server communication
+must be defined here to make sure it get covered by our API compatbility tests.
+
+Compatibility note:
+- Adding a new body for new API is compatible as long as the SDK method using
+  the new API is properly decorated with `versions.minimal_api_version`.
+- Adding a new field with default value to an existing body is compatible at
+  API level, but the business logic must handle the case where the field is
+  not proccessed by an old version of remote client/server. This can usually
+  be done by checking `versions.get_remote_api_version()`.
+- Other changes are not compatible at API level, so must be handled specially.
+  A common pattern is to keep both the old and new version of the body and
+  checking `versions.get_remote_api_version()` to decide which body to use. For
+  example, say we refactor the `LaunchBody`, the original `LaunchBody` must be
+  kept in the codebase and the new body should be added via `LaunchBodyV2`.
+  Then if the remote runs in an old version, the local code should still send
+  `LaunchBody` to keep the backward compatibility. `LaunchBody` can be removed
+  later when constants.MIN_COMPATIBLE_API_VERSION is updated to a version that
+  supports `LaunchBodyV2`
+
+Also refer to sky.server.constants.MIN_COMPATIBLE_API_VERSION and the
+sky.server.versions module for more details.
 """
 import os
 import typing
@@ -94,7 +112,18 @@ def get_override_skypilot_config_path_from_client() -> Optional[str]:
     return skypilot_config.loaded_config_path_serialized()
 
 
-class RequestBody(pydantic.BaseModel):
+class BasePayload(pydantic.BaseModel):
+    """The base payload for the SkyPilot API."""
+    # Ignore extra fields in the request body, which is useful for backward
+    # compatibility. The difference with `allow` is that `ignore` will not
+    # include the unknown fields when dump the model, i.e., we can add new
+    # fields to the request body without breaking the existing old API server
+    # where the handler function does not accept the new field in function
+    # signature.
+    model_config = pydantic.ConfigDict(extra='ignore')
+
+
+class RequestBody(BasePayload):
     """The request body for the SkyPilot API."""
     env_vars: Dict[str, str] = {}
     entrypoint: str = ''
@@ -103,11 +132,6 @@ class RequestBody(pydantic.BaseModel):
     override_skypilot_config: Optional[Dict[str, Any]] = {}
     override_skypilot_config_path: Optional[str] = None
 
-    # Allow extra fields in the request body, which is useful for backward
-    # compatibility, i.e., we can add new fields to the request body without
-    # breaking the existing old API server.
-    model_config = pydantic.ConfigDict(extra='allow')
-
     def __init__(self, **data):
         data['env_vars'] = data.get('env_vars', request_body_env_vars())
         usage_lib_entrypoint = usage_lib.messages.usage.entrypoint
@@ -665,3 +689,24 @@ class GetConfigBody(RequestBody):
 class CostReportBody(RequestBody):
     """The request body for the cost report endpoint."""
     days: Optional[int] = 30
+
+
+class RequestPayload(BasePayload):
+    """The payload for the requests."""
+
+    request_id: str
+    name: str
+    entrypoint: str
+    request_body: str
+    status: str
+    created_at: float
+    user_id: str
+    return_value: str
+    error: str
+    pid: Optional[int]
+    schedule_type: str
+    user_name: Optional[str] = None
+    # Resources the request operates on.
+    cluster_name: Optional[str] = None
+    status_msg: Optional[str] = None
+    should_retry: bool = False

sky/server/requests/requests.py

@@ -98,28 +98,6 @@ class ScheduleType(enum.Enum):
     SHORT = 'short'
 
 
-@dataclasses.dataclass
-class RequestPayload:
-    """The payload for the requests."""
-
-    request_id: str
-    name: str
-    entrypoint: str
-    request_body: str
-    status: str
-    created_at: float
-    user_id: str
-    return_value: str
-    error: str
-    pid: Optional[int]
-    schedule_type: str
-    user_name: Optional[str] = None
-    # Resources the request operates on.
-    cluster_name: Optional[str] = None
-    status_msg: Optional[str] = None
-    should_retry: bool = False
-
-
 @dataclasses.dataclass
 class Request:
     """A SkyPilot API request."""
@@ -185,7 +163,7 @@ class Request:
     @classmethod
     def from_row(cls, row: Tuple[Any, ...]) -> 'Request':
         content = dict(zip(REQUEST_COLUMNS, row))
-        return cls.decode(RequestPayload(**content))
+        return cls.decode(payloads.RequestPayload(**content))
 
     def to_row(self) -> Tuple[Any, ...]:
         payload = self.encode()
@@ -194,7 +172,7 @@ class Request:
             row.append(getattr(payload, k))
         return tuple(row)
 
-    def readable_encode(self) -> RequestPayload:
+    def readable_encode(self) -> payloads.RequestPayload:
         """Serialize the SkyPilot API request for display purposes.
 
         This function should be called on the server side to serialize the
@@ -212,7 +190,7 @@ class Request:
                           payloads.RequestBody), (self.name, self.request_body)
         user = global_user_state.get_user(self.user_id)
         user_name = user.name if user is not None else None
-        return RequestPayload(
+        return payloads.RequestPayload(
             request_id=self.request_id,
             name=self.name,
             entrypoint=self.entrypoint.__name__,
@@ -230,12 +208,12 @@ class Request:
             should_retry=self.should_retry,
         )
 
-    def encode(self) -> RequestPayload:
+    def encode(self) -> payloads.RequestPayload:
         """Serialize the SkyPilot API request."""
         assert isinstance(self.request_body,
                           payloads.RequestBody), (self.name, self.request_body)
         try:
-            return RequestPayload(
+            return payloads.RequestPayload(
                 request_id=self.request_id,
                 name=self.name,
                 entrypoint=encoders.pickle_and_encode(self.entrypoint),
@@ -264,7 +242,7 @@ class Request:
             raise
 
     @classmethod
-    def decode(cls, payload: RequestPayload) -> 'Request':
+    def decode(cls, payload: payloads.RequestPayload) -> 'Request':
         """Deserialize the SkyPilot API request."""
         try:
             return cls(

sky/server/rest.py

@@ -12,6 +12,8 @@ import colorama
 from sky import exceptions
 from sky import sky_logging
 from sky.adaptors import common as adaptors_common
+from sky.server import constants
+from sky.server import versions
 from sky.utils import common_utils
 from sky.utils import rich_utils
 from sky.utils import ux_utils
@@ -28,6 +30,11 @@ F = TypeVar('F', bound=Callable[..., Any])
 
 _RETRY_CONTEXT = contextvars.ContextVar('retry_context', default=None)
 
+_session = requests.Session()
+_session.headers[constants.API_VERSION_HEADER] = str(constants.API_VERSION)
+_session.headers[constants.VERSION_HEADER] = (
+    versions.get_local_readable_version())
+
 
 class RetryContext:
 
@@ -132,13 +139,17 @@ def handle_server_unavailable(response: 'requests.Response') -> None:
 def request(method, url, **kwargs) -> 'requests.Response':
     """Send a request to the API server, retry on server temporarily
     unavailable."""
-    response = requests.request(method, url, **kwargs)
-    handle_server_unavailable(response)
-    return response
+    return request_without_retry(method, url, **kwargs)
 
 
 def request_without_retry(method, url, **kwargs) -> 'requests.Response':
     """Send a request to the API server without retry."""
-    response = requests.request(method, url, **kwargs)
+    response = _session.request(method, url, **kwargs)
     handle_server_unavailable(response)
+    remote_api_version = response.headers.get(constants.API_VERSION_HEADER)
+    remote_version = response.headers.get(constants.VERSION_HEADER)
+    if remote_api_version is not None:
+        versions.set_remote_api_version(int(remote_api_version))
+    if remote_version is not None:
+        versions.set_remote_version(remote_version)
     return response