skypilot-nightly 1.0.0.dev20250717__py3-none-any.whl → 1.0.0.dev20250720__py3-none-any.whl

sky/global_user_state.py

@@ -12,12 +12,12 @@ import os
 import pathlib
 import pickle
 import re
-import threading
 import time
 import typing
 from typing import Any, Dict, List, Optional, Set, Tuple
 import uuid
 
+from alembic import command as alembic_command
 import sqlalchemy
 from sqlalchemy import exc as sqlalchemy_exc
 from sqlalchemy import orm
@@ -32,9 +32,10 @@ from sky import skypilot_config
 from sky.skylet import constants
 from sky.utils import common_utils
 from sky.utils import context_utils
-from sky.utils import db_utils
 from sky.utils import registry
 from sky.utils import status_lib
+from sky.utils.db import db_utils
+from sky.utils.db import migration_utils
 
 if typing.TYPE_CHECKING:
     from sky import backends
@@ -48,7 +49,6 @@ _ENABLED_CLOUDS_KEY_PREFIX = 'enabled_clouds_'
 _ALLOWED_CLOUDS_KEY_PREFIX = 'allowed_clouds_'
 
 _SQLALCHEMY_ENGINE: Optional[sqlalchemy.engine.Engine] = None
-_DB_INIT_LOCK = threading.Lock()
 
 Base = declarative.declarative_base()
 
@@ -238,152 +238,20 @@ def create_table(engine: sqlalchemy.engine.Engine):
             # 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.
 
-    # Create tables if they don't exist
-    db_utils.add_tables_to_db_sqlalchemy(Base.metadata, engine)
-
-    # For backward compatibility.
-    # TODO(zhwu): Remove this function after all users have migrated to
-    # the latest version of SkyPilot.
-    with orm.Session(engine) as session:
-        # Add autostop column to clusters table
-        db_utils.add_column_to_table_sqlalchemy(session,
-                                                'clusters',
-                                                'autostop',
-                                                sqlalchemy.Integer(),
-                                                default_statement='DEFAULT -1')
-
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'clusters',
-            'metadata',
-            sqlalchemy.Text(),
-            default_statement='DEFAULT \'{}\'')
-
-        db_utils.add_column_to_table_sqlalchemy(session,
-                                                'clusters',
-                                                'to_down',
-                                                sqlalchemy.Integer(),
-                                                default_statement='DEFAULT 0')
-
-        # The cloud identity that created the cluster.
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'clusters',
-            'owner',
-            sqlalchemy.Text(),
-            default_statement='DEFAULT NULL')
-
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'clusters',
-            'cluster_hash',
-            sqlalchemy.Text(),
-            default_statement='DEFAULT NULL')
-
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'clusters',
-            'storage_mounts_metadata',
-            sqlalchemy.LargeBinary(),
-            default_statement='DEFAULT NULL')
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'clusters',
-            'cluster_ever_up',
-            sqlalchemy.Integer(),
-            default_statement='DEFAULT 0',
-            # Set the value to 1 so that all the existing clusters before #2977
-            # are considered as ever up, i.e:
-            #   existing cluster's default (null) -> 1;
-            #   new cluster's default -> 0;
-            # This is conservative for the existing clusters: even if some INIT
-            # clusters were never really UP, setting it to 1 means they won't be
-            # auto-deleted during any failover.
-            value_to_replace_existing_entries=1)
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'clusters',
-            'status_updated_at',
-            sqlalchemy.Integer(),
-            default_statement='DEFAULT NULL')
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'clusters',
-            'user_hash',
-            sqlalchemy.Text(),
-            default_statement='DEFAULT NULL',
-            value_to_replace_existing_entries=common_utils.get_current_user(
-            ).id)
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'clusters',
-            'config_hash',
-            sqlalchemy.Text(),
-            default_statement='DEFAULT NULL')
-
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'cluster_history',
-            'user_hash',
-            sqlalchemy.Text(),
-            default_statement='DEFAULT NULL')
-
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'clusters',
-            'workspace',
-            sqlalchemy.Text(),
-            default_statement='DEFAULT \'default\'',
-            value_to_replace_existing_entries=constants.
-            SKYPILOT_DEFAULT_WORKSPACE)
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'clusters',
-            'last_creation_yaml',
-            sqlalchemy.Text(),
-            default_statement='DEFAULT NULL',
-        )
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'clusters',
-            'last_creation_command',
-            sqlalchemy.Text(),
-            default_statement='DEFAULT NULL')
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'users',
-            'password',
-            sqlalchemy.Text(),
-            default_statement='DEFAULT NULL')
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'users',
-            'created_at',
-            sqlalchemy.Integer(),
-            default_statement='DEFAULT NULL')
-
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'cluster_history',
-            'last_creation_yaml',
-            sqlalchemy.Text(),
-            default_statement='DEFAULT NULL')
-
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'cluster_history',
-            'last_creation_command',
-            sqlalchemy.Text(),
-            default_statement='DEFAULT NULL')
-
-        session.commit()
+    # Get alembic config for state db and run migrations
+    alembic_config = migration_utils.get_alembic_config(
+        engine, migration_utils.GLOBAL_USER_STATE_DB_NAME)
+    # pylint: disable=line-too-long
+    alembic_config.config_ini_section = migration_utils.GLOBAL_USER_STATE_DB_NAME
+    alembic_command.upgrade(alembic_config,
+                            migration_utils.GLOBAL_USER_STATE_VERSION)
 
 
 def initialize_and_get_db() -> sqlalchemy.engine.Engine:
     global _SQLALCHEMY_ENGINE
     if _SQLALCHEMY_ENGINE is not None:
         return _SQLALCHEMY_ENGINE
