skypilot-nightly 1.0.0.dev20250828__py3-none-any.whl → 1.0.0.dev20250831__py3-none-any.whl

sky/server/server.py

@@ -21,6 +21,7 @@ import uuid
 import zipfile
 
 import aiofiles
+import anyio
 import fastapi
 from fastapi.middleware import cors
 import starlette.middleware.base
@@ -847,7 +848,7 @@ async def upload_zip_file(request: fastapi.Request, user_hash: str,
     client_file_mounts_dir = (
         common.API_SERVER_CLIENT_DIR.expanduser().resolve() / user_id /
         'file_mounts')
-    client_file_mounts_dir.mkdir(parents=True, exist_ok=True)
+    await anyio.Path(client_file_mounts_dir).mkdir(parents=True, exist_ok=True)
 
     # Check upload_id to be a valid SkyPilot run_timestamp appended with 8 hex
     # characters, e.g. 'sky-2025-01-17-09-10-13-933602-35d31c22'.
@@ -870,7 +871,7 @@ async def upload_zip_file(request: fastapi.Request, user_hash: str,
         zip_file_path = client_file_mounts_dir / f'{upload_id}.zip'
     else:
         chunk_dir = client_file_mounts_dir / upload_id
-        chunk_dir.mkdir(parents=True, exist_ok=True)
+        await anyio.Path(chunk_dir).mkdir(parents=True, exist_ok=True)
         zip_file_path = chunk_dir / f'part{chunk_index}.incomplete'
 
     try:
@@ -916,9 +917,9 @@ async def upload_zip_file(request: fastapi.Request, user_hash: str,
                         await zip_file.write(data)
 
     logger.info(f'Uploaded zip file: {zip_file_path}')
-    unzip_file(zip_file_path, client_file_mounts_dir)
+    await unzip_file(zip_file_path, client_file_mounts_dir)
     if total_chunks > 1:
-        shutil.rmtree(chunk_dir)
+        await context_utils.to_thread(shutil.rmtree, chunk_dir)
     return payloads.UploadZipFileResponse(
         status=responses.UploadStatus.COMPLETED.value)
 
@@ -933,61 +934,69 @@ def _is_relative_to(path: pathlib.Path, parent: pathlib.Path) -> bool:
         return False
 
 
-def unzip_file(zip_file_path: pathlib.Path,
-               client_file_mounts_dir: pathlib.Path) -> None:
-    """Unzips a zip file."""
-    try:
-        with zipfile.ZipFile(zip_file_path, 'r') as zipf:
-            for member in zipf.infolist():
-                # Determine the new path
-                original_path = os.path.normpath(member.filename)
-                new_path = client_file_mounts_dir / original_path.lstrip('/')
-
-                if (member.external_attr >> 28) == 0xA:
-                    # Symlink. Read the target path and create a symlink.
-                    new_path.parent.mkdir(parents=True, exist_ok=True)
-                    target = zipf.read(member).decode()
-                    assert not os.path.isabs(target), target
-                    # Since target is a relative path, we need to check that it
-                    # is under `client_file_mounts_dir` for security.
-                    full_target_path = (new_path.parent / target).resolve()
-                    if not _is_relative_to(full_target_path,
-                                           client_file_mounts_dir):
-                        raise ValueError(f'Symlink target {target} leads to a '
-                                         'file not in userspace. Aborted.')
-
-                    if new_path.exists() or new_path.is_symlink():
-                        new_path.unlink(missing_ok=True)
-                    new_path.symlink_to(
-                        target,
-                        target_is_directory=member.filename.endswith('/'))
-                    continue
-
-                # Handle directories
-                if member.filename.endswith('/'):
-                    new_path.mkdir(parents=True, exist_ok=True)
-                    continue
-
-                # Handle files
-                new_path.parent.mkdir(parents=True, exist_ok=True)
-                with zipf.open(member) as member_file, new_path.open('wb') as f:
-                    # Use shutil.copyfileobj to copy files in chunks, so it does
-                    # not load the entire file into memory.
-                    shutil.copyfileobj(member_file, f)
-    except zipfile.BadZipFile as e:
-        logger.error(f'Bad zip file: {zip_file_path}')
-        raise fastapi.HTTPException(
-            status_code=400,
-            detail=f'Invalid zip file: {common_utils.format_exception(e)}')
-    except Exception as e:
-        logger.error(f'Error unzipping file: {zip_file_path}')
-        raise fastapi.HTTPException(
-            status_code=500,
-            detail=(f'Error unzipping file: '
-                    f'{common_utils.format_exception(e)}'))
+async def unzip_file(zip_file_path: pathlib.Path,
+                     client_file_mounts_dir: pathlib.Path) -> None:
+    """Unzips a zip file without blocking the event loop."""
 
-    # Cleanup the temporary file
-    zip_file_path.unlink()
+    def _do_unzip() -> None:
+        try:
+            with zipfile.ZipFile(zip_file_path, 'r') as zipf:
+                for member in zipf.infolist():
+                    # Determine the new path
+                    original_path = os.path.normpath(member.filename)
+                    new_path = client_file_mounts_dir / original_path.lstrip(
+                        '/')
+
+                    if (member.external_attr >> 28) == 0xA:
+                        # Symlink. Read the target path and create a symlink.
+                        new_path.parent.mkdir(parents=True, exist_ok=True)
+                        target = zipf.read(member).decode()
+                        assert not os.path.isabs(target), target
+                        # Since target is a relative path, we need to check that
+                        # it is under `client_file_mounts_dir` for security.
+                        full_target_path = (new_path.parent / target).resolve()
+                        if not _is_relative_to(full_target_path,
+                                               client_file_mounts_dir):
+                            raise ValueError(
+                                f'Symlink target {target} leads to a '
+                                'file not in userspace. Aborted.')
+
+                        if new_path.exists() or new_path.is_symlink():
+                            new_path.unlink(missing_ok=True)
+                        new_path.symlink_to(
+                            target,
+                            target_is_directory=member.filename.endswith('/'))
+                        continue
+
+                    # Handle directories
+                    if member.filename.endswith('/'):
+                        new_path.mkdir(parents=True, exist_ok=True)
+                        continue
+
+                    # Handle files
+                    new_path.parent.mkdir(parents=True, exist_ok=True)
+                    with zipf.open(member) as member_file, new_path.open(
+                            'wb') as f:
+                        # Use shutil.copyfileobj to copy files in chunks,
+                        # so it does not load the entire file into memory.
+                        shutil.copyfileobj(member_file, f)
+        except zipfile.BadZipFile as e:
+            logger.error(f'Bad zip file: {zip_file_path}')
+            raise fastapi.HTTPException(
+                status_code=400,
+                detail=f'Invalid zip file: {common_utils.format_exception(e)}')
+        except Exception as e:
+            logger.error(f'Error unzipping file: {zip_file_path}')
+            raise fastapi.HTTPException(
+                status_code=500,
+                detail=(f'Error unzipping file: '
+                        f'{common_utils.format_exception(e)}'))
+        finally:
+            # Cleanup the temporary file regardless of
+            # success/failure handling above
+            zip_file_path.unlink(missing_ok=True)
+
+    await context_utils.to_thread(_do_unzip)
 
 
 @app.post('/launch')

sky/setup_files/dependencies.py

@@ -8,8 +8,12 @@ This file is imported by setup.py, so:
 import sys
 from typing import Dict, List
 
+clouds_with_ray = ['ibm', 'docker', 'scp']
+
 install_requires = [
     'wheel<0.46.0',  # https://github.com/skypilot-org/skypilot/issues/5153
+    'setuptools',  # TODO: match version to pyproject.toml once #5153 is fixed
+    'pip',
     'cachetools',
     # NOTE: ray requires click>=7.0.
     # click 8.2.0 has a bug in parsing the command line arguments:
@@ -71,6 +75,7 @@ install_requires = [
     'types-paramiko',
     'alembic',
     'aiohttp',
+    'anyio',
 ]
 
 # See requirements-dev.txt for the version of grpc and protobuf
@@ -92,6 +97,7 @@ server_dependencies = [
     'passlib',
     'pyjwt',
     'aiohttp',
+    'anyio',
     GRPC,
     PROTOBUF,
 ]
@@ -143,7 +149,7 @@ extras_require: Dict[str, List[str]] = {
         'azure-storage-blob>=12.23.1',
         'msgraph-sdk',
         'msrestazure',
-    ] + local_ray,
+    ],
     # We need google-api-python-client>=2.69.0 to enable 'discardLocalSsd'
     # parameter for stopping instances. Reference:
     # https://github.com/googleapis/google-api-python-client/commit/f6e9d3869ed605b06f7cbf2e8cf2db25108506e6
@@ -164,7 +170,7 @@ extras_require: Dict[str, List[str]] = {
     'lambda': [],  # No dependencies needed for lambda
     'cloudflare': aws_dependencies,
     'scp': local_ray,
-    'oci': ['oci'] + local_ray,
+    'oci': ['oci'],
     # Kubernetes 32.0.0 has an authentication bug: https://github.com/kubernetes-client/python/issues/2333 # pylint: disable=line-too-long
     'kubernetes': [
         'kubernetes>=20.0.0,!=32.0.0', 'websockets', 'python-dateutil'
@@ -195,10 +201,21 @@ extras_require: Dict[str, List[str]] = {
     'server': server_dependencies,
 }
 
-# Nebius needs python3.10. If python 3.9 [all] will not install nebius
+# Calculate which clouds should be included in the [all] installation.
+clouds_for_all = set(extras_require)
+clouds_for_all.remove('remote')
+
 if sys.version_info < (3, 10):
-    filtered_keys = [k for k in extras_require if k != 'nebius']
-    extras_require['all'] = sum(
-        [v for k, v in extras_require.items() if k != 'nebius'], [])
-else:
-    extras_require['all'] = sum(extras_require.values(), [])
+    # Nebius needs python3.10. If python 3.9 [all] will not install nebius
+    clouds_for_all.remove('nebius')
+
+if sys.version_info >= (3, 12):
+    # The version of ray we use does not work with >= 3.12, so avoid clouds
+    # that require ray.
+    clouds_for_all -= set(clouds_with_ray)
+    # vast requires setuptools==51.1.1 which will not work with python >= 3.12
+    # TODO: Remove once https://github.com/vast-ai/vast-sdk/pull/6 is released
+    clouds_for_all.remove('vast')
+
+extras_require['all'] = list(
+    set().union(*[extras_require[cloud] for cloud in clouds_for_all]))

sky/setup_files/setup.py

@@ -178,6 +178,8 @@ setuptools.setup(
         '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',

sky/sky_logging.py

@@ -19,6 +19,9 @@ _FORMAT = '%(levelname).1s %(asctime)s %(filename)s:%(lineno)d] %(message)s'
 _DATE_FORMAT = '%m-%d %H:%M:%S'
 _SENSITIVE_LOGGER = ['sky.provisioner', 'sky.optimizer']
 
+_DEBUG_LOG_DIR = os.path.expanduser(
+    os.path.join(constants.SKY_LOGS_DIRECTORY, 'request_debug'))
+
 DEBUG = logging.DEBUG
 INFO = logging.INFO
 WARNING = logging.WARNING
@@ -254,3 +257,28 @@ def generate_tmp_logging_file_path(file_name: str) -> str:
     log_path = os.path.expanduser(os.path.join(log_dir, file_name))
 
     return log_path
+
+
+@contextlib.contextmanager
+def add_debug_log_handler(request_id: str):
+    if os.getenv(constants.ENV_VAR_ENABLE_REQUEST_DEBUG_LOGGING) != 'true':
+        yield
+        return
+
+    os.makedirs(_DEBUG_LOG_DIR, exist_ok=True)
+    log_path = os.path.join(_DEBUG_LOG_DIR, f'{request_id}.log')
+    try:
+        debug_log_handler = logging.FileHandler(log_path)
+        debug_log_handler.setFormatter(FORMATTER)
+        debug_log_handler.setLevel(logging.DEBUG)
+        _root_logger.addHandler(debug_log_handler)
+        # sky.provision sets up its own logger/handler with propogate=False,
+        # so add it there too.
+        provision_logger = logging.getLogger('sky.provision')
+        provision_logger.addHandler(debug_log_handler)
+        provision_logger.setLevel(logging.DEBUG)
+        yield
+    finally:
+        _root_logger.removeHandler(debug_log_handler)
+        provision_logger.removeHandler(debug_log_handler)
+        debug_log_handler.close()

sky/skylet/constants.py

@@ -70,6 +70,7 @@ DEACTIVATE_SKY_REMOTE_PYTHON_ENV = (
 
 # Prefix for SkyPilot environment variables
 SKYPILOT_ENV_VAR_PREFIX = 'SKYPILOT_'
+SKYPILOT_SERVER_ENV_VAR_PREFIX = 'SKYPILOT_SERVER_'
 
 # The name for the environment variable that stores the unique ID of the
 # current task. This will stay the same across multiple recoveries of the
@@ -417,6 +418,7 @@ LOCAL_SKYPILOT_CONFIG_PATH_PLACEHOLDER = 'skypilot:local_skypilot_config_path'
 # Path to the generated cluster config yamls and ssh configs.
 SKY_USER_FILE_PATH = '~/.sky/generated'
 
+# TODO(cooperc): Update all env vars to begin with SKYPILOT_ or SKYPILOT_SERVER_
 # Environment variable that is set to 'true' if this is a skypilot server.
 ENV_VAR_IS_SKYPILOT_SERVER = 'IS_SKYPILOT_SERVER'
 
@@ -436,6 +438,10 @@ ENV_VAR_ENABLE_BASIC_AUTH = 'ENABLE_BASIC_AUTH'
 SKYPILOT_INITIAL_BASIC_AUTH = 'SKYPILOT_INITIAL_BASIC_AUTH'
 ENV_VAR_ENABLE_SERVICE_ACCOUNTS = 'ENABLE_SERVICE_ACCOUNTS'
 
+# Enable debug logging for requests.
+ENV_VAR_ENABLE_REQUEST_DEBUG_LOGGING = (
+    f'{SKYPILOT_SERVER_ENV_VAR_PREFIX}ENABLE_REQUEST_DEBUG_LOGGING')
+
 SKYPILOT_DEFAULT_WORKSPACE = 'default'
 
 # BEGIN constants used for service catalog.

sky/templates/aws-ray.yml.j2

@@ -48,6 +48,7 @@ provider:
   # The upper-level SkyPilot code has make sure there will not be resource
   # leakage.
   disable_launch_config_check: true
+  max_efa_interfaces: {{max_efa_interfaces}}
 
 auth:
   ssh_user: {{ssh_user}}

sky/utils/annotations.py

@@ -7,7 +7,7 @@ from typing_extensions import ParamSpec
 
 # Whether the current process is a SkyPilot API server process.
 is_on_api_server = True
-FUNCTIONS_NEED_RELOAD_CACHE = []
+_FUNCTIONS_NEED_RELOAD_CACHE = []
 
 T = TypeVar('T')
 P = ParamSpec('P')
@@ -50,7 +50,13 @@ def lru_cache(scope: Literal['global', 'request'], *lru_cache_args,
         else:
             cached_func = functools.lru_cache(*lru_cache_args,
                                               **lru_cache_kwargs)(func)
-            FUNCTIONS_NEED_RELOAD_CACHE.append(cached_func)
+            _FUNCTIONS_NEED_RELOAD_CACHE.append(cached_func)
             return cached_func
 
     return decorator
+
+
+def clear_request_level_cache():
+    """Clear the request-level cache."""
+    for func in _FUNCTIONS_NEED_RELOAD_CACHE:
+        func.cache_clear()

sky/utils/cluster_utils.py

@@ -11,7 +11,7 @@ import uuid
 from sky.skylet import constants
 from sky.utils import command_runner
 from sky.utils import common_utils
-from sky.utils import timeline
+from sky.utils import lock_events
 
 # The cluster yaml used to create the current cluster where the module is
 # called.
@@ -107,7 +107,7 @@ class SSHConfigHelper(object):
         return auth_config['ssh_private_key']
 
     @classmethod
-    @timeline.FileLockEvent(ssh_conf_lock_path)
+    @lock_events.FileLockEvent(ssh_conf_lock_path)
     def add_cluster(
         cls,
         cluster_name: str,
@@ -334,7 +334,7 @@ class SSHConfigHelper(object):
             cluster_name: Cluster name.
         """
 
-        with timeline.FileLockEvent(
+        with lock_events.FileLockEvent(
                 cls.ssh_conf_per_cluster_lock_path.format(cluster_name)):
             cluster_config_path = os.path.expanduser(
                 cls.ssh_cluster_path.format(cluster_name))

sky/utils/db/migration_utils.py

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

sky/utils/kubernetes_enums.py

@@ -42,4 +42,5 @@ class KubernetesAutoscalerType(enum.Enum):
     """Enum for the different types of cluster autoscalers for Kubernetes."""
     GKE = 'gke'
     KARPENTER = 'karpenter'
+    COREWEAVE = 'coreweave'
     GENERIC = 'generic'

sky/utils/lock_events.py

@@ -0,0 +1,94 @@
+"""Lock events."""
+
+import functools
+import os
+from typing import Optional, Union
+
+import filelock
+
+from sky.utils import locks
+from sky.utils import timeline
+
+
+class DistributedLockEvent:
+    """Serve both as a distributed lock and event for the lock."""
+
+    def __init__(self, lock_id: str, timeout: Optional[float] = None):
+        self._lock_id = lock_id
+        self._lock = locks.get_lock(lock_id, timeout)
+        self._hold_lock_event = timeline.Event(
+            f'[DistributedLock.hold]:{lock_id}')
+
+    def acquire(self):
+        was_locked = self._lock.is_locked
+        with timeline.Event(f'[DistributedLock.acquire]:{self._lock_id}'):
+            self._lock.acquire()
+        if not was_locked and self._lock.is_locked:
+            # start holding the lock after initial acquiring
+            self._hold_lock_event.begin()
+
+    def release(self):
+        was_locked = self._lock.is_locked
+        self._lock.release()
+        if was_locked and not self._lock.is_locked:
+            # stop holding the lock after initial releasing
+            self._hold_lock_event.end()
+
+    def __enter__(self):
+        self.acquire()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.release()
+
+    def __call__(self, f):
+
+        @functools.wraps(f)
+        def wrapper(*args, **kwargs):
+            with self:
+                return f(*args, **kwargs)
+
+        return wrapper
+
+
+class FileLockEvent:
+    """Serve both as a file lock and event for the lock."""
+
+    def __init__(self, lockfile: Union[str, os.PathLike], timeout: float = -1):
+        self._lockfile = lockfile
+        os.makedirs(os.path.dirname(os.path.abspath(self._lockfile)),
+                    exist_ok=True)
+        self._lock = filelock.FileLock(self._lockfile, timeout)
+        self._hold_lock_event = timeline.Event(
+            f'[FileLock.hold]:{self._lockfile}')
+
+    def acquire(self):
+        was_locked = self._lock.is_locked
+        with timeline.Event(f'[FileLock.acquire]:{self._lockfile}'):
+            self._lock.acquire()
+        if not was_locked and self._lock.is_locked:
+            # start holding the lock after initial acquiring
+            self._hold_lock_event.begin()
+
+    def release(self):
+        was_locked = self._lock.is_locked
+        self._lock.release()
+        if was_locked and not self._lock.is_locked:
+            # stop holding the lock after initial releasing
+            self._hold_lock_event.end()
+
+    def __enter__(self):
+        self.acquire()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.release()
+
+    def __call__(self, f):
+        # Make this class callable as a decorator.
+        @functools.wraps(f)
+        def wrapper(*args, **kwargs):
+            with self:
+                return f(*args, **kwargs)
+
+        return wrapper

sky/utils/schemas.py

@@ -1410,6 +1410,9 @@ def get_config_schema():
                 **_NETWORK_CONFIG_SCHEMA, 'tenant_id': {
                     'type': 'string',
                 },
+                'domain': {
+                    'type': 'string',
+                },
                 'region_configs': {
                     'type': 'object',
                     'required': [],
@@ -1668,6 +1671,9 @@ def get_config_schema():
                         'tenant_id': {
                             'type': 'string',
                         },
+                        'domain': {
+                            'type': 'string',
+                        },
                         'disabled': {
                             'type': 'boolean'
                         },

sky/utils/timeline.py

@@ -4,7 +4,6 @@ The timeline follows the trace event format defined here:
 https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview
 """  # pylint: disable=line-too-long
 import atexit
-import functools
 import json
 import os
 import threading
@@ -12,14 +11,15 @@ import time
 import traceback
 from typing import Callable, Optional, Union
 
-import filelock
-
 from sky.utils import common_utils
-from sky.utils import locks
 
 _events = []
 
 
+def _get_events_file_path():
+    return os.environ.get('SKYPILOT_TIMELINE_FILE_PATH')
+
+
 class Event:
     """Record an event.
 
@@ -29,6 +29,10 @@ class Event:
     """
 
     def __init__(self, name: str, message: Optional[str] = None):
+        self._skipped = False
+        if not _get_events_file_path():
+            self._skipped = True
+            return
         self._name = name
         self._message = message
         # See the module doc for the event format.
@@ -45,6 +49,8 @@ class Event:
             self._event['args'] = {'message': self._message}
 
     def begin(self):
+        if self._skipped:
+            return
         event_begin = self._event.copy()
         event_begin.update({
             'ph': 'B',
@@ -56,6 +62,8 @@ class Event:
         _events.append(event_begin)
 
     def end(self):
+        if self._skipped:
+            return
         event_end = self._event.copy()
         event_end.update({
             'ph': 'E',
@@ -77,103 +85,26 @@ def event(name_or_fn: Union[str, Callable], message: Optional[str] = None):
     return common_utils.make_decorator(Event, name_or_fn, message=message)
 
 
-class DistributedLockEvent:
-    """Serve both as a distributed lock and event for the lock."""
-
-    def __init__(self, lock_id: str, timeout: Optional[float] = None):
-        self._lock_id = lock_id
-        self._lock = locks.get_lock(lock_id, timeout)
-        self._hold_lock_event = Event(f'[DistributedLock.hold]:{lock_id}')
-
-    def acquire(self):
-        was_locked = self._lock.is_locked
-        with Event(f'[DistributedLock.acquire]:{self._lock_id}'):
-            self._lock.acquire()
-        if not was_locked and self._lock.is_locked:
-            # start holding the lock after initial acquiring
-            self._hold_lock_event.begin()
-
-    def release(self):
-        was_locked = self._lock.is_locked
-        self._lock.release()
-        if was_locked and not self._lock.is_locked:
-            # stop holding the lock after initial releasing
-            self._hold_lock_event.end()
-
-    def __enter__(self):
-        self.acquire()
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.release()
-
-    def __call__(self, f):
-
-        @functools.wraps(f)
-        def wrapper(*args, **kwargs):
-            with self:
-                return f(*args, **kwargs)
-
-        return wrapper
-
-
-class FileLockEvent:
-    """Serve both as a file lock and event for the lock."""
-
-    def __init__(self, lockfile: Union[str, os.PathLike], timeout: float = -1):
-        self._lockfile = lockfile
-        os.makedirs(os.path.dirname(os.path.abspath(self._lockfile)),
-                    exist_ok=True)
-        self._lock = filelock.FileLock(self._lockfile, timeout)
-        self._hold_lock_event = Event(f'[FileLock.hold]:{self._lockfile}')
-
-    def acquire(self):
-        was_locked = self._lock.is_locked
-        with Event(f'[FileLock.acquire]:{self._lockfile}'):
-            self._lock.acquire()
-        if not was_locked and self._lock.is_locked:
-            # start holding the lock after initial acquiring
-            self._hold_lock_event.begin()
-
-    def release(self):
-        was_locked = self._lock.is_locked
-        self._lock.release()
-        if was_locked and not self._lock.is_locked:
-            # stop holding the lock after initial releasing
-            self._hold_lock_event.end()
-
-    def __enter__(self):
-        self.acquire()
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.release()
-
-    def __call__(self, f):
-        # Make this class callable as a decorator.
-        @functools.wraps(f)
-        def wrapper(*args, **kwargs):
-            with self:
-                return f(*args, **kwargs)
-
-        return wrapper
-
-
 def save_timeline():
-    file_path = os.environ.get('SKYPILOT_TIMELINE_FILE_PATH')
-    if not file_path:
+    events_file_path = _get_events_file_path()
+    if not events_file_path:
         return
+    global _events
+    events_to_write = _events
+    _events = []
     json_output = {
-        'traceEvents': _events,
+        'traceEvents': events_to_write,
         'displayTimeUnit': 'ms',
         'otherData': {
-            'log_dir': os.path.dirname(os.path.abspath(file_path)),
+            'log_dir': os.path.dirname(os.path.abspath(events_file_path)),
         }
     }
-    os.makedirs(os.path.dirname(os.path.abspath(file_path)), exist_ok=True)
-    with open(file_path, 'w', encoding='utf-8') as f:
+    os.makedirs(os.path.dirname(os.path.abspath(events_file_path)),
+                exist_ok=True)
+    with open(events_file_path, 'w', encoding='utf-8') as f:
         json.dump(json_output, f)
+    del events_to_write
 
 
-if os.environ.get('SKYPILOT_TIMELINE_FILE_PATH'):
+if _get_events_file_path():
     atexit.register(save_timeline)