skypilot-nightly 1.0.0.dev20251009__py3-none-any.whl → 1.0.0.dev20251107__py3-none-any.whl

sky/task.py

@@ -1,6 +1,5 @@
 """Task: a coarse-grained stage in an application."""
 import collections
-import inspect
 import json
 import os
 import re
@@ -29,10 +28,6 @@ from sky.utils import yaml_utils
 
 logger = sky_logging.init_logger(__name__)
 
-# A lambda generating commands (node rank_i, node addrs -> cmd_i).
-CommandGen = Callable[[int, List[str]], Optional[str]]
-CommandOrCommandGen = Union[str, CommandGen]
-
 _VALID_NAME_REGEX = '[a-zA-Z0-9]+(?:[._-]{1,2}[a-zA-Z0-9]+)*'
 _VALID_NAME_DESCR = ('ASCII characters and may contain lowercase and'
                      ' uppercase letters, digits, underscores, periods,'
@@ -236,7 +231,7 @@ class Task:
         name: Optional[str] = None,
         *,
         setup: Optional[Union[str, List[str]]] = None,
-        run: Optional[Union[CommandOrCommandGen, List[str]]] = None,
+        run: Optional[Union[str, List[str]]] = None,
         envs: Optional[Dict[str, str]] = None,
         secrets: Optional[Dict[str, str]] = None,
         workdir: Optional[Union[str, Dict[str, Any]]] = None,
@@ -349,7 +344,7 @@ class Task:
         self._volumes = volumes or {}
 
         # concatenate commands if given as list
-        def _concat(commands):
+        def _concat(commands: Optional[Union[str, List[str]]]) -> Optional[str]:
             if isinstance(commands, list):
                 return '\n'.join(commands)
             return commands
@@ -447,42 +442,9 @@ class Task:
 
     def validate_run(self):
         """Validates if the run command is valid."""
-        if callable(self.run):
-            run_sig = inspect.signature(self.run)
-            # Check that run is a function with 2 arguments.
-            if len(run_sig.parameters) != 2:
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(_RUN_FN_CHECK_FAIL_MSG.format(run_sig))
-
-            type_list = [int, List[str]]
-            # Check annotations, if exists
-            for i, param in enumerate(run_sig.parameters.values()):
-                if param.annotation != inspect.Parameter.empty:
-                    if param.annotation != type_list[i]:
-                        with ux_utils.print_exception_no_traceback():
-                            raise ValueError(
-                                _RUN_FN_CHECK_FAIL_MSG.format(run_sig))
-
-            # Check self containedness.
-            run_closure = inspect.getclosurevars(self.run)
-            if run_closure.nonlocals:
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(
-                        'run command generator must be self contained. '
-                        f'Found nonlocals: {run_closure.nonlocals}')
-            if run_closure.globals:
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(
-                        'run command generator must be self contained. '
-                        f'Found globals: {run_closure.globals}')
-            if run_closure.unbound:
-                # Do not raise an error here. Import statements, which are
-                # allowed, will be considered as unbounded.
-                pass
-        elif self.run is not None and not isinstance(self.run, str):
+        if self.run is not None and not isinstance(self.run, str):
             with ux_utils.print_exception_no_traceback():
-                raise ValueError('run must be either a shell script (str) or '
-                                 f'a command generator ({CommandGen}). '
+                raise ValueError('run must be a shell script (str). '
                                  f'Got {type(self.run)}')
 
     def expand_and_validate_file_mounts(self):
@@ -1130,7 +1092,7 @@ class Task:
     def set_resources(
         self, resources: Union['resources_lib.Resources',
                                List['resources_lib.Resources'],
-                               Set['resources_lib.Resources']]
+                               Set['resources_lib.Resources'], Dict[str, Any]]
     ) -> 'Task':
         """Sets the required resources to execute this task.
 
@@ -1144,7 +1106,9 @@ class Task:
         Returns:
           self: The current task, with resources set.
         """
-        if isinstance(resources, resources_lib.Resources):
+        if isinstance(resources, dict):
+            resources = resources_lib.Resources.from_yaml_config(resources)
+        elif isinstance(resources, resources_lib.Resources):
             resources = {resources}
         # TODO(woosuk): Check if the resources are None.
         self.resources = _with_docker_login_config(resources, self.envs,
@@ -1172,6 +1136,10 @@ class Task:
         self.set_resources(type(self.resources)(new_resources_list))
         return self
 
+    def get_resource_config(self) -> Dict[str, Any]:
+        return _resources_to_config(self.resources,
+                                    factor_out_common_fields=True)
+
     @property
     def service(self) -> Optional[service_spec.SkyServiceSpec]:
         return self._service
@@ -1552,6 +1520,16 @@ class Task:
                     self.update_file_mounts({
                         mnt_path: blob_path,
                     })
+                elif store_type is storage_lib.StoreType.COREWEAVE:
+                    if storage.source is not None and not isinstance(
+                            storage.source,
+                            list) and storage.source.startswith('cw://'):
+                        blob_path = storage.source
+                    else:
+                        blob_path = 'cw://' + storage.name
+                    self.update_file_mounts({
+                        mnt_path: blob_path,
+                    })
                 else:
                     with ux_utils.print_exception_no_traceback():
                         raise ValueError(f'Storage Type {store_type} '
@@ -1688,16 +1666,7 @@ class Task:
 
         add_if_not_none('name', self.name)
 
-        tmp_resource_config: Union[Dict[str, Union[str, int]],
-                                   Dict[str, List[Dict[str, Union[str, int]]]]]
-        if len(self.resources) > 1:
-            resource_list = []
-            for r in self.resources:
-                resource_list.append(r.to_yaml_config())
-            key = 'ordered' if isinstance(self.resources, list) else 'any_of'
-            tmp_resource_config = {key: resource_list}
-        else:
-            tmp_resource_config = list(self.resources)[0].to_yaml_config()
+        tmp_resource_config = _resources_to_config(self.resources)
 
         add_if_not_none('resources', tmp_resource_config)
 
@@ -1810,3 +1779,47 @@ class Task:
         else:
             s += '\n  resources: default instances'
         return s
+
+
+def _resources_to_config(
+        resources: Union[List['resources_lib.Resources'],
+                         Set['resources_lib.Resources']],
+        factor_out_common_fields: bool = False) -> Dict[str, Any]:
+    if len(resources) > 1:
+        resource_list: List[Dict[str, Union[str, int]]] = []
+        for r in resources:
+            resource_list.append(r.to_yaml_config())
+        group_key = 'ordered' if isinstance(resources, list) else 'any_of'
+        if factor_out_common_fields:
+            return _factor_out_common_resource_fields(resource_list, group_key)
+        return {group_key: resource_list}
+    else:
+        return list(resources)[0].to_yaml_config()
+
+
+def _factor_out_common_resource_fields(configs: List[Dict[str, Union[str,
+                                                                     int]]],
+                                       group_key: str) -> Dict[str, Any]:
+    """Factors out the fields that are common to all resources."""
+    return_config: Dict[str, Any] = configs[0].copy()
+    if len(configs) > 1:
+        for config in configs[1:]:
+            for key, value in config.items():
+                if key in return_config and return_config[key] != value:
+                    del return_config[key]
+    num_empty_configs = 0
+    for config in configs:
+        keys_to_delete = []
+        for key, value in config.items():
+            if key in return_config:
+                keys_to_delete.append(key)
+        for key in keys_to_delete:
+            del config[key]
+        if not config:
+            num_empty_configs += 1
+
+    if num_empty_configs == len(configs):
+        return return_config
+    if len(configs) > 0:
+        return_config[group_key] = configs
+    return return_config

sky/templates/kubernetes-ray.yml.j2

@@ -209,7 +209,9 @@ provider:
         metadata:
           labels:
             parent: skypilot
+            # TODO (kyuds): remove this label for v0.12.0, as skypilot-cluster label is deprecated in favor of skypilot-cluster-name.
             skypilot-cluster: {{cluster_name_on_cloud}}
+            skypilot-cluster-name: {{cluster_name_on_cloud}}
             skypilot-user: {{ user }}
           name: {{cluster_name_on_cloud}}-head-ssh
         spec:
@@ -227,7 +229,9 @@ provider:
         metadata:
             labels:
               parent: skypilot
+              # TODO (kyuds): remove this label for v0.12.0, as skypilot-cluster label is deprecated in favor of skypilot-cluster-name.
               skypilot-cluster: {{cluster_name_on_cloud}}
+              skypilot-cluster-name: {{cluster_name_on_cloud}}
               skypilot-user: {{ user }}
             # NOTE: If you're running multiple Ray clusters with services
             # on one Kubernetes cluster, they must have unique service
@@ -247,7 +251,9 @@ provider:
         metadata:
           labels:
             parent: skypilot
+            # TODO (kyuds): remove this label for v0.12.0, as skypilot-cluster label is deprecated in favor of skypilot-cluster-name.
             skypilot-cluster: {{cluster_name_on_cloud}}
+            skypilot-cluster-name: {{cluster_name_on_cloud}}
             skypilot-user: {{ user }}
           name: {{cluster_name_on_cloud}}-worker{{ worker_id }}
         spec:
@@ -272,6 +278,7 @@ available_node_types:
         labels:
             parent: skypilot
             # component will be set for the head node pod to be the same as the head node service selector above if a
+            # TODO (kyuds): remove this label for v0.11.0, as skypilot-cluster label is deprecated in favor of skypilot-cluster-name.
             skypilot-cluster: {{cluster_name_on_cloud}}
             skypilot-user: {{ user }}
             # Custom tags for the pods
@@ -1059,7 +1066,7 @@ available_node_types:
                 # Also, skip the jobs that are waiting to be scheduled as those does not have a controller process running.
                 # For SkyServe, this will be None and every service will be recovered. This is because SkyServe
                 # will delete the service from the database after it is terminated so everything in the database is running.
-                ALL_IN_PROGRESS_JOBS=$({{sky_python_cmd}} -c "from sky.jobs import state; jobs = state.get_managed_jobs(); print(' '.join({str(job['job_id']) for job in jobs if job['schedule_state'] not in [state.ManagedJobScheduleState.DONE, state.ManagedJobScheduleState.WAITING]}) if jobs else None)")
+                ALL_IN_PROGRESS_JOBS=$({{sky_python_cmd}} -c "from sky.jobs import state; jobs, _ = state.get_managed_jobs_with_filters(fields=['job_id', 'schedule_state']); print(' '.join({str(job['job_id']) for job in jobs if job['schedule_state'] not in [state.ManagedJobScheduleState.DONE, state.ManagedJobScheduleState.WAITING]}) if jobs else None)")
                 if [ "$ALL_IN_PROGRESS_JOBS" != "None" ]; then
                   read -ra ALL_IN_PROGRESS_JOBS_SEQ <<< "$ALL_IN_PROGRESS_JOBS"
                 fi

sky/templates/nebius-ray.yml.j2

@@ -156,6 +156,7 @@ setup_commands:
     echo '{{env_var}}={{env_value}}' | sudo tee -a /etc/environment;
     {%- endfor %}
   {%- endif %}
+    IP=$(hostname -I | awk '{print $1}'); echo "$IP $(hostname)" | sudo tee -a /etc/hosts;
     sudo bash -c 'rm -rf /etc/security/limits.d; echo "* soft nofile 1048576" >> /etc/security/limits.conf; echo "* hard nofile 1048576" >> /etc/security/limits.conf';
     sudo grep -e '^DefaultTasksMax' /etc/systemd/system.conf || (sudo bash -c 'echo "DefaultTasksMax=infinity" >> /etc/systemd/system.conf'); sudo systemctl set-property user-$(id -u $(whoami)).slice TasksMax=infinity; sudo systemctl daemon-reload;
     mkdir -p ~/.ssh; (grep -Pzo -q "Host \*\n  StrictHostKeyChecking no\n  IdentityFile ~/.ssh/sky-cluster-key\n  IdentityFile ~/.ssh/id_rsa" ~/.ssh/config) || printf "Host *\n  StrictHostKeyChecking no\n  IdentityFile ~/.ssh/sky-cluster-key\n  IdentityFile ~/.ssh/id_rsa\n" >> ~/.ssh/config;

sky/templates/shadeform-ray.yml.j2

@@ -0,0 +1,72 @@
+cluster_name: {{cluster_name_on_cloud}}
+
+# The maximum number of workers nodes to launch in addition to the head node.
+max_workers: {{num_nodes - 1}}
+upscaling_speed: {{num_nodes - 1}}
+idle_timeout_minutes: 60
+
+provider:
+  type: external
+  module: sky.provision.shadeform
+  region: "{{region}}"
+  disable_launch_config_check: true
+
+auth:
+  ssh_user: shadeform
+  ssh_private_key: {{ssh_private_key}}
+  ssh_key_id: {{ssh_key_id}}
+
+available_node_types:
+  ray_head_default:
+{%- if custom_resources %}
+    resources: {{custom_resources}}
+{%- else %}
+    resources: {}
+{%- endif %}
+    node_config:
+      InstanceType: {{instance_type}}
+      PublicKey: |-
+        skypilot:ssh_public_key_content
+
+head_node_type: ray_head_default
+
+# Format: `REMOTE_PATH : LOCAL_PATH`
+file_mounts: {
+  "{{sky_ray_yaml_remote_path}}": "{{sky_ray_yaml_local_path}}",
+  "{{sky_remote_path}}/{{sky_wheel_hash}}": "{{sky_local_path}}",
+{%- for remote_path, local_path in credentials.items() %}
+  "{{remote_path}}": "{{local_path}}",
+{%- endfor %}
+}
+
+rsync_exclude: []
+
+initialization_commands: []
+
+# List of shell commands to run to set up nodes.
+# NOTE: these are very performance-sensitive. Each new item opens/closes an SSH
+# connection, which is expensive. Try your best to co-locate commands into fewer
+# items!
+#
+# Increment the following for catching performance bugs easier:
+#   current num items (num SSH connections): 1
+setup_commands:
+  # Create ~/.ssh/config file in case the file does not exist in the image.
+  # Line 'rm ..': there is another installation of pip.
+  # Line 'sudo bash ..': set the ulimit as suggested by ray docs for performance. https://docs.ray.io/en/latest/cluster/vms/user-guides/large-cluster-best-practices.html#system-configuration
+  # Line 'sudo grep ..': set the number of threads per process to unlimited to avoid ray job submit stucking issue when the number of running ray jobs increase.
+  # Line 'mkdir -p ..': disable host key check
+  # Line 'python3 -c ..': patch the buggy ray files and enable `-o allow_other` option for `goofys`
+  - {%- for initial_setup_command in initial_setup_commands %}
+    {{ initial_setup_command }}
+    {%- endfor %}
+    mkdir -p ~/.ssh; touch ~/.ssh/config; which patch > /dev/null || sudo apt install -y patch;
+    {{ conda_installation_commands }}
+    {{ ray_skypilot_installation_commands }}
+    sudo bash -c 'rm -rf /etc/security/limits.d; echo "* soft nofile 1048576" >> /etc/security/limits.conf; echo "* hard nofile 1048576" >> /etc/security/limits.conf';
+    sudo grep -e '^DefaultTasksMax' /etc/systemd/system.conf || (sudo bash -c 'echo "DefaultTasksMax=infinity" >> /etc/systemd/system.conf'); sudo systemctl set-property user-$(id -u $(whoami)).slice TasksMax=infinity; sudo systemctl daemon-reload;
+    (grep -Pzo -q "Host \*\n  StrictHostKeyChecking no" ~/.ssh/config) || printf "Host *\n  StrictHostKeyChecking no\n" >> ~/.ssh/config;
+    {{ ssh_max_sessions_config }}
+
+# Command to start ray clusters are now placed in `sky.provision.instance_setup`.
+# We do not need to list it here anymore.

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,27 @@ if __name__ == '__main__':
         # TODO(aylei): Remove this after 0.10.0
         server_url = f'http://{server_url}'
 
+    disable_latency_measurement = os.environ.get(
+        skylet_constants.SSH_DISABLE_LATENCY_MEASUREMENT_ENV_VAR, '0') == '1'
+    if disable_latency_measurement:
+        timestamps_are_supported = False
+    else:
+        health_url = f'{server_url}/api/health'
+        cookie_hdr = _get_cookie_header(health_url)
+        health_response = requests.get(health_url, headers=cookie_hdr)
+        health_data = health_response.json()
+        timestamps_are_supported = int(health_data.get('api_version', 0)) > 21
+
     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

@@ -14,6 +14,7 @@ from sky import models
 from sky import sky_logging
 from sky.skylet import constants
 from sky.users import rbac
+from sky.utils import annotations
 from sky.utils import common_utils
 from sky.utils.db import db_utils
 
@@ -42,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)
@@ -52,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:
@@ -254,6 +258,9 @@ class PermissionService:
         with _policy_lock():
             self._load_policy_no_lock()
 
+    # Right now, not a lot of users are using multiple workspaces,
+    # so 5 should be more than enough.
+    @annotations.lru_cache(scope='request', maxsize=5)
     def check_workspace_permission(self, user_id: str,
                                    workspace_name: str) -> bool:
         """Check workspace permission.

sky/utils/admin_policy_utils.py

@@ -2,6 +2,7 @@
 import contextlib
 import copy
 import importlib
+import typing
 from typing import Iterator, Optional, Tuple, Union
 import urllib.parse
 
@@ -13,12 +14,16 @@ from sky import exceptions
 from sky import sky_logging
 from sky import skypilot_config
 from sky import task as task_lib
+from sky.server.requests import request_names
 from sky.utils import common_utils
 from sky.utils import config_utils
 from sky.utils import ux_utils
 
 logger = sky_logging.init_logger(__name__)
 
+if typing.TYPE_CHECKING:
+    from sky import models
+
 
 def _is_url(policy_string: str) -> bool:
     """Check if the policy string is a URL."""
@@ -73,6 +78,7 @@ def _get_policy_impl(
 @contextlib.contextmanager
 def apply_and_use_config_in_current_request(
     entrypoint: Union['dag_lib.Dag', 'task_lib.Task'],
+    request_name: request_names.AdminPolicyRequestName,
     request_options: Optional[admin_policy.RequestOptions] = None,
     at_client_side: bool = False,
 ) -> Iterator['dag_lib.Dag']:
@@ -86,7 +92,8 @@ def apply_and_use_config_in_current_request(
     Refer to `apply()` for more details.
     """
     original_config = skypilot_config.to_dict()
-    dag, mutated_config = apply(entrypoint, request_options, at_client_side)
+    dag, mutated_config = apply(entrypoint, request_name, request_options,
+                                at_client_side)
     if mutated_config != original_config:
         with skypilot_config.replace_skypilot_config(mutated_config):
             yield dag
@@ -96,6 +103,7 @@ def apply_and_use_config_in_current_request(
 
 def apply(
     entrypoint: Union['dag_lib.Dag', 'task_lib.Task'],
+    request_name: request_names.AdminPolicyRequestName,
     request_options: Optional[admin_policy.RequestOptions] = None,
     at_client_side: bool = False,
 ) -> Tuple['dag_lib.Dag', config_utils.Config]:
@@ -126,9 +134,13 @@ def apply(
     if policy is None:
         return dag, skypilot_config.to_dict()
 
+    user = None
     if at_client_side:
         logger.info(f'Applying client admin policy: {policy}')
     else:
+        # When being called by the server, the middleware has set the
+        # current user and this information is available at this point.
+        user = common_utils.get_current_user()
         logger.info(f'Applying server admin policy: {policy}')
     config = copy.deepcopy(skypilot_config.to_dict())
     mutated_dag = dag_lib.Dag()
@@ -136,8 +148,9 @@ def apply(
 
     mutated_config = None
     for task in dag.tasks:
-        user_request = admin_policy.UserRequest(task, config, request_options,
-                                                at_client_side)
+        user_request = admin_policy.UserRequest(task, config, request_name,
+                                                request_options, at_client_side,
+                                                user)
         try:
             mutated_user_request = policy.apply(user_request)
         # Avoid duplicate exception wrapping.