-    with _DB_INIT_LOCK:
+    with migration_utils.db_lock(migration_utils.GLOBAL_USER_STATE_DB_NAME):
         if _SQLALCHEMY_ENGINE is None:
             conn_string = None
             if os.environ.get(constants.ENV_VAR_IS_SKYPILOT_SERVER) is not None:
@@ -520,6 +388,7 @@ def get_user(user_id: str) -> Optional[models.User]:
                        created_at=row.created_at)
 
 
+@_init_db
 def get_user_by_name(username: str) -> List[models.User]:
     with orm.Session(_SQLALCHEMY_ENGINE) as session:
         rows = session.query(user_table).filter_by(name=username).all()
@@ -533,6 +402,7 @@ def get_user_by_name(username: str) -> List[models.User]:
     ]
 
 
+@_init_db
 def delete_user(user_id: str) -> None:
     with orm.Session(_SQLALCHEMY_ENGINE) as session:
         session.query(user_table).filter_by(id=user_id).delete()

sky/jobs/client/sdk.py

@@ -192,7 +192,7 @@ def cancel(
 
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
-@rest.retry_on_server_unavailable()
+@rest.retry_transient_errors()
 def tail_logs(name: Optional[str] = None,
               job_id: Optional[int] = None,
               follow: bool = True,

sky/jobs/server/core.py

@@ -147,6 +147,18 @@ def launch(
         None if dryrun.
     """
     entrypoint = task
+    # using hasattr instead of isinstance to avoid importing sky
+    if hasattr(task, 'metadata'):
+        metadata = task.metadata
+    else:
+        # we are a Dag, not a Task
+        if len(task.tasks) == 1:
+            metadata = task.tasks[0].metadata
+        else:
+            # doesn't make sense to have a git commit since there might be
+            # different metadatas for each task
+            metadata = {}
+
     dag_uuid = str(uuid.uuid4().hex[:4])
     dag = dag_utils.convert_entrypoint_to_dag(entrypoint)
     dag.resolve_and_validate_volumes()
@@ -311,6 +323,8 @@ def launch(
         controller_task.set_resources(controller_resources)
 
         controller_task.managed_job_dag = dag
+        # pylint: disable=protected-access
+        controller_task._metadata = metadata
 
         logger.info(
             f'{colorama.Fore.YELLOW}'

sky/jobs/state.py

@@ -6,11 +6,11 @@ import functools
 import json
 import os
 import pathlib
-import threading
 import time
 import typing
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
+from alembic import command as alembic_command
 import colorama
 import sqlalchemy
 from sqlalchemy import exc as sqlalchemy_exc
@@ -24,7 +24,8 @@ from sky import sky_logging
 from sky import skypilot_config
 from sky.skylet import constants
 from sky.utils import common_utils
-from sky.utils import db_utils
+from sky.utils.db import db_utils
+from sky.utils.db import migration_utils
 
 if typing.TYPE_CHECKING:
     from sqlalchemy.engine import row
@@ -36,7 +37,6 @@ CallbackType = Callable[[str], None]
 logger = sky_logging.init_logger(__name__)
 
 _SQLALCHEMY_ENGINE: Optional[sqlalchemy.engine.Engine] = None
-_DB_INIT_LOCK = threading.Lock()
 
 Base = declarative.declarative_base()
 
@@ -130,97 +130,18 @@ def create_table(engine: sqlalchemy.engine.Engine):
             # 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.
 
-    # Create tables if they don't exist
-    db_utils.add_tables_to_db_sqlalchemy(Base.metadata, engine)
-
-    # Backward compatibility: add columns that not exist in older databases
-    with orm.Session(engine) as session:
-        db_utils.add_column_to_table_sqlalchemy(session, 'spot',
-                                                'failure_reason',
-                                                sqlalchemy.Text())
-        db_utils.add_column_to_table_sqlalchemy(session,
-                                                'spot',
-                                                'spot_job_id',
-                                                sqlalchemy.Integer(),
-                                                copy_from='job_id')
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'spot',
-            'task_id',
-            sqlalchemy.Integer(),
-            default_statement='DEFAULT 0',
-            value_to_replace_existing_entries=0)
-        db_utils.add_column_to_table_sqlalchemy(session,
-                                                'spot',
-                                                'task_name',
-                                                sqlalchemy.Text(),
-                                                copy_from='job_name')
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'spot',
-            'specs',
-            sqlalchemy.Text(),
-            value_to_replace_existing_entries=json.dumps({
-                'max_restarts_on_errors': 0,
-            }))
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'spot',
-            'local_log_file',
-            sqlalchemy.Text(),
-            default_statement='DEFAULT NULL')
-
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'spot',
-            'metadata',
-            sqlalchemy.Text(),
-            default_statement='DEFAULT \'{}\'',
-            value_to_replace_existing_entries='{}')
-
-        db_utils.add_column_to_table_sqlalchemy(session, 'job_info',
-                                                'schedule_state',
-                                                sqlalchemy.Text())
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'job_info',
-            'controller_pid',
-            sqlalchemy.Integer(),
-            default_statement='DEFAULT NULL')
-        db_utils.add_column_to_table_sqlalchemy(session, 'job_info',
-                                                'dag_yaml_path',
-                                                sqlalchemy.Text())
-        db_utils.add_column_to_table_sqlalchemy(session, 'job_info',
-                                                'env_file_path',
-                                                sqlalchemy.Text())
-        db_utils.add_column_to_table_sqlalchemy(session, 'job_info',
-                                                'user_hash', sqlalchemy.Text())
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'job_info',
-            'workspace',
-            sqlalchemy.Text(),
-            default_statement='DEFAULT NULL',
-            value_to_replace_existing_entries='default')
-        db_utils.add_column_to_table_sqlalchemy(
-            session,
-            'job_info',
-            'priority',
-            sqlalchemy.Integer(),
-            value_to_replace_existing_entries=constants.DEFAULT_PRIORITY)
-        db_utils.add_column_to_table_sqlalchemy(session, 'job_info',
-                                                'entrypoint', sqlalchemy.Text())
-        db_utils.add_column_to_table_sqlalchemy(session, 'job_info',
-                                                'original_user_yaml_path',
-                                                sqlalchemy.Text())
-        session.commit()
+    # Get alembic config for spot jobs db and run migrations
+    alembic_config = migration_utils.get_alembic_config(
+        engine, migration_utils.SPOT_JOBS_DB_NAME)
+    alembic_config.config_ini_section = migration_utils.SPOT_JOBS_DB_NAME
+    alembic_command.upgrade(alembic_config, migration_utils.SPOT_JOBS_VERSION)
 
 
 def initialize_and_get_db() -> sqlalchemy.engine.Engine:
     global _SQLALCHEMY_ENGINE
     if _SQLALCHEMY_ENGINE is not None:
         return _SQLALCHEMY_ENGINE
-    with _DB_INIT_LOCK:
+    with migration_utils.db_lock(migration_utils.SPOT_JOBS_DB_NAME):
         if _SQLALCHEMY_ENGINE is None:
             conn_string = None
             if os.environ.get(constants.ENV_VAR_IS_SKYPILOT_SERVER) is not None:

sky/jobs/utils.py

@@ -67,6 +67,9 @@ JOB_STARTED_STATUS_CHECK_GAP_SECONDS = 5
 
 _LOG_STREAM_CHECK_CONTROLLER_GAP_SECONDS = 5
 
+_JOB_STATUS_FETCH_MAX_RETRIES = 3
+_JOB_K8S_TRANSIENT_NW_MSG = 'Unable to connect to the server: dial tcp'
+
 _JOB_WAITING_STATUS_MESSAGE = ux_utils.spinner_message(
     'Waiting for task to start[/]'
     '{status_str}. It may take a few minutes.\n'
@@ -250,19 +253,31 @@ def get_job_status(backend: 'backends.CloudVmRayBackend',
         logger.info(f'Cluster {cluster_name} not found.')
         return None
     assert isinstance(handle, backends.CloudVmRayResourceHandle), handle
-    status = None
-    try:
-        logger.info('=== Checking the job status... ===')
-        statuses = backend.get_job_status(handle, stream_logs=False)
-        status = list(statuses.values())[0]
-        if status is None:
-            logger.info('No job found.')
-        else:
-            logger.info(f'Job status: {status}')
-    except exceptions.CommandError:
-        logger.info('Failed to connect to the cluster.')
-    logger.info('=' * 34)
-    return status
+    for i in range(_JOB_STATUS_FETCH_MAX_RETRIES):
+        try:
+            logger.info('=== Checking the job status... ===')
+            statuses = backend.get_job_status(handle, stream_logs=False)
+            status = list(statuses.values())[0]
+            if status is None:
+                logger.info('No job found.')
+            else:
+                logger.info(f'Job status: {status}')
+            logger.info('=' * 34)
+            return status
+        except exceptions.CommandError as e:
+            # Retry on k8s transient network errors. This is useful when using
+            # coreweave which may have transient network issue sometimes.
+            if (e.detailed_reason is not None and
+                    _JOB_K8S_TRANSIENT_NW_MSG in e.detailed_reason):
+                logger.info('Failed to connect to the cluster. Retrying '
+                            f'({i + 1}/{_JOB_STATUS_FETCH_MAX_RETRIES})...')
+                logger.info('=' * 34)
+                time.sleep(1)
+            else:
+                logger.info(f'Failed to get job status: {e.detailed_reason}')
+                logger.info('=' * 34)
+                return None
+    return None
 
 
 def _controller_process_alive(pid: int, job_id: int) -> bool:

sky/schemas/db/README

@@ -0,0 +1,4 @@
+Migrations for sqlalchemy databases. Currently includes:
+global_user_state
+spot_jobs (managed jobs state)
+skypilot_config

sky/schemas/db/env.py

@@ -0,0 +1,90 @@
+"""Alembic environment configuration for state database migrations."""
+from logging.config import fileConfig
+
+from alembic import context
+from sqlalchemy import engine_from_config
+from sqlalchemy import pool
+
+# this is the Alembic Config object, which provides
+# access to the values within the .ini file in use.
+config = context.config
+
+# NOTE: We intentionally disable Alembic's logging configuration to prevent
+# it from overriding SkyPilot's logging setup. Alembic's fileConfig() call
+# globally reconfigures Python's logging system, which can suppress SkyPilot's
+# output messages that tests expect to see.
+#
+# Original code (now disabled):
+if config.config_file_name is not None:
+    fileConfig(config.config_file_name, disable_existing_loggers=False)
+
+# add your model's MetaData object here
+# for 'autogenerate' support
+# from myapp import mymodel
+# target_metadata = mymodel.Base.metadata
+target_metadata = None
+
+# other values from the config, defined by the needs of env.py,
+# can be acquired:
+# my_important_option = config.get_main_option("my_important_option")
+# ... etc.
+
+
+def run_migrations_offline() -> None:
+    """Run migrations in 'offline' mode.
+
+    This configures the context with just a URL
+    and not an Engine, though an Engine is acceptable
+    here as well.  By skipping the Engine creation
+    we don't even need a DBAPI to be available.
+
+    Calls to context.execute() here emit the given string to the
+    script output.
+
+    """
+    url = config.get_main_option('sqlalchemy.url')
+    version_table = config.get_section_option(config.config_ini_section,
+                                              'version_table',
+                                              'alembic_version')
+    context.configure(
+        url=url,
+        target_metadata=target_metadata,
+        literal_binds=True,
+        dialect_opts={'paramstyle': 'named'},
+        version_table=version_table,
+    )
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+def run_migrations_online() -> None:
+    """Run migrations in 'online' mode.
+
+    In this scenario we need to create an Engine
+    and associate a connection with the context.
+
+    """
+    connectable = engine_from_config(
+        config.get_section(config.config_ini_section, {}),
+        prefix='sqlalchemy.',
+        poolclass=pool.NullPool,
+    )
+    version_table = config.get_section_option(config.config_ini_section,
+                                              'version_table',
+                                              'alembic_version')
+    with connectable.connect() as connection:
+        context.configure(
+            connection=connection,
+            target_metadata=target_metadata,
+            version_table=version_table,
+        )
+
+        with context.begin_transaction():
+            context.run_migrations()
+
+
+if context.is_offline_mode():
+    run_migrations_offline()
+else:
+    run_migrations_online()

sky/schemas/db/global_user_state/001_initial_schema.py

@@ -0,0 +1,124 @@
+"""Initial schema for state database with backwards compatibility columns
+
+Revision ID: 001
+Revises:
+Create Date: 2024-01-01 12:00:00.000000
+
+"""
+# pylint: disable=invalid-name
+from alembic import op
+import sqlalchemy as sa
+
+from sky.global_user_state import Base
+from sky.utils.db import db_utils
+
+# revision identifiers, used by Alembic.
+revision = '001'
+down_revision = None
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    with op.get_context().autocommit_block():
+        # Create any missing tables with current schema first
+        db_utils.add_tables_to_db_sqlalchemy(Base.metadata, op.get_bind())
+
+        # Add all missing columns to clusters table
+        # This allows each column addition to fail independently without rolling
+        # back the entire migration, which is needed for backwards compatibility
+
+        # Add all missing columns to clusters table
+        db_utils.add_column_to_table_alembic('clusters',
+                                             'autostop',
+                                             sa.Integer(),
+                                             server_default='-1')
+        db_utils.add_column_to_table_alembic('clusters',
+                                             'metadata',
+                                             sa.Text(),
+                                             server_default='{}')
+        db_utils.add_column_to_table_alembic('clusters',
+                                             'to_down',
+                                             sa.Integer(),
+                                             server_default='0')
+        db_utils.add_column_to_table_alembic('clusters',
+                                             'owner',
+                                             sa.Text(),
+                                             server_default=None)
+        db_utils.add_column_to_table_alembic('clusters',
+                                             'cluster_hash',
+                                             sa.Text(),
+                                             server_default=None)
+        db_utils.add_column_to_table_alembic('clusters',
+                                             'launched_nodes',
+                                             sa.Integer(),
+                                             server_default='0')
+        db_utils.add_column_to_table_alembic('clusters',
+                                             'disk_tier',
+                                             sa.Text(),
+                                             server_default=None)
+        db_utils.add_column_to_table_alembic('clusters',
+                                             'config_hash',
+                                             sa.Text(),
+                                             server_default=None)
+        db_utils.add_column_to_table_alembic('clusters',
+                                             'user_hash',
+                                             sa.Text(),
+                                             server_default=None)
+        db_utils.add_column_to_table_alembic('clusters',
+                                             'workspace',
+                                             sa.Text(),
+                                             server_default='default')
+        db_utils.add_column_to_table_alembic('clusters',
+                                             'last_creation_yaml',
+                                             sa.Text(),
+                                             server_default=None)
+        db_utils.add_column_to_table_alembic('clusters',
+                                             'last_creation_command',
+                                             sa.Text(),
+                                             server_default=None)
+        db_utils.add_column_to_table_alembic('clusters',
+                                             'config_hash_locked',
+                                             sa.Boolean(),
+                                             server_default='FALSE')
+        db_utils.add_column_to_table_alembic('clusters',
+                                             'handle_locked',
+                                             sa.Boolean(),
+                                             server_default='FALSE')
+        db_utils.add_column_to_table_alembic('clusters',
+                                             'num_failures',
+                                             sa.Integer(),
+                                             server_default='0')
+        db_utils.add_column_to_table_alembic('clusters',
+                                             'configs',
+                                             sa.Text(),
+                                             server_default='[]')
+
+        # Add all missing columns to cluster_history table
+        db_utils.add_column_to_table_alembic('cluster_history',
+                                             'user_hash',
+                                             sa.Text(),
+                                             server_default=None)
+        db_utils.add_column_to_table_alembic('cluster_history',
+                                             'last_creation_yaml',
+                                             sa.Text(),
+                                             server_default=None)
+        db_utils.add_column_to_table_alembic('cluster_history',
+                                             'last_creation_command',
+                                             sa.Text(),
+                                             server_default=None)
+
+        # Add all missing columns to users table
+        db_utils.add_column_to_table_alembic('users',
+                                             'password',
+                                             sa.Text(),
+                                             server_default=None)
+        db_utils.add_column_to_table_alembic('users',
+                                             'created_at',
+                                             sa.Integer(),
+                                             server_default=None)
+
+
+def downgrade():
+    # Drop all tables
+    Base.metadata.drop_all(bind=op.get_bind())