skypilot-nightly 1.0.0.dev20250610__py3-none-any.whl → 1.0.0.dev20250611__py3-none-any.whl

sky/skylet/job_lib.py

@@ -3,6 +3,7 @@
 This is a remote utility module that provides job queue functionality.
 """
 import enum
+import functools
 import getpass
 import json
 import os
@@ -10,6 +11,7 @@ import pathlib
 import shlex
 import signal
 import sqlite3
+import threading
 import time
 import typing
 from typing import Any, Dict, List, Optional, Sequence
@@ -119,9 +121,25 @@ def create_table(cursor, conn):
     conn.commit()
 
 
-_DB = db_utils.SQLiteConn(_DB_PATH, create_table)
-_CURSOR = _DB.cursor
-_CONN = _DB.conn
+_DB = None
+_db_init_lock = threading.Lock()
+
+
+def init_db(func):
+    """Initialize the database."""
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        global _DB
+        if _DB is not None:
+            return func(*args, **kwargs)
+
+        with _db_init_lock:
+            if _DB is None:
+                _DB = db_utils.SQLiteConn(_DB_PATH, create_table)
+        return func(*args, **kwargs)
+
+    return wrapper
 
 
 class JobStatus(enum.Enum):
@@ -210,30 +228,37 @@ _PRE_RESOURCE_STATUSES = [JobStatus.PENDING]
 class JobScheduler:
     """Base class for job scheduler"""
 
+    @init_db
     def queue(self, job_id: int, cmd: str) -> None:
-        _CURSOR.execute('INSERT INTO pending_jobs VALUES (?,?,?,?)',
-                        (job_id, cmd, 0, int(time.time())))
-        _CONN.commit()
+        assert _DB is not None
+        _DB.cursor.execute('INSERT INTO pending_jobs VALUES (?,?,?,?)',
+                           (job_id, cmd, 0, int(time.time())))
+        _DB.conn.commit()
         set_status(job_id, JobStatus.PENDING)
         self.schedule_step()
 
+    @init_db
     def remove_job_no_lock(self, job_id: int) -> None:
-        _CURSOR.execute(f'DELETE FROM pending_jobs WHERE job_id={job_id!r}')
-        _CONN.commit()
+        assert _DB is not None
+        _DB.cursor.execute(f'DELETE FROM pending_jobs WHERE job_id={job_id!r}')
+        _DB.conn.commit()
 
+    @init_db
     def _run_job(self, job_id: int, run_cmd: str):
-        _CURSOR.execute((f'UPDATE pending_jobs SET submit={int(time.time())} '
-                         f'WHERE job_id={job_id!r}'))
-        _CONN.commit()
+        assert _DB is not None
+        _DB.cursor.execute(
+            (f'UPDATE pending_jobs SET submit={int(time.time())} '
+             f'WHERE job_id={job_id!r}'))
+        _DB.conn.commit()
         pid = subprocess_utils.launch_new_process_tree(run_cmd)
         # TODO(zhwu): Backward compatibility, remove this check after 0.10.0.
         # This is for the case where the job is submitted with SkyPilot older
         # than #4318, using ray job submit.
         if 'job submit' in run_cmd:
             pid = -1
-        _CURSOR.execute((f'UPDATE jobs SET pid={pid} '
-                         f'WHERE job_id={job_id!r}'))
-        _CONN.commit()
+        _DB.cursor.execute((f'UPDATE jobs SET pid={pid} '
+                            f'WHERE job_id={job_id!r}'))
+        _DB.conn.commit()
 
     def schedule_step(self, force_update_jobs: bool = False) -> None:
         if force_update_jobs:
@@ -282,8 +307,10 @@ class JobScheduler:
 class FIFOScheduler(JobScheduler):
     """First in first out job scheduler"""
 
+    @init_db
     def _get_pending_job_ids(self) -> List[int]:
-        rows = _CURSOR.execute(
+        assert _DB is not None
+        rows = _DB.cursor.execute(
             'SELECT job_id FROM pending_jobs ORDER BY job_id').fetchall()
         return [row[0] for row in rows]
 
@@ -308,26 +335,30 @@ def make_job_command_with_user_switching(username: str,
     return ['sudo', '-H', 'su', '--login', username, '-c', command]
 
 
+@init_db
 def add_job(job_name: str, username: str, run_timestamp: str,
             resources_str: str) -> int:
     """Atomically reserve the next available job id for the user."""
+    assert _DB is not None
     job_submitted_at = time.time()
     # job_id will autoincrement with the null value
-    _CURSOR.execute(
+    _DB.cursor.execute(
         'INSERT INTO jobs VALUES (null, ?, ?, ?, ?, ?, ?, null, ?, 0)',
         (job_name, username, job_submitted_at, JobStatus.INIT.value,
          run_timestamp, None, resources_str))
-    _CONN.commit()
-    rows = _CURSOR.execute('SELECT job_id FROM jobs WHERE run_timestamp=(?)',
-                           (run_timestamp,))
+    _DB.conn.commit()
+    rows = _DB.cursor.execute('SELECT job_id FROM jobs WHERE run_timestamp=(?)',
+                              (run_timestamp,))
     for row in rows:
         job_id = row[0]
     assert job_id is not None
     return job_id
 
 
+@init_db
 def _set_status_no_lock(job_id: int, status: JobStatus) -> None:
     """Setting the status of the job in the database."""
+    assert _DB is not None
     assert status != JobStatus.RUNNING, (
         'Please use set_job_started() to set job status to RUNNING')
     if status.is_terminal():
@@ -339,15 +370,15 @@ def _set_status_no_lock(job_id: int, status: JobStatus) -> None:
         check_end_at_str = ' AND end_at IS NULL'
         if status != JobStatus.FAILED_SETUP:
             check_end_at_str = ''
-        _CURSOR.execute(
+        _DB.cursor.execute(
             'UPDATE jobs SET status=(?), end_at=(?) '
             f'WHERE job_id=(?) {check_end_at_str}',
             (status.value, end_at, job_id))
     else:
-        _CURSOR.execute(
+        _DB.cursor.execute(
             'UPDATE jobs SET status=(?), end_at=NULL '
             'WHERE job_id=(?)', (status.value, job_id))
-    _CONN.commit()
+    _DB.conn.commit()
 
 
 def set_status(job_id: int, status: JobStatus) -> None:
@@ -357,16 +388,19 @@ def set_status(job_id: int, status: JobStatus) -> None:
         _set_status_no_lock(job_id, status)
 
 
+@init_db
 def set_job_started(job_id: int) -> None:
     # TODO(mraheja): remove pylint disabling when filelock version updated.
     # pylint: disable=abstract-class-instantiated
+    assert _DB is not None
     with filelock.FileLock(_get_lock_path(job_id)):
-        _CURSOR.execute(
+        _DB.cursor.execute(
             'UPDATE jobs SET status=(?), start_at=(?), end_at=NULL '
             'WHERE job_id=(?)', (JobStatus.RUNNING.value, time.time(), job_id))
-        _CONN.commit()
+        _DB.conn.commit()
 
 
+@init_db
 def get_status_no_lock(job_id: int) -> Optional[JobStatus]:
     """Get the status of the job with the given id.
 
