skypilot-nightly 1.0.0.dev20250718__py3-none-any.whl → 1.0.0.dev20250723__py3-none-any.whl

sky/utils/common_utils.py

@@ -369,6 +369,83 @@ def get_pretty_entrypoint_cmd() -> str:
     return ' '.join(argv)
 
 
+def read_last_n_lines(file_path: str,
+                      n: int,
+                      chunk_size: int = 8192,
+                      encoding: str = 'utf-8',
+                      errors: str = 'replace') -> List[str]:
+    """Read the last N lines of a file.
+
+    Args:
+        file_path: Path to the file to read.
+        n: Number of lines to read from the end of the file.
+        chunk_size: Size of chunks in bytes.
+        encoding: Encoding to use when decoding binary chunks.
+        errors: Error handling for decode errors (e.g., 'replace', 'ignore').
+
+    Returns:
+        A list of the last N lines, preserving newlines where applicable.
+    """
+
+    assert n >= 0, f'n must be non-negative. Got {n}'
+    assert chunk_size > 0, f'chunk_size must be positive. Got {chunk_size}'
+    assert os.path.exists(file_path), f'File not found: {file_path}'
+
+    if n == 0:
+        return []
+
+    try:
+        with open(file_path, 'rb') as f:
+            # Start reading from the end of the file
+            f.seek(0, os.SEEK_END)
+            file_size = f.tell()
+            if file_size == 0:
+                return []
+
+            pos = file_size
+            lines_found = 0
+            chunks = []
+
+            # Read backwards in chunks until we've found at least n newlines
+            while pos > 0 and lines_found <= n:
+                read_size = min(chunk_size, pos)
+                pos -= read_size
+                f.seek(pos)
+                chunk = f.read(read_size)
+                chunks.append(chunk)
+                lines_found += chunk.count(b'\n')
+
+            # Combine all chunks in reverse order since we read backwards
+            full_bytes = b''.join(reversed(chunks))
+
+            # Split by newline byte. Note: this handles '\n' endings.
+            all_lines = full_bytes.split(b'\n')
+
+            # Handle edge case: if file ends with a newline, last element is b''
+            if all_lines and all_lines[-1] == b'':
+                result_bytes = all_lines[-n - 1:-1]
+            else:
+                result_bytes = all_lines[-n:]
+
+            # Decode each line and normalize CR/LF endings
+            decoded_lines = [
+                line.decode(encoding, errors=errors).rstrip('\r') + '\n'
+                for line in result_bytes[:-1]
+            ]
+
+            # Decode the final line — only add newline if it was present
+            last_line = result_bytes[-1].decode(encoding,
+                                                errors=errors).rstrip('\r')
+            decoded_lines.append(last_line)
+
+            return decoded_lines
+
+    except OSError as e:
+        with ux_utils.print_exception_no_traceback():
+            raise RuntimeError(
+                f'Failed to read last {n} lines from {file_path}: {e}') from e
+
+
 def _redact_secrets_values(argv: List[str]) -> List[str]:
     """Redact sensitive values from --secret arguments.
 
@@ -485,8 +562,9 @@ def read_yaml_all(path: str) -> List[Dict[str, Any]]:
         return read_yaml_all_str(f.read())
 
 
-def dump_yaml(path: str, config: Union[List[Dict[str, Any]],
-                                       Dict[str, Any]]) -> None:
+def dump_yaml(path: str,
+              config: Union[List[Dict[str, Any]], Dict[str, Any]],
+              blank: bool = False) -> None:
     """Dumps a YAML file.
 
     Args:
