skypilot-nightly 1.0.0.dev20250616__py3-none-any.whl → 1.0.0.dev20250618__py3-none-any.whl

sky/templates/kubernetes-ray.yml.j2

@@ -273,6 +273,15 @@ available_node_types:
             {% if (k8s_acc_label_key is not none and k8s_acc_label_values is not none) %}
             skypilot-binpack: "gpu"
             {% endif %}
+            {% if k8s_kueue_local_queue_name %}
+            kueue.x-k8s.io/queue-name: {{k8s_kueue_local_queue_name}}
+            kueue.x-k8s.io/pod-group-name: {{cluster_name_on_cloud}}
+            {% endif %}
+        {% if k8s_kueue_local_queue_name %}
+        annotations:
+            kueue.x-k8s.io/retriable-in-group: "false"
+            kueue.x-k8s.io/pod-group-total-count: "{{ num_nodes|string }}"
+        {% endif %}
       spec:
         # serviceAccountName: skypilot-service-account
         serviceAccountName: {{k8s_service_account_name}}
@@ -632,19 +641,66 @@ available_node_types:
               {% if high_availability %}
               mkdir -p {{k8s_high_availability_deployment_run_script_dir}}
               if [ -f {{k8s_high_availability_deployment_volume_mount_path}}/k8s_container_ready ]; then
+                SKYPILOT_HA_RECOVERY_LOG="/tmp/ha_recovery.log"
+                echo "Starting HA recovery at $(date)" >> $SKYPILOT_HA_RECOVERY_LOG
+                start_time=$SECONDS
+                retry_count=0
+
+                # Wait for Ray to be ready, as the following commands is depending on Ray.
+                GET_RAY_STATUS_CMD=$({{sky_python_cmd}} -c 'from sky.provision import instance_setup; print(instance_setup.RAY_STATUS_WITH_SKY_RAY_PORT_COMMAND)')
+                while true; do
+                  retry_count=$((retry_count + 1))
+                  current_duration=$(( SECONDS - start_time ))
+                  echo "Attempt $retry_count to get Ray status after $current_duration seconds..." >> $SKYPILOT_HA_RECOVERY_LOG
+
+                  bash --login -c "$GET_RAY_STATUS_CMD"
+                  if [ $? -eq 0 ]; then
+                    wait_duration=$(( SECONDS - start_time ))
+                    echo "Ray ready after waiting $wait_duration seconds (took $retry_count attempts)" >> $SKYPILOT_HA_RECOVERY_LOG
+                    break
+                  fi
+                  echo "Waiting for Ray to be ready..." >> $SKYPILOT_HA_RECOVERY_LOG
+                  sleep 2
+                done
+
                 # ! Keep this aligned with `CloudVmRayBackend._setup()`
-                # Suppose all `task.setup` are the same for skyserve controller task.
+                # Suppose all `task.setup` are the same for sky serve / managed jobs controller task.
                 # So be careful for compatibility issue once you change it.
                 chmod +x {{k8s_high_availability_deployment_setup_script_path}}
                 /bin/bash --login -c "true && export OMP_NUM_THREADS=1 PYTHONWARNINGS='ignore' && {{k8s_high_availability_deployment_setup_script_path}} > /tmp/controller_recovery_setup_commands.log 2>&1"
