skypilot-nightly 1.0.0.dev20250502__py3-none-any.whl → 1.0.0.dev20251203__py3-none-any.whl

sky/utils/db/kv_cache.py

@@ -0,0 +1,149 @@
+"""Persistent KV cache, backed by a sqlite or postgres database."""
+import functools
+import threading
+import time
+from typing import Optional
+
+import sqlalchemy
+from sqlalchemy import exc as sqlalchemy_exc
+from sqlalchemy import orm
+from sqlalchemy.dialects import postgresql
+from sqlalchemy.dialects import sqlite
+from sqlalchemy.ext import declarative
+
+from sky import sky_logging
+from sky.metrics import utils as metrics_lib
+from sky.utils import common_utils
+from sky.utils.db import db_utils
+from sky.utils.db import migration_utils
+
+logger = sky_logging.init_logger(__name__)
+
+_SQLALCHEMY_ENGINE: Optional[sqlalchemy.engine.Engine] = None
+_SQLALCHEMY_ENGINE_LOCK = threading.Lock()
+
+Base = declarative.declarative_base()
+
+kv_cache_table = sqlalchemy.Table(
+    'kv_cache',
+    Base.metadata,
+    sqlalchemy.Column('key', sqlalchemy.Text, primary_key=True),
+    sqlalchemy.Column('value', sqlalchemy.Text),
+    sqlalchemy.Column('expires_at', sqlalchemy.Float),
+)
+
+
+def create_table(engine: sqlalchemy.engine.Engine):
+    # Enable WAL mode to avoid locking issues.
+    # See: issue #1441 and PR #1509
+    # https://github.com/microsoft/WSL/issues/2395
+    # TODO(romilb): We do not enable WAL for WSL because of known issue in WSL.
+    #  This may cause the database locked problem from WSL issue #1441.
+    if (engine.dialect.name == db_utils.SQLAlchemyDialect.SQLITE.value and
+            not common_utils.is_wsl()):
+        try:
+            with orm.Session(engine) as session:
+                session.execute(sqlalchemy.text('PRAGMA journal_mode=WAL'))
+                session.commit()
+        except sqlalchemy_exc.OperationalError as e:
+            if 'database is locked' not in str(e):
+                raise
+            # If the database is locked, it is OK to continue, as the WAL mode
+            # is not critical and is likely to be enabled by other processes.
+
+    migration_utils.safe_alembic_upgrade(engine,
+                                         migration_utils.KV_CACHE_DB_NAME,
+                                         migration_utils.KV_CACHE_VERSION)
+
+
+# We wrap the sqlalchemy engine initialization in a thread
+# lock to ensure that multiple threads do not initialize the
+# engine which could result in a rare race condition where
+# a session has already been created with _SQLALCHEMY_ENGINE = e1,
+# and then another thread overwrites _SQLALCHEMY_ENGINE = e2
+# which could result in e1 being garbage collected unexpectedly.
+def initialize_and_get_db() -> sqlalchemy.engine.Engine:
+    global _SQLALCHEMY_ENGINE
+
+    if _SQLALCHEMY_ENGINE is not None:
+        return _SQLALCHEMY_ENGINE
+    with _SQLALCHEMY_ENGINE_LOCK:
+        if _SQLALCHEMY_ENGINE is not None:
+            return _SQLALCHEMY_ENGINE
+        # get an engine to the db
+        engine = db_utils.get_engine('kv_cache')
+
+        # run migrations if needed
+        create_table(engine)
+        _SQLALCHEMY_ENGINE = engine
+        return _SQLALCHEMY_ENGINE
+
+
+def _init_db(func):
+    """Initialize the database."""
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        initialize_and_get_db()
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+initialize_and_get_db()
+
+
+@_init_db
+@metrics_lib.time_me
+def add_or_update_cache_entry(
+    key: str,
+    value: str,
+    expires_at: float,
+) -> None:
+    """Store the mapping from user hash to user name for display purposes.
+
+    Args:
+        key: The key of the cache entry.
+        value: The value of the cache entry.
+        expires_at: The timestamp when the cache entry expires.
+    """
+    assert _SQLALCHEMY_ENGINE is not None
+    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')
+
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        insert_stmt = insert_func(kv_cache_table).values(key=key,
+                                                         value=value,
+                                                         expires_at=expires_at)
+        do_update_stmt = insert_stmt.on_conflict_do_update(
+            index_elements=[kv_cache_table.c.key],
+            set_={
+                kv_cache_table.c.value: value,
+                kv_cache_table.c.expires_at: expires_at
+            })
+        session.execute(do_update_stmt)
+
+        session.commit()
+
+
+@_init_db
+@metrics_lib.time_me
+def get_cache_entry(key: str) -> Optional[str]:
+    """Get the value of the cache entry.
+
+    Args:
+        key: The key of the cache entry.
+    """
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        result = session.execute(
+            sqlalchemy.select(kv_cache_table.c.value).where(
+                kv_cache_table.c.key == key).where(
+                    kv_cache_table.c.expires_at > time.time()))
+        return result.scalar()

