skypilot-nightly 1.0.0.dev20250604__py3-none-any.whl → 1.0.0.dev20250606__py3-none-any.whl

sky/global_user_state.py

@@ -44,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(
@@ -183,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:
@@ -197,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',
@@ -309,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)
@@ -344,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]
 
@@ -421,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(
@@ -448,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:
@@ -529,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()
@@ -537,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()
@@ -548,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
@@ -581,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
@@ -590,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')(
@@ -608,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,
@@ -622,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,
@@ -635,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
@@ -643,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
@@ -651,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)})
@@ -663,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
@@ -672,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:
@@ -689,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:
@@ -726,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:
@@ -743,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})
@@ -754,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
@@ -763,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:
@@ -808,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
@@ -841,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 = []
@@ -876,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 ==
@@ -912,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]
@@ -920,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 = []
@@ -944,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:
@@ -981,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:
@@ -1010,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()
@@ -1028,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]
@@ -1037,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)})
@@ -1051,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)
@@ -1060,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')(
@@ -1076,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:
@@ -1099,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:
@@ -1108,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:
@@ -1142,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:
@@ -1172,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:
@@ -1191,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

@@ -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/resources.py

@@ -45,7 +45,9 @@ class AutostopConfig:
     # to be complete.
     enabled: bool
     # If enabled is False, these values are ignored.
-    idle_minutes: int = 5
+    # Keep the default value to 0 to make the behavior consistent with the CLI
+    # flags.
+    idle_minutes: int = 0
     down: bool = False
 
     def to_yaml_config(self) -> Union[Literal[False], Dict[str, Any]]:
@@ -883,6 +885,25 @@ class Resources:
             valid_volumes.append(volume)
         self._volumes = valid_volumes
 
+    def override_autostop_config(self,
+                                 down: bool = False,
+                                 idle_minutes: Optional[int] = None) -> None:
+        """Override autostop config to the resource.
+
+        Args:
+            down: If true, override the autostop config to use autodown.
+            idle_minutes: If not None, override the idle minutes to autostop or
+                autodown.
+        """
+        if not down and idle_minutes is None:
+            return
+        if self._autostop_config is None:
+            self._autostop_config = AutostopConfig(enabled=True,)
+        if down:
+            self._autostop_config.down = down
+        if idle_minutes is not None:
+            self._autostop_config.idle_minutes = idle_minutes
+
     def is_launchable(self) -> bool:
         """Returns whether the resource is launchable."""
         return self.cloud is not None and self._instance_type is not None