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

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

sky/schemas/db/script.py.mako

@@ -0,0 +1,28 @@
+"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+${imports if imports else ""}
+
+# revision identifiers, used by Alembic.
+revision: str = ${repr(up_revision)}
+down_revision: Union[str, Sequence[str], None] = ${repr(down_revision)}
+branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
+depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    ${upgrades if upgrades else "pass"}
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    ${downgrades if downgrades else "pass"}

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

@@ -0,0 +1,97 @@
+"""Initial schema for spot jobs database with backwards compatibility columns
+
+Revision ID: 001
+Revises:
+Create Date: 2024-01-01 12:00:00.000000
+
+"""
+# pylint: disable=invalid-name
+import json
+
+from alembic import op
+import sqlalchemy as sa
+
+from sky.jobs.state import Base
+from sky.skylet import constants
+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():
+    """Create initial schema and add all backwards compatibility columns"""
+    with op.get_context().autocommit_block():
+        # Create all tables with their current schema
+        db_utils.add_tables_to_db_sqlalchemy(Base.metadata, op.get_bind())
+
+        # Add backwards compatibility columns using helper function that matches
+        # original add_column_to_table_sqlalchemy behavior exactly
+
+        # Spot table columns
+        db_utils.add_column_to_table_alembic('spot', 'failure_reason',
+                                             sa.Text())
+        db_utils.add_column_to_table_alembic('spot',
+                                             'spot_job_id',
+                                             sa.Integer(),
+                                             copy_from='job_id')
+        db_utils.add_column_to_table_alembic(
+            'spot',
+            'task_id',
+            sa.Integer(),
+            server_default='0',
+            value_to_replace_existing_entries=0)
+        db_utils.add_column_to_table_alembic('spot',
+                                             'task_name',
+                                             sa.Text(),
+                                             copy_from='job_name')
+        db_utils.add_column_to_table_alembic(
+            'spot',
+            'specs',
+            sa.Text(),
+            value_to_replace_existing_entries=json.dumps(
+                {'max_restarts_on_errors': 0}))
+        db_utils.add_column_to_table_alembic('spot', 'local_log_file',
+                                             sa.Text())
+        db_utils.add_column_to_table_alembic(
+            'spot',
+            'metadata',
+            sa.Text(),
+            server_default='{}',
+            value_to_replace_existing_entries='{}')
+
+        # Job info table columns
+        db_utils.add_column_to_table_alembic('job_info', 'schedule_state',
+                                             sa.Text())
+        db_utils.add_column_to_table_alembic('job_info', 'controller_pid',
+                                             sa.Integer())
+        db_utils.add_column_to_table_alembic('job_info', 'dag_yaml_path',
+                                             sa.Text())
+        db_utils.add_column_to_table_alembic('job_info', 'env_file_path',
+                                             sa.Text())
+        db_utils.add_column_to_table_alembic('job_info', 'user_hash', sa.Text())
+        db_utils.add_column_to_table_alembic(
+            'job_info',
+            'workspace',
+            sa.Text(),
+            value_to_replace_existing_entries=constants.
+            SKYPILOT_DEFAULT_WORKSPACE)
+        db_utils.add_column_to_table_alembic(
+            'job_info',
+            'priority',
+            sa.Integer(),
+            server_default=str(constants.DEFAULT_PRIORITY),
+            value_to_replace_existing_entries=constants.DEFAULT_PRIORITY)
+        db_utils.add_column_to_table_alembic('job_info', 'entrypoint',
+                                             sa.Text())
+        db_utils.add_column_to_table_alembic('job_info',
+                                             'original_user_yaml_path',
+                                             sa.Text())
+
+
+def downgrade():
+    """Drop all tables"""
+    Base.metadata.drop_all(bind=op.get_bind())

sky/serve/client/sdk.py

@@ -297,7 +297,8 @@ def tail_logs(service_name: str,
               target: Union[str, 'serve_utils.ServiceComponent'],
               replica_id: Optional[int] = None,
               follow: bool = True,
-              output_stream: Optional['io.TextIOBase'] = None) -> None:
+              output_stream: Optional['io.TextIOBase'] = None,
+              tail: Optional[int] = None) -> None:
     """Tails logs for a service.
 
     Usage:
@@ -367,6 +368,7 @@ def tail_logs(service_name: str,
         target=target,
         replica_id=replica_id,
         follow=follow,
+        tail=tail,
     )
     response = server_common.make_authenticated_request(
         'POST',
@@ -390,7 +392,8 @@ def sync_down_logs(service_name: str,
                        str, 'serve_utils.ServiceComponent',
                        List[Union[str,
                                   'serve_utils.ServiceComponent']]]] = None,
-                   replica_ids: Optional[List[int]] = None) -> None:
+                   replica_ids: Optional[List[int]] = None,
+                   tail: Optional[int] = None) -> None:
     """Sync down logs from the service components to a local directory.
 
     This function syncs logs from the specified service components (controller,
@@ -429,6 +432,7 @@ def sync_down_logs(service_name: str,
         local_dir=local_dir,
         targets=targets,
         replica_ids=replica_ids,
+        tail=tail,
     )
     response = server_common.make_authenticated_request(
         'POST',

sky/serve/controller.py

@@ -156,9 +156,13 @@ class SkyServeController:
                 return responses.JSONResponse(content={'message': 'Success'},
                                               status_code=200)
             except Exception as e:  # pylint: disable=broad-except
-                logger.error(f'Error in update_service: '
-                             f'{common_utils.format_exception(e)}')
-                return responses.JSONResponse(content={'message': 'Error'},
+                exception_str = common_utils.format_exception(e)
+                logger.error(f'Error in update_service: {exception_str}')
+                return responses.JSONResponse(content={
+                    'message': 'Error',
+                    'exception': exception_str,
+                    'traceback': traceback.format_exc()
+                },
                                               status_code=500)
 
         @self._app.post('/controller/terminate_replica')

sky/serve/serve_state.py

@@ -13,7 +13,7 @@ from typing import Any, Dict, List, Optional, Tuple
 import colorama
 
 from sky.serve import constants
-from sky.utils import db_utils
+from sky.utils.db import db_utils
 
 if typing.TYPE_CHECKING:
     from sky.serve import replica_managers