skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/global_user_state.py

@@ -16,15 +16,20 @@ import typing
 from typing import Any, Dict, List, Optional, Set, Tuple
 import uuid
 
-from sky import clouds
-from sky import status_lib
+from sky import models
+from sky import sky_logging
 from sky.utils import common_utils
 from sky.utils import db_utils
+from sky.utils import registry
+from sky.utils import status_lib
 
 if typing.TYPE_CHECKING:
     from sky import backends
+    from sky import clouds
     from sky.data import Storage
 
+logger = sky_logging.init_logger(__name__)
+
 _ENABLED_CLOUDS_KEY = 'enabled_clouds'
 
 _DB_PATH = os.path.expanduser('~/.sky/state.db')
@@ -55,12 +60,15 @@ def create_table(cursor, conn):
         last_use TEXT,
         status TEXT,
         autostop INTEGER DEFAULT -1,
-        metadata TEXT DEFAULT "{}",
+        metadata TEXT DEFAULT '{}',
         to_down INTEGER DEFAULT 0,
         owner TEXT DEFAULT null,
         cluster_hash TEXT DEFAULT null,
         storage_mounts_metadata BLOB DEFAULT null,
-        cluster_ever_up INTEGER DEFAULT 0)""")
+        cluster_ever_up INTEGER DEFAULT 0,
+        status_updated_at INTEGER DEFAULT null,
+        config_hash TEXT DEFAULT null,
+        user_hash TEXT DEFAULT null)""")
 
     # Table for Cluster History
     # usage_intervals: List[Tuple[int, int]]
@@ -83,7 +91,8 @@ def create_table(cursor, conn):
         num_nodes int,
         requested_resources BLOB,
         launched_resources BLOB,
-        usage_intervals BLOB)""")
+        usage_intervals BLOB,
+        user_hash TEXT)""")
     # Table for configs (e.g. enabled clouds)
     cursor.execute("""\
         CREATE TABLE IF NOT EXISTS config (
@@ -96,6 +105,11 @@ def create_table(cursor, conn):
         handle BLOB,
         last_use TEXT,
         status TEXT)""")
+    # Table for User
+    cursor.execute("""\
+        CREATE TABLE IF NOT EXISTS users (
+        id TEXT PRIMARY KEY,
+        name TEXT)""")
     # For backward compatibility.
     # TODO(zhwu): Remove this function after all users have migrated to
     # the latest version of SkyPilot.
@@ -104,11 +118,12 @@ def create_table(cursor, conn):
                                  'INTEGER DEFAULT -1')
 
     db_utils.add_column_to_table(cursor, conn, 'clusters', 'metadata',
-                                 'TEXT DEFAULT "{}"')
+                                 'TEXT DEFAULT \'{}\'')
 
     db_utils.add_column_to_table(cursor, conn, 'clusters', 'to_down',
                                  'INTEGER DEFAULT 0')
 
+    # The cloud identity that created the cluster.
     db_utils.add_column_to_table(cursor, conn, 'clusters', 'owner', 'TEXT')
 
     db_utils.add_column_to_table(cursor, conn, 'clusters', 'cluster_hash',
@@ -130,17 +145,52 @@ def create_table(cursor, conn):
         # 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(cursor, conn, 'clusters', 'status_updated_at',
+                                 'INTEGER DEFAULT null')
+    db_utils.add_column_to_table(
+        cursor,
+        conn,
+        'clusters',
+        'user_hash',
+        'TEXT DEFAULT null',
+        value_to_replace_existing_entries=common_utils.get_user_hash())
+    db_utils.add_column_to_table(cursor, conn, 'clusters', 'config_hash',
+                                 'TEXT DEFAULT null')
+
+    db_utils.add_column_to_table(cursor, conn, 'clusters', 'config_hash',
+                                 'TEXT DEFAULT null')
+
+    db_utils.add_column_to_table(cursor, conn, 'cluster_history', 'user_hash',
+                                 'TEXT DEFAULT null')
     conn.commit()
 
 
 _DB = db_utils.SQLiteConn(_DB_PATH, create_table)
 
 
