skypilot-nightly 1.0.0.dev20250905__py3-none-any.whl → 1.0.0.dev20251210__py3-none-any.whl

sky/server/requests/payloads.py

@@ -60,6 +60,11 @@ EXTERNAL_LOCAL_ENV_VARS = [
     'AWS_ACCESS_KEY_ID',
     'AWS_SECRET_ACCESS_KEY',
     'AWS_SESSION_TOKEN',
+    # Allow overriding the Azure authentication.
+    'AZURE_CLIENT_ID',
+    'AZURE_CLIENT_SECRET',
+    'AZURE_TENANT_ID',
+    'AZURE_SUBSCRIPTION_ID',
     # Allow overriding the GCP authentication.
     'GOOGLE_APPLICATION_CREDENTIALS',
     # Allow overriding the kubeconfig.
@@ -67,7 +72,6 @@ EXTERNAL_LOCAL_ENV_VARS = [
 ]
 
 
-@annotations.lru_cache(scope='global')
 def request_body_env_vars() -> dict:
     env_vars = {}
     for env_var in os.environ:
@@ -78,14 +82,22 @@ def request_body_env_vars() -> dict:
         if common.is_api_server_local() and env_var in EXTERNAL_LOCAL_ENV_VARS:
             env_vars[env_var] = os.environ[env_var]
     env_vars[constants.USER_ID_ENV_VAR] = common_utils.get_user_hash()
-    env_vars[constants.USER_ENV_VAR] = common_utils.get_current_user_name()
+    env_vars[constants.USER_ENV_VAR] = common_utils.get_local_user_name()
     env_vars[
         usage_constants.USAGE_RUN_ID_ENV_VAR] = usage_lib.messages.usage.run_id
+    if not common.is_api_server_local():
+        # Used in job controller, for local API server, keep the
+        # SKYPILOT_CONFIG env var to use the config for the managed job.
+        env_vars.pop(skypilot_config.ENV_VAR_SKYPILOT_CONFIG, None)
     # Remove the path to config file, as the config content is included in the
     # request body and will be merged with the config on the server side.
-    env_vars.pop(skypilot_config.ENV_VAR_SKYPILOT_CONFIG, None)
     env_vars.pop(skypilot_config.ENV_VAR_GLOBAL_CONFIG, None)
     env_vars.pop(skypilot_config.ENV_VAR_PROJECT_CONFIG, None)
+    # Remove the config related env vars, as the client config override
+    # should be passed in the request body.
+    # Any new environment variables that are server-specific should
+    # use SKYPILOT_SERVER_ENV_VAR_PREFIX.
+    env_vars.pop(constants.ENV_VAR_DB_CONNECTION_URI, None)
     return env_vars
 
 
@@ -309,8 +321,13 @@ class StatusBody(RequestBody):
     cluster_names: Optional[List[str]] = None
     refresh: common_lib.StatusRefreshMode = common_lib.StatusRefreshMode.NONE
     all_users: bool = True
-    # TODO (kyuds): default to False post 0.10.5
+    # TODO (kyuds): default to False post 0.12.0
     include_credentials: bool = True
+    # Only return fields that are needed for the
+    # dashboard / CLI summary response
+    summary_response: bool = False
+    # Include the cluster handle in the response
+    include_handle: bool = True
 
 
 class StartBody(RequestBody):
@@ -355,9 +372,10 @@ class CancelBody(RequestBody):
         return kwargs
 
 
-class ClusterNameBody(RequestBody):
+class ProvisionLogsBody(RequestBody):
     """Cluster node."""
     cluster_name: str
+    worker: Optional[int] = None
 
 
 class ClusterJobBody(RequestBody):
@@ -458,6 +476,7 @@ class VolumeApplyBody(RequestBody):
     size: Optional[str] = None
     config: Optional[Dict[str, Any]] = None
     labels: Optional[Dict[str, str]] = None
+    use_existing: Optional[bool] = None
 
 
 class VolumeDeleteBody(RequestBody):
@@ -470,6 +489,17 @@ class VolumeListBody(RequestBody):
     pass
 
 
+class VolumeValidateBody(RequestBody):
+    """The request body for the volume validate endpoint."""
+    name: Optional[str] = None
+    volume_type: Optional[str] = None
+    infra: Optional[str] = None
+    size: Optional[str] = None
+    labels: Optional[Dict[str, str]] = None
+    config: Optional[Dict[str, Any]] = None
+    use_existing: Optional[bool] = None
+
+
 class EndpointsBody(RequestBody):
     """The request body for the endpoint."""
     cluster: str
@@ -507,6 +537,14 @@ class JobsQueueBody(RequestBody):
     skip_finished: bool = False
     all_users: bool = False
     job_ids: Optional[List[int]] = None
+
+
+class JobsQueueV2Body(RequestBody):
+    """The request body for the jobs queue endpoint."""
+    refresh: bool = False
+    skip_finished: bool = False
+    all_users: bool = False
+    job_ids: Optional[List[int]] = None
     user_match: Optional[str] = None
     workspace_match: Optional[str] = None
     name_match: Optional[str] = None
@@ -514,6 +552,9 @@ class JobsQueueBody(RequestBody):
     page: Optional[int] = None
     limit: Optional[int] = None
     statuses: Optional[List[str]] = None
+    # The fields to return in the response.
+    # Refer to the fields in the `class ManagedJobRecord` in `response.py`
+    fields: Optional[List[str]] = None
 
 
 class JobsCancelBody(RequestBody):
@@ -546,6 +587,8 @@ class RequestStatusBody(pydantic.BaseModel):
     """The request body for the API request status endpoint."""
     request_ids: Optional[List[str]] = None
     all_status: bool = False
+    limit: Optional[int] = None
+    fields: Optional[List[str]] = None
 
 
 class ServeUpBody(RequestBody):
@@ -627,6 +670,11 @@ class KubernetesNodeInfoRequestBody(RequestBody):
     context: Optional[str] = None
 
 
+class SlurmNodeInfoRequestBody(RequestBody):
+    """The request body for the slurm node info endpoint."""
+    slurm_cluster_name: Optional[str] = None
+
+
 class ListAcceleratorsBody(RequestBody):
     """The request body for the list accelerators endpoint."""
     gpus_only: bool = True
@@ -651,12 +699,13 @@ class ListAcceleratorCountsBody(RequestBody):
 class LocalUpBody(RequestBody):
     """The request body for the local up endpoint."""
     gpus: bool = True
-    ips: Optional[List[str]] = None
-    ssh_user: Optional[str] = None
-    ssh_key: Optional[str] = None
-    cleanup: bool = False
-    context_name: Optional[str] = None
-    password: Optional[str] = None
+    name: Optional[str] = None
+    port_start: Optional[int] = None
+
+
+class LocalDownBody(RequestBody):
+    """The request body for the local down endpoint."""
+    name: Optional[str] = None
 
 
 class SSHUpBody(RequestBody):
@@ -696,19 +745,22 @@ class JobsDownloadLogsBody(RequestBody):
 
 class JobsPoolApplyBody(RequestBody):
     """The request body for the jobs pool apply endpoint."""
-    task: str
+    task: Optional[str] = None
+    workers: Optional[int] = None
     pool_name: str
     mode: serve.UpdateMode
 
     def to_kwargs(self) -> Dict[str, Any]:
         kwargs = super().to_kwargs()
-        dag = common.process_mounts_in_task_on_api_server(self.task,
-                                                          self.env_vars,
-                                                          workdir_only=False)
-        assert len(
-            dag.tasks) == 1, ('Must only specify one task in the DAG for '
-                              'a pool.', dag)
-        kwargs['task'] = dag.tasks[0]
+        if self.task is not None:
+            dag = common.process_mounts_in_task_on_api_server(
+                self.task, self.env_vars, workdir_only=False)
+            assert len(
+                dag.tasks) == 1, ('Must only specify one task in the DAG for '
+                                  'a pool.', dag)
+            kwargs['task'] = dag.tasks[0]
+        else:
+            kwargs['task'] = None
         return kwargs
 
 
@@ -779,6 +831,12 @@ class GetConfigBody(RequestBody):
 class CostReportBody(RequestBody):
     """The request body for the cost report endpoint."""
     days: Optional[int] = 30
+    # we use hashes instead of names to avoid the case where
+    # the name is not unique
+    cluster_hashes: Optional[List[str]] = None
+    # Only return fields that are needed for the dashboard
+    # summary page
+    dashboard_summary_response: bool = False
 
 
 class RequestPayload(BasePayload):
@@ -801,3 +859,9 @@ class RequestPayload(BasePayload):
     status_msg: Optional[str] = None
     should_retry: bool = False
     finished_at: Optional[float] = None
+
+
+class SlurmGpuAvailabilityRequestBody(RequestBody):
+    """Request body for getting Slurm real-time GPU availability."""
+    name_filter: Optional[str] = None
+    quantity_filter: Optional[int] = None

sky/server/requests/preconditions.py

@@ -90,7 +90,7 @@ class Precondition(abc.ABC):
         while True:
             if self.timeout > 0 and time.time() - start_time > self.timeout:
                 # Cancel the request on timeout.
-                api_requests.set_request_failed(
+                await api_requests.set_request_failed_async(
                     self.request_id,
                     exceptions.RequestCancelled(
                         f'Request {self.request_id} precondition wait timed '
@@ -98,13 +98,15 @@ class Precondition(abc.ABC):
                 return False
 
             # Check if the request has been cancelled
-            request = await api_requests.get_request_async(self.request_id)
+            request = await api_requests.get_request_async(self.request_id,
+                                                           fields=['status'])
             if request is None:
                 logger.error(f'Request {self.request_id} not found')
                 return False
             if request.status == api_requests.RequestStatus.CANCELLED:
                 logger.debug(f'Request {self.request_id} cancelled')
                 return False
+            del request
 
             try:
                 met, status_msg = await self.check()
@@ -112,13 +114,11 @@ class Precondition(abc.ABC):
                     return True
                 if status_msg is not None and status_msg != last_status_msg:
                     # Update the status message if it has changed.
-                    async with api_requests.update_request_async(
-                            self.request_id) as req:
-                        assert req is not None, self.request_id
-                        req.status_msg = status_msg
+                    await api_requests.update_status_msg_async(
+                        self.request_id, status_msg)
                     last_status_msg = status_msg
             except (Exception, SystemExit, KeyboardInterrupt) as e:  # pylint: disable=broad-except
-                api_requests.set_request_failed(self.request_id, e)
+                await api_requests.set_request_failed_async(self.request_id, e)
                 logger.info(f'Request {self.request_id} failed due to '
                             f'{common_utils.format_exception(e)}')
                 return False
@@ -146,10 +146,9 @@ class ClusterStartCompletePrecondition(Precondition):
         self.cluster_name = cluster_name
 
     async def check(self) -> Tuple[bool, Optional[str]]:
-        cluster_record = global_user_state.get_cluster_from_name(
+        cluster_status = global_user_state.get_status_from_cluster_name(
             self.cluster_name)
-        if (cluster_record and
-                cluster_record['status'] is status_lib.ClusterStatus.UP):
+        if cluster_status is status_lib.ClusterStatus.UP:
             # Shortcut for started clusters, ignore cluster not found
             # since the cluster record might not yet be created by the
             # launch task.
@@ -165,11 +164,14 @@ class ClusterStartCompletePrecondition(Precondition):
         requests = await api_requests.get_request_tasks_async(
             req_filter=api_requests.RequestTaskFilter(
                 status=[
-                    api_requests.RequestStatus.RUNNING,
-                    api_requests.RequestStatus.PENDING
+                    api_requests.RequestStatus.PENDING,
+                    api_requests.RequestStatus.RUNNING
                 ],
                 include_request_names=['sky.launch', 'sky.start'],
-                cluster_names=[self.cluster_name]))
+                cluster_names=[self.cluster_name],
+                # Only get the request ID to avoid fetching the whole request.
+                # We're only interested in the count, not the whole request.
+                fields=['request_id']))
         if len(requests) == 0:
             # No running or pending tasks, the start process is done.
             return True, None

sky/server/requests/request_names.py

@@ -0,0 +1,123 @@
+"""Request names."""
+import enum
+
+
+class RequestName(str, enum.Enum):
+    """Enum of all the request names."""
+    # General requests
+    CHECK = 'check'
+    ENABLED_CLOUDS = 'enabled_clouds'
+    REALTIME_KUBERNETES_GPU_AVAILABILITY = (
+        'realtime_kubernetes_gpu_availability')
+    KUBERNETES_NODE_INFO = 'kubernetes_node_info'
+    REALTIME_SLURM_GPU_AVAILABILITY = 'realtime_slurm_gpu_availability'
+    SLURM_NODE_INFO = 'slurm_node_info'
+    STATUS_KUBERNETES = 'status_kubernetes'
+    LIST_ACCELERATORS = 'list_accelerators'
+    LIST_ACCELERATOR_COUNTS = 'list_accelerator_counts'
+    OPTIMIZE = 'optimize'
+    # Cluster requests
+    CLUSTER_LAUNCH = 'launch'
+    CLUSTER_EXEC = 'exec'
+    CLUSTER_STOP = 'stop'
+    CLUSTER_STATUS = 'status'
+    CLUSTER_ENDPOINTS = 'endpoints'
+    CLUSTER_DOWN = 'down'
+    CLUSTER_START = 'start'
+    CLUSTER_AUTOSTOP = 'autostop'
+    CLUSTER_QUEUE = 'queue'
+    CLUSTER_JOB_STATUS = 'job_status'
+    CLUSTER_JOB_CANCEL = 'cancel'
+    CLUSTER_JOB_LOGS = 'logs'
+    CLUSTER_JOB_DOWNLOAD_LOGS = 'download_logs'
+    CLUSTER_COST_REPORT = 'cost_report'
+    # Storage requests
+    STORAGE_LS = 'storage_ls'
+    STORAGE_DELETE = 'storage_delete'
+    # Local requests
+    LOCAL_UP = 'local_up'
+    LOCAL_DOWN = 'local_down'
+    # API requests
+    API_CANCEL = 'api_cancel'
+    ALL_CONTEXTS = 'all_contexts'
+    # Managed jobs requests
+    JOBS_LAUNCH = 'jobs.launch'
+    JOBS_QUEUE = 'jobs.queue'
+    JOBS_QUEUE_V2 = 'jobs.queue_v2'
+    JOBS_CANCEL = 'jobs.cancel'
+    JOBS_LOGS = 'jobs.logs'
+    JOBS_DOWNLOAD_LOGS = 'jobs.download_logs'
+    JOBS_POOL_APPLY = 'jobs.pool_apply'
+    JOBS_POOL_DOWN = 'jobs.pool_down'
+    JOBS_POOL_STATUS = 'jobs.pool_status'
+    JOBS_POOL_LOGS = 'jobs.pool_logs'
+    JOBS_POOL_SYNC_DOWN_LOGS = 'jobs.pool_sync_down_logs'
+    # Serve requests
+    SERVE_UP = 'serve.up'
+    SERVE_UPDATE = 'serve.update'
+    SERVE_DOWN = 'serve.down'
+    SERVE_TERMINATE_REPLICA = 'serve.terminate_replica'
+    SERVE_STATUS = 'serve.status'
+    SERVE_LOGS = 'serve.logs'
+    SERVE_SYNC_DOWN_LOGS = 'serve.sync_down_logs'
+    # Volumes requests
+    VOLUME_LIST = 'volume_list'
+    VOLUME_DELETE = 'volume_delete'
+    VOLUME_APPLY = 'volume_apply'
+    # Workspaces requests
+    WORKSPACES_GET = 'workspaces.get'
+    WORKSPACES_UPDATE = 'workspaces.update'
+    WORKSPACES_CREATE = 'workspaces.create'
+    WORKSPACES_DELETE = 'workspaces.delete'
+    WORKSPACES_GET_CONFIG = 'workspaces.get_config'
+    WORKSPACES_UPDATE_CONFIG = 'workspaces.update_config'
+    # SSH node pools requests
+    SSH_NODE_POOLS_UP = 'ssh_node_pools.up'
+    SSH_NODE_POOLS_DOWN = 'ssh_node_pools.down'
+    # Internal request daemons
+    REQUEST_DAEMON_STATUS_REFRESH = 'status-refresh'
+    REQUEST_DAEMON_VOLUME_REFRESH = 'volume-refresh'
+    REQUEST_DAEMON_MANAGED_JOB_STATUS_REFRESH = 'managed-job-status-refresh'
+    REQUEST_DAEMON_SKY_SERVE_STATUS_REFRESH = 'sky-serve-status-refresh'
+    REQUEST_DAEMON_POOL_STATUS_REFRESH = 'pool-status-refresh'
+
+    def __repr__(self):
+        return self.value
+
+    def __str__(self):
+        return self.value
+
+
+class AdminPolicyRequestName(str, enum.Enum):
+    """Enum of all the request names that are
+    used for admin policy application."""
+    # General requests
+    # validate call is not stored in the request db,
+    # but is defined here for admin policy application
+    VALIDATE = 'validate'
+    OPTIMIZE = RequestName.OPTIMIZE.value
+    # Cluster requests
+    CLUSTER_LAUNCH = RequestName.CLUSTER_LAUNCH.value
+    CLUSTER_EXEC = RequestName.CLUSTER_EXEC.value
+    # Jobs requests
+    JOBS_LAUNCH = RequestName.JOBS_LAUNCH.value
+    # jobs launch controller request is not stored in the request db,
+    # but is defined here for admin policy application
+    JOBS_LAUNCH_CONTROLLER = 'jobs.launch_controller'
+    JOBS_POOL_APPLY = RequestName.JOBS_POOL_APPLY.value
+    JOBS_LAUNCH_CLUSTER = 'jobs.launch_cluster'
+    # Serve requests
+    SERVE_UP = RequestName.SERVE_UP.value
+    # serve launch controller request is not stored in the request db,
+    # but is defined here for admin policy application
+    SERVE_LAUNCH_CONTROLLER = 'serve.launch_controller'
+    SERVE_UPDATE = RequestName.SERVE_UPDATE.value
+    # serve launch replica request is not stored in the request db,
+    # but is defined here for admin policy application
+    SERVE_LAUNCH_REPLICA = 'serve.launch_replica'
+
+    def __repr__(self):
+        return self.value
+
+    def __str__(self):
+        return self.value