skypilot-nightly 1.0.0.dev20250623__py3-none-any.whl → 1.0.0.dev20250625__py3-none-any.whl

sky/jobs/server/core.py

@@ -1,5 +1,6 @@
 """SDK functions for managed jobs."""
 import os
+import pathlib
 import tempfile
 import typing
 from typing import Any, Dict, List, Optional, Tuple, Union
@@ -20,6 +21,7 @@ from sky.backends import backend_utils
 from sky.catalog import common as service_catalog_common
 from sky.data import storage as storage_lib
 from sky.jobs import constants as managed_job_constants
+from sky.jobs import state as managed_job_state
 from sky.jobs import utils as managed_job_utils
 from sky.provision import common as provision_common
 from sky.skylet import constants as skylet_constants
@@ -43,6 +45,72 @@ if typing.TYPE_CHECKING:
 logger = sky_logging.init_logger(__name__)
 
 
+def _maybe_upload_files_to_controller(dag: 'sky.Dag') -> Dict[str, str]:
+    """Maybe upload files to the controller.
+
+    In consolidation mode, we don't need to upload files to the controller as
+    the API server and the controller are colocated.
+    """
+    local_to_controller_file_mounts: Dict[str, str] = {}
+
+    if managed_job_utils.is_consolidation_mode():
+        return local_to_controller_file_mounts
+
+    if storage_lib.get_cached_enabled_storage_cloud_names_or_refresh():
+        for task_ in dag.tasks:
+            controller_utils.maybe_translate_local_file_mounts_and_sync_up(
+                task_, task_type='jobs')
+    else:
+        # We do not have any cloud storage available, so fall back to
+        # two-hop file_mount uploading.
+        # Note: we can't easily hack sync_storage_mounts() to upload
+        # directly to the controller, because the controller may not
+        # even be up yet.
+        for task_ in dag.tasks:
+            if task_.storage_mounts:
+                # Technically, we could convert COPY storage_mounts that
+                # have a local source and do not specify `store`, but we
+                # will not do that for now. Only plain file_mounts are
+                # supported.
+                raise exceptions.NotSupportedError(
+                    'Cloud-based file_mounts are specified, but no cloud '
+                    'storage is available. Please specify local '
+                    'file_mounts only.')
+
+            # Merge file mounts from all tasks.
+            local_to_controller_file_mounts.update(
+                controller_utils.translate_local_file_mounts_to_two_hop(task_))
+
+    return local_to_controller_file_mounts
+
+
+def _maybe_submit_job_locally(prefix: str, dag: 'sky.Dag') -> Optional[int]:
+    """Submit the managed job locally if in consolidation mode.
+
+    In normal mode the managed job submission is done in the ray job submission.
+    For consolidation mode, we need to manually submit it. Check the following
+    function for the normal mode submission:
+    sky/backends/cloud_vm_ray_backend.py::CloudVmRayBackend,
+    _exec_code_on_head::_maybe_add_managed_job_code
+    """
+    if not managed_job_utils.is_consolidation_mode():
+        return None
+
+    # Create local directory for the managed job.
+    pathlib.Path(prefix).expanduser().mkdir(parents=True, exist_ok=True)
+    consolidation_mode_job_id = managed_job_state.set_job_info_without_job_id(
+        dag.name,
+        workspace=skypilot_config.get_active_workspace(
+            force_user_workspace=True),
+        entrypoint=common_utils.get_current_command())
+    for task_id, task in enumerate(dag.tasks):
+        resources_str = backend_utils.get_task_resources_str(
+            task, is_managed_job=True)
+        managed_job_state.set_pending(consolidation_mode_job_id, task_id,
+                                      task.name, resources_str)
+    return consolidation_mode_job_id
+
+
 @timeline.event
 @usage_lib.entrypoint
 def launch(
@@ -77,6 +145,7 @@ def launch(
     entrypoint = task
     dag_uuid = str(uuid.uuid4().hex[:4])
     dag = dag_utils.convert_entrypoint_to_dag(entrypoint)
+    dag.resolve_and_validate_volumes()
     # Always apply the policy again here, even though it might have been applied
     # in the CLI. This is to ensure that we apply the policy to the final DAG
     # and get the mutated config.
@@ -86,6 +155,9 @@ def launch(
             raise ValueError('Only single-task or chain DAG is '
                              f'allowed for job_launch. Dag: {dag}')
     dag.validate()
+    # TODO(aylei): use consolidated job controller instead of performing
+    # pre-mount operations when submitting jobs.
+    dag.pre_mount_volumes()
 
     user_dag_str = dag_utils.dump_chain_dag_to_yaml_str(dag)
 
@@ -103,7 +175,7 @@ def launch(
                     'will be auto-generated) .')
         task_names.add(task_.name)
 
-        # Check for priority in resources first, then fall back to job priority
+        # Check for priority in resources
         task_priority = None
         if task_.resources:
             # Convert set to list to access elements by index
@@ -121,20 +193,6 @@ def launch(
                             f'{resource.priority} but expected {task_priority}.'
                         )
 
-            # Check for conflict between resources priority and job
-            # priority
-            if task_.job_priority is not None:
-                with ux_utils.print_exception_no_traceback():
-                    raise ValueError(
-                        f'Task {task_.name!r}: Cannot specify both '
-                        f'resources.priority ({task_priority}) and '
-                        f'job.priority ({task_.job_priority}). Please use only '
-                        'one priority specification method.')
-
-        # Fall back to job priority if no resources priority found
-        if task_priority is None:
-            task_priority = task_.job_priority
-
         if task_priority is not None:
             if (priority is not None and priority != task_priority):
                 with ux_utils.print_exception_no_traceback():
@@ -183,34 +241,7 @@ def launch(
                         f'with:\n\n`sky down {cluster_name} --purge`\n\n'
                         f'Reason: {common_utils.format_exception(e)}')
 
-        local_to_controller_file_mounts = {}
-
-        if storage_lib.get_cached_enabled_storage_cloud_names_or_refresh():
-            for task_ in dag.tasks:
-                controller_utils.maybe_translate_local_file_mounts_and_sync_up(
-                    task_, task_type='jobs')
-
-        else:
-            # We do not have any cloud storage available, so fall back to
-            # two-hop file_mount uploading.
-            # Note: we can't easily hack sync_storage_mounts() to upload
-            # directly to the controller, because the controller may not
-            # even be up yet.
-            for task_ in dag.tasks:
-                if task_.storage_mounts:
-                    # Technically, we could convert COPY storage_mounts that
-                    # have a local source and do not specify `store`, but we
-                    # will not do that for now. Only plain file_mounts are
-                    # supported.
-                    raise exceptions.NotSupportedError(
-                        'Cloud-based file_mounts are specified, but no cloud '
-                        'storage is available. Please specify local '
-                        'file_mounts only.')
-
-                # Merge file mounts from all tasks.
-                local_to_controller_file_mounts.update(
-                    controller_utils.translate_local_file_mounts_to_two_hop(
-                        task_))
+    local_to_controller_file_mounts = _maybe_upload_files_to_controller(dag)
 
     # Has to use `\` to avoid yapf issue.
     with tempfile.NamedTemporaryFile(prefix=f'managed-dag-{dag.name}-',
@@ -233,6 +264,13 @@ def launch(
             controller=controller,
             task_resources=sum([list(t.resources) for t in dag.tasks], []))
 
+        consolidation_mode_job_id = _maybe_submit_job_locally(prefix, dag)
+
+        # This is only needed for non-consolidation mode. For consolidation
+        # mode, the controller uses the same catalog as API server.
+        modified_catalogs = {} if consolidation_mode_job_id is not None else (
+            service_catalog_common.get_modified_catalog_file_mounts())
+
         vars_to_fill = {
             'remote_original_user_yaml_path': remote_original_user_yaml_path,
             'original_user_dag_path': original_user_yaml_path.name,
@@ -244,9 +282,9 @@ def launch(
             'dag_name': dag.name,
             'remote_user_config_path': remote_user_config_path,
             'remote_env_file_path': remote_env_file_path,
-            'modified_catalogs':
-                service_catalog_common.get_modified_catalog_file_mounts(),
+            'modified_catalogs': modified_catalogs,
             'priority': priority,
+            'consolidation_mode_job_id': consolidation_mode_job_id,
             **controller_utils.shared_controller_vars_to_fill(
                 controller,
                 remote_user_config_path=remote_user_config_path,
@@ -285,12 +323,44 @@ def launch(
                 # workspace A, but the controller is in workspace B, the
                 # intermediate bucket and newly created bucket should be in
                 # workspace A.
-                return execution.launch(task=controller_task,
-                                        cluster_name=controller_name,
-                                        stream_logs=stream_logs,
-                                        retry_until_up=True,
-                                        fast=True,
-                                        _disable_controller_check=True)
+                if consolidation_mode_job_id is None:
+                    return execution.launch(task=controller_task,
+                                            cluster_name=controller_name,
+                                            stream_logs=stream_logs,
+                                            retry_until_up=True,
+                                            fast=True,
+                                            _disable_controller_check=True)
+                # Manually launch the scheduler process in consolidation mode.
+                local_handle = backend_utils.is_controller_accessible(
+                    controller=controller, stopped_message='')
+                backend = backend_utils.get_backend_from_handle(local_handle)
+                assert isinstance(backend, backends.CloudVmRayBackend)
+                backend.sync_file_mounts(
+                    handle=local_handle,
+                    all_file_mounts=controller_task.file_mounts,
+                    storage_mounts=controller_task.storage_mounts)
+                run_script = controller_task.run
+                assert isinstance(run_script, str)
+                # Manually add the env variables to the run script. Originally
+                # this is done in ray jobs submission but now we have to do it
+                # manually because there is no ray runtime on the API server.
+                env_cmds = [
+                    f'export {k}={v!r}'
+                    for k, v in controller_task.envs.items()
+                ]
+                run_script = '\n'.join(env_cmds + [run_script])
+                # Dump script for high availability recovery.
+                if controller_utils.high_availability_specified(
+                        controller_name):
+                    dump_script_path = (
+                        managed_job_utils.get_ha_dump_script_path(
+                            consolidation_mode_job_id))
+                    dump_script_path.parent.mkdir(parents=True, exist_ok=True)
+                    with open(dump_script_path, 'w',
+                              encoding='utf-8') as script_f:
+                        script_f.write(run_script)
+                backend.run_on_head(local_handle, run_script)
+                return consolidation_mode_job_id, local_handle
 
 
 def queue_from_kubernetes_pod(

sky/jobs/state.py

@@ -463,6 +463,21 @@ def set_job_info(job_id: int, name: str, workspace: str, entrypoint: str):
              entrypoint))
 
 
+@_init_db
+def set_job_info_without_job_id(name: str, workspace: str,
+                                entrypoint: str) -> int:
+    assert _DB_PATH is not None
+    with db_utils.safe_cursor(_DB_PATH) as cursor:
+        cursor.execute(
+            """\
+            INSERT INTO job_info
+            (name, schedule_state, workspace, entrypoint)
+            VALUES (?, ?, ?, ?)""",
+            (name, ManagedJobScheduleState.INACTIVE.value, workspace,
+             entrypoint))
+        return cursor.lastrowid
+
+
 @_init_db
 def set_pending(job_id: int, task_id: int, task_name: str, resources_str: str):
     """Set the task to pending state."""

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/models.py

@@ -6,6 +6,8 @@ import getpass
 import os
 from typing import Any, Dict, Optional
 
+import pydantic
+
 from sky.skylet import constants
 from sky.utils import common_utils
 
@@ -48,6 +50,8 @@ class KubernetesNodeInfo:
     # Resources available on the node. E.g., {'nvidia.com/gpu': '2'}
     total: Dict[str, int]
     free: Dict[str, int]
+    # IP address of the node (external IP preferred, fallback to internal IP)
+    ip_address: Optional[str] = None
 
 
 @dataclasses.dataclass
@@ -76,3 +80,15 @@ class KubernetesNodesInfo:
             },
             hint=data['hint'],
         )
+
+
+class VolumeConfig(pydantic.BaseModel):
+    """Configuration for creating a volume."""
+    name: str
+    type: str
+    cloud: str
+    region: Optional[str]
+    zone: Optional[str]
+    name_on_cloud: str
+    size: Optional[str]
+    config: Dict[str, Any] = {}

sky/provision/__init__.py

@@ -8,6 +8,7 @@ import inspect
 import typing
 from typing import Any, Dict, List, Optional, Type
 
+from sky import models
 from sky import sky_logging
 # These provision.<cloud> modules should never fail even if underlying cloud SDK
 # dependencies are not installed. This is ensured by using sky.adaptors inside
@@ -103,6 +104,31 @@ def bootstrap_instances(
     raise NotImplementedError
 
 
+@_route_to_cloud_impl
+def apply_volume(provider_name: str,
+                 config: models.VolumeConfig) -> models.VolumeConfig:
+    """Create or register a volume.
+
+    This function creates or registers a volume with the provided configuration,
+    and returns a VolumeConfig object with updated configuration.
+    """
+    raise NotImplementedError
+
+
+@_route_to_cloud_impl
+def delete_volume(provider_name: str,
+                  config: models.VolumeConfig) -> models.VolumeConfig:
+    """Delete a volume."""
+    raise NotImplementedError
+
+
+@_route_to_cloud_impl
+def get_volume_usedby(provider_name: str,
+                      config: models.VolumeConfig) -> List[str]:
+    """Get the usedby of a volume."""
+    raise NotImplementedError
+
+
 @_route_to_cloud_impl
 def run_instances(provider_name: str, region: str, cluster_name_on_cloud: str,
                   config: common.ProvisionConfig) -> common.ProvisionRecord:

sky/provision/kubernetes/__init__.py

@@ -11,3 +11,6 @@ from sky.provision.kubernetes.instance import wait_instances
 from sky.provision.kubernetes.network import cleanup_ports
 from sky.provision.kubernetes.network import open_ports
 from sky.provision.kubernetes.network import query_ports
+from sky.provision.kubernetes.volume import apply_volume
+from sky.provision.kubernetes.volume import delete_volume
+from sky.provision.kubernetes.volume import get_volume_usedby