skypilot-nightly 1.0.0.dev20250603__py3-none-any.whl → 1.0.0.dev20250605__py3-none-any.whl

sky/global_user_state.py

@@ -26,6 +26,7 @@ import yaml
 
 from sky import models
 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 context_utils
@@ -43,18 +44,6 @@ logger = sky_logging.init_logger(__name__)
 
 _ENABLED_CLOUDS_KEY_PREFIX = 'enabled_clouds_'
 
-_DB_PATH = os.path.expanduser('~/.sky/state.db')
-pathlib.Path(_DB_PATH).parents[0].mkdir(parents=True, exist_ok=True)
-
-if os.environ.get(constants.SKYPILOT_API_SERVER_DB_URL_ENV_VAR):
-    # If SKYPILOT_API_SERVER_DB_URL_ENV_VAR is set, use it as the database URI.
-    logger.debug(
-        f'using db URI from {constants.SKYPILOT_API_SERVER_DB_URL_ENV_VAR}')
-    _SQLALCHEMY_ENGINE = sqlalchemy.create_engine(
-        os.environ.get(constants.SKYPILOT_API_SERVER_DB_URL_ENV_VAR))
-else:
-    _SQLALCHEMY_ENGINE = sqlalchemy.create_engine('sqlite:///' + _DB_PATH)
-
 Base = declarative.declarative_base()
 
 config_table = sqlalchemy.Table(
@@ -182,11 +171,11 @@ def create_table():
     # 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 (_SQLALCHEMY_ENGINE.dialect.name
+    if (SQLALCHEMY_ENGINE.dialect.name
             == db_utils.SQLAlchemyDialect.SQLITE.value and
             not common_utils.is_wsl()):
         try:
-            with orm.Session(_SQLALCHEMY_ENGINE) as session:
+            with orm.Session(SQLALCHEMY_ENGINE) as session:
                 session.execute(sqlalchemy.text('PRAGMA journal_mode=WAL'))
                 session.commit()
         except sqlalchemy_exc.OperationalError as e:
@@ -196,12 +185,12 @@ def create_table():
             # is not critical and is likely to be enabled by other processes.
 
     # Create tables if they don't exist
-    Base.metadata.create_all(bind=_SQLALCHEMY_ENGINE)
+    Base.metadata.create_all(bind=SQLALCHEMY_ENGINE)
 
     # For backward compatibility.
     # TODO(zhwu): Remove this function after all users have migrated to
     # the latest version of SkyPilot.
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         # Add autostop column to clusters table
         db_utils.add_column_to_table_sqlalchemy(session,
                                                 'clusters',
@@ -308,34 +297,81 @@ def create_table():
         session.commit()
 
 
+conn_string = None
+if os.environ.get(constants.ENV_VAR_IS_SKYPILOT_SERVER) is not None:
+    conn_string = skypilot_config.get_nested(('db',), None)
+if conn_string:
+    logger.debug(f'using db URI from {conn_string}')
+    SQLALCHEMY_ENGINE = sqlalchemy.create_engine(conn_string)
+else:
+    _DB_PATH = os.path.expanduser('~/.sky/state.db')
+    pathlib.Path(_DB_PATH).parents[0].mkdir(parents=True, exist_ok=True)
+    SQLALCHEMY_ENGINE = sqlalchemy.create_engine('sqlite:///' + _DB_PATH)
 create_table()
 
 
-def add_or_update_user(user: models.User):
-    """Store the mapping from user hash to user name for display purposes."""
+def add_or_update_user(user: models.User) -> bool:
+    """Store the mapping from user hash to user name for display purposes.
+
+    Returns:
+        Boolean: whether the user is newly added
+    """
     if user.name is None:
-        return
+        return False
 
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        if (_SQLALCHEMY_ENGINE.dialect.name ==
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
+        if (SQLALCHEMY_ENGINE.dialect.name ==
                 db_utils.SQLAlchemyDialect.SQLITE.value):
+            # For SQLite, use INSERT OR IGNORE followed by UPDATE to detect new
+            # vs existing
             insert_func = sqlite.insert
-        elif (_SQLALCHEMY_ENGINE.dialect.name ==
+
+            # First try INSERT OR IGNORE - this won't fail if user exists
+            insert_stmnt = insert_func(user_table).prefix_with(
+                'OR IGNORE').values(id=user.id, name=user.name)
+            result = session.execute(insert_stmnt)
+
+            # Check if the INSERT actually inserted a row
+            was_inserted = result.rowcount > 0
+
+            if not was_inserted:
+                # User existed, so update it
+                session.query(user_table).filter_by(id=user.id).update(
+                    {user_table.c.name: user.name})
+
+            session.commit()
+            return was_inserted
+
+        elif (SQLALCHEMY_ENGINE.dialect.name ==
               db_utils.SQLAlchemyDialect.POSTGRESQL.value):
+            # For PostgreSQL, use INSERT ... ON CONFLICT with RETURNING to
+            # detect insert vs update
             insert_func = postgresql.insert
+            insert_stmnt = insert_func(user_table).values(id=user.id,
+                                                          name=user.name)
+
+            # Use a sentinel in the RETURNING clause to detect insert vs update
+            upsert_stmnt = insert_stmnt.on_conflict_do_update(
+                index_elements=[user_table.c.id],
+                set_={
+                    user_table.c.name: user.name
+                }).returning(
+                    user_table.c.id,
+                    # This will be True for INSERT, False for UPDATE
+                    sqlalchemy.literal_column('(xmax = 0)').label('was_inserted'
+                                                                 ))
+
+            result = session.execute(upsert_stmnt)
+            session.commit()
+
+            row = result.fetchone()
+            return bool(row.was_inserted) if row else False
         else:
             raise ValueError('Unsupported database dialect')
-        insert_stmnt = insert_func(user_table).values(id=user.id,
-                                                      name=user.name)
-        do_update_stmt = insert_stmnt.on_conflict_do_update(
-            index_elements=[user_table.c.id],
-            set_={user_table.c.name: user.name})
-        session.execute(do_update_stmt)
-        session.commit()
 
 
 def get_user(user_id: str) -> models.User:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         row = session.query(user_table).filter_by(id=user_id).first()
     if row is None:
         return models.User(id=user_id)
@@ -343,7 +379,7 @@ def get_user(user_id: str) -> models.User:
 
 
 def get_all_users() -> List[models.User]:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         rows = session.query(user_table).all()
     return [models.User(id=row.id, name=row.name) for row in rows]
 
@@ -368,9 +404,6 @@ def add_or_update_cluster(cluster_name: str,
         config_hash: Configuration hash for the cluster.
         task_config: The config of the task being launched.
     """
-    # TODO(zhwu): have to be imported here to avoid circular import.
-    from sky import skypilot_config  # pylint: disable=import-outside-toplevel
-
     # FIXME: launched_at will be changed when `sky launch -c` is called.
     handle = pickle.dumps(cluster_handle)
     cluster_launched_at = int(time.time()) if is_launch else None
@@ -423,7 +456,7 @@ def add_or_update_cluster(cluster_name: str,
             'config_hash': config_hash,
         })
 
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         # with_for_update() locks the row until commit() or rollback()
         # is called, or until the code escapes the with block.
         cluster_row = session.query(cluster_table).filter_by(
@@ -450,10 +483,10 @@ def add_or_update_cluster(cluster_name: str,
                 'last_creation_command': last_use,
             })
 
-        if (_SQLALCHEMY_ENGINE.dialect.name ==
+        if (SQLALCHEMY_ENGINE.dialect.name ==
                 db_utils.SQLAlchemyDialect.SQLITE.value):
             insert_func = sqlite.insert
-        elif (_SQLALCHEMY_ENGINE.dialect.name ==
+        elif (SQLALCHEMY_ENGINE.dialect.name ==
               db_utils.SQLAlchemyDialect.POSTGRESQL.value):
             insert_func = postgresql.insert
         else:
@@ -531,7 +564,7 @@ def _get_user_hash_or_current_user(user_hash: Optional[str]) -> str:
 def update_cluster_handle(cluster_name: str,
                           cluster_handle: 'backends.ResourceHandle'):
     handle = pickle.dumps(cluster_handle)
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         session.query(cluster_table).filter_by(name=cluster_name).update(
             {cluster_table.c.handle: handle})
         session.commit()
@@ -539,7 +572,7 @@ def update_cluster_handle(cluster_name: str,
 
 def update_last_use(cluster_name: str):
     """Updates the last used command for the cluster."""
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         session.query(cluster_table).filter_by(name=cluster_name).update(
             {cluster_table.c.last_use: common_utils.get_current_command()})
         session.commit()
@@ -550,7 +583,7 @@ def remove_cluster(cluster_name: str, terminate: bool) -> None:
     cluster_hash = _get_hash_for_existing_cluster(cluster_name)
     usage_intervals = _get_cluster_usage_intervals(cluster_hash)
 
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         # usage_intervals is not None and not empty
         if usage_intervals:
             assert cluster_hash is not None, cluster_name
@@ -583,7 +616,7 @@ def remove_cluster(cluster_name: str, terminate: bool) -> None:
 def get_handle_from_cluster_name(
         cluster_name: str) -> Optional['backends.ResourceHandle']:
     assert cluster_name is not None, 'cluster_name cannot be None'
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         row = session.query(cluster_table).filter_by(name=cluster_name).first()
     if row is None:
         return None
@@ -592,12 +625,12 @@ def get_handle_from_cluster_name(
 
 def get_glob_cluster_names(cluster_name: str) -> List[str]:
     assert cluster_name is not None, 'cluster_name cannot be None'
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        if (_SQLALCHEMY_ENGINE.dialect.name ==
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
+        if (SQLALCHEMY_ENGINE.dialect.name ==
                 db_utils.SQLAlchemyDialect.SQLITE.value):
             rows = session.query(cluster_table).filter(
                 cluster_table.c.name.op('GLOB')(cluster_name)).all()
-        elif (_SQLALCHEMY_ENGINE.dialect.name ==
+        elif (SQLALCHEMY_ENGINE.dialect.name ==
               db_utils.SQLAlchemyDialect.POSTGRESQL.value):
             rows = session.query(cluster_table).filter(
                 cluster_table.c.name.op('SIMILAR TO')(
@@ -610,7 +643,7 @@ def get_glob_cluster_names(cluster_name: str) -> List[str]:
 def set_cluster_status(cluster_name: str,
                        status: status_lib.ClusterStatus) -> None:
     current_time = int(time.time())
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         count = session.query(cluster_table).filter_by(
             name=cluster_name).update({
                 cluster_table.c.status: status.value,
@@ -624,7 +657,7 @@ def set_cluster_status(cluster_name: str,
 
 def set_cluster_autostop_value(cluster_name: str, idle_minutes: int,
                                to_down: bool) -> None:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         count = session.query(cluster_table).filter_by(
             name=cluster_name).update({
                 cluster_table.c.autostop: idle_minutes,
@@ -637,7 +670,7 @@ def set_cluster_autostop_value(cluster_name: str, idle_minutes: int,
 
 
 def get_cluster_launch_time(cluster_name: str) -> Optional[int]:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         row = session.query(cluster_table).filter_by(name=cluster_name).first()
     if row is None or row.launched_at is None:
         return None
@@ -645,7 +678,7 @@ def get_cluster_launch_time(cluster_name: str) -> Optional[int]:
 
 
 def get_cluster_info(cluster_name: str) -> Optional[Dict[str, Any]]:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         row = session.query(cluster_table).filter_by(name=cluster_name).first()
     if row is None or row.metadata is None:
         return None
@@ -653,7 +686,7 @@ def get_cluster_info(cluster_name: str) -> Optional[Dict[str, Any]]:
 
 
 def set_cluster_info(cluster_name: str, metadata: Dict[str, Any]) -> None:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         count = session.query(cluster_table).filter_by(
             name=cluster_name).update(
                 {cluster_table.c.metadata: json.dumps(metadata)})
@@ -665,7 +698,7 @@ def set_cluster_info(cluster_name: str, metadata: Dict[str, Any]) -> None:
 
 def get_cluster_storage_mounts_metadata(
         cluster_name: str) -> Optional[Dict[str, Any]]:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         row = session.query(cluster_table).filter_by(name=cluster_name).first()
     if row is None or row.storage_mounts_metadata is None:
         return None
@@ -674,7 +707,7 @@ def get_cluster_storage_mounts_metadata(
 
 def set_cluster_storage_mounts_metadata(
         cluster_name: str, storage_mounts_metadata: Dict[str, Any]) -> None:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         count = session.query(cluster_table).filter_by(
             name=cluster_name).update({
                 cluster_table.c.storage_mounts_metadata:
@@ -691,7 +724,7 @@ def _get_cluster_usage_intervals(
 ) -> Optional[List[Tuple[int, Optional[int]]]]:
     if cluster_hash is None:
         return None
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         row = session.query(cluster_history_table).filter_by(
             cluster_hash=cluster_hash).first()
     if row is None or row.usage_intervals is None:
@@ -728,7 +761,7 @@ def _get_cluster_duration(cluster_hash: str) -> int:
 def _set_cluster_usage_intervals(
         cluster_hash: str, usage_intervals: List[Tuple[int,
                                                        Optional[int]]]) -> None:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         count = session.query(cluster_history_table).filter_by(
             cluster_hash=cluster_hash).update({
                 cluster_history_table.c.usage_intervals:
@@ -745,7 +778,7 @@ def set_owner_identity_for_cluster(cluster_name: str,
     if owner_identity is None:
         return
     owner_identity_str = json.dumps(owner_identity)
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         count = session.query(cluster_table).filter_by(
             name=cluster_name).update(
                 {cluster_table.c.owner: owner_identity_str})
@@ -756,7 +789,7 @@ def set_owner_identity_for_cluster(cluster_name: str,
 
 
 def _get_hash_for_existing_cluster(cluster_name: str) -> Optional[str]:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         row = session.query(cluster_table).filter_by(name=cluster_name).first()
     if row is None or row.cluster_hash is None:
         return None
@@ -765,7 +798,7 @@ def _get_hash_for_existing_cluster(cluster_name: str) -> Optional[str]:
 
 def get_launched_resources_from_cluster_hash(
         cluster_hash: str) -> Optional[Tuple[int, Any]]:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         row = session.query(cluster_history_table).filter_by(
             cluster_hash=cluster_hash).first()
     if row is None:
@@ -810,7 +843,7 @@ def _load_storage_mounts_metadata(
 @context_utils.cancellation_guard
 def get_cluster_from_name(
         cluster_name: Optional[str]) -> Optional[Dict[str, Any]]:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         row = session.query(cluster_table).filter_by(name=cluster_name).first()
     if row is None:
         return None
@@ -843,7 +876,7 @@ def get_cluster_from_name(
 
 
 def get_clusters() -> List[Dict[str, Any]]:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         rows = session.query(cluster_table).order_by(
             sqlalchemy.desc(cluster_table.c.launched_at)).all()
     records = []
@@ -878,7 +911,7 @@ def get_clusters() -> List[Dict[str, Any]]:
 
 
 def get_clusters_from_history() -> List[Dict[str, Any]]:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         rows = session.query(
             cluster_history_table.join(cluster_table,
                                        cluster_history_table.c.cluster_hash ==
@@ -914,7 +947,7 @@ def get_clusters_from_history() -> List[Dict[str, Any]]:
 
 
 def get_cluster_names_start_with(starts_with: str) -> List[str]:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         rows = session.query(cluster_table).filter(
             cluster_table.c.name.like(f'{starts_with}%')).all()
     return [row.name for row in rows]
@@ -922,7 +955,7 @@ def get_cluster_names_start_with(starts_with: str) -> List[str]:
 
 def get_cached_enabled_clouds(cloud_capability: 'cloud.CloudCapability',
                               workspace: str) -> List['clouds.Cloud']:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         row = session.query(config_table).filter_by(
             key=_get_enabled_clouds_key(cloud_capability, workspace)).first()
     ret = []
@@ -946,11 +979,11 @@ def get_cached_enabled_clouds(cloud_capability: 'cloud.CloudCapability',
 def set_enabled_clouds(enabled_clouds: List[str],
                        cloud_capability: 'cloud.CloudCapability',
                        workspace: str) -> None:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        if (_SQLALCHEMY_ENGINE.dialect.name ==
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
+        if (SQLALCHEMY_ENGINE.dialect.name ==
                 db_utils.SQLAlchemyDialect.SQLITE.value):
             insert_func = sqlite.insert
-        elif (_SQLALCHEMY_ENGINE.dialect.name ==
+        elif (SQLALCHEMY_ENGINE.dialect.name ==
               db_utils.SQLAlchemyDialect.POSTGRESQL.value):
             insert_func = postgresql.insert
         else:
@@ -983,11 +1016,11 @@ def add_or_update_storage(storage_name: str,
     if not status_check(storage_status):
         raise ValueError(f'Error in updating global state. Storage Status '
                          f'{storage_status} is passed in incorrectly')
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        if (_SQLALCHEMY_ENGINE.dialect.name ==
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
+        if (SQLALCHEMY_ENGINE.dialect.name ==
                 db_utils.SQLAlchemyDialect.SQLITE.value):
             insert_func = sqlite.insert
-        elif (_SQLALCHEMY_ENGINE.dialect.name ==
+        elif (SQLALCHEMY_ENGINE.dialect.name ==
               db_utils.SQLAlchemyDialect.POSTGRESQL.value):
             insert_func = postgresql.insert
         else:
@@ -1012,14 +1045,14 @@ def add_or_update_storage(storage_name: str,
 
 def remove_storage(storage_name: str):
     """Removes Storage from Database"""
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         session.query(storage_table).filter_by(name=storage_name).delete()
         session.commit()
 
 
 def set_storage_status(storage_name: str,
                        status: status_lib.StorageStatus) -> None:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         count = session.query(storage_table).filter_by(
             name=storage_name).update({storage_table.c.status: status.value})
         session.commit()
@@ -1030,7 +1063,7 @@ def set_storage_status(storage_name: str,
 
 def get_storage_status(storage_name: str) -> Optional[status_lib.StorageStatus]:
     assert storage_name is not None, 'storage_name cannot be None'
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         row = session.query(storage_table).filter_by(name=storage_name).first()
     if row:
         return status_lib.StorageStatus[row.status]
@@ -1039,7 +1072,7 @@ def get_storage_status(storage_name: str) -> Optional[status_lib.StorageStatus]:
 
 def set_storage_handle(storage_name: str,
                        handle: 'Storage.StorageMetadata') -> None:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         count = session.query(storage_table).filter_by(
             name=storage_name).update(
                 {storage_table.c.handle: pickle.dumps(handle)})
@@ -1053,7 +1086,7 @@ def get_handle_from_storage_name(
         storage_name: Optional[str]) -> Optional['Storage.StorageMetadata']:
     if storage_name is None:
         return None
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         row = session.query(storage_table).filter_by(name=storage_name).first()
     if row:
         return pickle.loads(row.handle)
@@ -1062,12 +1095,12 @@ def get_handle_from_storage_name(
 
 def get_glob_storage_name(storage_name: str) -> List[str]:
     assert storage_name is not None, 'storage_name cannot be None'
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        if (_SQLALCHEMY_ENGINE.dialect.name ==
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
+        if (SQLALCHEMY_ENGINE.dialect.name ==
                 db_utils.SQLAlchemyDialect.SQLITE.value):
             rows = session.query(storage_table).filter(
                 storage_table.c.name.op('GLOB')(storage_name)).all()
-        elif (_SQLALCHEMY_ENGINE.dialect.name ==
+        elif (SQLALCHEMY_ENGINE.dialect.name ==
               db_utils.SQLAlchemyDialect.POSTGRESQL.value):
             rows = session.query(storage_table).filter(
                 storage_table.c.name.op('SIMILAR TO')(
@@ -1078,14 +1111,14 @@ def get_glob_storage_name(storage_name: str) -> List[str]:
 
 
 def get_storage_names_start_with(starts_with: str) -> List[str]:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         rows = session.query(storage_table).filter(
             storage_table.c.name.like(f'{starts_with}%')).all()
     return [row.name for row in rows]
 
 
 def get_storage() -> List[Dict[str, Any]]:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         rows = session.query(storage_table).all()
     records = []
     for row in rows:
@@ -1101,7 +1134,7 @@ def get_storage() -> List[Dict[str, Any]]:
 
 
 def get_ssh_keys(user_hash: str) -> Tuple[str, str, bool]:
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         row = session.query(ssh_key_table).filter_by(
             user_hash=user_hash).first()
     if row:
@@ -1110,11 +1143,11 @@ def get_ssh_keys(user_hash: str) -> Tuple[str, str, bool]:
 
 
 def set_ssh_keys(user_hash: str, ssh_public_key: str, ssh_private_key: str):
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        if (_SQLALCHEMY_ENGINE.dialect.name ==
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
+        if (SQLALCHEMY_ENGINE.dialect.name ==
                 db_utils.SQLAlchemyDialect.SQLITE.value):
             insert_func = sqlite.insert
-        elif (_SQLALCHEMY_ENGINE.dialect.name ==
+        elif (SQLALCHEMY_ENGINE.dialect.name ==
               db_utils.SQLAlchemyDialect.POSTGRESQL.value):
             insert_func = postgresql.insert
         else:
@@ -1144,7 +1177,7 @@ def get_cluster_yaml_str(cluster_yaml_path: Optional[str]) -> Optional[str]:
         raise ValueError('Attempted to read a None YAML.')
     cluster_file_name = os.path.basename(cluster_yaml_path)
     cluster_name, _ = os.path.splitext(cluster_file_name)
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         row = session.query(cluster_yaml_table).filter_by(
             cluster_name=cluster_name).first()
     if row is None:
@@ -1174,11 +1207,11 @@ def get_cluster_yaml_dict(cluster_yaml_path: Optional[str]) -> Dict[str, Any]:
 
 def set_cluster_yaml(cluster_name: str, yaml_str: str) -> None:
     """Set the cluster yaml in the database."""
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        if (_SQLALCHEMY_ENGINE.dialect.name ==
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
+        if (SQLALCHEMY_ENGINE.dialect.name ==
                 db_utils.SQLAlchemyDialect.SQLITE.value):
             insert_func = sqlite.insert
-        elif (_SQLALCHEMY_ENGINE.dialect.name ==
+        elif (SQLALCHEMY_ENGINE.dialect.name ==
               db_utils.SQLAlchemyDialect.POSTGRESQL.value):
             insert_func = postgresql.insert
         else:
@@ -1193,7 +1226,7 @@ def set_cluster_yaml(cluster_name: str, yaml_str: str) -> None:
 
 
 def remove_cluster_yaml(cluster_name: str):
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(SQLALCHEMY_ENGINE) as session:
         session.query(cluster_yaml_table).filter_by(
             cluster_name=cluster_name).delete()
         session.commit()

sky/jobs/client/sdk.py

@@ -92,7 +92,8 @@ def launch(
 @server_common.check_server_healthy_or_start
 def queue(refresh: bool,
           skip_finished: bool = False,
-          all_users: bool = False) -> server_common.RequestId:
+          all_users: bool = False,
+          job_ids: Optional[List[int]] = None) -> server_common.RequestId:
     """Gets statuses of managed jobs.
 
     Please refer to sky.cli.job_queue for documentation.
@@ -101,6 +102,7 @@ def queue(refresh: bool,
         refresh: Whether to restart the jobs controller if it is stopped.
         skip_finished: Whether to skip finished jobs.
         all_users: Whether to show all users' jobs.
+        job_ids: IDs of the managed jobs to show.
 
     Returns:
         The request ID of the queue request.
@@ -135,6 +137,7 @@ def queue(refresh: bool,
         refresh=refresh,
         skip_finished=skip_finished,
         all_users=all_users,
+        job_ids=job_ids,
     )
     response = requests.post(
         f'{server_common.get_server_url()}/jobs/queue',

sky/jobs/server/core.py

@@ -20,7 +20,7 @@ from sky import sky_logging
 from sky import skypilot_config
 from sky import task as task_lib
 from sky.backends import backend_utils
-from sky.clouds.service_catalog import common as service_catalog_common
+from sky.catalog import common as service_catalog_common
 from sky.data import storage as storage_lib
 from sky.jobs import constants as managed_job_constants
 from sky.jobs import utils as managed_job_utils
@@ -378,7 +378,8 @@ def _maybe_restart_controller(
 @usage_lib.entrypoint
 def queue(refresh: bool,
           skip_finished: bool = False,
-          all_users: bool = False) -> List[Dict[str, Any]]:
+          all_users: bool = False,
+          job_ids: Optional[List[int]] = None) -> List[Dict[str, Any]]:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
     """Gets statuses of managed jobs.
 
@@ -450,6 +451,9 @@ def queue(refresh: bool,
         jobs = list(
             filter(lambda job: job['job_id'] in non_finished_job_ids, jobs))
 
+    if job_ids:
+        jobs = [job for job in jobs if job['job_id'] in job_ids]
+
     return jobs
 
 

sky/models.py

@@ -7,6 +7,7 @@ from typing import Any, Dict, Optional
 
 @dataclasses.dataclass
 class User:
+    """Dataclass to store user information."""
     # User hash
     id: str
     # Display name of the user

sky/optimizer.py

@@ -1313,7 +1313,7 @@ def _fill_in_launchable_resources(
             if feasible_resources.resources_list:
                 # Assume feasible_resources is sorted by prices. Guaranteed by
                 # the implementation of get_feasible_launchable_resources and
-                # the underlying service_catalog filtering
+                # the underlying catalog filtering
                 cheapest = feasible_resources.resources_list[0]
                 # Generate region/zone-specified resources.
                 launchable[resources].extend(

sky/provision/cudo/cudo_machine_type.py

@@ -1,7 +1,7 @@
 """Cudo Compute VM spec helper for SkyPilot."""
 import csv
 
-from sky.clouds.service_catalog.common import get_catalog_path
+from sky.catalog.common import get_catalog_path
 
 VMS_CSV = 'cudo/vms.csv'