+def add_or_update_user(user: models.User):
+    """Store the mapping from user hash to user name for display purposes."""
+    if user.name is None:
+        return
+    _DB.cursor.execute('INSERT OR REPLACE INTO users (id, name) VALUES (?, ?)',
+                       (user.id, user.name))
+    _DB.conn.commit()
+
+
+def get_user(user_id: str) -> models.User:
+    row = _DB.cursor.execute('SELECT id, name FROM users WHERE id=?',
+                             (user_id,)).fetchone()
+    if row is None:
+        return models.User(id=user_id)
+    return models.User(id=row[0], name=row[1])
+
+
 def add_or_update_cluster(cluster_name: str,
                           cluster_handle: 'backends.ResourceHandle',
                           requested_resources: Optional[Set[Any]],
                           ready: bool,
-                          is_launch: bool = True):
+                          is_launch: bool = True,
+                          config_hash: Optional[str] = None):
     """Adds or updates cluster_name -> cluster_handle mapping.
 
     Args:
@@ -155,10 +205,11 @@ def add_or_update_cluster(cluster_name: str,
     # 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
-    last_use = common_utils.get_pretty_entry_point() if is_launch else None
+    last_use = common_utils.get_current_command() if is_launch else None
     status = status_lib.ClusterStatus.INIT
     if ready:
         status = status_lib.ClusterStatus.UP
+    status_updated_at = int(time.time())
 
     # TODO (sumanth): Cluster history table will have multiple entries
     # when the cluster failover through multiple regions (one entry per region).
@@ -183,6 +234,8 @@ def add_or_update_cluster(cluster_name: str,
             cluster_launched_at = int(time.time())
         usage_intervals.append((cluster_launched_at, None))
 
+    user_hash = common_utils.get_user_hash()
+
     _DB.cursor.execute(
         'INSERT or REPLACE INTO clusters'
         # All the fields need to exist here, even if they don't need
@@ -191,7 +244,8 @@ def add_or_update_cluster(cluster_name: str,
         # specified.
         '(name, launched_at, handle, last_use, status, '
         'autostop, to_down, metadata, owner, cluster_hash, '
-        'storage_mounts_metadata, cluster_ever_up) '
+        'storage_mounts_metadata, cluster_ever_up, status_updated_at, '
+        'config_hash, user_hash) '
         'VALUES ('
         # name
         '?, '
@@ -217,7 +271,7 @@ def add_or_update_cluster(cluster_name: str,
         # Keep the old metadata value if it exists, otherwise set it to
         # default {}.
         'COALESCE('
-        '(SELECT metadata FROM clusters WHERE name=?), "{}"),'
+        '(SELECT metadata FROM clusters WHERE name=?), \'{}\'),'
         # Keep the old owner value if it exists, otherwise set it to
         # default null.
         'COALESCE('
@@ -228,7 +282,14 @@ def add_or_update_cluster(cluster_name: str,
         'COALESCE('
         '(SELECT storage_mounts_metadata FROM clusters WHERE name=?), null), '
         # cluster_ever_up
-        '((SELECT cluster_ever_up FROM clusters WHERE name=?) OR ?)'
+        '((SELECT cluster_ever_up FROM clusters WHERE name=?) OR ?), '
+        # status_updated_at
+        '?,'
+        # config_hash
+        'COALESCE(?, (SELECT config_hash FROM clusters WHERE name=?)),'
+        # user_hash: keep original user_hash if it exists
+        'COALESCE('
+        '(SELECT user_hash FROM clusters WHERE name=?), ?)'
         ')',
         (
             # name
@@ -260,6 +321,14 @@ def add_or_update_cluster(cluster_name: str,
             # cluster_ever_up
             cluster_name,
             int(ready),
+            # status_updated_at
+            status_updated_at,
+            # config_hash
+            config_hash,
+            cluster_name,
+            # user_hash
+            cluster_name,
+            user_hash,
         ))
 
     launched_nodes = getattr(cluster_handle, 'launched_nodes', None)
@@ -267,7 +336,7 @@ def add_or_update_cluster(cluster_name: str,
     _DB.cursor.execute(
         'INSERT or REPLACE INTO cluster_history'
         '(cluster_hash, name, num_nodes, requested_resources, '
-        'launched_resources, usage_intervals) '
+        'launched_resources, usage_intervals, user_hash) '
         'VALUES ('
         # hash
         '?, '
@@ -280,7 +349,10 @@ def add_or_update_cluster(cluster_name: str,
         # number of nodes
         '?, '
         # usage intervals
-        '?)',
+        '?, '
+        # user_hash
+        '?'
+        ')',
         (
             # hash
             cluster_hash,
@@ -294,15 +366,37 @@ def add_or_update_cluster(cluster_name: str,
             pickle.dumps(launched_resources),
             # usage intervals
             pickle.dumps(usage_intervals),
+            # user_hash
+            user_hash,
         ))
 
     _DB.conn.commit()
 
 
+def _get_user_hash_or_current_user(user_hash: Optional[str]) -> str:
+    """Returns the user hash or the current user hash, if user_hash is None.
+
+    This is to ensure that the clusters created before the client-server
+    architecture (no user hash info previously) are associated with the current
+    user.
+    """
+    if user_hash is not None:
+        return user_hash
+    return common_utils.get_user_hash()
+
+
+def update_cluster_handle(cluster_name: str,
+                          cluster_handle: 'backends.ResourceHandle'):
+    handle = pickle.dumps(cluster_handle)
+    _DB.cursor.execute('UPDATE clusters SET handle=(?) WHERE name=(?)',
+                       (handle, cluster_name))
+    _DB.conn.commit()
+
+
 def update_last_use(cluster_name: str):
     """Updates the last used command for the cluster."""
     _DB.cursor.execute('UPDATE clusters SET last_use=(?) WHERE name=(?)',
-                       (common_utils.get_pretty_entry_point(), cluster_name))
+                       (common_utils.get_current_command(), cluster_name))
     _DB.conn.commit()
 
 
@@ -330,11 +424,13 @@ def remove_cluster(cluster_name: str, terminate: bool) -> None:
         # stopped VM, which leads to timeout.
         if hasattr(handle, 'stable_internal_external_ips'):
             handle.stable_internal_external_ips = None
+        current_time = int(time.time())
         _DB.cursor.execute(
-            'UPDATE clusters SET handle=(?), status=(?) '
-            'WHERE name=(?)', (
+            'UPDATE clusters SET handle=(?), status=(?), '
+            'status_updated_at=(?) WHERE name=(?)', (
                 pickle.dumps(handle),
                 status_lib.ClusterStatus.STOPPED.value,
+                current_time,
                 cluster_name,
             ))
     _DB.conn.commit()
@@ -359,10 +455,10 @@ def get_glob_cluster_names(cluster_name: str) -> List[str]:
 
 def set_cluster_status(cluster_name: str,
                        status: status_lib.ClusterStatus) -> None:
-    _DB.cursor.execute('UPDATE clusters SET status=(?) WHERE name=(?)', (
-        status.value,
-        cluster_name,
-    ))
+    current_time = int(time.time())
+    _DB.cursor.execute(
+        'UPDATE clusters SET status=(?), status_updated_at=(?) WHERE name=(?)',
+        (status.value, current_time, cluster_name))
     count = _DB.cursor.rowcount
     _DB.conn.commit()
     assert count <= 1, count
@@ -570,15 +666,19 @@ def _load_storage_mounts_metadata(
 
 def get_cluster_from_name(
         cluster_name: Optional[str]) -> Optional[Dict[str, Any]]:
-    rows = _DB.cursor.execute('SELECT * FROM clusters WHERE name=(?)',
-                              (cluster_name,)).fetchall()
+    rows = _DB.cursor.execute(
+        'SELECT name, launched_at, handle, last_use, status, autostop, '
+        'metadata, to_down, owner, cluster_hash, storage_mounts_metadata, '
+        'cluster_ever_up, status_updated_at, config_hash, user_hash '
+        'FROM clusters WHERE name=(?)', (cluster_name,)).fetchall()
     for row in rows:
         # Explicitly specify the number of fields to unpack, so that
         # we can add new fields to the database in the future without
         # breaking the previous code.
         (name, launched_at, handle, last_use, status, autostop, metadata,
-         to_down, owner, cluster_hash, storage_mounts_metadata,
-         cluster_ever_up) = row[:12]
+         to_down, owner, cluster_hash, storage_mounts_metadata, cluster_ever_up,
+         status_updated_at, config_hash, user_hash) = row
+        user_hash = _get_user_hash_or_current_user(user_hash)
         # TODO: use namedtuple instead of dict
         record = {
             'name': name,
@@ -594,6 +694,10 @@ def get_cluster_from_name(
             'storage_mounts_metadata':
                 _load_storage_mounts_metadata(storage_mounts_metadata),
             'cluster_ever_up': bool(cluster_ever_up),
+            'status_updated_at': status_updated_at,
+            'user_hash': user_hash,
+            'user_name': get_user(user_hash).name,
+            'config_hash': config_hash,
         }
         return record
     return None
@@ -601,12 +705,16 @@ def get_cluster_from_name(
 
 def get_clusters() -> List[Dict[str, Any]]:
     rows = _DB.cursor.execute(
-        'select * from clusters order by launched_at desc').fetchall()
+        'select name, launched_at, handle, last_use, status, autostop, '
+        'metadata, to_down, owner, cluster_hash, storage_mounts_metadata, '
+        'cluster_ever_up, status_updated_at, config_hash, user_hash '
+        'from clusters order by launched_at desc').fetchall()
     records = []
     for row in rows:
         (name, launched_at, handle, last_use, status, autostop, metadata,
-         to_down, owner, cluster_hash, storage_mounts_metadata,
-         cluster_ever_up) = row[:12]
+         to_down, owner, cluster_hash, storage_mounts_metadata, cluster_ever_up,
+         status_updated_at, config_hash, user_hash) = row
+        user_hash = _get_user_hash_or_current_user(user_hash)
         # TODO: use namedtuple instead of dict
         record = {
             'name': name,
@@ -622,6 +730,10 @@ def get_clusters() -> List[Dict[str, Any]]:
             'storage_mounts_metadata':
                 _load_storage_mounts_metadata(storage_mounts_metadata),
             'cluster_ever_up': bool(cluster_ever_up),
+            'status_updated_at': status_updated_at,
+            'user_hash': user_hash,
+            'user_name': get_user(user_hash).name,
+            'config_hash': config_hash,
         }
 
         records.append(record)
@@ -631,7 +743,8 @@ def get_clusters() -> List[Dict[str, Any]]:
 def get_clusters_from_history() -> List[Dict[str, Any]]:
     rows = _DB.cursor.execute(
         'SELECT ch.cluster_hash, ch.name, ch.num_nodes, '
-        'ch.launched_resources, ch.usage_intervals, clusters.status  '
+        'ch.launched_resources, ch.usage_intervals, clusters.status, '
+        'ch.user_hash  '
         'FROM cluster_history ch '
         'LEFT OUTER JOIN clusters '
         'ON ch.cluster_hash=clusters.cluster_hash ').fetchall()
@@ -650,7 +763,9 @@ def get_clusters_from_history() -> List[Dict[str, Any]]:
             launched_resources,
             usage_intervals,
             status,
-        ) = row[:6]
+            user_hash,
+        ) = row[:7]
+        user_hash = _get_user_hash_or_current_user(user_hash)
 
         if status is not None:
             status = status_lib.ClusterStatus[status]
@@ -664,6 +779,7 @@ def get_clusters_from_history() -> List[Dict[str, Any]]:
             'cluster_hash': cluster_hash,
             'usage_intervals': pickle.loads(usage_intervals),
             'status': status,
+            'user_hash': user_hash,
         }
 
         records.append(record)
@@ -679,17 +795,17 @@ def get_cluster_names_start_with(starts_with: str) -> List[str]:
     return [row[0] for row in rows]
 
 
-def get_cached_enabled_clouds() -> List[clouds.Cloud]:
+def get_cached_enabled_clouds() -> List['clouds.Cloud']:
     rows = _DB.cursor.execute('SELECT value FROM config WHERE key = ?',
                               (_ENABLED_CLOUDS_KEY,))
     ret = []
     for (value,) in rows:
         ret = json.loads(value)
         break
-    enabled_clouds: List[clouds.Cloud] = []
+    enabled_clouds: List['clouds.Cloud'] = []
     for c in ret:
         try:
-            cloud = clouds.CLOUD_REGISTRY.from_str(c)
+            cloud = registry.CLOUD_REGISTRY.from_str(c)
         except ValueError:
             # Handle the case for the clouds whose support has been removed from
             # SkyPilot, e.g., 'local' was a cloud in the past and may be stored
@@ -712,7 +828,7 @@ def add_or_update_storage(storage_name: str,
                           storage_status: status_lib.StorageStatus):
     storage_launched_at = int(time.time())
     handle = pickle.dumps(storage_handle)
-    last_use = common_utils.get_pretty_entry_point()
+    last_use = common_utils.get_current_command()
 
     def status_check(status):
         return status in status_lib.StorageStatus
@@ -794,7 +910,7 @@ def get_storage_names_start_with(starts_with: str) -> List[str]:
 
 
 def get_storage() -> List[Dict[str, Any]]:
-    rows = _DB.cursor.execute('select * from storage')
+    rows = _DB.cursor.execute('SELECT * FROM storage')
     records = []
     for name, launched_at, handle, last_use, status in rows:
         # TODO: use namedtuple instead of dict

sky/jobs/__init__.py

@@ -1,33 +1,32 @@
 """Managed jobs."""
 import pathlib
 
+from sky.jobs.client.sdk import cancel
+from sky.jobs.client.sdk import dashboard
+from sky.jobs.client.sdk import download_logs
+from sky.jobs.client.sdk import launch
+from sky.jobs.client.sdk import queue
+from sky.jobs.client.sdk import tail_logs
 from sky.jobs.constants import JOBS_CLUSTER_NAME_PREFIX_LENGTH
+from sky.jobs.constants import JOBS_CONTROLLER_LOGS_DIR
 from sky.jobs.constants import JOBS_CONTROLLER_TEMPLATE
 from sky.jobs.constants import JOBS_CONTROLLER_YAML_PREFIX
 from sky.jobs.constants import JOBS_TASK_YAML_PREFIX
-from sky.jobs.core import cancel
-from sky.jobs.core import launch
-from sky.jobs.core import queue
-from sky.jobs.core import tail_logs
-from sky.jobs.recovery_strategy import DEFAULT_RECOVERY_STRATEGY
-from sky.jobs.recovery_strategy import RECOVERY_STRATEGIES
+from sky.jobs.recovery_strategy import StrategyExecutor
 from sky.jobs.state import ManagedJobStatus
 from sky.jobs.utils import dump_managed_job_queue
 from sky.jobs.utils import format_job_table
-from sky.jobs.utils import JOB_CONTROLLER_NAME
 from sky.jobs.utils import load_managed_job_queue
 from sky.jobs.utils import ManagedJobCodeGen
 
 pathlib.Path(JOBS_TASK_YAML_PREFIX).expanduser().parent.mkdir(parents=True,
                                                               exist_ok=True)
 __all__ = [
-    'RECOVERY_STRATEGIES',
-    'DEFAULT_RECOVERY_STRATEGY',
-    'JOB_CONTROLLER_NAME',
     # Constants
     'JOBS_CONTROLLER_TEMPLATE',
     'JOBS_CONTROLLER_YAML_PREFIX',
     'JOBS_TASK_YAML_PREFIX',
+    'JOBS_CONTROLLER_LOGS_DIR',
     # Enums
     'ManagedJobStatus',
     # Core
@@ -35,9 +34,12 @@ __all__ = [
     'launch',
     'queue',
     'tail_logs',
+    'dashboard',
+    'download_logs',
     # utils
     'ManagedJobCodeGen',
     'format_job_table',
     'dump_managed_job_queue',
     'load_managed_job_queue',
+    'StrategyExecutor',
 ]