skypilot-nightly 1.0.0.dev20251029__py3-none-any.whl → 1.0.0.dev20251101__py3-none-any.whl

sky/templates/websocket_proxy.py

@@ -11,15 +11,23 @@ This script is useful for users who do not have local Kubernetes credentials.
 import asyncio
 from http.cookiejar import MozillaCookieJar
 import os
+import struct
 import sys
-from typing import Dict
+import time
+from typing import Dict, Optional
 from urllib.request import Request
 
+import requests
 import websockets
 from websockets.asyncio.client import ClientConnection
 from websockets.asyncio.client import connect
 
+from sky.server import constants
+from sky.server.server import KubernetesSSHMessageType
+from sky.skylet import constants as skylet_constants
+
 BUFFER_SIZE = 2**16  # 64KB
+HEARTBEAT_INTERVAL_SECONDS = 10
 
 # Environment variable for a file path to the API cookie file.
 # Keep in sync with server/constants.py
@@ -28,6 +36,8 @@ API_COOKIE_FILE_ENV_VAR = 'SKYPILOT_API_COOKIE_FILE'
 # Keep in sync with server/constants.py
 API_COOKIE_FILE_DEFAULT_LOCATION = '~/.sky/cookies.txt'
 
+MAX_UNANSWERED_PINGS = 100
+
 
 def _get_cookie_header(url: str) -> Dict[str, str]:
     """Extract Cookie header value from a cookie jar for a specific URL"""
@@ -49,7 +59,7 @@ def _get_cookie_header(url: str) -> Dict[str, str]:
     return {'Cookie': cookie_header}
 
 
-async def main(url: str) -> None:
+async def main(url: str, timestamps_supported: bool) -> None:
     cookie_header = _get_cookie_header(url)
     async with connect(url,
                        ping_interval=None,
@@ -75,45 +85,149 @@ async def main(url: str) -> None:
                 asyncio.streams.FlowControlMixin, sys.stdout)  # type: ignore
             stdout_writer = asyncio.StreamWriter(transport, protocol, None,
                                                  loop)
+            # Dictionary to store last ping time for latency measurement
+            last_ping_time_dict: Optional[Dict[int, float]] = None
+            if timestamps_supported:
+                last_ping_time_dict = {}
+
+            # Use an Event to signal when websocket is closed
+            websocket_closed_event = asyncio.Event()
+            websocket_lock = asyncio.Lock()
 
-            await asyncio.gather(stdin_to_websocket(stdin_reader, websocket),
-                                 websocket_to_stdout(websocket, stdout_writer))
+            await asyncio.gather(
+                stdin_to_websocket(stdin_reader, websocket,
+                                   timestamps_supported, websocket_closed_event,
+                                   websocket_lock),
+                websocket_to_stdout(websocket, stdout_writer,
+                                    timestamps_supported, last_ping_time_dict,
+                                    websocket_closed_event, websocket_lock),
+                latency_monitor(websocket, last_ping_time_dict,
+                                websocket_closed_event, websocket_lock),
+                return_exceptions=True)
         finally:
             if old_settings:
                 termios.tcsetattr(sys.stdin.fileno(), termios.TCSADRAIN,
                                   old_settings)
 
 