@@ -494,7 +572,11 @@ def dump_yaml(path: str, config: Union[List[Dict[str, Any]],
         config: the configuration to dump.
     """
     with open(path, 'w', encoding='utf-8') as f:
-        f.write(dump_yaml_str(config))
+        contents = dump_yaml_str(config)
+        if blank and isinstance(config, dict) and len(config) == 0:
+            # when dumping to yaml, an empty dict will go in as {}.
+            contents = ''
+        f.write(contents)
 
 
 def dump_yaml_str(config: Union[List[Dict[str, Any]], Dict[str, Any]]) -> str:

sky/utils/config_utils.py

@@ -248,6 +248,8 @@ def get_cloud_config_value_from_dict(
     region_key = None
     if cloud == 'kubernetes':
         region_key = 'context_configs'
+    if cloud == 'nebius':
+        region_key = 'region_configs'
 
     per_context_config = None
     if region is not None and region_key is not None:
@@ -255,6 +257,19 @@ def get_cloud_config_value_from_dict(
             keys=(cloud, region_key, region) + keys,
             default_value=None,
             override_configs=override_configs)
+        if not per_context_config and cloud == 'nebius':
+            # TODO (kyuds): Backward compatibility, remove after 0.11.0.
+            per_context_config = input_config.get_nested(
+                keys=(cloud, region) + keys,
+                default_value=None,
+                override_configs=override_configs)
+            if per_context_config is not None:
+                logger.info(
+                    'Nebius configuration is using the legacy format. \n'
+                    'This format will be deprecated after 0.11.0, refer to '
+                    '`https://docs.skypilot.co/en/latest/reference/config.html#nebius` '  # pylint: disable=line-too-long
+                    'for the new format. Please use `region_configs` to specify region specific configuration.'
+                )
     # if no override found for specified region
     general_config = input_config.get_nested(keys=(cloud,) + keys,
                                              default_value=default_value,

sky/utils/{db_utils.py → db/db_utils.py}

@@ -9,6 +9,9 @@ from typing import Any, Callable, Optional
 import sqlalchemy
 from sqlalchemy import exc as sqlalchemy_exc
 
+from sky import sky_logging
+
+logger = sky_logging.init_logger(__name__)
 if typing.TYPE_CHECKING:
     from sqlalchemy.orm import Session
 
@@ -146,6 +149,62 @@ def add_column_to_table_sqlalchemy(
     session.commit()
 
 
+def add_column_to_table_alembic(
+    table_name: str,
+    column_name: str,
+    column_type: sqlalchemy.types.TypeEngine,
+    server_default: Optional[str] = None,
+    copy_from: Optional[str] = None,
+    value_to_replace_existing_entries: Optional[Any] = None,
+):
+    """Add a column to a table using Alembic operations.
+
+    This provides the same interface as add_column_to_table_sqlalchemy but
+    uses Alembic's connection context for proper migration support.
+
+    Args:
+        table_name: Name of the table to add column to
+        column_name: Name of the new column
+        column_type: SQLAlchemy column type
+        server_default: Server-side default value for the column
+        copy_from: Column name to copy values from (for existing rows)
+        value_to_replace_existing_entries: Default value for existing NULL
+            entries
+    """
+    from alembic import op  # pylint: disable=import-outside-toplevel
+
+    try:
+        # Create the column with server_default if provided
+        column = sqlalchemy.Column(column_name,
+                                   column_type,
+                                   server_default=server_default)
+        op.add_column(table_name, column)
+
+        # Handle data migration
+        if copy_from is not None:
+            op.execute(
+                sqlalchemy.text(
+                    f'UPDATE {table_name} SET {column_name} = {copy_from}'))
+
+        if value_to_replace_existing_entries is not None:
+            # Use parameterized query for safety
+            op.get_bind().execute(
+                sqlalchemy.text(f'UPDATE {table_name} '
+                                f'SET {column_name} = :replacement_value '
+                                f'WHERE {column_name} IS NULL'),
+                {'replacement_value': value_to_replace_existing_entries})
+    except sqlalchemy_exc.ProgrammingError as e:
+        if 'already exists' in str(e).lower():
+            pass  # Column already exists, that's fine
+        else:
+            raise
+    except sqlalchemy_exc.OperationalError as e:
+        if 'duplicate column name' in str(e).lower():
+            pass  # Column already exists, that's fine
+        else:
+            raise
+
+
 class SQLiteConn(threading.local):
     """Thread-local connection to the sqlite3 database."""
 

sky/utils/db/migration_utils.py

@@ -0,0 +1,93 @@
+"""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
+
+DB_INIT_LOCK_TIMEOUT_SECONDS = 10
+
+GLOBAL_USER_STATE_DB_NAME = 'state_db'
+GLOBAL_USER_STATE_VERSION = '001'
+GLOBAL_USER_STATE_LOCK_PATH = '~/.sky/locks/.state_db.lock'
+
+SPOT_JOBS_DB_NAME = 'spot_jobs_db'
+SPOT_JOBS_VERSION = '001'
+SPOT_JOBS_LOCK_PATH = '~/.sky/locks/.spot_jobs_db.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"""
+    # Use the alembic.ini file from setup_files (included in wheel)
+    # 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)
+    alembic_cfg.set_section_option(section, 'sqlalchemy.url', url)
+
+    return alembic_cfg
+
+
+def safe_alembic_upgrade(engine: sqlalchemy.engine.Engine,
+                         alembic_config: Config, target_revision: str):
+    """Only upgrade if current version is older than target.
+
+    This handles the case where a database was created with a newer version of
+    the code and we're now running older code. Since our migrations are purely
+    additive, it's safe to run a newer database with older code.
+
+    Args:
+        engine: SQLAlchemy engine for the database
+        alembic_config: Alembic configuration object
+        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)
+
+    current_rev = None
+
+    # Get the current revision from the database
+    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()
+
+    if current_rev is None:
+        alembic_command.upgrade(alembic_config, target_revision)
+        return
+
+    # Compare revisions - assuming they are numeric strings like '001', '002'
+    current_rev_num = int(current_rev)
+    target_rev_num = int(target_revision)
+
+    # only upgrade if current revision is older than target revision
+    if current_rev_num < target_rev_num:
+        alembic_command.upgrade(alembic_config, target_revision)

sky/utils/locks.py

@@ -0,0 +1,319 @@
+"""Lock for SkyPilot.
+
+This module provides an abstraction for locking that can use
+either local file locks or database-based distributed locks.
+"""
+import abc
+import hashlib
+import logging
+import os
+import time
+from typing import Any, Optional
+
+import filelock
+import sqlalchemy
+
+from sky import global_user_state
+from sky.skylet import constants
+from sky.utils import common_utils
+from sky.utils.db import db_utils
+
+logger = logging.getLogger(__name__)
+
+
+class LockTimeout(RuntimeError):
+    """Raised when a lock acquisition times out."""
+    pass
+
+
+class AcquireReturnProxy:
+    """A context manager that releases the lock when exiting.
+
+    This proxy is returned by acquire() and ensures proper cleanup
+    when used in a with statement.
+    """
+
+    def __init__(self, lock: 'DistributedLock') -> None:
+        self.lock = lock
+
+    def __enter__(self) -> 'DistributedLock':
+        return self.lock
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.lock.release()
+
+
+class DistributedLock(abc.ABC):
+    """Abstract base class for a distributed lock.
+
+    Provides a context manager interface for acquiring and releasing locks
+    that can work across multiple processes and potentially multiple machines.
+    """
+
+    def __init__(self,
+                 lock_id: str,
+                 timeout: Optional[float] = None,
+                 poll_interval: float = 0.1):
+        """Initialize the lock.
+
+        Args:
+            lock_id: Unique identifier for the lock.
+            timeout: Maximum time to wait for lock acquisition.
+                If None, wait indefinitely.
+            poll_interval: Interval in seconds to poll for lock acquisition.
+        """
+        self.lock_id = lock_id
+        self.timeout = timeout
+        self.poll_interval = poll_interval
+
+    @abc.abstractmethod
+    def acquire(self, blocking: bool = True) -> AcquireReturnProxy:
+        """Acquire the lock.
+
+        Args:
+            blocking: If True, block until lock is acquired or timeout.
+                     If False, return immediately.
+
+        Returns:
+            AcquireReturnProxy that can be used as a context manager.
+
+        Raises:
+            LockTimeout: If lock cannot be acquired.
+        """
+        pass
+
+    @abc.abstractmethod
+    def release(self) -> None:
+        """Release the lock."""
+        pass
+
+    @abc.abstractmethod
+    def force_unlock(self) -> None:
+        """Force unlock the lock if it is acquired."""
+        pass
+
+    @abc.abstractmethod
+    def is_locked(self) -> bool:
+        """Check if the lock is acquired."""
+        pass
+
+    def __enter__(self) -> 'DistributedLock':
+        """Context manager entry."""
+        self.acquire()
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Context manager exit."""
+        self.release()
+
+
+class FileLock(DistributedLock):
+    """A wrapper around filelock.FileLock.
+
+    This implements a distributed lock that works across multiple processes
+    when they share the same filesystem.
+    """
+
+    def __init__(self,
+                 lock_id: str,
+                 timeout: Optional[float] = None,
+                 poll_interval: float = 0.1):
+        """Initialize the file lock.
+
+        Args:
+            lock_id: Unique identifier for the lock.
+            timeout: Maximum time to wait for lock acquisition.
+            poll_interval: Interval in seconds to poll for lock acquisition.
+        """
+        super().__init__(lock_id, timeout, poll_interval)
+        os.makedirs(constants.SKY_LOCKS_DIR, exist_ok=True)
+        self.lock_path = os.path.join(constants.SKY_LOCKS_DIR,
+                                      f'.{lock_id}.lock')
+        if timeout is None:
+            timeout = -1
+        self._filelock: filelock.FileLock = filelock.FileLock(self.lock_path,
+                                                              timeout=timeout)
+
+    def acquire(self, blocking: bool = True) -> AcquireReturnProxy:
+        """Acquire the file lock."""
+        try:
+            acquired = self._filelock.acquire(blocking=blocking)
+            if not acquired:
+                raise LockTimeout(f'Failed to acquire file lock {self.lock_id}')
+            return AcquireReturnProxy(self)
+        except filelock.Timeout as e:
+            raise LockTimeout(
+                f'Failed to acquire file lock {self.lock_id}') from e
+
+    def release(self) -> None:
+        """Release the file lock."""
+        self._filelock.release()
+
+    def force_unlock(self) -> None:
+        """Force unlock the file lock."""
+        common_utils.remove_file_if_exists(self.lock_path)
+
+    def is_locked(self) -> bool:
+        return self._filelock.is_locked()
+
+
+class PostgresLock(DistributedLock):
+    """PostgreSQL advisory lock implementation.
+
+    Uses PostgreSQL advisory locks to implement distributed locking
+    that works across multiple machines sharing the same database.
+    Reference:
+    https://www.postgresql.org/docs/current/explicit-locking.html
+    #ADVISORY-LOCKS
+    """
+
+    def __init__(self,
+                 lock_id: str,
+                 timeout: Optional[float] = None,
+                 poll_interval: float = 1):
+        """Initialize the postgres lock.
+
+        Args:
+            lock_id: Unique identifier for the lock.
+            timeout: Maximum time to wait for lock acquisition.
+            poll_interval: Interval in seconds to poll for lock acquisition,
+                default to 1 second to avoid storming the database.
+        """
+        super().__init__(lock_id, timeout, poll_interval)
+        # Convert string lock_id to integer for postgres advisory locks
+        self._lock_key = self._string_to_lock_key(lock_id)
+        self._acquired = False
+        self._connection: Optional[sqlalchemy.pool.PoolProxiedConnection] = None
+
+    def _string_to_lock_key(self, s: str) -> int:
+        """Convert string to a 64-bit integer for advisory lock key."""
+        hash_digest = hashlib.sha256(s.encode('utf-8')).digest()
+        # Take first 8 bytes and convert to int, ensure positive 64-bit
+        return int.from_bytes(hash_digest[:8], 'big') & ((1 << 63) - 1)
+
+    def _get_connection(self) -> sqlalchemy.pool.PoolProxiedConnection:
+        """Get database connection."""
+        engine = global_user_state.initialize_and_get_db()
+        if engine.dialect.name != db_utils.SQLAlchemyDialect.POSTGRESQL.value:
+            raise ValueError('PostgresLock requires PostgreSQL database. '
+                             f'Current dialect: {engine.dialect.name}')
+        return engine.raw_connection()
+
+    def acquire(self, blocking: bool = True) -> AcquireReturnProxy:
+        """Acquire the postgres advisory lock."""
+        if self._acquired:
+            return AcquireReturnProxy(self)
+
+        self._connection = self._get_connection()
+        cursor = self._connection.cursor()
+
+        start_time = time.time()
+
+        try:
+            while True:
+                cursor.execute('SELECT pg_try_advisory_lock(%s)',
+                               (self._lock_key,))
+                result = cursor.fetchone()[0]
+
+                if result:
+                    self._acquired = True
+                    return AcquireReturnProxy(self)
+
+                if not blocking:
+                    raise LockTimeout(
+                        f'Failed to immediately acquire postgres lock '
+                        f'{self.lock_id}')
+
+                if (self.timeout is not None and
+                        time.time() - start_time > self.timeout):
+                    raise LockTimeout(
+                        f'Failed to acquire postgres lock {self.lock_id} '
+                        f'within {self.timeout} seconds')
+
+                time.sleep(self.poll_interval)
+
+        except Exception:
+            if self._connection:
+                self._connection.close()
+                self._connection = None
+            raise
+
+    def release(self) -> None:
+        """Release the postgres advisory lock."""
+        if not self._acquired or not self._connection:
+            return
+
+        try:
+            cursor = self._connection.cursor()
+            cursor.execute('SELECT pg_advisory_unlock(%s)', (self._lock_key,))
+            self._connection.commit()
+            self._acquired = False
+        finally:
+            if self._connection:
+                self._connection.close()
+                self._connection = None
+
+    def force_unlock(self) -> None:
+        """Force unlock the postgres advisory lock."""
+        try:
+            if not self._connection:
+                self._connection = self._get_connection()
+            cursor = self._connection.cursor()
+            cursor.execute('SELECT pg_advisory_unlock(%s)', (self._lock_key,))
+            self._connection.commit()
+        except Exception as e:
+            raise RuntimeError(
+                f'Failed to force unlock postgres lock {self.lock_id}: {e}'
+            ) from e
+        finally:
+            if self._connection:
+                self._connection.close()
+                self._connection = None
+
+    def is_locked(self) -> bool:
+        """Check if the postgres advisory lock is acquired."""
+        return self._acquired
+
+
+def get_lock(lock_id: str,
+             timeout: Optional[float] = None,
+             lock_type: Optional[str] = None,
+             poll_interval: Optional[float] = None) -> DistributedLock:
+    """Create a distributed lock instance.
+
+    Args:
+        lock_id: Unique identifier for the lock.
+        timeout: Maximum time seconds to wait for lock acquisition,
+                 None means wait indefinitely.
+        lock_type: Type of lock to create ('filelock' or 'postgres').
+                   If None, auto-detect based on database configuration.
+
+    Returns:
+        DistributedLock instance.
+    """
+    if lock_type is None:
+        lock_type = _detect_lock_type()
+
+    if lock_type == 'postgres':
+        if poll_interval is None:
+            return PostgresLock(lock_id, timeout)
+        return PostgresLock(lock_id, timeout, poll_interval)
+    elif lock_type == 'filelock':
+        if poll_interval is None:
+            return FileLock(lock_id, timeout)
+        return FileLock(lock_id, timeout, poll_interval)
+    else:
+        raise ValueError(f'Unknown lock type: {lock_type}')
+
+
+def _detect_lock_type() -> str:
+    """Auto-detect the appropriate lock type based on configuration."""
+    try:
+        engine = global_user_state.initialize_and_get_db()
+        if engine.dialect.name == db_utils.SQLAlchemyDialect.POSTGRESQL.value:
+            return 'postgres'
+    except Exception:  # pylint: disable=broad-except
+        # Fall back to filelock if database detection fails
+        pass
+
+    return 'filelock'