sky/utils/db/migration_utils.py

@@ -0,0 +1,137 @@
+"""Constants for the database schemas."""
+
+import contextlib
+import logging
+import os
+
+from alembic import command as alembic_command
+from alembic.config import Config
+from alembic.runtime import migration
+import filelock
+import sqlalchemy
+
+from sky import sky_logging
+from sky.skylet import constants
+
+logger = sky_logging.init_logger(__name__)
+
+DB_INIT_LOCK_TIMEOUT_SECONDS = 10
+
+GLOBAL_USER_STATE_DB_NAME = 'state_db'
+GLOBAL_USER_STATE_VERSION = '011'
+GLOBAL_USER_STATE_LOCK_PATH = f'~/.sky/locks/.{GLOBAL_USER_STATE_DB_NAME}.lock'
+
+SPOT_JOBS_DB_NAME = 'spot_jobs_db'
+SPOT_JOBS_VERSION = '007'
+SPOT_JOBS_LOCK_PATH = f'~/.sky/locks/.{SPOT_JOBS_DB_NAME}.lock'
+
+SERVE_DB_NAME = 'serve_db'
+SERVE_VERSION = '002'
+SERVE_LOCK_PATH = f'~/.sky/locks/.{SERVE_DB_NAME}.lock'
+
+SKYPILOT_CONFIG_DB_NAME = 'sky_config_db'
+SKYPILOT_CONFIG_VERSION = '001'
+SKYPILOT_CONFIG_LOCK_PATH = f'~/.sky/locks/.{SKYPILOT_CONFIG_DB_NAME}.lock'
+
+KV_CACHE_DB_NAME = 'kv_cache_db'
+KV_CACHE_VERSION = '001'
+KV_CACHE_LOCK_PATH = f'~/.sky/locks/.{KV_CACHE_DB_NAME}.lock'
+
+
+@contextlib.contextmanager
+def db_lock(db_name: str):
+    lock_path = os.path.expanduser(f'~/.sky/locks/.{db_name}.lock')
+    try:
+        with filelock.FileLock(lock_path, timeout=DB_INIT_LOCK_TIMEOUT_SECONDS):
+            yield
+    except filelock.Timeout as e:
+        raise RuntimeError(f'Failed to initialize database due to a timeout '
+                           f'when trying to acquire the lock at '
+                           f'{lock_path}. '
+                           'Please try again or manually remove the lock '
+                           f'file if you believe it is stale.') from e
+
+
+def get_alembic_config(engine: sqlalchemy.engine.Engine, section: str):
+    """Get Alembic configuration for the given section"""
+    # From sky/utils/db/migration_utils.py -> sky/setup_files/alembic.ini
+    alembic_ini_path = os.path.join(
+        os.path.dirname(os.path.dirname(os.path.dirname(__file__))),
+        'setup_files', 'alembic.ini')
+    alembic_cfg = Config(alembic_ini_path, ini_section=section)
+
+    # Override the database URL to match SkyPilot's current connection
+    # Use render_as_string to get the full URL with password
+    url = engine.url.render_as_string(hide_password=False)
+    # Replace % with %% to escape the % character in the URL
+    # set_section_option uses variable interpolation, which treats % as a
+    # special character.
+    # any '%' symbol not used for interpolation needs to be escaped.
+    url = url.replace('%', '%%')
+    alembic_cfg.set_section_option(section, 'sqlalchemy.url', url)
+
+    return alembic_cfg
+
+
+def needs_upgrade(engine: sqlalchemy.engine.Engine, section: str,
+                  target_revision: str):
+    """Check if the database needs to be upgraded.
+
+    Args:
+        engine: SQLAlchemy engine for the database
+        section: Alembic section to upgrade (e.g., 'state_db' or 'spot_jobs_db')
+        target_revision: Target revision to upgrade to (e.g., '001')
+    """
+    current_rev = None
+
+    # get alembic config for the given section
+    alembic_config = get_alembic_config(engine, section)
+    version_table = alembic_config.get_section_option(
+        alembic_config.config_ini_section, 'version_table', 'alembic_version')
+
+    with engine.connect() as connection:
+        context = migration.MigrationContext.configure(
+            connection, opts={'version_table': version_table})
+        current_rev = context.get_current_revision()
+
+    target_rev_num = int(target_revision)
+    if current_rev is None:
+        if os.environ.get(constants.ENV_VAR_IS_SKYPILOT_SERVER) is not None:
+            logger.debug(f'{section} database currently uninitialized, '
+                         f'targeting revision {target_rev_num}')
+        return True
+
+    # Compare revisions - assuming they are numeric strings like '001', '002'
+    current_rev_num = int(current_rev)
+    if (current_rev_num < target_rev_num and
+            os.environ.get(constants.ENV_VAR_IS_SKYPILOT_SERVER) is not None):
+        logger.debug(
+            f'{section} database currently at revision {current_rev_num}, '
+            f'targeting revision {target_rev_num}')
+
+    return current_rev_num < target_rev_num
+
+
+def safe_alembic_upgrade(engine: sqlalchemy.engine.Engine, section: str,
+                         target_revision: str):
+    """Upgrade the database if needed. Uses a file lock to ensure
+    that only one process tries to upgrade the database at a time.
+
+    Args:
+        engine: SQLAlchemy engine for the database
+        section: Alembic section to upgrade (e.g., 'state_db' or 'spot_jobs_db')
+        target_revision: Target revision to upgrade to (e.g., '001')
+    """
+    # set alembic logger to warning level
+    alembic_logger = logging.getLogger('alembic')
+    alembic_logger.setLevel(logging.WARNING)
+
+    alembic_config = get_alembic_config(engine, section)
+
+    # only acquire lock if db needs upgrade
+    if needs_upgrade(engine, section, target_revision):
+        with db_lock(section):
+            # check again if db needs upgrade in case another
+            # process upgraded it while we were waiting for the lock
+            if needs_upgrade(engine, section, target_revision):
+                alembic_command.upgrade(alembic_config, target_revision)