@@ -375,8 +409,9 @@ def get_status_no_lock(job_id: int) -> Optional[JobStatus]:
     the status in a while loop as in `log_lib._follow_job_logs`. Otherwise, use
     `get_status`.
     """
-    rows = _CURSOR.execute('SELECT status FROM jobs WHERE job_id=(?)',
-                           (job_id,))
+    assert _DB is not None
+    rows = _DB.cursor.execute('SELECT status FROM jobs WHERE job_id=(?)',
+                              (job_id,))
     for (status,) in rows:
         if status is None:
             return None
@@ -391,11 +426,13 @@ def get_status(job_id: int) -> Optional[JobStatus]:
         return get_status_no_lock(job_id)
 
 
+@init_db
 def get_statuses_payload(job_ids: List[Optional[int]]) -> str:
+    assert _DB is not None
     # Per-job lock is not required here, since the staled job status will not
     # affect the caller.
     query_str = ','.join(['?'] * len(job_ids))
-    rows = _CURSOR.execute(
+    rows = _DB.cursor.execute(
         f'SELECT job_id, status FROM jobs WHERE job_id IN ({query_str})',
         job_ids)
     statuses = {job_id: None for job_id in job_ids}
@@ -419,14 +456,17 @@ def load_statuses_payload(
     return statuses
 
 
+@init_db
 def get_latest_job_id() -> Optional[int]:
-    rows = _CURSOR.execute(
+    assert _DB is not None
+    rows = _DB.cursor.execute(
         'SELECT job_id FROM jobs ORDER BY job_id DESC LIMIT 1')
     for (job_id,) in rows:
         return job_id
     return None
 
 
+@init_db
 def get_job_submitted_or_ended_timestamp_payload(job_id: int,
                                                  get_ended_time: bool) -> str:
     """Get the job submitted/ended timestamp.
