skypilot-nightly 1.0.0.dev20250922__py3-none-any.whl → 1.0.0.dev20250926__py3-none-any.whl

sky/setup_files/dependencies.py

@@ -79,7 +79,7 @@ install_requires = [
     # Required for API server metrics
     'prometheus_client>=0.8.0',
     'passlib',
-    'bcrypt',
+    'bcrypt==4.0.1',
     'pyjwt',
     'gitpython',
     'types-paramiko',
@@ -112,6 +112,7 @@ server_dependencies = [
     GRPC,
     PROTOBUF,
     'aiosqlite',
+    'greenlet',
 ]
 
 local_ray = [
@@ -192,7 +193,9 @@ extras_require: Dict[str, List[str]] = {
     'remote': remote,
     # For the container registry auth api. Reference:
     # https://github.com/runpod/runpod-python/releases/tag/1.6.1
-    'runpod': ['runpod>=1.6.1'],
+    # RunPod needs a TOML parser to read ~/.runpod/config.toml. On Python 3.11+
+    # stdlib provides tomllib; on lower versions we depend on tomli explicitly.
+    'runpod': ['runpod>=1.6.1', 'tomli; python_version < "3.11"'],
     'fluidstack': [],  # No dependencies needed for fluidstack
     'cudo': ['cudo-compute>=0.1.10'],
     'paperspace': [],  # No dependencies needed for paperspace

sky/setup_files/setup.py

@@ -148,47 +148,47 @@ if os.path.exists(readme_filepath):
     long_description = io.open(readme_filepath, 'r', encoding='utf-8').read()
     long_description = parse_readme(long_description)
 
-atexit.register(revert_commit_hash)
-replace_commit_hash()
-
-setuptools.setup(
-    # NOTE: this affects the package.whl wheel name. When changing this (if
-    # ever), you must grep for '.whl' and change all corresponding wheel paths
-    # (templates/*.j2 and wheel_utils.py).
-    name='skypilot-nightly',
-    version=find_version(),
-    packages=setuptools.find_packages(),
-    author='SkyPilot Team',
-    license='Apache 2.0',
-    readme='README.md',
-    description='SkyPilot: Run AI on Any Infra — Unified, Faster, Cheaper.',
-    long_description=long_description,
-    long_description_content_type='text/markdown',
-    setup_requires=['wheel'],
-    requires_python='>=3.7',
-    install_requires=dependencies['install_requires'],
-    extras_require=dependencies['extras_require'],
-    entry_points={
-        'console_scripts': ['sky = sky.cli:cli'],
-    },
-    include_package_data=True,
-    classifiers=[
-        'Programming Language :: Python :: 3.7',
-        'Programming Language :: Python :: 3.8',
-        'Programming Language :: Python :: 3.9',
-        'Programming Language :: Python :: 3.10',
-        'Programming Language :: Python :: 3.11',
-        'Programming Language :: Python :: 3.12',
-        'Programming Language :: Python :: 3.13',
-        'License :: OSI Approved :: Apache Software License',
-        'Operating System :: OS Independent',
-        'Topic :: Software Development :: Libraries :: Python Modules',
-        'Topic :: System :: Distributed Computing',
-    ],
-    project_urls={
-        'Homepage': 'https://github.com/skypilot-org/skypilot',
-        'Issues': 'https://github.com/skypilot-org/skypilot/issues',
-        'Discussion': 'https://github.com/skypilot-org/skypilot/discussions',
-        'Documentation': 'https://docs.skypilot.co/',
-    },
-)
+if __name__ == '__main__':
+    atexit.register(revert_commit_hash)
+    replace_commit_hash()
+    setuptools.setup(
+        # NOTE: this affects the package.whl wheel name. When changing this (if
+        # ever), you must grep for '.whl' and change all corresponding wheel paths
+        # (templates/*.j2 and wheel_utils.py).
+        name='skypilot-nightly',
+        version=find_version(),
+        packages=setuptools.find_packages(),
+        author='SkyPilot Team',
+        license='Apache 2.0',
+        readme='README.md',
+        description='SkyPilot: Run AI on Any Infra — Unified, Faster, Cheaper.',
+        long_description=long_description,
+        long_description_content_type='text/markdown',
+        setup_requires=['wheel'],
+        requires_python='>=3.7',
+        install_requires=dependencies['install_requires'],
+        extras_require=dependencies['extras_require'],
+        entry_points={
+            'console_scripts': ['sky = sky.cli:cli'],
+        },
+        include_package_data=True,
+        classifiers=[
+            'Programming Language :: Python :: 3.7',
+            'Programming Language :: Python :: 3.8',
+            'Programming Language :: Python :: 3.9',
+            'Programming Language :: Python :: 3.10',
+            'Programming Language :: Python :: 3.11',
+            'Programming Language :: Python :: 3.12',
+            'Programming Language :: Python :: 3.13',
+            'License :: OSI Approved :: Apache Software License',
+            'Operating System :: OS Independent',
+            'Topic :: Software Development :: Libraries :: Python Modules',
+            'Topic :: System :: Distributed Computing',
+        ],
+        project_urls={
+            'Homepage': 'https://github.com/skypilot-org/skypilot',
+            'Issues': 'https://github.com/skypilot-org/skypilot/issues',
+            'Discussion': 'https://github.com/skypilot-org/skypilot/discussions',
+            'Documentation': 'https://docs.skypilot.co/',
+        },
+    )

sky/skylet/constants.py

@@ -57,6 +57,9 @@ SKY_REMOTE_PYTHON_ENV: str = f'~/{SKY_REMOTE_PYTHON_ENV_NAME}'
 ACTIVATE_SKY_REMOTE_PYTHON_ENV = f'source {SKY_REMOTE_PYTHON_ENV}/bin/activate'
 # uv is used for venv and pip, much faster than python implementations.
 SKY_UV_INSTALL_DIR = '"$HOME/.local/bin"'
+# set UV_SYSTEM_PYTHON to false in case the
+# user provided docker image set it to true.
+# unset PYTHONPATH in case the user provided docker image set it.
 SKY_UV_CMD = ('UV_SYSTEM_PYTHON=false '
               f'{SKY_UNSET_PYTHONPATH} {SKY_UV_INSTALL_DIR}/uv')
 # This won't reinstall uv if it's already installed, so it's safe to re-run.
@@ -97,7 +100,7 @@ TASK_ID_LIST_ENV_VAR = f'{SKYPILOT_ENV_VAR_PREFIX}TASK_IDS'
 # cluster yaml is updated.
 #
 # TODO(zongheng,zhanghao): make the upgrading of skylet automatic?
-SKYLET_VERSION = '18'
+SKYLET_VERSION = '19'
 # The version of the lib files that skylet/jobs use. Whenever there is an API
 # change for the job_lib or log_lib, we need to bump this version, so that the
 # user can be notified to update their SkyPilot version on the remote cluster.

sky/skylet/events.py

@@ -11,6 +11,7 @@ import psutil
 from sky import clouds
 from sky import sky_logging
 from sky.backends import cloud_vm_ray_backend
+from sky.jobs import constants as managed_job_constants
 from sky.jobs import scheduler
 from sky.jobs import state as managed_job_state
 from sky.jobs import utils as managed_job_utils
@@ -21,6 +22,7 @@ from sky.skylet import job_lib
 from sky.usage import usage_lib
 from sky.utils import cluster_utils
 from sky.utils import registry
+from sky.utils import subprocess_utils
 from sky.utils import ux_utils
 from sky.utils import yaml_utils
 
@@ -74,6 +76,46 @@ class ManagedJobEvent(SkyletEvent):
     EVENT_INTERVAL_SECONDS = 300
 
     def _run(self):
+        if not os.path.exists(
+                os.path.expanduser(
+                    managed_job_constants.JOB_CONTROLLER_INDICATOR_FILE)):
+            # Note: since the skylet is started before the user setup (in
+            # jobs-controller.yaml.j2) runs, it's possible that we hit this
+            # before the indicator file is written. However, since we will wait
+            # EVENT_INTERVAL_SECONDS before the first run, this should be very
+            # unlikely.
+            logger.info('No jobs controller indicator file found.')
+            all_job_ids = managed_job_state.get_all_job_ids_by_name(None)
+            if not all_job_ids:
+                logger.info('No jobs running. Stopping controllers.')
+                # TODO(cooperc): Move this to a shared function also called by
+                # sdk.api_stop(). (#7229)
+                try:
+                    with open(os.path.expanduser(
+                            scheduler.JOB_CONTROLLER_PID_PATH),
+                              'r',
+                              encoding='utf-8') as f:
+                        pids = f.read().split('\n')[:-1]
+                        for pid in pids:
+                            if subprocess_utils.is_process_alive(
+                                    int(pid.strip())):
+                                subprocess_utils.kill_children_processes(
+                                    parent_pids=[int(pid.strip())], force=True)
+                    os.remove(
+                        os.path.expanduser(scheduler.JOB_CONTROLLER_PID_PATH))
+                except FileNotFoundError:
+                    # its fine we will create it
+                    pass
+                except Exception as e:  # pylint: disable=broad-except
+                    # in case we get perm issues or something is messed up, just
+                    # ignore it and assume the process is dead
+                    logger.error(
+                        f'Error looking at job controller pid file: {e}')
+                    pass
+            logger.info(f'{len(all_job_ids)} jobs running. Assuming the '
+                        'indicator file hasn\'t been written yet.')
+            return
+
         logger.info('=== Updating managed job status ===')
         managed_job_utils.update_managed_jobs_statuses()
         scheduler.maybe_start_controllers()

sky/templates/jobs-controller.yaml.j2

@@ -36,6 +36,9 @@ setup: |
   grep -q 'alias sky-env=' ~/.bashrc || echo 'alias sky-env="{{ sky_activate_python_env }}"' >> ~/.bashrc
   {% endif %}
 
+  # This is used by the skylet events to check if we are a jobs controller.
+  touch {{job_controller_indicator_file}}
+
 run: |
   {%- if consolidation_mode_job_id is none %}
   {{ sky_activate_python_env }}

sky/templates/kubernetes-ray.yml.j2

@@ -510,6 +510,16 @@ available_node_types:
               valueFrom:
                 fieldRef:
                   fieldPath: metadata.labels['ray-node-type']
+            - name: SKYPILOT_POD_CPU_CORE_LIMIT
+              valueFrom:
+                resourceFieldRef:
+                  containerName: ray-node
+                  resource: requests.cpu
+            - name: SKYPILOT_POD_MEMORY_BYTES_LIMIT
+              valueFrom:
+                resourceFieldRef:
+                  containerName: ray-node
+                  resource: requests.memory
             {% for key, value in k8s_env_vars.items() if k8s_env_vars is not none %}
             - name: {{ key }}
               value: {{ value }}
@@ -630,13 +640,6 @@ available_node_types:
           command: ["/bin/bash", "-c", "--"]
           args:
             - |
-              # For backwards compatibility, we put a marker file in the pod
-              # to indicate that the pod is running with the changes introduced
-              # in project nimbus: https://github.com/skypilot-org/skypilot/pull/4393
-              # TODO: Remove this marker file and it's usage in setup_commands
-              # after v0.10.0 release.
-              touch /tmp/skypilot_is_nimbus
-
               # Helper function to conditionally use sudo
               # TODO(zhwu): consolidate the two prefix_cmd and sudo replacements
               prefix_cmd() { if [ $(id -u) -ne 0 ]; then echo "sudo"; else echo ""; fi; }
@@ -898,15 +901,20 @@ available_node_types:
                 {{ conda_installation_commands }}
                 {{ ray_installation_commands }}
 
-                VIRTUAL_ENV=~/skypilot-runtime ~/.local/bin/uv pip install skypilot[kubernetes,remote]
+                # set UV_SYSTEM_PYTHON to false in case the user provided docker image set it to true.
+                # unset PYTHONPATH in case the user provided docker image set it.
+                VIRTUAL_ENV=~/skypilot-runtime UV_SYSTEM_PYTHON=false env -u PYTHONPATH ~/.local/bin/uv pip install skypilot[kubernetes,remote]
                 # Wait for `patch` package to be installed before applying ray patches
                 until dpkg -l | grep -q "^ii  patch "; do
                   sleep 0.1
                   echo "Waiting for patch package to be installed..."
                 done
                 # Apply Ray patches for progress bar fix
-                ~/.local/bin/uv pip list | grep "ray " | grep 2.9.3 2>&1 > /dev/null && {
-                  VIRTUAL_ENV=~/skypilot-runtime python -c "from sky.skylet.ray_patches import patch; patch()" || exit 1;
+                # set UV_SYSTEM_PYTHON to false in case the user provided docker image set it to true.
+                # unset PYTHONPATH in case the user provided docker image set it.
+                # ~/.sky/python_path is seeded by conda_installation_commands
+                VIRTUAL_ENV=~/skypilot-runtime UV_SYSTEM_PYTHON=false env -u PYTHONPATH ~/.local/bin/uv pip list | grep "ray " | grep 2.9.3 2>&1 > /dev/null && {
+                  $(cat ~/.sky/python_path) -c "from sky.skylet.ray_patches import patch; patch()" || exit 1;
                 }
                 touch /tmp/ray_skypilot_installation_complete
                 echo "=== Ray and skypilot installation completed ==="
@@ -1333,18 +1341,16 @@ setup_commands:
     # Wait for SSH setup to complete before proceeding
     if [ -f /tmp/apt_ssh_setup_started ]; then
       echo "=== Logs for asynchronous SSH setup ===";
-      [ -f /tmp/apt_ssh_setup_complete ] && cat /tmp/${STEPS[0]}.log ||
-      { tail -f -n +1 /tmp/${STEPS[0]}.log & TAIL_PID=$!; echo "Tail PID: $TAIL_PID"; until [ -f /tmp/apt_ssh_setup_complete ] || [ -f /tmp/${STEPS[0]}.failed ]; do sleep 0.5; done; kill $TAIL_PID || true; };
+      ([ -f /tmp/apt_ssh_setup_complete ]|| [ -f /tmp/${STEPS[0]}.failed ]) && cat /tmp/${STEPS[0]}.log ||
+      { tail -f -n +1 /tmp/${STEPS[0]}.log & TAIL_PID=$!; echo "Tail PID: $TAIL_PID"; sleep 0.5; until [ -f /tmp/apt_ssh_setup_complete ] || [ -f /tmp/${STEPS[0]}.failed ]; do sleep 0.5; done; kill $TAIL_PID || true; };
       [ -f /tmp/${STEPS[0]}.failed ] && { echo "Error: ${STEPS[0]} failed. Exiting."; exit 1; } || true;
     fi
 
     echo "=== Logs for asynchronous ray and skypilot installation ===";
-    if [ -f /tmp/skypilot_is_nimbus ]; then
-      echo "=== Logs for asynchronous ray and skypilot installation ===";
-      [ -f /tmp/ray_skypilot_installation_complete ] && cat /tmp/${STEPS[1]}.log ||
-      { tail -f -n +1 /tmp/${STEPS[1]}.log & TAIL_PID=$!; echo "Tail PID: $TAIL_PID"; until [ -f /tmp/ray_skypilot_installation_complete ] || [ -f /tmp/${STEPS[1]}.failed ]; do sleep 0.5; done; kill $TAIL_PID || true; };
-      [ -f /tmp/${STEPS[1]}.failed ] && { echo "Error: ${STEPS[1]} failed. Exiting."; exit 1; } || true;
-    fi
+    ([ -f /tmp/ray_skypilot_installation_complete ]|| [ -f /tmp/${STEPS[1]}.failed ]) && cat /tmp/${STEPS[1]}.log ||
+    { tail -f -n +1 /tmp/${STEPS[1]}.log & TAIL_PID=$!; echo "Tail PID: $TAIL_PID"; sleep 0.5; until [ -f /tmp/ray_skypilot_installation_complete ] || [ -f /tmp/${STEPS[1]}.failed ]; do sleep 0.5; done; kill $TAIL_PID || true; };
+    [ -f /tmp/${STEPS[1]}.failed ] && { echo "Error: ${STEPS[1]} failed. Exiting."; exit 1; } || true;
+
     end_epoch=$(date +%s);
     echo "=== Ray and skypilot dependencies installation completed in $(($end_epoch - $start_epoch)) secs ===";
     start_epoch=$(date +%s);

sky/usage/usage_lib.py

@@ -14,6 +14,7 @@ from typing_extensions import ParamSpec
 
 import sky
 from sky import sky_logging
+from sky import skypilot_config
 from sky.adaptors import common as adaptors_common
 from sky.usage import constants
 from sky.utils import common_utils
@@ -167,6 +168,7 @@ class UsageMessageToReport(MessageToReport):
         self.runtimes: Dict[str, float] = {}  # update_runtime
         self.exception: Optional[str] = None  # entrypoint_context
         self.stacktrace: Optional[str] = None  # entrypoint_context
+        self.skypilot_config: Optional[Dict[str, Any]] = None
 
         # Whether API server is deployed remotely.
         self.using_remote_api_server: bool = (
@@ -177,6 +179,7 @@ class UsageMessageToReport(MessageToReport):
             self.client_entrypoint = common_utils.get_current_client_entrypoint(
                 msg)
         self.entrypoint = msg
+        self.skypilot_config = dict(skypilot_config.to_dict())
 
     def set_internal(self):
         self.internal = True

sky/utils/cli_utils/status_utils.py

@@ -11,6 +11,7 @@ from sky.utils import common_utils
 from sky.utils import log_utils
 from sky.utils import resources_utils
 from sky.utils import status_lib
+from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
     from sky.provision.kubernetes import utils as kubernetes_utils
@@ -105,11 +106,9 @@ def show_status_table(cluster_records: List[responses.StatusResponse],
 
     if query_clusters:
         cluster_names = {record['name'] for record in cluster_records}
-        not_found_clusters = [
-            repr(cluster)
-            for cluster in query_clusters
-            if cluster not in cluster_names
-        ]
+        not_found_clusters = ux_utils.get_non_matched_query(
+            query_clusters, cluster_names)
+        not_found_clusters = [repr(cluster) for cluster in not_found_clusters]
         if not_found_clusters:
             cluster_str = 'Cluster'
             if len(not_found_clusters) > 1:

sky/utils/context.py

@@ -2,15 +2,21 @@
 
 import asyncio
 from collections.abc import Mapping
-from collections.abc import MutableMapping
 import contextvars
+import copy
 import functools
+import inspect
 import os
 import pathlib
 import subprocess
 import sys
-import typing
-from typing import Any, Callable, Dict, Optional, TextIO, TypeVar
+from typing import (Callable, Dict, Iterator, MutableMapping, Optional, TextIO,
+                    TYPE_CHECKING, TypeVar)
+
+from typing_extensions import ParamSpec
+
+if TYPE_CHECKING:
+    from sky.skypilot_config import ConfigContext
 
 
 class Context(object):
@@ -88,7 +94,7 @@ class Context(object):
         else:
             self._log_file_handle = open(log_file, 'a', encoding='utf-8')
         self._log_file = log_file
-        if original_log_file is not None:
+        if original_log_handle is not None:
             original_log_handle.close()
         return original_log_file
 
@@ -102,8 +108,30 @@ class Context(object):
         for k, v in envs.items():
             self.env_overrides[k] = v
 
+    def cleanup(self):
+        """Clean up the context."""
+        if self._log_file_handle is not None:
+            self._log_file_handle.close()
+            self._log_file_handle = None
+
+    def copy(self) -> 'Context':
+        """Create a copy of the context.
+
+        Changes to the current context after this call will not affect the copy.
+        The new context will get its own handle/fd for the log file.
+        The new context will get an independent copy of the env var overrides.
+        The new context will get an independent copy of the config context.
+        Cancellation of the current context will not be propagated to the copy.
+        """
+        new_context = Context()
+        new_context.redirect_log(self._log_file)
+        new_context.env_overrides = self.env_overrides.copy()
+        new_context.config_context = copy.deepcopy(self.config_context)
+        return new_context
 
-_CONTEXT = contextvars.ContextVar('sky_context', default=None)
+
+_CONTEXT = contextvars.ContextVar[Optional[Context]]('sky_context',
+                                                     default=None)
 
 
 def get() -> Optional[Context]:
@@ -116,7 +144,7 @@ def get() -> Optional[Context]:
     return _CONTEXT.get()
 
 
-class ContextualEnviron(MutableMapping):
+class ContextualEnviron(MutableMapping[str, str]):
     """Environment variables wrapper with contextual overrides.
 
     An instance of ContextualEnviron will typically be used to replace
@@ -155,10 +183,10 @@ class ContextualEnviron(MutableMapping):
        assert os.environ['FOO'] == 'BAR1'
     """
 
-    def __init__(self, environ):
+    def __init__(self, environ: 'os._Environ[str]') -> None:
         self._environ = environ
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: str) -> str:
         ctx = get()
         if ctx is not None:
             if key in ctx.env_overrides:
@@ -170,10 +198,10 @@ class ContextualEnviron(MutableMapping):
                 return value
         return self._environ[key]
 
-    def __iter__(self):
-        ctx = get()
-        deleted_keys = set()
-        if ctx is not None:
+    def __iter__(self) -> Iterator[str]:
+
+        def iter_from_context(ctx: Context) -> Iterator[str]:
+            deleted_keys = set()
             for key, value in ctx.env_overrides.items():
                 if value is None:
                     deleted_keys.add(key)
@@ -182,20 +210,24 @@ class ContextualEnviron(MutableMapping):
                 # Deduplicate the keys
                 if key not in ctx.env_overrides and key not in deleted_keys:
                     yield key
+
+        ctx = get()
+        if ctx is not None:
+            return iter_from_context(ctx)
         else:
             return self._environ.__iter__()
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(dict(self))
 
-    def __setitem__(self, key, value):
+    def __setitem__(self, key: str, value: str) -> None:
         ctx = get()
         if ctx is not None:
             ctx.env_overrides[key] = value
         else:
             self._environ.__setitem__(key, value)
 
-    def __delitem__(self, key):
+    def __delitem__(self, key: str) -> None:
         ctx = get()
         if ctx is not None:
             if key in ctx.env_overrides:
@@ -211,10 +243,13 @@ class ContextualEnviron(MutableMapping):
         else:
             self._environ.__delitem__(key)
 
-    def __repr__(self):
-        return self._environ.__repr__()
+    def __repr__(self) -> str:
+        # Adapted from os._Environ.__repr__
+        formatted_items = ', '.join(
+            f'{key!r}: {value!r}' for key, value in self.items())
+        return f'ctx_environ({{{formatted_items}}})'
 
-    def copy(self):
+    def copy(self) -> Dict[str, str]:
         copied = self._environ.copy()
         ctx = get()
         if ctx is not None:
@@ -225,7 +260,7 @@ class ContextualEnviron(MutableMapping):
                     copied[key] = ctx.env_overrides[key]
         return copied
 
-    def setdefault(self, key, default=None):
+    def setdefault(self, key: str, default: str) -> str:
         return self._environ.setdefault(key, default)
 
     def __ior__(self, other):
@@ -260,27 +295,67 @@ class Popen(subprocess.Popen):
         super().__init__(*args, env=env, **kwargs)
 
 
-F = TypeVar('F', bound=Callable[..., Any])
+P = ParamSpec('P')
+T = TypeVar('T')
 
 
-def contextual(func: F) -> F:
+def contextual(func: Callable[P, T]) -> Callable[P, T]:
     """Decorator to initialize a context before executing the function.
 
-    If a context is already initialized, this decorator will reset the context,
-    i.e. all contextual variables set previously will be cleared.
+    If a context is already initialized, this decorator will create a new
+    context that inherits the values from the existing context.
     """
 
     @functools.wraps(func)
-    def wrapper(*args, **kwargs):
-        initialize()
-        return func(*args, **kwargs)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+        original_ctx = get()
+        initialize(original_ctx)
+        ctx = get()
+        cleanup_after_await = False
+
+        def cleanup():
+            try:
+                if ctx is not None:
+                    ctx.cleanup()
+            finally:
+                # Note: _CONTEXT.reset() is not reliable - may fail with
+                # ValueError: <Token ... at ...> was created in a different
+                # Context
+                # We must make sure this happens because otherwise we may try to
+                # write to the wrong log.
+                _CONTEXT.set(original_ctx)
+
+        # There are two cases:
+        # 1. The function is synchronous (that is, return type is not awaitable)
+        #    In this case, we use a finally block to cleanup the context.
+        # 2. The function is asynchronous (that is, return type is awaitable)
+        #    In this case, we need to construct an async def wrapper and await
+        #    the value, then call the cleanup function in the finally block.
+
+        async def await_with_cleanup(awaitable):
+            try:
+                return await awaitable
+            finally:
+                cleanup()
+
+        try:
+            ret = func(*args, **kwargs)
+            if inspect.isawaitable(ret):
+                cleanup_after_await = True
+                return await_with_cleanup(ret)
+            else:
+                return ret
+        finally:
+            if not cleanup_after_await:
+                cleanup()
 
-    return typing.cast(F, wrapper)
+    return wrapper
 
 
-def initialize():
+def initialize(base_context: Optional[Context] = None) -> None:
     """Initialize the current SkyPilot context."""
-    _CONTEXT.set(Context())
+    new_context = base_context.copy() if base_context is not None else Context()
+    _CONTEXT.set(new_context)
 
 
 class _ContextualStream:

sky/utils/controller_utils.py

@@ -620,15 +620,16 @@ def get_controller_resources(
     controller_resources_to_use: resources.Resources = list(
         controller_resources)[0]
 
-    controller_record = global_user_state.get_cluster_from_name(
+    controller_handle = global_user_state.get_handle_from_cluster_name(
         controller.value.cluster_name)
-    if controller_record is not None:
-        handle = controller_record.get('handle', None)
-        if handle is not None:
+    if controller_handle is not None:
+        if controller_handle is not None:
             # Use the existing resources, but override the autostop config with
             # the one currently specified in the config.
-            controller_resources_to_use = handle.launched_resources.copy(
-                autostop=controller_resources_config_copied.get('autostop'))
+            controller_resources_to_use = (
+                controller_handle.launched_resources.copy(
+                    autostop=controller_resources_config_copied.get('autostop'))
+            )
 
     # If the controller and replicas are from the same cloud (and region/zone),
     # it should provide better connectivity. We will let the controller choose

sky/utils/db/db_utils.py

@@ -201,6 +201,7 @@ def add_column_to_table_alembic(
     server_default: Optional[str] = None,
     copy_from: Optional[str] = None,
     value_to_replace_existing_entries: Optional[Any] = None,
+    index: Optional[bool] = None,
 ):
     """Add a column to a table using Alembic operations.
 
@@ -215,6 +216,8 @@ def add_column_to_table_alembic(
         copy_from: Column name to copy values from (for existing rows)
         value_to_replace_existing_entries: Default value for existing NULL
             entries
+        index: If True, create an index on this column. If None, no index
+            is created.
     """
     from alembic import op  # pylint: disable=import-outside-toplevel
 
@@ -222,7 +225,8 @@ def add_column_to_table_alembic(
         # Create the column with server_default if provided
         column = sqlalchemy.Column(column_name,
                                    column_type,
-                                   server_default=server_default)
+                                   server_default=server_default,
+                                   index=index)
         op.add_column(table_name, column)
 
         # Handle data migration

sky/utils/db/migration_utils.py

@@ -17,7 +17,7 @@ logger = sky_logging.init_logger(__name__)
 DB_INIT_LOCK_TIMEOUT_SECONDS = 10
 
 GLOBAL_USER_STATE_DB_NAME = 'state_db'
-GLOBAL_USER_STATE_VERSION = '008'
+GLOBAL_USER_STATE_VERSION = '009'
 GLOBAL_USER_STATE_LOCK_PATH = '~/.sky/locks/.state_db.lock'
 
 SPOT_JOBS_DB_NAME = 'spot_jobs_db'