sky/utils/directory_utils.py

@@ -0,0 +1,12 @@
+"""Directory utilities."""
+
+import os
+
+# This file is in '<project_root>/sky/utils/directory_utils.py'
+# So we need to go up 2 levels to get to the '<project_root>/sky' directory
+SKY_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+
+def get_sky_dir():
+    """Get the sky root directory."""
+    return SKY_DIR

sky/utils/env_options.py

@@ -19,8 +19,21 @@ class Options(enum.Enum):
     # will not be multiple identities, and skipping the check can increase
     # robustness.
     SKIP_CLOUD_IDENTITY_CHECK = ('SKYPILOT_SKIP_CLOUD_IDENTITY_CHECK', False)
+    # Internal: This environment variable is set to "true" by Buildkite
+    # agent when running tests. It is used to identify when SkyPilot is
+    # running in a Buildkite container environment, which requires special
+    # handling for networking between containers.
+    RUNNING_IN_BUILDKITE = ('BUILDKITE', False)
+    # Internal: This is used for testing to enable grpc for communication
+    # between the API server and the Skylet.
+    ENABLE_GRPC = ('SKYPILOT_ENABLE_GRPC', False)
+    # Allow all contexts for Kubernetes if allowed_contexts is not set in
+    # config.
+    ALLOW_ALL_KUBERNETES_CONTEXTS = ('SKYPILOT_ALLOW_ALL_KUBERNETES_CONTEXTS',
+                                     False)
 
     def __init__(self, env_var: str, default: bool) -> None:
+        super().__init__()
         self.env_var = env_var
         self.default = default