@@ -440,9 +480,10 @@ def get_job_submitted_or_ended_timestamp_payload(job_id: int,
     `format_job_queue()`), because the job may stay in PENDING if the cluster is
     busy.
     """
+    assert _DB is not None
     field = 'end_at' if get_ended_time else 'submitted_at'
-    rows = _CURSOR.execute(f'SELECT {field} FROM jobs WHERE job_id=(?)',
-                           (job_id,))
+    rows = _DB.cursor.execute(f'SELECT {field} FROM jobs WHERE job_id=(?)',
+                              (job_id,))
     for (timestamp,) in rows:
         return message_utils.encode_payload(timestamp)
     return message_utils.encode_payload(None)
@@ -496,10 +537,12 @@ def _get_records_from_rows(rows) -> List[Dict[str, Any]]:
     return records
 
 
+@init_db
 def _get_jobs(
         user_hash: Optional[str],
         status_list: Optional[List[JobStatus]] = None) -> List[Dict[str, Any]]:
     """Returns jobs with the given fields, sorted by job_id, descending."""
+    assert _DB is not None
     if status_list is None:
         status_list = list(JobStatus)
     status_str_list = [repr(status.value) for status in status_list]
@@ -509,14 +552,16 @@ def _get_jobs(
         # We use the old username field for compatibility.
         filter_str += ' AND username=(?)'
         params.append(user_hash)
-    rows = _CURSOR.execute(
+    rows = _DB.cursor.execute(
         f'SELECT * FROM jobs {filter_str} ORDER BY job_id DESC', params)
     records = _get_records_from_rows(rows)
     return records
 
 
+@init_db
 def _get_jobs_by_ids(job_ids: List[int]) -> List[Dict[str, Any]]:
-    rows = _CURSOR.execute(
+    assert _DB is not None
+    rows = _DB.cursor.execute(
         f"""\
         SELECT * FROM jobs
         WHERE job_id IN ({','.join(['?'] * len(job_ids))})
@@ -527,8 +572,10 @@ def _get_jobs_by_ids(job_ids: List[int]) -> List[Dict[str, Any]]:
     return records
 
 
+@init_db
 def _get_pending_job(job_id: int) -> Optional[Dict[str, Any]]:
-    rows = _CURSOR.execute(
+    assert _DB is not None
+    rows = _DB.cursor.execute(
         'SELECT created_time, submit, run_cmd FROM pending_jobs '
         f'WHERE job_id={job_id!r}')
     for row in rows:
@@ -698,16 +745,18 @@ def update_job_status(job_ids: List[int],
     return statuses
 
 
+@init_db
 def fail_all_jobs_in_progress() -> None:
+    assert _DB is not None
     in_progress_status = [
         status.value for status in JobStatus.nonterminal_statuses()
     ]
-    _CURSOR.execute(
+    _DB.cursor.execute(
         f"""\
         UPDATE jobs SET status=(?)
         WHERE status IN ({','.join(['?'] * len(in_progress_status))})
         """, (JobStatus.FAILED_DRIVER.value, *in_progress_status))
-    _CONN.commit()
+    _DB.conn.commit()
 
 
 def update_status() -> None:
@@ -720,12 +769,14 @@ def update_status() -> None:
     update_job_status(nonterminal_job_ids)
 
 
+@init_db
 def is_cluster_idle() -> bool:
     """Returns if the cluster is idle (no in-flight jobs)."""
+    assert _DB is not None
     in_progress_status = [
         status.value for status in JobStatus.nonterminal_statuses()
     ]
-    rows = _CURSOR.execute(
+    rows = _DB.cursor.execute(
         f"""\
         SELECT COUNT(*) FROM jobs
         WHERE status IN ({','.join(['?'] * len(in_progress_status))})
@@ -905,27 +956,31 @@ def cancel_jobs_encoded_results(jobs: Optional[List[int]],
     return message_utils.encode_payload(cancelled_ids)
 
 
+@init_db
 def get_run_timestamp(job_id: Optional[int]) -> Optional[str]:
     """Returns the relative path to the log file for a job."""
-    _CURSOR.execute(
+    assert _DB is not None
+    _DB.cursor.execute(
         """\
             SELECT * FROM jobs
             WHERE job_id=(?)""", (job_id,))
-    row = _CURSOR.fetchone()
+    row = _DB.cursor.fetchone()
     if row is None:
         return None
     run_timestamp = row[JobInfoLoc.RUN_TIMESTAMP.value]
     return run_timestamp
 
 
+@init_db
 def run_timestamp_with_globbing_payload(job_ids: List[Optional[str]]) -> str:
     """Returns the relative paths to the log files for job with globbing."""
+    assert _DB is not None
     query_str = ' OR '.join(['job_id GLOB (?)'] * len(job_ids))
-    _CURSOR.execute(
+    _DB.cursor.execute(
         f"""\
             SELECT * FROM jobs
             WHERE {query_str}""", job_ids)
-    rows = _CURSOR.fetchall()
+    rows = _DB.cursor.fetchall()
     run_timestamps = {}
     for row in rows:
         job_id = row[JobInfoLoc.JOB_ID.value]

sky/users/permission.py

@@ -30,26 +30,34 @@ class PermissionService:
     """Permission service for SkyPilot API Server."""
 
     def __init__(self):
-        global _enforcer_instance
-        if _enforcer_instance is None:
-            # For different threads, we share the same enforcer instance.
-            with _lock:
-                if _enforcer_instance is None:
-                    _enforcer_instance = self
-                    engine = global_user_state.initialize_and_get_db()
-                    adapter = sqlalchemy_adapter.Adapter(engine)
-                    model_path = os.path.join(os.path.dirname(__file__),
-                                              'model.conf')
-                    enforcer = casbin.Enforcer(model_path, adapter)
-                    self.enforcer = enforcer
-        else:
-            self.enforcer = _enforcer_instance.enforcer
-        with _policy_lock():
-            self._maybe_initialize_policies()
+        self.enforcer = None
+        self.init_lock = threading.Lock()
+
+    def _lazy_initialize(self):
+        if self.enforcer is not None:
+            return
+        with self.init_lock:
+            if self.enforcer is not None:
+                return
+            global _enforcer_instance
+            if _enforcer_instance is None:
+                # For different threads, we share the same enforcer instance.
+                with _lock:
+                    if _enforcer_instance is None:
+                        _enforcer_instance = self
+                        engine = global_user_state.initialize_and_get_db()
+                        adapter = sqlalchemy_adapter.Adapter(engine)
+                        model_path = os.path.join(os.path.dirname(__file__),
+                                                  'model.conf')
+                        enforcer = casbin.Enforcer(model_path, adapter)
+                        self.enforcer = enforcer
+            else:
+                self.enforcer = _enforcer_instance.enforcer
+            with _policy_lock():
+                self._maybe_initialize_policies()
 
     def _maybe_initialize_policies(self) -> None:
         """Initialize policies if they don't already exist."""
-        # TODO(zhwu): we should avoid running this on client side.
         logger.debug(f'Initializing policies in process: {os.getpid()}')
         self._load_policy_no_lock()
 
@@ -128,6 +136,7 @@ class PermissionService:
 
     def add_user_if_not_exists(self, user_id: str) -> None:
         """Add user role relationship."""
+        self._lazy_initialize()
         with _policy_lock():
             self._add_user_if_not_exists_no_lock(user_id)
 
@@ -147,6 +156,7 @@ class PermissionService:
 
     def update_role(self, user_id: str, new_role: str) -> None:
         """Update user role relationship."""
+        self._lazy_initialize()
         with _policy_lock():
             # Get current roles
             self._load_policy_no_lock()
@@ -179,6 +189,7 @@ class PermissionService:
         Returns:
             A list of role names that the user has.
         """
+        self._lazy_initialize()
         self._load_policy_no_lock()
         return self.enforcer.get_roles_for_user(user_id)
 
@@ -191,6 +202,7 @@ class PermissionService:
         # it is a hot path in every request. It is ok to have a stale policy,
         # as long as it is eventually consistent.
         # self._load_policy_no_lock()
+        self._lazy_initialize()
         return self.enforcer.enforce(user_id, path, method)
 
     def _load_policy_no_lock(self):
@@ -199,6 +211,7 @@ class PermissionService:
 
     def load_policy(self):
         """Load policy from storage with lock."""
+        self._lazy_initialize()
         with _policy_lock():
             self._load_policy_no_lock()
 
@@ -214,6 +227,7 @@ class PermissionService:
         For public workspaces, the permission is granted via a wildcard policy
         ('*').
         """
+        self._lazy_initialize()
         if os.getenv(constants.ENV_VAR_IS_SKYPILOT_SERVER) is None:
             # When it is not on API server, we allow all users to access all
             # workspaces, as the workspace check has been done on API server.
@@ -241,6 +255,7 @@ class PermissionService:
                    For public workspaces, this should be ['*'].
                    For private workspaces, this should be specific user IDs.
         """
+        self._lazy_initialize()
         with _policy_lock():
             for user in users:
                 logger.debug(f'Adding workspace policy: user={user}, '
@@ -258,6 +273,7 @@ class PermissionService:
                    For public workspaces, this should be ['*'].
                    For private workspaces, this should be specific user IDs.
         """
+        self._lazy_initialize()
         with _policy_lock():
             self._load_policy_no_lock()
             # Remove all existing policies for this workspace
@@ -271,6 +287,7 @@ class PermissionService:
 
     def remove_workspace_policy(self, workspace_name: str) -> None:
         """Remove workspace policy."""
+        self._lazy_initialize()
         with _policy_lock():
             self.enforcer.remove_filtered_policy(1, workspace_name)
             self.enforcer.save_policy()

sky/utils/admin_policy_utils.py

@@ -3,6 +3,7 @@ import contextlib
 import copy
 import importlib
 from typing import Iterator, Optional, Tuple, Union
+import urllib.parse
 
 import colorama
 
@@ -19,18 +20,34 @@ from sky.utils import ux_utils
 logger = sky_logging.init_logger(__name__)
 
 
-def _get_policy_cls(
-        policy: Optional[str]) -> Optional[admin_policy.AdminPolicy]:
+def _is_url(policy_string: str) -> bool:
+    """Check if the policy string is a URL."""
+    try:
+        parsed = urllib.parse.urlparse(policy_string)
+        return parsed.scheme in ('http', 'https')
+    except Exception:  # pylint: disable=broad-except
+        return False
+
+
+def _get_policy_impl(
+        policy_location: Optional[str]
+) -> Optional[admin_policy.PolicyInterface]:
     """Gets admin-defined policy."""
-    if policy is None:
+    if policy_location is None:
         return None
+
+    if _is_url(policy_location):
+        # Use the built-in URL policy class when an URL is specified.
+        return admin_policy.RestfulAdminPolicy(policy_location)
+
+    # Handle module path format
     try:
-        module_path, class_name = policy.rsplit('.', 1)
+        module_path, class_name = policy_location.rsplit('.', 1)
         module = importlib.import_module(module_path)
     except ImportError as e:
         with ux_utils.print_exception_no_traceback():
             raise ImportError(
-                f'Failed to import policy module: {policy}. '
+                f'Failed to import policy module: {policy_location}. '
                 'Please check if the module is installed in your Python '
                 'environment.') from e
 
@@ -42,13 +59,15 @@ def _get_policy_cls(
                 f'Could not find {class_name} class in module {module_path}. '
                 'Please check with your policy admin for details.') from e
 
-    # Check if the module implements the AdminPolicy interface.
+    # Currently we only allow users to define subclass of AdminPolicy
+    # instead of inheriting from PolicyInterface or PolicyTemplate.
     if not issubclass(policy_cls, admin_policy.AdminPolicy):
         with ux_utils.print_exception_no_traceback():
             raise ValueError(
-                f'Policy class {policy!r} does not implement the AdminPolicy '
-                'interface. Please check with your policy admin for details.')
-    return policy_cls
+                f'Policy class {policy_cls!r} does not implement the '
+                'AdminPolicy interface. Please check with your policy admin '
+                'for details.')
+    return policy_cls()
 
 
 @contextlib.contextmanager
@@ -102,9 +121,9 @@ def apply(
     else:
         dag = entrypoint
 
-    policy = skypilot_config.get_nested(('admin_policy',), None)
-    policy_cls = _get_policy_cls(policy)
-    if policy_cls is None:
+    policy_location = skypilot_config.get_nested(('admin_policy',), None)
+    policy = _get_policy_impl(policy_location)
+    if policy is None:
         return dag, skypilot_config.to_dict()
 
     if at_client_side:
@@ -120,7 +139,7 @@ def apply(
         user_request = admin_policy.UserRequest(task, config, request_options,
                                                 at_client_side)
         try:
-            mutated_user_request = policy_cls.validate_and_mutate(user_request)
+            mutated_user_request = policy.apply(user_request)
         except Exception as e:  # pylint: disable=broad-except
             with ux_utils.print_exception_no_traceback():
                 raise exceptions.UserRequestRejectedByPolicy(

sky/utils/schemas.py

@@ -1149,9 +1149,17 @@ def get_config_schema():
 
     admin_policy_schema = {
         'type': 'string',
-        # Check regex to be a valid python module path
-        'pattern': (r'^[a-zA-Z_][a-zA-Z0-9_]*'
-                    r'(\.[a-zA-Z_][a-zA-Z0-9_]*)+$'),
+        'anyOf': [
+            {
+                # Check regex to be a valid python module path
+                'pattern': (r'^[a-zA-Z_][a-zA-Z0-9_]*'
+                            r'(\.[a-zA-Z_][a-zA-Z0-9_]*)+$'),
+            },
+            {
+                # Check for valid HTTP/HTTPS URL
+                'pattern': r'^https?://.*$',
+            }
+        ]
     }
 
     allowed_clouds = {

{skypilot_nightly-1.0.0.dev20250610.dist-info → skypilot_nightly-1.0.0.dev20250611.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: skypilot-nightly
-Version: 1.0.0.dev20250610
+Version: 1.0.0.dev20250611
 Summary: SkyPilot: Run AI on Any Infra — Unified, Faster, Cheaper.
 Author: SkyPilot Team
 License: Apache 2.0