+async def latency_monitor(websocket: ClientConnection,
+                          last_ping_time_dict: Optional[dict],
+                          websocket_closed_event: asyncio.Event,
+                          websocket_lock: asyncio.Lock):
+    """Periodically send PING messages (type 1) to measure latency."""
+    if last_ping_time_dict is None:
+        return
+    next_id = 0
+    while not websocket_closed_event.is_set():
+        try:
+            await asyncio.sleep(HEARTBEAT_INTERVAL_SECONDS)
+            if len(last_ping_time_dict) >= MAX_UNANSWERED_PINGS:
+                # We are not getting responses, clear the dictionary so
+                # as not to grow unbounded.
+                last_ping_time_dict.clear()
+            ping_time = time.time()
+            next_id += 1
+            last_ping_time_dict[next_id] = ping_time
+            message_header_bytes = struct.pack(
+                '!BI', KubernetesSSHMessageType.PINGPONG.value, next_id)
+            try:
+                async with websocket_lock:
+                    await websocket.send(message_header_bytes)
+            except websockets.exceptions.ConnectionClosed as e:
+                # Websocket is already closed.
+                print(f'Failed to send PING message: {e}', file=sys.stderr)
+                break
+        except Exception as e:
+            print(f'Error in latency_monitor: {e}', file=sys.stderr)
+            websocket_closed_event.set()
+            raise e
+
+
 async def stdin_to_websocket(reader: asyncio.StreamReader,
-                             websocket: ClientConnection):
+                             websocket: ClientConnection,
+                             timestamps_supported: bool,
+                             websocket_closed_event: asyncio.Event,
+                             websocket_lock: asyncio.Lock):
     try:
-        while True:
+        while not websocket_closed_event.is_set():
             # Read at most BUFFER_SIZE bytes, this not affect
             # responsiveness since it will return as soon as
             # there is at least one byte.
             # The BUFFER_SIZE is chosen to be large enough to improve
             # throughput.
             data = await reader.read(BUFFER_SIZE)
+
             if not data:
                 break
-            await websocket.send(data)
+            if timestamps_supported:
+                # Send message with type 0 to indicate data.
+                message_type_bytes = struct.pack(
+                    '!B', KubernetesSSHMessageType.REGULAR_DATA.value)
+                data = message_type_bytes + data
+            async with websocket_lock:
+                await websocket.send(data)
+
     except Exception as e:  # pylint: disable=broad-except
         print(f'Error in stdin_to_websocket: {e}', file=sys.stderr)
     finally:
-        await websocket.close()
+        async with websocket_lock:
+            await websocket.close()
+        websocket_closed_event.set()
 
 
 async def websocket_to_stdout(websocket: ClientConnection,
-                              writer: asyncio.StreamWriter):
+                              writer: asyncio.StreamWriter,
+                              timestamps_supported: bool,
+                              last_ping_time_dict: Optional[dict],
+                              websocket_closed_event: asyncio.Event,
+                              websocket_lock: asyncio.Lock):
     try:
-        while True:
+        while not websocket_closed_event.is_set():
             message = await websocket.recv()
+            if (timestamps_supported and len(message) > 0 and
+                    last_ping_time_dict is not None):
+                message_type = struct.unpack('!B', message[:1])[0]
+                if message_type == KubernetesSSHMessageType.REGULAR_DATA.value:
+                    # Regular data - strip type byte and write to stdout
+                    message = message[1:]
+                elif message_type == KubernetesSSHMessageType.PINGPONG.value:
+                    # PONG response - calculate latency and send measurement
+                    if not len(message) == struct.calcsize('!BI'):
+                        raise ValueError(
+                            f'Invalid PONG message length: {len(message)}')
+                    pong_id = struct.unpack('!I', message[1:5])[0]
+                    pong_time = time.time()
+
+                    ping_time = last_ping_time_dict.pop(pong_id, None)
+
+                    if ping_time is None:
+                        continue
+
+                    latency_seconds = pong_time - ping_time
+                    latency_ms = int(latency_seconds * 1000)
+
+                    # Send latency measurement (type 2)
+                    message_type_bytes = struct.pack(
+                        '!B',
+                        KubernetesSSHMessageType.LATENCY_MEASUREMENT.value)
+                    latency_bytes = struct.pack('!Q', latency_ms)
+                    message = message_type_bytes + latency_bytes
+                    # Send to server.
+                    async with websocket_lock:
+                        await websocket.send(message)
+                    continue
+            # No timestamps support, write directly
             writer.write(message)
             await writer.drain()
     except websockets.exceptions.ConnectionClosed:
         print('WebSocket connection closed', file=sys.stderr)
     except Exception as e:  # pylint: disable=broad-except
         print(f'Error in websocket_to_stdout: {e}', file=sys.stderr)
