skypilot-nightly 1.0.0.dev20250609__py3-none-any.whl → 1.0.0.dev20250611__py3-none-any.whl

sky/provision/kubernetes/utils.py

@@ -1,6 +1,7 @@
 """Kubernetes utilities for SkyPilot."""
 import dataclasses
 import functools
+import hashlib
 import json
 import math
 import os
@@ -1555,11 +1556,11 @@ def is_kubeconfig_exec_auth(
             == schemas.RemoteIdentityOptions.LOCAL_CREDENTIALS.value):
         ctx_name = context_obj['name']
         exec_msg = ('exec-based authentication is used for '
-                    f'Kubernetes context {ctx_name!r}.'
-                    ' This may cause issues with autodown or when running '
-                    'Managed Jobs or SkyServe controller on Kubernetes. '
-                    'To fix, configure SkyPilot to create a service account '
-                    'for running pods by setting the following in '
+                    f'Kubernetes context {ctx_name!r}. '
+                    'Make sure that the corresponding cloud provider is '
+                    'also enabled through `sky check` (e.g.: GCP for GKE). '
+                    'Alternatively, configure SkyPilot to create a service '
+                    'account for running pods by setting the following in '
                     '~/.sky/config.yaml:\n'
                     '    kubernetes:\n'
                     '      remote_identity: SERVICE_ACCOUNT\n'
@@ -2877,8 +2878,8 @@ def get_context_from_config(provider_config: Dict[str, Any]) -> Optional[str]:
     context = provider_config.get('context',
                                   get_current_kube_config_context_name())
     if context == kubernetes.in_cluster_context_name():
-        # If the context (also used as the region) is in-cluster, we need to
-        # we need to use in-cluster auth by setting the context to None.
+        # If the context (also used as the region) is in-cluster, we need
+        # to use in-cluster auth by setting the context to None.
         context = None
     return context
 
@@ -3135,3 +3136,101 @@ def get_kubeconfig_paths() -> List[str]:
     for path in paths.split(kubernetes.ENV_KUBECONFIG_PATH_SEPARATOR):
         expanded.append(os.path.expanduser(path))
     return expanded
+
+
+def format_kubeconfig_exec_auth(config: Any,
+                                output_path: str,
+                                inject_wrapper: bool = True) -> bool:
+    """Reformat the kubeconfig so that exec-based authentication can be used
+    with SkyPilot. Will create a new kubeconfig file under <output_path>
+    regardless of whether a change has been made.
+
+    kubectl internally strips all environment variables except for system
+    defaults. If `inject_wrapper` is true, a wrapper executable is applied
+    to inject the relevant PATH information before exec-auth is executed.
+
+    Contents of sky-kube-exec-wrapper:
+
+    #!/bin/bash
+    export PATH="$HOME/skypilot-runtime/bin:$HOME/google-cloud-sdk:$PATH"
+    exec "$@"
+
+    refer to `skylet/constants.py` for more information.
+
+    Args:
+        config (dict): kubeconfig parsed by yaml.safe_load
+        output_path (str): Path where the potentially modified kubeconfig file
+          will be saved
+        inject_wrapper (bool): Whether to inject the wrapper script
+    Returns: whether config was updated, for logging purposes
+    """
+    updated = False
+    for user in config.get('users', []):
+        exec_info = user.get('user', {}).get('exec', {})
+        current_command = exec_info.get('command', '')
+
+        if current_command:
+            # Strip the path and keep only the executable name
+            executable = os.path.basename(current_command)
+            if executable == kubernetes_constants.SKY_K8S_EXEC_AUTH_WRAPPER:
+                # we don't want this happening recursively.
+                continue
+
+            if inject_wrapper:
+                exec_info[
+                    'command'] = kubernetes_constants.SKY_K8S_EXEC_AUTH_WRAPPER
+                if exec_info.get('args') is None:
+                    exec_info['args'] = []
+                exec_info['args'].insert(0, executable)
+                updated = True
+            elif executable != current_command:
+                exec_info['command'] = executable
+                updated = True
+
+            # Handle Nebius kubeconfigs: change --profile to 'sky'
+            if executable == 'nebius':
+                args = exec_info.get('args', [])
+                if args and '--profile' in args:
+                    try:
+                        profile_index = args.index('--profile')
+                        if profile_index + 1 < len(args):
+                            old_profile = args[profile_index + 1]
+                            if old_profile != 'sky':
+                                args[profile_index + 1] = 'sky'
+                                updated = True
+                    except ValueError:
+                        pass
+
+    os.makedirs(os.path.dirname(os.path.expanduser(output_path)), exist_ok=True)
+    with open(output_path, 'w', encoding='utf-8') as file:
+        yaml.safe_dump(config, file)
+
+    return updated
+
+
+def format_kubeconfig_exec_auth_with_cache(kubeconfig_path: str) -> str:
+    """Reformat the kubeconfig file or retrieve it from cache if it has already
+    been formatted before. Store it in the cache directory if necessary.
+
+    Having a cache for this is good if users spawn an extreme number of jobs
+    concurrently.
+
+    Args:
+        kubeconfig_path (str): kubeconfig path
+    Returns: updated kubeconfig path
+    """
+    # TODO(kyuds): GC cache files
+    with open(kubeconfig_path, 'r', encoding='utf-8') as file:
+        config = yaml.safe_load(file)
+    normalized = yaml.dump(config, sort_keys=True)
+    hashed = hashlib.sha1(normalized.encode('utf-8')).hexdigest()
+    path = os.path.expanduser(
+        f'{kubernetes_constants.SKY_K8S_EXEC_AUTH_KUBECONFIG_CACHE}/{hashed}.yaml'
+    )
+
+    # If we have already converted the same kubeconfig before, just return.
+    if os.path.isfile(path):
+        return path
+
+    format_kubeconfig_exec_auth(config, path)
+    return path

sky/serve/client/sdk.py

@@ -10,6 +10,8 @@ from sky.client import common as client_common
 from sky.server import common as server_common
 from sky.server.requests import payloads
 from sky.usage import usage_lib
+from sky.utils import admin_policy_utils
+from sky.utils import context
 from sky.utils import dag_utils
 
 if typing.TYPE_CHECKING:
@@ -23,6 +25,7 @@ else:
     requests = adaptors_common.LazyImport('requests')
 
 
+@context.contextual
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 def up(
@@ -55,30 +58,36 @@ def up(
     from sky.client import sdk  # pylint: disable=import-outside-toplevel
 
     dag = dag_utils.convert_entrypoint_to_dag(task)
-    sdk.validate(dag)
-    request_id = sdk.optimize(dag)
-    sdk.stream_and_get(request_id)
-    if _need_confirmation:
-        prompt = f'Launching a new service {service_name!r}. Proceed?'
-        if prompt is not None:
-            click.confirm(prompt, default=True, abort=True, show_default=True)
-
-    dag = client_common.upload_mounts_to_api_server(dag)
-    dag_str = dag_utils.dump_chain_dag_to_yaml_str(dag)
-
-    body = payloads.ServeUpBody(
-        task=dag_str,
-        service_name=service_name,
-    )
-    response = requests.post(
-        f'{server_common.get_server_url()}/serve/up',
-        json=json.loads(body.model_dump_json()),
-        timeout=(5, None),
-        cookies=server_common.get_api_cookie_jar(),
-    )
-    return server_common.get_request_id(response)
+    with admin_policy_utils.apply_and_use_config_in_current_request(
+            dag, at_client_side=True) as dag:
+        sdk.validate(dag)
+        request_id = sdk.optimize(dag)
+        sdk.stream_and_get(request_id)
+        if _need_confirmation:
+            prompt = f'Launching a new service {service_name!r}. Proceed?'
+            if prompt is not None:
+                click.confirm(prompt,
+                              default=True,
+                              abort=True,
+                              show_default=True)
+
+        dag = client_common.upload_mounts_to_api_server(dag)
+        dag_str = dag_utils.dump_chain_dag_to_yaml_str(dag)
+
+        body = payloads.ServeUpBody(
+            task=dag_str,
+            service_name=service_name,
+        )
+        response = requests.post(
+            f'{server_common.get_server_url()}/serve/up',
+            json=json.loads(body.model_dump_json()),
+            timeout=(5, None),
+            cookies=server_common.get_api_cookie_jar(),
+        )
+        return server_common.get_request_id(response)
 
 
+@context.contextual
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 def update(
@@ -112,30 +121,32 @@ def update(
     from sky.client import sdk  # pylint: disable=import-outside-toplevel
 
     dag = dag_utils.convert_entrypoint_to_dag(task)
-    sdk.validate(dag)
-    request_id = sdk.optimize(dag)
-    sdk.stream_and_get(request_id)
-    if _need_confirmation:
-        click.confirm(f'Updating service {service_name!r}. Proceed?',
-                      default=True,
-                      abort=True,
-                      show_default=True)
-
-    dag = client_common.upload_mounts_to_api_server(dag)
-    dag_str = dag_utils.dump_chain_dag_to_yaml_str(dag)
-    body = payloads.ServeUpdateBody(
-        task=dag_str,
-        service_name=service_name,
-        mode=mode,
-    )
+    with admin_policy_utils.apply_and_use_config_in_current_request(
+            dag, at_client_side=True) as dag:
+        sdk.validate(dag)
+        request_id = sdk.optimize(dag)
+        sdk.stream_and_get(request_id)
+        if _need_confirmation:
+            click.confirm(f'Updating service {service_name!r}. Proceed?',
+                          default=True,
+                          abort=True,
+                          show_default=True)
+
+        dag = client_common.upload_mounts_to_api_server(dag)
+        dag_str = dag_utils.dump_chain_dag_to_yaml_str(dag)
+        body = payloads.ServeUpdateBody(
+            task=dag_str,
+            service_name=service_name,
+            mode=mode,
+        )
 
-    response = requests.post(
-        f'{server_common.get_server_url()}/serve/update',
-        json=json.loads(body.model_dump_json()),
-        timeout=(5, None),
-        cookies=server_common.get_api_cookie_jar(),
-    )
-    return server_common.get_request_id(response)
+        response = requests.post(
+            f'{server_common.get_server_url()}/serve/update',
+            json=json.loads(body.model_dump_json()),
+            timeout=(5, None),
+            cookies=server_common.get_api_cookie_jar(),
+        )
+        return server_common.get_request_id(response)
 
 
 @usage_lib.entrypoint

sky/server/common.py

@@ -420,11 +420,7 @@ def _start_api_server(deploy: bool = False,
                 dashboard_msg += (
                     'Dashboard may be stale when installed from source, '
                     'to rebuild: npm --prefix sky/dashboard install '
-                    '&& npm --prefix sky/dashboard run build\n')
-            dashboard_msg += (
-                f'{ux_utils.INDENT_LAST_SYMBOL}{colorama.Fore.GREEN}'
-                f'Dashboard: {get_dashboard_url(server_url)}')
-            dashboard_msg += f'{colorama.Style.RESET_ALL}'
+                    '&& npm --prefix sky/dashboard run build')
         logger.info(
             ux_utils.finishing_message(
                 f'SkyPilot API server started. {dashboard_msg}'))

sky/server/requests/executor.py

@@ -19,6 +19,7 @@ The number of the workers is determined by the system resources.
 See the [README.md](../README.md) for detailed architecture of the executor.
 """
 import asyncio
+import concurrent.futures
 import contextlib
 import multiprocessing
 import os
@@ -93,21 +94,21 @@ class RequestQueue:
         else:
             raise RuntimeError(f'Invalid queue backend: {backend}')
 
-    def put(self, request: Tuple[str, bool]) -> None:
+    def put(self, request: Tuple[str, bool, bool]) -> None:
         """Put and request to the queue.
 
         Args:
-            request: A tuple of request_id and ignore_return_value.
+            request: A tuple of request_id, ignore_return_value, and retryable.
         """
         self.queue.put(request)  # type: ignore
 
-    def get(self) -> Optional[Tuple[str, bool]]:
+    def get(self) -> Optional[Tuple[str, bool, bool]]:
         """Get a request from the queue.
 
         It is non-blocking if the queue is empty, and returns None.
 
         Returns:
-            A tuple of request_id and ignore_return_value.
+            A tuple of request_id, ignore_return_value, and retryable.
         """
         try:
             return self.queue.get(block=False)
@@ -159,7 +160,7 @@ class RequestWorker:
             if request_element is None:
                 time.sleep(0.1)
                 return
-            request_id, ignore_return_value = request_element
+            request_id, ignore_return_value, retryable = request_element
             request = api_requests.get_request(request_id)
             assert request is not None, f'Request with ID {request_id} is None'
             if request.status == api_requests.RequestStatus.CANCELLED:
@@ -171,8 +172,14 @@ class RequestWorker:
             # multiple requests can share the same process pid, which may cause
             # issues with SkyPilot core functions if they rely on the exit of
             # the process, such as subprocess_daemon.py.
-            executor.submit_until_success(_request_execution_wrapper,
-                                          request_id, ignore_return_value)
+            fut = executor.submit_until_success(_request_execution_wrapper,
+                                                request_id, ignore_return_value)
+            if retryable:
+                # If the task might fail and be retried, start a thread to
+                # monitor the future and process retry.
+                threading.Thread(target=self.handle_task_result,
+                                 args=(fut, request_element),
+                                 daemon=True).start()
 
             logger.info(f'[{self}] Submitted request: {request_id}')
         except (Exception, SystemExit) as e:  # pylint: disable=broad-except
@@ -182,6 +189,16 @@ class RequestWorker:
                 f'{request_id if "request_id" in locals() else ""} '
                 f'{common_utils.format_exception(e, use_bracket=True)}')
 
+    def handle_task_result(self, fut: concurrent.futures.Future,
+                           request_element: Tuple[str, bool, bool]) -> None:
+        try:
+            fut.result()
+        except exceptions.ExecutionRetryableError as e:
+            time.sleep(e.retry_wait_seconds)
+            # Reschedule the request.
+            queue = _get_queue(self.schedule_type)
+            queue.put(request_element)
+
     def run(self) -> None:
         # Handle the SIGTERM signal to abort the executor process gracefully.
         proc_group = f'{self.schedule_type.value}'
@@ -316,7 +333,9 @@ def _request_execution_wrapper(request_id: str,
         func = request_task.entrypoint
         request_body = request_task.request_body
 
-    with log_path.open('w', encoding='utf-8') as f:
+    # Append to the log file instead of overwriting it since there might be
+    # logs from previous retries.
+    with log_path.open('a', encoding='utf-8') as f:
         # Store copies of the original stdout and stderr file descriptors
         original_stdout, original_stderr = _redirect_output(f)
         # Redirect the stdout/stderr before overriding the environment and
@@ -340,6 +359,17 @@ def _request_execution_wrapper(request_id: str,
             subprocess_utils.kill_children_processes()
             _restore_output(original_stdout, original_stderr)
             return
+        except exceptions.ExecutionRetryableError as e:
+            logger.error(e)
+            logger.info(e.hint)
+            with api_requests.update_request(request_id) as request_task:
+                assert request_task is not None, request_id
+                # Retried request will undergo rescheduling and a new execution,
+                # clear the pid of the request.
+                request_task.pid = None
+            # Yield control to the scheduler for uniform handling of retries.
+            _restore_output(original_stdout, original_stderr)
+            raise
         except (Exception, SystemExit) as e:  # pylint: disable=broad-except
             api_requests.set_request_failed(request_id, e)
             _restore_output(original_stdout, original_stderr)
@@ -463,17 +493,17 @@ def prepare_request(
     return request
 
 
-def schedule_request(
-        request_id: str,
-        request_name: str,
-        request_body: payloads.RequestBody,
-        func: Callable[P, Any],
-        request_cluster_name: Optional[str] = None,
-        ignore_return_value: bool = False,
-        schedule_type: api_requests.ScheduleType = (
-            api_requests.ScheduleType.LONG),
-        is_skypilot_system: bool = False,
-        precondition: Optional[preconditions.Precondition] = None) -> None:
+def schedule_request(request_id: str,
+                     request_name: str,
+                     request_body: payloads.RequestBody,
+                     func: Callable[P, Any],
+                     request_cluster_name: Optional[str] = None,
+                     ignore_return_value: bool = False,
+                     schedule_type: api_requests.ScheduleType = (
+                         api_requests.ScheduleType.LONG),
+                     is_skypilot_system: bool = False,
+                     precondition: Optional[preconditions.Precondition] = None,
+                     retryable: bool = False) -> None:
     """Enqueue a request to the request queue.
 
     Args:
@@ -498,7 +528,7 @@ def schedule_request(
                     request_cluster_name, schedule_type, is_skypilot_system)
 
     def enqueue():
-        input_tuple = (request_id, ignore_return_value)
+        input_tuple = (request_id, ignore_return_value, retryable)
         logger.info(f'Queuing request: {request_id}')
         _get_queue(schedule_type).put(input_tuple)
 

sky/server/requests/payloads.py

@@ -79,6 +79,9 @@ def get_override_skypilot_config_from_client() -> Dict[str, Any]:
     # server endpoint on the server side. This avoids the warning at
     # server-side.
     config.pop_nested(('api_server',), default_value=None)
+    # Remove the admin policy, as the policy has been applied on the client
+    # side.
+    config.pop_nested(('admin_policy',), default_value=None)
     return config
 
 

sky/server/requests/process.py

@@ -6,6 +6,7 @@ import threading
 import time
 from typing import Callable, Dict, Optional, Tuple
 
+from sky import exceptions
 from sky.utils import atomic
 from sky.utils import subprocess_utils
 
@@ -67,14 +68,24 @@ class PoolExecutor(concurrent.futures.ProcessPoolExecutor):
 
 
 # Define the worker function outside of the class to avoid pickling self
-def _disposable_worker(fn, initializer: Optional[Callable], initargs: Tuple,
-                       args, kwargs):
+def _disposable_worker(fn, initializer, initargs, result_queue, args, kwargs):
+    """The worker function that is used to run the task.
+
+    Args:
+        fn: The function to run.
+        initializer: The initializer function to run before running the task.
+        initargs: The arguments to pass to the initializer function.
+        result_queue: The queue to put the result and exception into.
+        args: The arguments to pass to the function.
+        kwargs: The keyword arguments to pass to the function.
+    """
     try:
         if initializer is not None:
             initializer(*initargs)
-        fn(*args, **kwargs)
+        result = fn(*args, **kwargs)
+        result_queue.put(result)
     except BaseException as e:  # pylint: disable=broad-except
-        return e
+        result_queue.put(e)
 
 
 class DisposableExecutor:
@@ -98,28 +109,52 @@ class DisposableExecutor:
         self._initializer: Optional[Callable] = initializer
         self._initargs: Tuple = initargs
 
-    def _monitor_worker(self, process: multiprocessing.Process) -> None:
+    def _monitor_worker(self, process: multiprocessing.Process,
+                        future: concurrent.futures.Future,
+                        result_queue: multiprocessing.Queue) -> None:
         """Monitor the worker process and cleanup when it's done."""
-        process.join()
-        if process.pid:
-            with self._lock:
-                if process.pid in self.workers:
-                    del self.workers[process.pid]
-
-    # Submit is not compatible with ProcessPoolExecutor because we does not
-    # bother to return a Future. Can be improved if needed.
-    def submit(self, fn, *args, **kwargs) -> bool:
-        """Submit a task for execution."""
+        try:
+            process.join()
+            if not future.cancelled():
+                try:
+                    # Get result from the queue if process completed
+                    if not result_queue.empty():
+                        result = result_queue.get(block=False)
+                        if isinstance(result, BaseException):
+                            future.set_exception(result)
+                        else:
+                            future.set_result(result)
+                    else:
+                        # Process ended but no result
+                        future.set_result(None)
+                except (multiprocessing.TimeoutError, BrokenPipeError,
+                        EOFError) as e:
+                    future.set_exception(e)
+        finally:
+            if process.pid:
+                with self._lock:
+                    if process.pid in self.workers:
+                        del self.workers[process.pid]
+
+    def submit(self, fn, *args, **kwargs) -> concurrent.futures.Future:
+        """Submit a task for execution and return a Future."""
+        future: concurrent.futures.Future = concurrent.futures.Future()
+
         if self._shutdown:
-            return False
+            raise RuntimeError('Cannot submit task after executor is shutdown')
+
         with self._lock:
             if (self.max_workers is not None and
                     len(self.workers) >= self.max_workers):
-                return False
+                raise exceptions.ExecutionPoolFullError(
+                    'Maximum workers reached')
 
+        result_queue: multiprocessing.Queue = multiprocessing.Queue()
         process = multiprocessing.Process(target=_disposable_worker,
                                           args=(fn, self._initializer,
-                                                self._initargs, args, kwargs))
+                                                self._initargs, result_queue,
+                                                args, kwargs))
+        process.daemon = True
         process.start()
 
         with self._lock:
@@ -128,13 +163,13 @@ class DisposableExecutor:
                 raise RuntimeError('Failed to start process')
             self.workers[pid] = process
 
-        # Start monitor thread to cleanup the worker process when it's done.
+        # Start monitor thread to cleanup the worker process when it's done
         monitor_thread = threading.Thread(target=self._monitor_worker,
-                                          args=(process,),
+                                          args=(process, future, result_queue),
                                           daemon=True)
         monitor_thread.start()
 
-        return True
+        return future
 
     def has_idle_workers(self) -> bool:
         """Check if there are any idle workers."""
@@ -173,12 +208,14 @@ class BurstableExecutor:
             self._burst_executor = DisposableExecutor(max_workers=burst_workers,
                                                       **kwargs)
 
-    def submit_until_success(self, fn, *args, **kwargs):
+    def submit_until_success(self, fn, *args,
+                             **kwargs) -> concurrent.futures.Future:
         """Submit a task for execution until success.
 
         Prioritizes submitting to the guaranteed pool. If no idle workers
         are available in the guaranteed pool, it will submit to the burst
-        pool.
+        pool. If the burst pool is full, it will retry the whole process until
+        the task is submitted successfully.
         TODO(aylei): this is coupled with executor.RequestWorker since we
         know the worker is dedicated to request scheduling and it either
         blocks on request polling or request submitting. So it is no harm
@@ -188,17 +225,20 @@ class BurstableExecutor:
 
         while True:
             if self._executor is not None and self._executor.has_idle_workers():
-                self._executor.submit(fn, *args, **kwargs)
-                break
+                logger.info('Submitting to the guaranteed pool')
+                return self._executor.submit(fn, *args, **kwargs)
             if (self._burst_executor is not None and
                     self._burst_executor.has_idle_workers()):
-                self._burst_executor.submit(fn, *args, **kwargs)
-                break
+                try:
+                    fut = self._burst_executor.submit(fn, *args, **kwargs)
+                    return fut
+                except exceptions.ExecutionPoolFullError:
+                    # The burst pool is full, try the next candidate.
+                    pass
             if self._executor is not None:
                 # No idle workers in either pool, still queue the request
                 # to the guaranteed pool to keep behavior consistent.
-                self._executor.submit(fn, *args, **kwargs)
-                break
+                return self._executor.submit(fn, *args, **kwargs)
             logger.debug('No guaranteed pool set and the burst pool is full, '
                          'retry later.')
             time.sleep(0.1)

sky/server/server.py

@@ -693,6 +693,7 @@ async def launch(launch_body: payloads.LaunchBody,
         func=execution.launch,
         schedule_type=requests_lib.ScheduleType.LONG,
         request_cluster_name=launch_body.cluster_name,
+        retryable=launch_body.retry_until_up,
     )