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

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,
     )
 
 

sky/server/stream_utils.py

@@ -3,7 +3,7 @@
 import asyncio
 import collections
 import pathlib
-from typing import AsyncGenerator, Deque, Optional
+from typing import AsyncGenerator, Deque, List, Optional
 
 import aiofiles
 import fastapi
@@ -15,6 +15,12 @@ from sky.utils import rich_utils
 
 logger = sky_logging.init_logger(__name__)
 
+# When streaming log lines, buffer the lines in memory and flush them in chunks
+# to improve log tailing throughput. Buffer size is the max size bytes of each
+# chunk and the timeout threshold for flushing the buffer to ensure
+# responsiveness.
+_BUFFER_SIZE = 8 * 1024  # 8KB
+_BUFFER_TIMEOUT = 0.02  # 20ms
 _HEARTBEAT_INTERVAL = 30
 
 
@@ -36,7 +42,16 @@ async def log_streamer(request_id: Optional[str],
                        plain_logs: bool = False,
                        tail: Optional[int] = None,
                        follow: bool = True) -> AsyncGenerator[str, None]:
-    """Streams the logs of a request."""
+    """Streams the logs of a request.
+
+    Args:
+        request_id: The request ID to check whether the log tailing process
+            should be stopped.
+        log_path: The path to the log file.
+        plain_logs: Whether to show plain logs.
+        tail: The number of lines to tail. If None, tail the whole file.
+        follow: Whether to follow the log file.
+    """
 
     if request_id is not None:
         status_msg = rich_utils.EncodedStatusMessage(
@@ -80,65 +95,106 @@ async def log_streamer(request_id: Optional[str],
         if show_request_waiting_spinner:
             yield status_msg.stop()
 
-    # Find last n lines of the log file. Do not read the whole file into memory.
     async with aiofiles.open(log_path, 'rb') as f:
-        if tail is not None:
-            # TODO(zhwu): this will include the control lines for rich status,
-            # which may not lead to exact tail lines when showing on the client
-            # side.
-            lines: Deque[str] = collections.deque(maxlen=tail)
-            async for line_str in _yield_log_file_with_payloads_skipped(f):
-                lines.append(line_str)
-            for line_str in lines:
-                yield line_str
-
-        last_heartbeat_time = asyncio.get_event_loop().time()
+        async for chunk in _tail_log_file(f, request_id, plain_logs, tail,
+                                          follow):
+            yield chunk
+
+
+async def _tail_log_file(f: aiofiles.threadpool.binary.AsyncBufferedReader,
+                         request_id: Optional[str] = None,
+                         plain_logs: bool = False,
+                         tail: Optional[int] = None,
+                         follow: bool = True) -> AsyncGenerator[str, None]:
+    """Tail the opened log file, buffer the lines and flush in chunks."""
+
+    if tail is not None:
+        # Find last n lines of the log file. Do not read the whole file into
+        # memory.
+        # TODO(zhwu): this will include the control lines for rich status,
+        # which may not lead to exact tail lines when showing on the client
+        # side.
+        lines: Deque[str] = collections.deque(maxlen=tail)
+        async for line_str in _yield_log_file_with_payloads_skipped(f):
+            lines.append(line_str)
+        for line_str in lines:
+            yield line_str
 
-        while True:
-            # Sleep 0 to yield control to allow other coroutines to run,
-            # while keeps the loop tight to make log stream responsive.
-            await asyncio.sleep(0)
-            line: Optional[bytes] = await f.readline()
-            if not line:
-                if request_id is not None:
-                    request_task = requests_lib.get_request(request_id)
-                    if request_task.status > requests_lib.RequestStatus.RUNNING:
-                        if (request_task.status ==
-                                requests_lib.RequestStatus.CANCELLED):
-                            yield (f'{request_task.name!r} request {request_id}'
-                                   ' cancelled\n')
-                        break
-                if not follow:
+    last_heartbeat_time = asyncio.get_event_loop().time()
+
+    # Buffer the lines in memory and flush them in chunks to improve log
+    # tailing throughput.
+    buffer: List[str] = []
+    buffer_bytes = 0
+    last_flush_time = asyncio.get_event_loop().time()
+
+    async def flush_buffer() -> AsyncGenerator[str, None]:
+        nonlocal buffer, buffer_bytes, last_flush_time
+        if buffer:
+            yield ''.join(buffer)
+            buffer.clear()
+            buffer_bytes = 0
+            last_flush_time = asyncio.get_event_loop().time()
+
+    while True:
+        # Sleep 0 to yield control to allow other coroutines to run,
+        # while keeps the loop tight to make log stream responsive.
+        await asyncio.sleep(0)
+        current_time = asyncio.get_event_loop().time()
+        # Flush the buffer when it is not empty and the buffer is full or the
+        # flush timeout is reached.
+        if buffer and (buffer_bytes >= _BUFFER_SIZE or
+                       (current_time - last_flush_time) >= _BUFFER_TIMEOUT):
+            async for chunk in flush_buffer():
+                yield chunk
+
+        line: Optional[bytes] = await f.readline()
+        if not line:
+            if request_id is not None:
+                request_task = requests_lib.get_request(request_id)
+                if request_task.status > requests_lib.RequestStatus.RUNNING:
+                    if (request_task.status ==
+                            requests_lib.RequestStatus.CANCELLED):
+                        buffer.append(
+                            f'{request_task.name!r} request {request_id}'
+                            ' cancelled\n')
                     break
+            if not follow:
+                break
+
+            if current_time - last_heartbeat_time >= _HEARTBEAT_INTERVAL:
+                # Currently just used to keep the connection busy, refer to
+                # https://github.com/skypilot-org/skypilot/issues/5750 for
+                # more details.
+                buffer.append(
+                    message_utils.encode_payload(
+                        rich_utils.Control.HEARTBEAT.encode('')))
+                last_heartbeat_time = current_time
+
+            # Sleep shortly to avoid storming the DB and CPU, this has
+            # little impact on the responsivness here since we are waiting
+            # for a new line to come in.
+            await asyncio.sleep(0.1)
+            continue
 
-                current_time = asyncio.get_event_loop().time()
-                if current_time - last_heartbeat_time >= _HEARTBEAT_INTERVAL:
-                    # Currently just used to keep the connection busy, refer to
-                    # https://github.com/skypilot-org/skypilot/issues/5750 for
-                    # more details.
-                    yield message_utils.encode_payload(
-                        rich_utils.Control.HEARTBEAT.encode(''))
-                    last_heartbeat_time = current_time
-
-                # Sleep shortly to avoid storming the DB and CPU, this has
-                # little impact on the responsivness here since we are waiting
-                # for a new line to come in.
-                await asyncio.sleep(0.1)
+        # Refresh the heartbeat time, this is a trivial optimization for
+        # performance but it helps avoid unnecessary heartbeat strings
+        # being printed when the client runs in an old version.
+        last_heartbeat_time = asyncio.get_event_loop().time()
+        line_str = line.decode('utf-8')
+        if plain_logs:
+            is_payload, line_str = message_utils.decode_payload(
+                line_str, raise_for_mismatch=False)
+            # TODO(aylei): implement heartbeat mechanism for plain logs,
+            # sending invisible characters might be okay.
+            if is_payload:
                 continue
+        buffer.append(line_str)
+        buffer_bytes += len(line_str.encode('utf-8'))
 
-            # Refresh the heartbeat time, this is a trivial optimization for
-            # performance but it helps avoid unnecessary heartbeat strings
-            # being printed when the client runs in an old version.
-            last_heartbeat_time = asyncio.get_event_loop().time()
-            line_str = line.decode('utf-8')
-            if plain_logs:
-                is_payload, line_str = message_utils.decode_payload(
-                    line_str, raise_for_mismatch=False)
-                # TODO(aylei): implement heartbeat mechanism for plain logs,
-                # sending invisible characters might be okay.
-                if is_payload:
-                    continue
-            yield line_str
+    # Flush remaining lines in the buffer.
+    async for chunk in flush_buffer():
+        yield chunk
 
 
 def stream_response(

sky/skylet/constants.py

@@ -377,8 +377,7 @@ OVERRIDEABLE_CONFIG_KEYS_IN_TASK: List[Tuple[str, ...]] = [
 ]
 # When overriding the SkyPilot configs on the API server with the client one,
 # we skip the following keys because they are meant to be client-side configs.
-SKIPPED_CLIENT_OVERRIDE_KEYS: List[Tuple[str, ...]] = [('admin_policy',),
-                                                       ('api_server',),
+SKIPPED_CLIENT_OVERRIDE_KEYS: List[Tuple[str, ...]] = [('api_server',),
                                                        ('allowed_clouds',),
                                                        ('workspaces',), ('db',)]
 

sky/skypilot_config.py

@@ -58,6 +58,11 @@ import typing
 from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
 import filelock
+import sqlalchemy
+from sqlalchemy import orm
+from sqlalchemy.dialects import postgresql
+from sqlalchemy.dialects import sqlite
+from sqlalchemy.ext import declarative
 
 from sky import exceptions
 from sky import sky_logging
@@ -66,6 +71,7 @@ from sky.skylet import constants
 from sky.utils import common_utils
 from sky.utils import config_utils
 from sky.utils import context
+from sky.utils import db_utils
 from sky.utils import schemas
 from sky.utils import ux_utils
 from sky.utils.kubernetes import config_map_utils
@@ -110,6 +116,56 @@ ENV_VAR_PROJECT_CONFIG = f'{constants.SKYPILOT_ENV_VAR_PREFIX}PROJECT_CONFIG'
 _GLOBAL_CONFIG_PATH = '~/.sky/config.yaml'
 _PROJECT_CONFIG_PATH = '.sky.yaml'
 
+_SQLALCHEMY_ENGINE: Optional[sqlalchemy.engine.Engine] = None
+API_SERVER_CONFIG_KEY = 'api_server_config'
+
+Base = declarative.declarative_base()
+
+config_yaml_table = sqlalchemy.Table(
+    'config_yaml',
+    Base.metadata,
+    sqlalchemy.Column('key', sqlalchemy.Text, primary_key=True),
+    sqlalchemy.Column('value', sqlalchemy.Text),
+)
+
+
+def create_table():
+    # Create tables if they don't exist
+    Base.metadata.create_all(bind=_SQLALCHEMY_ENGINE)
+
+
+def _get_config_yaml_from_db(key: str) -> Optional[config_utils.Config]:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        row = session.query(config_yaml_table).filter_by(key=key).first()
+    if row:
+        db_config = config_utils.Config(yaml.safe_load(row.value))
+        db_config.pop_nested(('db',), None)
+        return db_config
+    return None
+
+
+def _set_config_yaml_to_db(key: str, config: config_utils.Config):
+    assert _SQLALCHEMY_ENGINE is not None
+    config.pop_nested(('db',), None)
+    config_str = common_utils.dump_yaml_str(dict(config))
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        if (_SQLALCHEMY_ENGINE.dialect.name ==
+                db_utils.SQLAlchemyDialect.SQLITE.value):
+            insert_func = sqlite.insert
+        elif (_SQLALCHEMY_ENGINE.dialect.name ==
+              db_utils.SQLAlchemyDialect.POSTGRESQL.value):
+            insert_func = postgresql.insert
+        else:
+            raise ValueError('Unsupported database dialect')
+        insert_stmnt = insert_func(config_yaml_table).values(key=key,
+                                                             value=config_str)
+        do_update_stmt = insert_stmnt.on_conflict_do_update(
+            index_elements=[config_yaml_table.c.key],
+            set_={config_yaml_table.c.value: config_str})
+        session.execute(do_update_stmt)
+        session.commit()
+
 
 class ConfigContext:
 
@@ -257,11 +313,6 @@ def _resolve_project_config_path() -> Optional[str]:
     return None
 
 
-def _get_project_config() -> config_utils.Config:
-    """Returns the project config."""
-    return _get_config_from_path(_resolve_project_config_path())
-
-
 def _resolve_server_config_path() -> Optional[str]:
     # find the server config file
     server_config_path = _get_config_file_path(ENV_VAR_GLOBAL_CONFIG)
@@ -507,26 +558,35 @@ def _reload_config_from_internal_file(internal_config_path: str) -> None:
 
 
 def _reload_config_as_server() -> None:
+    global _SQLALCHEMY_ENGINE
     # Reset the global variables, to avoid using stale values.
     _set_loaded_config(config_utils.Config())
     _set_loaded_config_path(None)
 
-    overrides: List[config_utils.Config] = []
     server_config_path = _resolve_server_config_path()
     server_config = _get_config_from_path(server_config_path)
-    if server_config:
-        overrides.append(server_config)
 
-    # layer the configs on top of each other based on priority
-    overlaid_server_config: config_utils.Config = config_utils.Config()
-    for override in overrides:
-        overlaid_server_config = overlay_skypilot_config(
-            original_config=overlaid_server_config, override_configs=override)
     if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
-        logger.debug(
-            f'server config: \n'
-            f'{common_utils.dump_yaml_str(dict(overlaid_server_config))}')
-    _set_loaded_config(overlaid_server_config)
+        logger.debug(f'server config: \n'
+                     f'{common_utils.dump_yaml_str(dict(server_config))}')
+
+    db_url = server_config.get_nested(('db',), None)
+    if db_url and len(server_config.keys()) > 1:
+        raise ValueError(
+            'if db config is specified, no other config is allowed')
+
+    if db_url:
+        if _SQLALCHEMY_ENGINE is None:
+            _SQLALCHEMY_ENGINE = sqlalchemy.create_engine(db_url)
+            create_table()
+        db_config = _get_config_yaml_from_db(API_SERVER_CONFIG_KEY)
+        if db_config:
+            if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
+                logger.debug(f'Config loaded from db:\n'
+                             f'{common_utils.dump_yaml_str(dict(db_config))}')
+            server_config = overlay_skypilot_config(server_config, db_config)
+
+    _set_loaded_config(server_config)
     _set_loaded_config_path(server_config_path)
 
 
@@ -778,13 +838,27 @@ def update_api_server_config_no_lock(config: config_utils.Config) -> None:
     if global_config_path is None:
         global_config_path = get_user_config_path()
 
-    # Always save to the local file (PVC in Kubernetes, local file otherwise)
-    common_utils.dump_yaml(global_config_path, dict(config))
-
-    if config_map_utils.is_running_in_kubernetes():
-        # In Kubernetes, sync the PVC config to ConfigMap for user convenience
-        # PVC file is the source of truth, ConfigMap is just a mirror for easy
-        # access
-        config_map_utils.patch_configmap_with_config(config, global_config_path)
+    db_updated = False
+    if os.environ.get(constants.ENV_VAR_IS_SKYPILOT_SERVER) is not None:
+        existing_db_url = get_nested(('db',), None)
+        if existing_db_url:
+            new_db_url = config.get_nested(('db',), None)
+            if new_db_url and new_db_url != existing_db_url:
+                raise ValueError('Cannot change db url while server is running')
+            logger.debug('saving api_server config to db')
+            _set_config_yaml_to_db(API_SERVER_CONFIG_KEY, config)
+            db_updated = True
+
+    if not db_updated:
+        # save to the local file (PVC in Kubernetes, local file otherwise)
+        common_utils.dump_yaml(global_config_path, dict(config))
+
+        if config_map_utils.is_running_in_kubernetes():
+            # In Kubernetes, sync the PVC config to ConfigMap for user
+            # convenience.
+            # PVC file is the source of truth, ConfigMap is just a mirror for
+            # easy access.
+            config_map_utils.patch_configmap_with_config(
+                config, global_config_path)
 
     _reload_config()

sky/users/permission.py

@@ -36,7 +36,7 @@ class PermissionService:
             with _lock:
                 if _enforcer_instance is None:
                     _enforcer_instance = self
-                    engine = global_user_state.SQLALCHEMY_ENGINE
+                    engine = global_user_state.initialize_and_get_db()
                     adapter = sqlalchemy_adapter.Adapter(engine)
                     model_path = os.path.join(os.path.dirname(__file__),
                                               'model.conf')

sky/utils/admin_policy_utils.py

@@ -55,6 +55,7 @@ def _get_policy_cls(
 def apply_and_use_config_in_current_request(
     entrypoint: Union['dag_lib.Dag', 'task_lib.Task'],
     request_options: Optional[admin_policy.RequestOptions] = None,
+    at_client_side: bool = False,
 ) -> Iterator['dag_lib.Dag']:
     """Applies an admin policy and override SkyPilot config for current request
 
@@ -66,7 +67,7 @@ def apply_and_use_config_in_current_request(
     Refer to `apply()` for more details.
     """
     original_config = skypilot_config.to_dict()
-    dag, mutated_config = apply(entrypoint, request_options)
+    dag, mutated_config = apply(entrypoint, request_options, at_client_side)
     if mutated_config != original_config:
         with skypilot_config.replace_skypilot_config(mutated_config):
             yield dag
@@ -77,6 +78,7 @@ def apply_and_use_config_in_current_request(
 def apply(
     entrypoint: Union['dag_lib.Dag', 'task_lib.Task'],
     request_options: Optional[admin_policy.RequestOptions] = None,
+    at_client_side: bool = False,
 ) -> Tuple['dag_lib.Dag', config_utils.Config]:
     """Applies an admin policy (if registered) to a DAG or a task.
 
@@ -105,14 +107,18 @@ def apply(
     if policy_cls is None:
         return dag, skypilot_config.to_dict()
 
-    logger.info(f'Applying policy: {policy}')
+    if at_client_side:
+        logger.info(f'Applying client admin policy: {policy}')
+    else:
+        logger.info(f'Applying server admin policy: {policy}')
     config = copy.deepcopy(skypilot_config.to_dict())
     mutated_dag = dag_lib.Dag()
     mutated_dag.name = dag.name
 
     mutated_config = None
     for task in dag.tasks:
-        user_request = admin_policy.UserRequest(task, config, request_options)
+        user_request = admin_policy.UserRequest(task, config, request_options,
+                                                at_client_side)
         try:
             mutated_user_request = policy_cls.validate_and_mutate(user_request)
         except Exception as e:  # pylint: disable=broad-except

sky/utils/context.py

@@ -4,11 +4,13 @@ import asyncio
 from collections.abc import Mapping
 from collections.abc import MutableMapping
 import contextvars
+import functools
 import os
 import pathlib
 import subprocess
 import sys
-from typing import Dict, Optional, TextIO
+import typing
+from typing import Any, Callable, Dict, Optional, TextIO, TypeVar
 
 
 class Context(object):
@@ -256,6 +258,24 @@ class Popen(subprocess.Popen):
         super().__init__(*args, env=env, **kwargs)
 
 
+F = TypeVar('F', bound=Callable[..., Any])
+
+
+def contextual(func: F) -> F:
+    """Decorator to intiailize 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.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        initialize()
+        return func(*args, **kwargs)
+
+    return typing.cast(F, wrapper)
+
+
 def initialize():
     """Initialize the current SkyPilot context."""
     _CONTEXT.set(Context())