+        raise e
+    finally:
+        async with websocket_lock:
+            await websocket.close()
+        websocket_closed_event.set()
 
 
 if __name__ == '__main__':
@@ -123,11 +237,25 @@ if __name__ == '__main__':
         # TODO(aylei): Remove this after 0.10.0
         server_url = f'http://{server_url}'
 
+    health_url = f'{server_url}/api/health'
+    health_response = requests.get(health_url)
+    health_data = health_response.json()
+    timestamps_are_supported = int(health_data['api_version']) > 21
+    disable_latency_measurement = os.environ.get(
+        skylet_constants.SSH_DISABLE_LATENCY_MEASUREMENT_ENV_VAR, '0') == '1'
+    timestamps_are_supported = (timestamps_are_supported and
+                                not disable_latency_measurement)
+
     server_proto, server_fqdn = server_url.split('://')
     websocket_proto = 'ws'
     if server_proto == 'https':
         websocket_proto = 'wss'
     server_url = f'{websocket_proto}://{server_fqdn}'
+
+    client_version_str = (f'&client_version={constants.API_VERSION}'
+                          if timestamps_are_supported else '')
+
     websocket_url = (f'{server_url}/kubernetes-pod-ssh-proxy'
-                     f'?cluster_name={sys.argv[2]}')
-    asyncio.run(main(websocket_url))
+                     f'?cluster_name={sys.argv[2]}'
+                     f'{client_version_str}')
+    asyncio.run(main(websocket_url, timestamps_are_supported))

sky/users/permission.py

@@ -43,7 +43,6 @@ class PermissionService:
         with _policy_lock():
             global _enforcer_instance
             if _enforcer_instance is None:
-                _enforcer_instance = self
                 engine = global_user_state.initialize_and_get_db()
                 db_utils.add_all_tables_to_db_sqlalchemy(
                     sqlalchemy_adapter.Base.metadata, engine)
@@ -53,6 +52,10 @@ class PermissionService:
                                           'model.conf')
                 enforcer = casbin.Enforcer(model_path, adapter)
                 self.enforcer = enforcer
+                # Only set the enforcer instance once the enforcer
+                # is successfully initialized, if we change it and then fail
+                # we will set it to None and all subsequent calls will fail.
+                _enforcer_instance = self
                 self._maybe_initialize_policies()
                 self._maybe_initialize_basic_auth_user()
             else:

sky/utils/db/migration_utils.py

@@ -22,7 +22,7 @@ GLOBAL_USER_STATE_VERSION = '010'
 GLOBAL_USER_STATE_LOCK_PATH = f'~/.sky/locks/.{GLOBAL_USER_STATE_DB_NAME}.lock'
 
 SPOT_JOBS_DB_NAME = 'spot_jobs_db'
-SPOT_JOBS_VERSION = '003'
+SPOT_JOBS_VERSION = '005'
 SPOT_JOBS_LOCK_PATH = f'~/.sky/locks/.{SPOT_JOBS_DB_NAME}.lock'
 
 SERVE_DB_NAME = 'serve_db'

sky/utils/resource_checker.py

@@ -278,7 +278,10 @@ def _get_active_resources(
         from sky.jobs.server import core as managed_jobs_core
         try:
             filtered_jobs, _, _, _ = managed_jobs_core.queue_v2(
-                refresh=False, skip_finished=True, all_users=True)
+                refresh=False,
+                skip_finished=True,
+                all_users=True,
+                fields=['job_id', 'user_hash', 'workspace'])
             return filtered_jobs
         except exceptions.ClusterNotUpError:
             logger.warning('All jobs should be finished.')

sky/utils/schemas.py

