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

sky/jobs/client/sdk.py

@@ -14,7 +14,9 @@ from sky.server import common as server_common
 from sky.server.requests import payloads
 from sky.skylet import constants
 from sky.usage import usage_lib
+from sky.utils import admin_policy_utils
 from sky.utils import common_utils
+from sky.utils import context
 from sky.utils import dag_utils
 
 if typing.TYPE_CHECKING:
@@ -29,6 +31,7 @@ else:
 logger = sky_logging.init_logger(__name__)
 
 
+@context.contextual
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 def launch(
@@ -65,27 +68,32 @@ def launch(
     """
 
     dag = dag_utils.convert_entrypoint_to_dag(task)
-    sdk.validate(dag)
-    if _need_confirmation:
-        request_id = sdk.optimize(dag)
-        sdk.stream_and_get(request_id)
-        prompt = f'Launching a managed job {dag.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.JobsLaunchBody(
-        task=dag_str,
-        name=name,
-    )
-    response = requests.post(
-        f'{server_common.get_server_url()}/jobs/launch',
-        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)
+        if _need_confirmation:
+            request_id = sdk.optimize(dag)
+            sdk.stream_and_get(request_id)
+            prompt = f'Launching a managed job {dag.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.JobsLaunchBody(
+            task=dag_str,
+            name=name,
+        )
+        response = requests.post(
+            f'{server_common.get_server_url()}/jobs/launch',
+            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/provision/kubernetes/constants.py

@@ -6,3 +6,12 @@ NO_GPU_HELP_MESSAGE = ('If your cluster contains GPUs, make sure '
                        '(e.g., skypilot.co/accelerator) are setup correctly. ')
 
 KUBERNETES_IN_CLUSTER_NAMESPACE_ENV_VAR = 'SKYPILOT_IN_CLUSTER_NAMESPACE'
+
+# Name of kubernetes exec auth wrapper script
+SKY_K8S_EXEC_AUTH_WRAPPER = 'sky-kube-exec-wrapper'
+
+# PATH envvar for kubectl exec auth execve
+SKY_K8S_EXEC_AUTH_PATH = '$HOME/skypilot-runtime/bin:$HOME/google-cloud-sdk/bin:$PATH'  # pylint: disable=line-too-long
+
+# cache directory for kubeconfig with modified exec auth
+SKY_K8S_EXEC_AUTH_KUBECONFIG_CACHE = '~/.sky/generated/kubeconfigs'

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