skypilot-nightly 1.0.0.dev20250621__py3-none-any.whl → 1.0.0.dev20250624__py3-none-any.whl

sky/jobs/utils.py

@@ -5,6 +5,7 @@ jobs.constants.MANAGED_JOBS_VERSION and handle the API change in the
 ManagedJobCodeGen.
 """
 import collections
+import datetime
 import enum
 import os
 import pathlib
@@ -33,7 +34,10 @@ from sky.skylet import constants
 from sky.skylet import job_lib
 from sky.skylet import log_lib
 from sky.usage import usage_lib
+from sky.utils import annotations
+from sky.utils import command_runner
 from sky.utils import common_utils
+from sky.utils import controller_utils
 from sky.utils import infra_utils
 from sky.utils import log_utils
 from sky.utils import message_utils
@@ -124,6 +128,114 @@ def terminate_cluster(cluster_name: str, max_retry: int = 6) -> None:
             time.sleep(backoff.current_backoff())
 
 
+def _check_consolidation_mode_consistency(
+        current_is_consolidation_mode: bool) -> None:
+    """Check the consistency of the consolidation mode."""
+    # Check whether the consolidation mode config is changed.
+    if current_is_consolidation_mode:
+        controller_cn = (
+            controller_utils.Controllers.JOBS_CONTROLLER.value.cluster_name)
+        if global_user_state.get_cluster_from_name(controller_cn) is not None:
+            with ux_utils.print_exception_no_traceback():
+                raise exceptions.InconsistentConsolidationModeError(
+                    f'{colorama.Fore.RED}Consolidation mode is '
+                    f'enabled, but the controller cluster '
+                    f'{controller_cn} is still running. Please '
+                    'terminate the controller cluster first.'
+                    f'{colorama.Style.RESET_ALL}')
+    else:
+        all_jobs = managed_job_state.get_managed_jobs()
+        if all_jobs:
+            nonterminal_jobs = (
+                managed_job_state.get_nonterminal_job_ids_by_name(
+                    None, all_users=True))
+            if nonterminal_jobs:
+                with ux_utils.print_exception_no_traceback():
+                    raise exceptions.InconsistentConsolidationModeError(
+                        f'{colorama.Fore.RED}Consolidation mode '
+                        'is disabled, but there are still '
+                        f'{len(nonterminal_jobs)} managed jobs '
+                        'running. Please terminate those jobs '
+                        f'first.{colorama.Style.RESET_ALL}')
+            else:
+                logger.warning(
+                    f'{colorama.Fore.YELLOW}Consolidation mode is disabled, '
+                    f'but there are {len(all_jobs)} jobs from previous '
+                    'consolidation mode. Reset the `jobs.controller.'
+                    'consolidation_mode` to `true` and run `sky jobs queue` '
+                    'to see those jobs. Switching to normal mode will '
+                    f'lose the job history.{colorama.Style.RESET_ALL}')
+
+
+# Whether to use consolidation mode or not. When this is enabled, the managed
+# jobs controller will not be running on a separate cluster, but locally on the
+# API Server. Under the hood, we submit the job monitoring logic as processes
+# directly in the API Server.
+# Use LRU Cache so that the check is only done once.
+@annotations.lru_cache(scope='request', maxsize=1)
+def is_consolidation_mode() -> bool:
+    consolidation_mode = skypilot_config.get_nested(
+        ('jobs', 'controller', 'consolidation_mode'), default_value=False)
+    _check_consolidation_mode_consistency(consolidation_mode)
+    return consolidation_mode
+
+
+def get_ha_dump_script_path(job_id: int) -> pathlib.Path:
+    """Get the path to the HA dump script for a job."""
+    return pathlib.Path(constants.PERSISTENT_RUN_SCRIPT_DIR).expanduser(
+    ).resolve() / f'sky_job_{job_id}'
+
+
+def ha_recovery_for_consolidation_mode():
+    """Recovery logic for HA mode."""
+    # No setup recovery is needed in consolidation mode, as the API server
+    # already has all runtime installed. Directly start jobs recovery here.
+    # Refers to sky/templates/kubernetes-ray.yml.j2 for more details.
+    runner = command_runner.LocalProcessCommandRunner()
+    with open(constants.HA_PERSISTENT_RECOVERY_LOG_PATH, 'w',
+              encoding='utf-8') as f:
+        start = time.time()
+        f.write(f'Starting HA recovery at {datetime.datetime.now()}\n')
+        for job in managed_job_state.get_managed_jobs():
+            job_id = job['job_id']
+            controller_pid = job['controller_pid']
+
+            # In consolidation mode, it is possible that only the API server
+            # process is restarted, and the controller process is not. In such
+            # case, we don't need to do anything and the controller process will
+            # just keep running.
+            if controller_pid is not None:
+                try:
+                    if _controller_process_alive(controller_pid, job_id):
+                        f.write(f'Controller pid {controller_pid} for '
+                                f'job {job_id} is still running. '
+                                'Skipping recovery.\n')
+                        continue
+                except Exception:  # pylint: disable=broad-except
+                    # _controller_process_alive may raise if psutil fails; we
+                    # should not crash the recovery logic because of this.
+                    f.write('Error checking controller pid '
+                            f'{controller_pid} for job {job_id}\n')
+
+            if job['schedule_state'] not in [
+                    managed_job_state.ManagedJobScheduleState.DONE,
+                    managed_job_state.ManagedJobScheduleState.WAITING
+            ]:
+                dump_script_path = get_ha_dump_script_path(job_id)
+                if not dump_script_path.exists():
+                    f.write(f'Job {job_id}\'s recovery file ({dump_script_path}'
+                            ') does not exist. Skipping recovery. Job '
+                            f'schedule state: {job["schedule_state"]}\n')
+                    continue
+                with open(dump_script_path, 'r', encoding='utf-8') as script_f:
+                    script = script_f.read()
+                runner.run(script)
+                f.write(f'Job {job_id} (file: {dump_script_path}) completed '
+                        f'recovery at {datetime.datetime.now()}\n')
+        f.write(f'HA recovery completed at {datetime.datetime.now()}\n')
+        f.write(f'Total recovery time: {time.time() - start} seconds\n')
+
+
 def get_job_status(backend: 'backends.CloudVmRayBackend',
                    cluster_name: str) -> Optional['job_lib.JobStatus']:
     """Check the status of the job running on a managed job cluster.