-                echo "=== Controller setup commands completed for recovery ==="
-
+                echo "=== Controller setup commands completed for recovery at $(date) ===" >> $SKYPILOT_HA_RECOVERY_LOG
+
+                touch {{k8s_high_availability_restarting_signal_file}}
+                # Get all in-progress jobs from managed jobs controller. We skip any jobs that are already done.
+                # 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)")
+                if [ "$ALL_IN_PROGRESS_JOBS" != "None" ]; then
+                  read -ra ALL_IN_PROGRESS_JOBS_SEQ <<< "$ALL_IN_PROGRESS_JOBS"
+                fi
                 for file in {{k8s_high_availability_deployment_run_script_dir}}/*; do
+                  # This is the cluster job id on managed jobs controller, but it is guaranteed to be the same as the managed job id,
+                  # so we directly use it here. See `CloudVmRayBackend._exec_code_on_head::_dump_code_to_file` for more details.
+                  JOB_ID=$(basename $file | sed 's/sky_job_//')
+                  # If the list of in-progress jobs is not None (meaning this is a managed job HA controller) and job is not in-progress, skip.
+                  if [ "$ALL_IN_PROGRESS_JOBS" != "None" ]; then
+                    if [[ ! " ${ALL_IN_PROGRESS_JOBS_SEQ[@]} " =~ " ${JOB_ID} " ]]; then
+                      continue
+                    fi
+                  fi
                   # ! Keep this aligned with `CloudVmRayBackend._execute()`
                   chmod +x $file
+                  # TODO(tian): This logic may run a lot of things if the jobs controller previously had many jobs.
+                  # We should do more tests and make sure it will scale well.
                   /bin/bash --login -c "true && export OMP_NUM_THREADS=1 PYTHONWARNINGS='ignore' && $file > /tmp/task_run_$(basename $file).log 2>&1"
-                  echo "=== Controller task run for service (file: $file) completed for recovery ==="
+                  echo "=== Controller task run for service / job (file: $file) completed for recovery at $(date) ===" >> $SKYPILOT_HA_RECOVERY_LOG
                 done
+                rm {{k8s_high_availability_restarting_signal_file}}
+
+                duration=$(( SECONDS - start_time ))
+                echo "HA recovery completed at $(date)" >> $SKYPILOT_HA_RECOVERY_LOG
+                echo "Total recovery time: $duration seconds" >> $SKYPILOT_HA_RECOVERY_LOG
               fi
 
               touch {{k8s_high_availability_deployment_volume_mount_path}}/k8s_container_ready

sky/templates/scp-ray.yml.j2

@@ -7,7 +7,7 @@ idle_timeout_minutes: 60
 
 provider:
   type: external
-  module: sky.skylet.providers.scp.SCPNodeProvider
+  module: sky.provision.scp
   region: {{region}}
   cache_stopped_nodes: True
 
@@ -24,19 +24,6 @@ available_node_types:
       InstanceType: {{instance_type}}
       imageId: {{image_id}}
       diskSize: {{disk_size}}
-{% if num_nodes > 1 %}
-  ray_worker_default:
-    min_workers: {{num_nodes - 1}}
-    max_workers: {{num_nodes - 1}}
-    resources: {}
-    node_config:
-      AuthorizedKey: |
-        skypilot:ssh_public_key_content
-      InstanceType: {{instance_type}}
-      imageId: {{image_id}}
-      diskSize: {{disk_size}}
-
-{%- endif %}
 
 head_node_type: ray_head_default
 
@@ -50,10 +37,6 @@ file_mounts: {
 {%- 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
@@ -77,36 +60,6 @@ setup_commands:
     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;
     [ -f /etc/fuse.conf ] && sudo sed -i 's/#user_allow_other/user_allow_other/g' /etc/fuse.conf || (sudo sh -c 'echo "user_allow_other" > /etc/fuse.conf'); # This is needed for `-o allow_other` option for `goofys`;
-    {{ ssh_max_sessions_config }}
-
-# Command to start ray on the head node. You don't need to change this.
-# 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! The same comment applies for worker_start_ray_commands.
-#
-# Increment the following for catching performance bugs easier:
-#   current num items (num SSH connections): 1
-head_start_ray_commands:
-  # NOTE: --disable-usage-stats in `ray start` saves 10 seconds of idle wait.
-  # Line "which prlimit ..": increase the limit of the number of open files for the raylet process, as the `ulimit` may not take effect at this point, because it requires
-  # all the sessions to be reloaded. This is a workaround.
-  - {{ sky_activate_python_env }}; {{ sky_ray_cmd }} stop; RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 {{ sky_ray_cmd }} start --disable-usage-stats --head --port={{ray_port}} --dashboard-port={{ray_dashboard_port}} --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} --temp-dir {{ray_temp_dir}} || exit 1;
-    which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
-    {{dump_port_command}}; {{ray_head_wait_initialized_command}}
-
-{%- if num_nodes > 1 %}
-worker_start_ray_commands:
-  - {{ sky_activate_python_env }}; {{ sky_ray_cmd }} stop; RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 {{ sky_ray_cmd }} start --disable-usage-stats --address=$RAY_HEAD_IP:{{ray_port}} --object-manager-port=8076 {{"--resources='%s'" % custom_resources if custom_resources}} --temp-dir {{ray_temp_dir}} || exit 1;
-    which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
-{%- else %}
-worker_start_ray_commands: []
-{%- endif %}
-
-head_node: {}
-worker_nodes: {}
 
-# These fields are required for external cloud providers.
-head_setup_commands: []
-worker_setup_commands: []
-cluster_synced_files: []
-file_mounts_sync_continuously: False
+# Command to start ray clusters are now placed in `sky.provision.instance_setup`.
+# We do not need to list it here anymore.

sky/users/permission.py

@@ -1,8 +1,8 @@
 """Permission service for SkyPilot API Server."""
 import contextlib
+import hashlib
 import logging
 import os
-import threading
 from typing import Generator, List
 
 import casbin
@@ -10,9 +10,11 @@ import filelock
 import sqlalchemy_adapter
 
 from sky import global_user_state
+from sky import models
 from sky import sky_logging
 from sky.skylet import constants
 from sky.users import rbac
+from sky.utils import common_utils
 
 logging.getLogger('casbin.policy').setLevel(sky_logging.ERROR)
 logging.getLogger('casbin.role').setLevel(sky_logging.ERROR)
@@ -23,43 +25,50 @@ POLICY_UPDATE_LOCK_PATH = os.path.expanduser('~/.sky/.policy_update.lock')
 POLICY_UPDATE_LOCK_TIMEOUT_SECONDS = 20
 
 _enforcer_instance = None
-_lock = threading.Lock()
 
 
 class PermissionService:
     """Permission service for SkyPilot API Server."""
 
     def __init__(self):
-        self.enforcer = None
-        self.init_lock = threading.Lock()
-
-    def _lazy_initialize(self):
-        if self.enforcer is not None:
-            return
-        with self.init_lock:
-            if self.enforcer is not None:
-                return
+        with _policy_lock():
             global _enforcer_instance
             if _enforcer_instance is None:
-                # For different threads, we share the same enforcer instance.
-                with _lock:
-                    if _enforcer_instance is None:
-                        _enforcer_instance = self
-                        engine = global_user_state.initialize_and_get_db()
-                        adapter = sqlalchemy_adapter.Adapter(engine)
-                        model_path = os.path.join(os.path.dirname(__file__),
-                                                  'model.conf')
-                        enforcer = casbin.Enforcer(model_path, adapter)
-                        self.enforcer = enforcer
-                    else:
-                        self.enforcer = _enforcer_instance.enforcer
+                _enforcer_instance = self
+                engine = global_user_state.initialize_and_get_db()
+                adapter = sqlalchemy_adapter.Adapter(engine)
+                model_path = os.path.join(os.path.dirname(__file__),
+                                          'model.conf')
+                enforcer = casbin.Enforcer(model_path, adapter)
+                self.enforcer = enforcer
+                self._maybe_initialize_policies()
+                self._maybe_initialize_basic_auth_user()
             else:
                 self.enforcer = _enforcer_instance.enforcer
-            with _policy_lock():
-                self._maybe_initialize_policies()
+
+    def _maybe_initialize_basic_auth_user(self) -> None:
+        """Initialize basic auth user if it is enabled."""
+        basic_auth = os.environ.get(constants.SKYPILOT_INITIAL_BASIC_AUTH)
+        if not basic_auth:
+            return
+        username, password = basic_auth.split(':', 1)
+        if username and password:
+            user_hash = hashlib.md5(
+                username.encode()).hexdigest()[:common_utils.USER_HASH_LENGTH]
+            user_info = global_user_state.get_user(user_hash)
+            if user_info:
+                logger.info(f'Basic auth user {username} already exists')
+                return
+            global_user_state.add_or_update_user(
+                models.User(id=user_hash, name=username, password=password))
+            self.enforcer.add_grouping_policy(user_hash,
+                                              rbac.RoleName.ADMIN.value)
+            self.enforcer.save_policy()
+            logger.info(f'Basic auth user {username} initialized')
 
     def _maybe_initialize_policies(self) -> None:
         """Initialize policies if they don't already exist."""
+        # TODO(zhwu): we should avoid running this on client side.
         logger.debug(f'Initializing policies in process: {os.getpid()}')
         self._load_policy_no_lock()
 
@@ -138,7 +147,6 @@ class PermissionService:
 
     def add_user_if_not_exists(self, user_id: str) -> None:
         """Add user role relationship."""
-        self._lazy_initialize()
         with _policy_lock():
             self._add_user_if_not_exists_no_lock(user_id)
 
@@ -156,9 +164,21 @@ class PermissionService:
             return True
         return False
 
+    def delete_user(self, user_id: str) -> None:
+        """Delete user role relationship."""
+        with _policy_lock():
+            # Get current roles
+            self._load_policy_no_lock()
+            # Avoid calling get_user_roles, as it will require the lock.
+            current_roles = self.enforcer.get_roles_for_user(user_id)
+            if not current_roles:
+                logger.warning(f'User {user_id} has no roles')
+                return
+            self.enforcer.remove_grouping_policy(user_id, current_roles[0])
+            self.enforcer.save_policy()
+
     def update_role(self, user_id: str, new_role: str) -> None:
         """Update user role relationship."""
-        self._lazy_initialize()
         with _policy_lock():
             # Get current roles
             self._load_policy_no_lock()
@@ -191,7 +211,6 @@ class PermissionService:
         Returns:
             A list of role names that the user has.
         """
-        self._lazy_initialize()
         self._load_policy_no_lock()
         return self.enforcer.get_roles_for_user(user_id)
 
@@ -204,7 +223,6 @@ class PermissionService:
         # it is a hot path in every request. It is ok to have a stale policy,
         # as long as it is eventually consistent.
         # self._load_policy_no_lock()
-        self._lazy_initialize()
         return self.enforcer.enforce(user_id, path, method)
 
     def _load_policy_no_lock(self):
@@ -213,7 +231,6 @@ class PermissionService:
 
     def load_policy(self):
         """Load policy from storage with lock."""
-        self._lazy_initialize()
         with _policy_lock():
             self._load_policy_no_lock()
 
@@ -229,7 +246,6 @@ class PermissionService:
         For public workspaces, the permission is granted via a wildcard policy
         ('*').
         """
-        self._lazy_initialize()
         if os.getenv(constants.ENV_VAR_IS_SKYPILOT_SERVER) is None:
             # When it is not on API server, we allow all users to access all
             # workspaces, as the workspace check has been done on API server.
@@ -257,7 +273,6 @@ class PermissionService:
                    For public workspaces, this should be ['*'].
                    For private workspaces, this should be specific user IDs.
         """
-        self._lazy_initialize()
         with _policy_lock():
             for user in users:
                 logger.debug(f'Adding workspace policy: user={user}, '
@@ -275,7 +290,6 @@ class PermissionService:
                    For public workspaces, this should be ['*'].
                    For private workspaces, this should be specific user IDs.
         """
-        self._lazy_initialize()
         with _policy_lock():
             self._load_policy_no_lock()
             # Remove all existing policies for this workspace
@@ -289,7 +303,6 @@ class PermissionService:
 
     def remove_workspace_policy(self, workspace_name: str) -> None:
         """Remove workspace policy."""
-        self._lazy_initialize()
         with _policy_lock():
             self.enforcer.remove_filtered_policy(1, workspace_name)
             self.enforcer.save_policy()

sky/users/rbac.py

@@ -25,8 +25,17 @@ _DEFAULT_USER_BLOCKLIST = [{
     'path': '/workspaces/delete',
     'method': 'POST'
 }, {
-    'path': '/users/update',
+    'path': '/users/delete',
     'method': 'POST'
+}, {
+    'path': '/users/create',
+    'method': 'POST'
+}, {
+    'path': '/users/import',
+    'method': 'POST'
+}, {
+    'path': '/users/export',
+    'method': 'GET'
 }]