@@ -1190,7 +1190,13 @@ def get_config_schema():
                         'consolidation_mode': {
                             'type': 'boolean',
                             'default': False,
-                        }
+                        },
+                        'controller_logs_gc_retention_hours': {
+                            'type': 'integer',
+                        },
+                        'task_logs_gc_retention_hours': {
+                            'type': 'integer',
+                        },
                     },
                 },
                 'bucket': {
@@ -1592,10 +1598,10 @@ def get_config_schema():
 
     allowed_workspace_cloud_names = list(constants.ALL_CLOUDS) + ['cloudflare']
     # Create pattern for not supported clouds, i.e.
-    # all clouds except gcp, kubernetes, ssh
+    # all clouds except aws, gcp, kubernetes, ssh, nebius
     not_supported_clouds = [
         cloud for cloud in allowed_workspace_cloud_names
-        if cloud.lower() not in ['gcp', 'kubernetes', 'ssh', 'nebius']
+        if cloud.lower() not in ['aws', 'gcp', 'kubernetes', 'ssh', 'nebius']
     ]
     not_supported_cloud_regex = '|'.join(not_supported_clouds)
     workspaces_schema = {
@@ -1606,7 +1612,8 @@ def get_config_schema():
             'type': 'object',
             'additionalProperties': False,
             'patternProperties': {
-                # Pattern for non-GCP clouds - only allows 'disabled' property
+                # Pattern for clouds with no workspace-specific config -
+                # only allow 'disabled' property.
                 f'^({not_supported_cloud_regex})$': {
                     'type': 'object',
                     'additionalProperties': False,
@@ -1641,6 +1648,18 @@ def get_config_schema():
                     },
                     'additionalProperties': False,
                 },
+                'aws': {
+                    'type': 'object',
+                    'properties': {
+                        'profile': {
+                            'type': 'string'
+                        },
+                        'disabled': {
+                            'type': 'boolean'
+                        },
+                    },
+                    'additionalProperties': False,
+                },
                 'ssh': {
                     'type': 'object',
                     'required': [],

sky/volumes/server/server.py

@@ -7,6 +7,7 @@ from sky import exceptions
 from sky import sky_logging
 from sky.server.requests import executor
 from sky.server.requests import payloads
+from sky.server.requests import request_names
 from sky.server.requests import requests as requests_lib
 from sky.utils import registry
 from sky.utils import volume as volume_utils
@@ -27,7 +28,7 @@ async def volume_list(request: fastapi.Request) -> None:
     request_body = payloads.RequestBody(**auth_user_env_vars_kwargs)
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='volume_list',
+        request_name=request_names.RequestName.VOLUME_LIST,
         request_body=request_body,
         func=core.volume_list,
         schedule_type=requests_lib.ScheduleType.SHORT,
@@ -40,7 +41,7 @@ async def volume_delete(request: fastapi.Request,
     """Deletes a volume."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='volume_delete',
+        request_name=request_names.RequestName.VOLUME_DELETE,
         request_body=volume_delete_body,
         func=core.volume_delete,
         schedule_type=requests_lib.ScheduleType.LONG,
@@ -114,7 +115,7 @@ async def volume_apply(request: fastapi.Request,
                 detail='Runpod network volume is only supported on Runpod')
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='volume_apply',
+        request_name=request_names.RequestName.VOLUME_APPLY,
         request_body=volume_apply_body,
         func=core.volume_apply,
         schedule_type=requests_lib.ScheduleType.LONG,

sky/workspaces/server.py

@@ -4,6 +4,7 @@ import fastapi
 
 from sky.server.requests import executor
 from sky.server.requests import payloads
+from sky.server.requests import request_names
 from sky.server.requests import requests as api_requests
 from sky.workspaces import core
 
@@ -24,7 +25,7 @@ async def get(request: fastapi.Request) -> None:
 
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='workspaces.get',
+        request_name=request_names.RequestName.WORKSPACES_GET,
         request_body=request_body,
         func=core.get_workspaces,
         schedule_type=api_requests.ScheduleType.SHORT,
@@ -37,7 +38,7 @@ async def update(request: fastapi.Request,
     """Updates a specific workspace configuration."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='workspaces.update',
+        request_name=request_names.RequestName.WORKSPACES_UPDATE,
         request_body=update_workspace_body,
         func=core.update_workspace,
         schedule_type=api_requests.ScheduleType.SHORT,
@@ -50,7 +51,7 @@ async def create(request: fastapi.Request,
     """Creates a new workspace configuration."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='workspaces.create',
+        request_name=request_names.RequestName.WORKSPACES_CREATE,
         request_body=create_workspace_body,
         func=core.create_workspace,
         schedule_type=api_requests.ScheduleType.SHORT,
@@ -63,7 +64,7 @@ async def delete(request: fastapi.Request,
     """Deletes a workspace configuration."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='workspaces.delete',
+        request_name=request_names.RequestName.WORKSPACES_DELETE,
         request_body=delete_workspace_body,
         func=core.delete_workspace,
         schedule_type=api_requests.ScheduleType.SHORT,
@@ -80,7 +81,7 @@ async def get_config(request: fastapi.Request) -> None:
     get_config_body = payloads.GetConfigBody(**auth_user_env_vars_kwargs)
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='workspaces.get_config',
+        request_name=request_names.RequestName.WORKSPACES_GET_CONFIG,
         request_body=get_config_body,
         func=core.get_config,
         schedule_type=api_requests.ScheduleType.SHORT,
@@ -93,7 +94,7 @@ async def update_config(request: fastapi.Request,
     """Updates the entire SkyPilot configuration."""
     await executor.schedule_request_async(
         request_id=request.state.request_id,
-        request_name='workspaces.update_config',
+        request_name=request_names.RequestName.WORKSPACES_UPDATE_CONFIG,
         request_body=update_config_body,
         func=core.update_config,
         schedule_type=api_requests.ScheduleType.SHORT,

{skypilot_nightly-1.0.0.dev20251029.dist-info → skypilot_nightly-1.0.0.dev20251101.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: skypilot-nightly
-Version: 1.0.0.dev20251029
+Version: 1.0.0.dev20251101
 Summary: SkyPilot: Run AI on Any Infra — Unified, Faster, Cheaper.
 Author: SkyPilot Team
 License: Apache 2.0
@@ -387,51 +387,51 @@ Requires-Dist: protobuf<7.0.0,>=5.26.1; extra == "shadeform"
 Requires-Dist: aiosqlite; extra == "shadeform"
 Requires-Dist: greenlet; extra == "shadeform"
 Provides-Extra: all
-Requires-Dist: msgraph-sdk; extra == "all"
-Requires-Dist: ibm-cloud-sdk-core; extra == "all"
-Requires-Dist: ibm-vpc; extra == "all"
+Requires-Dist: azure-cli>=2.65.0; extra == "all"
+Requires-Dist: aiohttp; extra == "all"
+Requires-Dist: google-api-python-client>=2.69.0; extra == "all"
 Requires-Dist: azure-core>=1.31.0; extra == "all"
-Requires-Dist: awscli>=1.27.10; extra == "all"
-Requires-Dist: sqlalchemy_adapter; extra == "all"
+Requires-Dist: botocore>=1.29.10; extra == "all"
+Requires-Dist: colorama<0.4.5; extra == "all"
+Requires-Dist: docker; extra == "all"
+Requires-Dist: protobuf<7.0.0,>=5.26.1; extra == "all"
+Requires-Dist: runpod>=1.6.1; extra == "all"
 Requires-Dist: azure-mgmt-network>=27.0.0; extra == "all"
-Requires-Dist: azure-identity>=1.19.0; extra == "all"
-Requires-Dist: ray[default]>=2.6.1; extra == "all"
+Requires-Dist: ibm-vpc; extra == "all"
+Requires-Dist: sqlalchemy_adapter; extra == "all"
+Requires-Dist: awscli>=1.27.10; extra == "all"
+Requires-Dist: anyio; extra == "all"
+Requires-Dist: kubernetes!=32.0.0,>=20.0.0; extra == "all"
+Requires-Dist: pydo>=0.3.0; extra == "all"
+Requires-Dist: grpcio>=1.63.0; extra == "all"
+Requires-Dist: tomli; python_version < "3.11" and extra == "all"
+Requires-Dist: passlib; extra == "all"
 Requires-Dist: nebius>=0.2.47; extra == "all"
-Requires-Dist: google-api-python-client>=2.69.0; extra == "all"
+Requires-Dist: azure-core>=1.24.0; extra == "all"
+Requires-Dist: ibm-cloud-sdk-core; extra == "all"
 Requires-Dist: vastai-sdk>=0.1.12; extra == "all"
-Requires-Dist: casbin; extra == "all"
-Requires-Dist: grpcio>=1.63.0; extra == "all"
+Requires-Dist: ibm-platform-services>=0.48.0; extra == "all"
+Requires-Dist: pyjwt; extra == "all"
+Requires-Dist: cudo-compute>=0.1.10; extra == "all"
+Requires-Dist: ray[default]>=2.6.1; extra == "all"
 Requires-Dist: aiosqlite; extra == "all"
+Requires-Dist: oci; extra == "all"
+Requires-Dist: ibm-cos-sdk; extra == "all"
 Requires-Dist: pyopenssl<24.3.0,>=23.2.0; extra == "all"
-Requires-Dist: runpod>=1.6.1; extra == "all"
+Requires-Dist: ecsapi>=0.2.0; extra == "all"
+Requires-Dist: msgraph-sdk; extra == "all"
+Requires-Dist: boto3>=1.26.1; extra == "all"
+Requires-Dist: casbin; extra == "all"
 Requires-Dist: greenlet; extra == "all"
-Requires-Dist: azure-common; extra == "all"
-Requires-Dist: colorama<0.4.5; extra == "all"
+Requires-Dist: pyvmomi==8.0.1.0.2; extra == "all"
+Requires-Dist: python-dateutil; extra == "all"
+Requires-Dist: azure-storage-blob>=12.23.1; extra == "all"
+Requires-Dist: azure-identity>=1.19.0; extra == "all"
 Requires-Dist: google-cloud-storage; extra == "all"
+Requires-Dist: msrestazure; extra == "all"
+Requires-Dist: azure-common; extra == "all"
 Requires-Dist: websockets; extra == "all"
 Requires-Dist: azure-mgmt-compute>=33.0.0; extra == "all"
-Requires-Dist: msrestazure; extra == "all"
-Requires-Dist: tomli; python_version < "3.11" and extra == "all"
-Requires-Dist: ecsapi>=0.2.0; extra == "all"
-Requires-Dist: python-dateutil; extra == "all"
-Requires-Dist: passlib; extra == "all"
-Requires-Dist: kubernetes!=32.0.0,>=20.0.0; extra == "all"
-Requires-Dist: docker; extra == "all"
-Requires-Dist: anyio; extra == "all"
-Requires-Dist: ibm-cos-sdk; extra == "all"
-Requires-Dist: pyvmomi==8.0.1.0.2; extra == "all"
-Requires-Dist: pyjwt; extra == "all"
-Requires-Dist: oci; extra == "all"
-Requires-Dist: azure-storage-blob>=12.23.1; extra == "all"
-Requires-Dist: cudo-compute>=0.1.10; extra == "all"
-Requires-Dist: azure-core>=1.24.0; extra == "all"
-Requires-Dist: aiohttp; extra == "all"
-Requires-Dist: protobuf<7.0.0,>=5.26.1; extra == "all"
-Requires-Dist: botocore>=1.29.10; extra == "all"
-Requires-Dist: azure-cli>=2.65.0; extra == "all"
-Requires-Dist: boto3>=1.26.1; extra == "all"
-Requires-Dist: pydo>=0.3.0; extra == "all"
-Requires-Dist: ibm-platform-services>=0.48.0; extra == "all"
 Provides-Extra: remote
 Requires-Dist: grpcio>=1.63.0; extra == "remote"
 Requires-Dist: protobuf<7.0.0,>=5.26.1; extra == "remote"