@@ -157,9 +269,8 @@ def _controller_process_alive(pid: int, job_id: int) -> bool:
     """Check if the controller process is alive."""
     try:
         process = psutil.Process(pid)
-        # The last two args of the command line should be --job-id <id>
-        job_args = process.cmdline()[-2:]
-        return process.is_running() and job_args == ['--job-id', str(job_id)]
+        cmd_str = ' '.join(process.cmdline())
+        return process.is_running() and f'--job-id {job_id}' in cmd_str
     except psutil.NoSuchProcess:
         return False
 
@@ -1136,7 +1247,6 @@ def format_job_table(
         'TASK',
         *(['WORKSPACE'] if show_workspace else []),
         'NAME',
-        'PRIORITY',
         *user_cols,
         'REQUESTED',
         'SUBMITTED',
@@ -1208,7 +1318,6 @@ def format_job_table(
             submitted_at = None
             end_at: Optional[int] = 0
             recovery_cnt = 0
-            priority = job_tasks[0].get('priority', '-')
             managed_job_status, current_task_id = _get_job_status_from_tasks(
                 job_tasks)
             for task in job_tasks:
@@ -1244,7 +1353,6 @@ def format_job_table(
                 '',
                 *([''] if show_workspace else []),
                 job_name,
-                str(priority),
                 *user_values,
                 '-',
                 submitted,
@@ -1275,13 +1383,11 @@ def format_job_table(
             submitted = log_utils.readable_time_duration(task['submitted_at'])
             user_values = get_user_column_values(task)
             task_workspace = '-' if len(job_tasks) > 1 else workspace
-            priority = task.get('priority', '-')
             values = [
                 task['job_id'] if len(job_tasks) == 1 else ' \u21B3',
                 task['task_id'] if len(job_tasks) > 1 else '-',
                 *([task_workspace] if show_workspace else []),
                 task['task_name'],
-                str(priority),
                 *user_values,
                 task['resources'],
                 # SUBMITTED

sky/resources.py

@@ -2296,7 +2296,7 @@ def parse_time_minutes(time: str) -> int:
 def parse_memory_resource(resource_qty_str: Union[str, int, float],
                           field_name: str,
                           ret_type: type = int,
-                          unit: str = 'g',
+                          unit: str = 'gb',
                           allow_plus: bool = False,
                           allow_x: bool = False,
                           allow_rounding: bool = False) -> str:

sky/server/requests/payloads.py

@@ -5,7 +5,6 @@ kwargs for the payloads, otherwise, we have to keep the default values the sync
 with the backend functions. The benefit of having the default values in the
 payloads is that a user can find the default values in the Restful API docs.
 """
-import getpass
 import os
 import typing
 from typing import Any, Dict, List, Optional, Tuple, Union
@@ -58,8 +57,7 @@ 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] = os.getenv(constants.USER_ENV_VAR,
-                                                 getpass.getuser())
+    env_vars[constants.USER_ENV_VAR] = common_utils.get_current_user_name()
     env_vars[
         usage_constants.USAGE_RUN_ID_ENV_VAR] = usage_lib.messages.usage.run_id
     # Remove the path to config file, as the config content is included in the
@@ -613,3 +611,8 @@ class UpdateConfigBody(RequestBody):
 class GetConfigBody(RequestBody):
     """The request body for getting the entire SkyPilot configuration."""
     pass
+
+
+class CostReportBody(RequestBody):
+    """The request body for the cost report endpoint."""
+    days: Optional[int] = 30

sky/server/requests/requests.py

@@ -327,6 +327,26 @@ def refresh_cluster_status_event():
         time.sleep(server_constants.CLUSTER_REFRESH_DAEMON_INTERVAL_SECONDS)
 
 
+def managed_job_status_refresh_event():
+    """Refresh the managed job status for controller consolidation mode."""
+    # pylint: disable=import-outside-toplevel
+    from sky.jobs import utils as managed_job_utils
+    if not managed_job_utils.is_consolidation_mode():
+        return
+    # We run the recovery logic before starting the event loop as those two are
+    # conflicting. Check PERSISTENT_RUN_RESTARTING_SIGNAL_FILE for details.
+    from sky.utils import controller_utils
+    if controller_utils.high_availability_specified(
+            controller_utils.Controllers.JOBS_CONTROLLER.value.cluster_name):
+        managed_job_utils.ha_recovery_for_consolidation_mode()
+    # After recovery, we start the event loop.
+    from sky.skylet import events
+    event = events.ManagedJobEvent()
+    while True:
+        time.sleep(events.EVENT_CHECKING_INTERVAL_SECONDS)
+        event.run()
+
+
 @dataclasses.dataclass
 class InternalRequestDaemon:
     id: str
@@ -341,7 +361,10 @@ INTERNAL_REQUEST_DAEMONS = [
     # cluster being stopped or down when `sky status -r` is called.
     InternalRequestDaemon(id='skypilot-status-refresh-daemon',
                           name='status',
-                          event_fn=refresh_cluster_status_event)
+                          event_fn=refresh_cluster_status_event),
+    InternalRequestDaemon(id='managed-job-status-refresh-daemon',
+                          name='managed-job-status',
+                          event_fn=managed_job_status_refresh_event),
 ]
 
 

sky/server/server.py

@@ -1044,13 +1044,14 @@ async def download(download_body: payloads.DownloadBody) -> None:
                                     detail=f'Error creating zip file: {str(e)}')
 
 
-@app.get('/cost_report')
-async def cost_report(request: fastapi.Request) -> None:
+@app.post('/cost_report')
+async def cost_report(request: fastapi.Request,
+                      cost_report_body: payloads.CostReportBody) -> None:
     """Gets the cost report of a cluster."""
     executor.schedule_request(
         request_id=request.state.request_id,
         request_name='cost_report',
-        request_body=payloads.RequestBody(),
+        request_body=cost_report_body,
         func=core.cost_report,
         schedule_type=requests_lib.ScheduleType.SHORT,
     )

sky/skylet/constants.py

@@ -401,6 +401,8 @@ PERSISTENT_RUN_SCRIPT_DIR = '~/.sky/.controller_recovery_task_run'
 PERSISTENT_RUN_RESTARTING_SIGNAL_FILE = (
     '~/.sky/.controller_recovery_restarting_signal')
 
+HA_PERSISTENT_RECOVERY_LOG_PATH = '/tmp/ha_recovery.log'
+
 # The placeholder for the local skypilot config path in file mounts for
 # controllers.
 LOCAL_SKYPILOT_CONFIG_PATH_PLACEHOLDER = 'skypilot:local_skypilot_config_path'
@@ -436,30 +438,20 @@ LOGGING_CONFIG_DIR = '~/.sky/logging'
 
 # Resources constants
 TIME_UNITS = {
-    's': 1 / 60,
-    'sec': 1 / 60,
     'm': 1,
-    'min': 1,
     'h': 60,
-    'hr': 60,
     'd': 24 * 60,
-    'day': 24 * 60,
+    'w': 7 * 24 * 60,
 }
 
 TIME_PATTERN: str = (
     f'^[0-9]+({"|".join([unit.lower() for unit in TIME_UNITS])})?$/i')
 
 MEMORY_SIZE_UNITS = {
-    'b': 1,
-    'k': 2**10,
     'kb': 2**10,
-    'm': 2**20,
     'mb': 2**20,
-    'g': 2**30,
     'gb': 2**30,
-    't': 2**40,
     'tb': 2**40,
-    'p': 2**50,
     'pb': 2**50,
 }
 
@@ -472,3 +464,5 @@ MEMORY_SIZE_PLUS_PATTERN = f'{MEMORY_SIZE_PATTERN[:-3]}+?$/i'
 MIN_PRIORITY = -1000
 MAX_PRIORITY = 1000
 DEFAULT_PRIORITY = 0
+
+COST_REPORT_DEFAULT_DAYS = 30

sky/task.py

@@ -342,8 +342,7 @@ class Task:
         self.resources: Union[List[sky.Resources],
                               Set[sky.Resources]] = {sky.Resources()}
         self._service: Optional[service_spec.SkyServiceSpec] = None
-        # The priority of the managed job running this task.
-        self._job_priority: Optional[int] = None
+
         # Resources that this task cannot run on.
         self.blocked_resources = blocked_resources
 
@@ -712,10 +711,6 @@ class Task:
             service = service_spec.SkyServiceSpec.from_yaml_config(service)
         task.set_service(service)
 
-        job = config.pop('job', None)
-        if job is not None and 'priority' in job:
-            task.set_job_priority(job['priority'])
-
         assert not config, f'Invalid task args: {config.keys()}'
         return task
 
@@ -976,23 +971,6 @@ class Task:
         self._service = service
         return self
 
-    @property
-    def job_priority(self) -> Optional[int]:
-        """The priority of the managed job running this task."""
-        return self._job_priority
-
-    def set_job_priority(self, priority: int) -> 'Task':
-        """Sets the job priority for this task.
-
-        Args:
-          priority: an integer between 0 and 1000.
-
-        Returns:
-          self: The current task, with job priority set.
-        """
-        self._job_priority = priority
-        return self
-
     def set_time_estimator(self, func: Callable[['sky.Resources'],
                                                 int]) -> 'Task':
         """Sets a func mapping resources to estimated time (secs).
@@ -1436,9 +1414,6 @@ class Task:
         if self.service is not None:
             add_if_not_none('service', self.service.to_yaml_config())
 
-        if self.job_priority is not None:
-            add_if_not_none('job', {'priority': self.job_priority})
-
         add_if_not_none('num_nodes', self.num_nodes)
 
         if self.inputs is not None:

sky/templates/jobs-controller.yaml.j2

@@ -31,7 +31,9 @@ setup: |
   {% endif %}
 
 run: |
+  {%- if consolidation_mode_job_id is none %}
   {{ sky_activate_python_env }}
+  {%- endif %}
 
   # Write env vars to a file
   {%- for env_name, env_value in controller_envs.items() %}
@@ -42,9 +44,18 @@ run: |
   # Note: The job is already in the `spot` table, marked as PENDING.
   # CloudVmRayBackend._exec_code_on_head() calls
   # managed_job_codegen.set_pending() before we get here.
-  python -u -m sky.jobs.scheduler {{remote_user_yaml_path}} \
+  {%- if consolidation_mode_job_id is not none %}
+  {{sky_python_cmd}} \
+  {%- else %}
+  python \
+  {%- endif %}
+    -u -m sky.jobs.scheduler {{remote_user_yaml_path}} \
     --user-yaml-path {{remote_original_user_yaml_path}} \
+    {%- if consolidation_mode_job_id is not none %}
+    --job-id {{consolidation_mode_job_id}} \
+    {%- else %}
     --job-id $SKYPILOT_INTERNAL_JOB_ID \
+    {%- endif %}
     --env-file {{remote_env_file_path}} \
     --priority {{priority}}
 

sky/templates/kubernetes-ray.yml.j2

@@ -641,7 +641,7 @@ available_node_types:
               {% if high_availability %}
               mkdir -p {{k8s_high_availability_deployment_run_script_dir}}
               if [ -f {{k8s_high_availability_deployment_volume_mount_path}}/k8s_container_ready ]; then
-                SKYPILOT_HA_RECOVERY_LOG="/tmp/ha_recovery.log"
+                SKYPILOT_HA_RECOVERY_LOG="{{ha_recovery_log_path}}"
                 echo "Starting HA recovery at $(date)" >> $SKYPILOT_HA_RECOVERY_LOG
                 start_time=$SECONDS
                 retry_count=0

sky/utils/admin_policy_utils.py

@@ -140,13 +140,17 @@ def apply(
                                                 at_client_side)
         try:
             mutated_user_request = policy.apply(user_request)
+        # Avoid duplicate exception wrapping.
+        except exceptions.UserRequestRejectedByPolicy as e:
+            with ux_utils.print_exception_no_traceback():
+                raise e
         except Exception as e:  # pylint: disable=broad-except
             with ux_utils.print_exception_no_traceback():
                 raise exceptions.UserRequestRejectedByPolicy(
                     f'{colorama.Fore.RED}User request rejected by policy '
                     f'{policy!r}{colorama.Fore.RESET}: '
                     f'{common_utils.format_exception(e, use_bracket=True)}'
-                ) from e
+                ) from None
         if mutated_config is None:
             mutated_config = mutated_user_request.skypilot_config
         else:

sky/utils/cli_utils/status_utils.py

@@ -7,7 +7,6 @@ import colorama
 
 from sky import backends
 from sky.utils import common_utils
-from sky.utils import controller_utils
 from sky.utils import log_utils
 from sky.utils import resources_utils
 from sky.utils import status_lib
@@ -137,7 +136,8 @@ def get_total_cost_of_displayed_records(
 
 def show_cost_report_table(cluster_records: List[_ClusterCostReportRecord],
                            show_all: bool,
-                           controller_name: Optional[str] = None):
+                           controller_name: Optional[str] = None,
+                           days: Optional[int] = None):
     """Compute cluster table values and display for cost report.
 
     For each cluster, this shows: cluster name, resources, launched time,
@@ -200,23 +200,21 @@ def show_cost_report_table(cluster_records: List[_ClusterCostReportRecord],
         cluster_table.add_row(row)
 
     if cluster_records:
+        controller_record = cluster_records[0]
         if controller_name is not None:
-            controller = controller_utils.Controllers.from_name(controller_name)
-            if controller is None:
-                raise ValueError(f'Controller {controller_name} not found.')
-            controller_handle: backends.CloudVmRayResourceHandle = (
-                cluster_records[0]['handle'])
-            autostop_config = (
-                controller_handle.launched_resources.autostop_config)
-            if autostop_config is not None:
+            autostop = controller_record.get('autostop', None)
+            autostop_str = ''
+            if autostop is not None:
                 autostop_str = (f'{colorama.Style.DIM} (will be autostopped if '
-                                f'idle for {autostop_config.idle_minutes}min)'
+                                f'idle for {autostop}min)'
                                 f'{colorama.Style.RESET_ALL}')
             click.echo(f'\n{colorama.Fore.CYAN}{colorama.Style.BRIGHT}'
                        f'{controller_name}{colorama.Style.RESET_ALL}'
                        f'{autostop_str}')
         else:
-            click.echo(f'{colorama.Fore.CYAN}{colorama.Style.BRIGHT}Clusters'
+            days_str = '' if days is None else f' (last {days} days)'
+            click.echo(f'{colorama.Fore.CYAN}{colorama.Style.BRIGHT}'
+                       f'Clusters{days_str}'
                        f'{colorama.Style.RESET_ALL}')
         click.echo(cluster_table)
 
@@ -345,7 +343,9 @@ def _get_infra(cluster_record: _ClusterRecord, truncate: bool = True) -> str:
 
 
 def _get_status_value_for_cost_report(
-        cluster_cost_report_record: _ClusterCostReportRecord) -> int:
+        cluster_cost_report_record: _ClusterCostReportRecord,
+        truncate: bool = True) -> int:
+    del truncate
     status = cluster_cost_report_record['status']
     if status is None:
         return -1
@@ -353,7 +353,9 @@ def _get_status_value_for_cost_report(
 
 
 def _get_status_for_cost_report(
-        cluster_cost_report_record: _ClusterCostReportRecord) -> str:
+        cluster_cost_report_record: _ClusterCostReportRecord,
+        truncate: bool = True) -> str:
+    del truncate
     status = cluster_cost_report_record['status']
     if status is None:
         return f'{colorama.Style.DIM}TERMINATED{colorama.Style.RESET_ALL}'
@@ -361,7 +363,9 @@ def _get_status_for_cost_report(
 
 
 def _get_resources_for_cost_report(
-        cluster_cost_report_record: _ClusterCostReportRecord) -> str:
+        cluster_cost_report_record: _ClusterCostReportRecord,
+        truncate: bool = True) -> str:
+    del truncate
     launched_nodes = cluster_cost_report_record['num_nodes']
     launched_resources = cluster_cost_report_record['resources']
 
@@ -373,7 +377,9 @@ def _get_resources_for_cost_report(
 
 
 def _get_price_for_cost_report(
-        cluster_cost_report_record: _ClusterCostReportRecord) -> str:
+        cluster_cost_report_record: _ClusterCostReportRecord,
+        truncate: bool = True) -> str:
+    del truncate
     launched_nodes = cluster_cost_report_record['num_nodes']
     launched_resources = cluster_cost_report_record['resources']
 
@@ -383,7 +389,9 @@ def _get_price_for_cost_report(
 
 
 def _get_estimated_cost_for_cost_report(
-        cluster_cost_report_record: _ClusterCostReportRecord) -> str:
+        cluster_cost_report_record: _ClusterCostReportRecord,
+        truncate: bool = True) -> str:
+    del truncate
     cost = cluster_cost_report_record['total_cost']
 
     if not cost: