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

sky/jobs/server/core.py

@@ -1,9 +1,13 @@
 """SDK functions for managed jobs."""
+import concurrent.futures
+import copy
+import ipaddress
 import os
 import pathlib
 import tempfile
 import typing
 from typing import Any, Dict, List, Optional, Tuple, Union
+from urllib import parse as urlparse
 import uuid
 
 import colorama
@@ -17,16 +21,21 @@ from sky import provision as provision_lib
 from sky import sky_logging
 from sky import skypilot_config
 from sky import task as task_lib
+from sky.adaptors import common as adaptors_common
 from sky.backends import backend_utils
+from sky.backends import cloud_vm_ray_backend
 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.metrics import utils as metrics_lib
 from sky.provision import common as provision_common
+from sky.schemas.api import responses
 from sky.serve import serve_state
 from sky.serve import serve_utils
 from sky.serve.server import impl
+from sky.server.requests import request_names
 from sky.skylet import constants as skylet_constants
 from sky.usage import usage_lib
 from sky.utils import admin_policy_utils
@@ -42,11 +51,47 @@ from sky.utils import ux_utils
 from sky.workspaces import core as workspaces_core
 
 if typing.TYPE_CHECKING:
+    from google.protobuf import json_format
+
     import sky
-    from sky.backends import cloud_vm_ray_backend
+    from sky.schemas.generated import managed_jobsv1_pb2
+else:
+    json_format = adaptors_common.LazyImport('google.protobuf.json_format')
+
+    managed_jobsv1_pb2 = adaptors_common.LazyImport(
+        'sky.schemas.generated.managed_jobsv1_pb2')
 
 logger = sky_logging.init_logger(__name__)
 
+_MANAGED_JOB_FIELDS_FOR_QUEUE_KUBERNETES = [
+    'job_id',
+    'task_id',
+    'workspace',
+    'job_name',
+    'task_name',
+    'resources',
+    'submitted_at',
+    'end_at',
+    'job_duration',
+    'recovery_count',
+    'status',
+    'pool',
+    'current_cluster_name',
+    'job_id_on_pool_cluster',
+    'start_at',
+    'infra',
+    'cloud',
+    'region',
+    'zone',
+    'cluster_resources',
+    'schedule_state',
+    'details',
+    'failure_reason',
+    'metadata',
+    'user_name',
+    'user_hash',
+]
+
 
 def _upload_files_to_controller(dag: 'sky.Dag') -> Dict[str, str]:
     """Upload files to the controller.
@@ -129,7 +174,8 @@ def _maybe_submit_job_locally(prefix: str, dag: 'sky.Dag',
                     force_user_workspace=True),
                 entrypoint=common_utils.get_current_command(),
                 pool=pool,
-                pool_hash=pool_hash))
+                pool_hash=pool_hash,
+                user_hash=common_utils.get_user_hash()))
         for task_id, task in enumerate(dag.tasks):
             resources_str = backend_utils.get_task_resources_str(
                 task, is_managed_job=True)
@@ -188,10 +234,12 @@ def launch(
 
     dag_uuid = str(uuid.uuid4().hex[:4])
     dag = dag_utils.convert_entrypoint_to_dag(entrypoint)
+
     # 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.
-    dag, mutated_user_config = admin_policy_utils.apply(dag)
+    dag, mutated_user_config = admin_policy_utils.apply(
+        dag, request_name=request_names.AdminPolicyRequestName.JOBS_LAUNCH)
     dag.resolve_and_validate_volumes()
     if not dag.is_chain():
         with ux_utils.print_exception_no_traceback():
@@ -202,6 +250,21 @@ def launch(
     # pre-mount operations when submitting jobs.
     dag.pre_mount_volumes()
 
+    # If there is a local postgres db, when the api server tries launching on
+    # the remote jobs controller it will fail. therefore, we should remove this
+    # before sending the config to the jobs controller.
+    # TODO(luca) there are a lot of potential problems with postgres being sent
+    # to the jobs controller. for example if the postgres is whitelisted to
+    # only the API server, this will then break. the simple solution to that is
+    # telling the user to add the jobs controller to the postgres whitelist.
+    if not managed_job_utils.is_consolidation_mode():
+        db_path = mutated_user_config.get('db', None)
+        if db_path is not None:
+            parsed = urlparse.urlparse(db_path)
+            if ((parsed.hostname == 'localhost' or
+                 ipaddress.ip_address(parsed.hostname).is_loopback)):
+                mutated_user_config.pop('db', None)
+
     user_dag_str_user_specified = dag_utils.dump_chain_dag_to_yaml_str(
         dag, use_user_specified_yaml=True)
 
@@ -263,15 +326,13 @@ def launch(
         # Check whether cached jobs controller cluster is accessible
         cluster_name = (
             controller_utils.Controllers.JOBS_CONTROLLER.value.cluster_name)
-        record = global_user_state.get_cluster_from_name(cluster_name)
-        if record is not None:
+        if global_user_state.cluster_with_name_exists(cluster_name):
             # there is a cached jobs controller cluster
             try:
                 # TODO: do something with returned status?
                 _, _ = backend_utils.refresh_cluster_status_handle(
                     cluster_name=cluster_name,
-                    force_refresh_statuses=set(status_lib.ClusterStatus),
-                    acquire_per_cluster_status_lock=False)
+                    force_refresh_statuses=set(status_lib.ClusterStatus))
             except (exceptions.ClusterOwnerIdentityMismatchError,
                     exceptions.CloudUserIdentityError,
                     exceptions.ClusterStatusFetchingError) as e:
@@ -309,6 +370,7 @@ def launch(
     def _submit_one(
         consolidation_mode_job_id: Optional[int] = None,
         job_rank: Optional[int] = None,
+        num_jobs: Optional[int] = None,
     ) -> Tuple[Optional[int], Optional[backends.ResourceHandle]]:
         rank_suffix = '' if job_rank is None else f'-{job_rank}'
         remote_original_user_yaml_path = (
@@ -328,11 +390,16 @@ def launch(
         ) as original_user_yaml_path:
             original_user_yaml_path.write(user_dag_str_user_specified)
             original_user_yaml_path.flush()
-            for task_ in dag.tasks:
+            # Copy tasks to avoid race conditions when multiple threads modify
+            # the same dag object concurrently. Each thread needs its own copy.
+            dag_copy = copy.deepcopy(dag)
+            for task_ in dag_copy.tasks:
                 if job_rank is not None:
                     task_.update_envs({'SKYPILOT_JOB_RANK': str(job_rank)})
+                if num_jobs is not None:
+                    task_.update_envs({'SKYPILOT_NUM_JOBS': str(num_jobs)})
 
-            dag_utils.dump_chain_dag_to_yaml(dag, f.name)
+            dag_utils.dump_chain_dag_to_yaml(dag_copy, f.name)
 
             vars_to_fill = {
                 'remote_original_user_yaml_path':
@@ -351,6 +418,8 @@ def launch(
                 'priority': priority,
                 'consolidation_mode_job_id': consolidation_mode_job_id,
                 'pool': pool,
+                'job_controller_indicator_file':
+                    managed_job_constants.JOB_CONTROLLER_INDICATOR_FILE,
                 **controller_utils.shared_controller_vars_to_fill(
                     controller,
                     remote_user_config_path=remote_user_config_path,
@@ -363,7 +432,8 @@ def launch(
 
             yaml_path = os.path.join(
                 managed_job_constants.JOBS_CONTROLLER_YAML_PREFIX,
-                f'{name}-{dag_uuid}-{consolidation_mode_job_id}.yaml')
+                f'{name}-{dag_uuid}-{consolidation_mode_job_id}-{job_rank}.yaml'
+            )
             common_utils.fill_template(
                 managed_job_constants.JOBS_CONTROLLER_TEMPLATE,
                 vars_to_fill,
@@ -371,16 +441,19 @@ def launch(
             controller_task = task_lib.Task.from_yaml(yaml_path)
             controller_task.set_resources(controller_resources)
 
-            controller_task.managed_job_dag = dag
+            controller_task.managed_job_dag = dag_copy
             # pylint: disable=protected-access
             controller_task._metadata = metadata
 
             job_identity = ''
             if job_rank is not None:
                 job_identity = f' (rank: {job_rank})'
-            logger.info(f'{colorama.Fore.YELLOW}'
-                        f'Launching managed job {dag.name!r}{job_identity} '
-                        f'from jobs controller...{colorama.Style.RESET_ALL}')
+            job_controller_postfix = (' from jobs controller' if
+                                      consolidation_mode_job_id is None else '')
+            logger.info(
+                f'{colorama.Fore.YELLOW}'
+                f'Launching managed job {dag.name!r}{job_identity}'
+                f'{job_controller_postfix}...{colorama.Style.RESET_ALL}')
 
             # Launch with the api server's user hash, so that sky status does
             # not show the owner of the controller as whatever user launched
@@ -395,19 +468,24 @@ def launch(
                     # intermediate bucket and newly created bucket should be in
                     # workspace A.
                     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)
+                        return execution.launch(
+                            task=controller_task,
+                            cluster_name=controller_name,
+                            stream_logs=stream_logs,
+                            retry_until_up=True,
+                            fast=True,
+                            _request_name=request_names.AdminPolicyRequestName.
+                            JOBS_LAUNCH_CONTROLLER,
+                            _disable_controller_check=True)
                     # Manually launch the scheduler 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)
-                    with sky_logging.silent():
+                    # Suppress file mount logs when submitting multiple jobs.
+                    should_silence = num_jobs is not None and num_jobs > 1
+                    with sky_logging.silent(should_silence):
                         backend.sync_file_mounts(
                             handle=local_handle,
                             all_file_mounts=controller_task.file_mounts,
@@ -423,12 +501,16 @@ def launch(
                         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):
-                        managed_job_state.set_ha_recovery_script(
-                            consolidation_mode_job_id, run_script)
-                    backend.run_on_head(local_handle, run_script)
+                    log_dir = os.path.join(skylet_constants.SKY_LOGS_DIRECTORY,
+                                           'managed_jobs')
+                    os.makedirs(log_dir, exist_ok=True)
+                    log_path = os.path.join(
+                        log_dir, f'submit-job-{consolidation_mode_job_id}.log')
+                    backend.run_on_head(local_handle,
+                                        run_script,
+                                        log_path=log_path)
+                    ux_utils.starting_message(
+                        f'Job submitted, ID: {consolidation_mode_job_id}')
                     return consolidation_mode_job_id, local_handle
 
     if pool is None:
@@ -437,15 +519,49 @@ def launch(
         assert len(consolidation_mode_job_ids) == 1
         return _submit_one(consolidation_mode_job_ids[0])
 
-    ids = []
-    all_handle = None
-    for job_rank in range(num_jobs):
-        job_id = (consolidation_mode_job_ids[job_rank]
+    ids: List[int] = []
+    all_handle: Optional[backends.ResourceHandle] = None
+
+    if num_jobs == 1:
+        job_id = (consolidation_mode_job_ids[0]
                   if consolidation_mode_job_ids is not None else None)
-        jid, handle = _submit_one(job_id, job_rank)
+        jid, handle = _submit_one(job_id, 0, num_jobs=num_jobs)
         assert jid is not None, (job_id, handle)
         ids.append(jid)
         all_handle = handle
+    else:
+        # Submit jobs in parallel using ThreadPoolExecutor
+        with concurrent.futures.ThreadPoolExecutor(
+                max_workers=min(num_jobs,
+                                os.cpu_count() or 1)) as executor:
+            # Submit jobs concurrently
+            future_to_rank = {}
+            for job_rank in range(num_jobs):
+                job_id = (consolidation_mode_job_ids[job_rank]
+                          if consolidation_mode_job_ids is not None else None)
+                future = executor.submit(_submit_one, job_id, job_rank,
+                                         num_jobs)
+                future_to_rank[future] = job_rank
+
+            # Collect results in order of job_rank to maintain consistent order.
+            results: List[Optional[Tuple[
+                int, Optional[backends.ResourceHandle]]]] = [None] * num_jobs
+            for future in concurrent.futures.as_completed(future_to_rank):
+                job_rank = future_to_rank[future]
+                try:
+                    jid, handle = future.result()
+                    assert jid is not None, (job_id, handle)
+                    results[job_rank] = (jid, handle)
+                    all_handle = handle  # Keep the last handle.
+                except Exception as e:
+                    logger.error(f'Error launching job {job_rank}: {e}')
+                    raise e
+
+            # Extract job IDs in order
+            for res in results:
+                if res is not None:
+                    ids.append(res[0])
+
     return ids, all_handle
 
 
@@ -498,7 +614,8 @@ def queue_from_kubernetes_pod(
         'kubernetes', cluster_info)[0]
 
     code = managed_job_utils.ManagedJobCodeGen.get_job_table(
-        skip_finished=skip_finished)
+        skip_finished=skip_finished,
+        fields=_MANAGED_JOB_FIELDS_FOR_QUEUE_KUBERNETES)
     returncode, job_table_payload, stderr = managed_jobs_runner.run(
         code,
         require_outputs=True,
@@ -575,8 +692,49 @@ def _maybe_restart_controller(
     return handle
 
 
+# For backwards compatibility
+# TODO(hailong): Remove before 0.12.0.
+@usage_lib.entrypoint
+def queue(refresh: bool,
+          skip_finished: bool = False,
+          all_users: bool = False,
+          job_ids: Optional[List[int]] = None) -> List[Dict[str, Any]]:
+    # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
+    """Gets statuses of managed jobs.
+
+    Please refer to sky.cli.job_queue for documentation.
+
+    Returns:
+        [
+            {
+                'job_id': int,
+                'job_name': str,
+                'resources': str,
+                'submitted_at': (float) timestamp of submission,
+                'end_at': (float) timestamp of end,
+                'job_duration': (float) duration in seconds,
+                'recovery_count': (int) Number of retries,
+                'status': (sky.jobs.ManagedJobStatus) of the job,
+                'cluster_resources': (str) resources of the cluster,
+                'region': (str) region of the cluster,
+                'user_name': (Optional[str]) job creator's user name,
+                'user_hash': (str) job creator's user hash,
+                'task_id': (int), set to 0 (except in pipelines, which may have multiple tasks), # pylint: disable=line-too-long
+                'task_name': (str), same as job_name (except in pipelines, which may have multiple tasks), # pylint: disable=line-too-long
+            }
+        ]
+    Raises:
+        sky.exceptions.ClusterNotUpError: the jobs controller is not up or
+            does not exist.
+        RuntimeError: if failed to get the managed jobs with ssh.
+    """
+    jobs, _, _, _ = queue_v2(refresh, skip_finished, all_users, job_ids)
+
+    return jobs
+
+
 @usage_lib.entrypoint
-def queue(
+def queue_v2_api(
     refresh: bool,
     skip_finished: bool = False,
     all_users: bool = False,
@@ -588,9 +746,34 @@ def queue(
     page: Optional[int] = None,
     limit: Optional[int] = None,
     statuses: Optional[List[str]] = None,
+    fields: Optional[List[str]] = None,
+) -> Tuple[List[responses.ManagedJobRecord], int, Dict[str, int], int]:
+    """Gets statuses of managed jobs and parse the
+    jobs to responses.ManagedJobRecord."""
+    jobs, total, status_counts, total_no_filter = queue_v2(
+        refresh, skip_finished, all_users, job_ids, user_match, workspace_match,
+        name_match, pool_match, page, limit, statuses, fields)
+    return [responses.ManagedJobRecord(**job) for job in jobs
+           ], total, status_counts, total_no_filter
+
+
+@metrics_lib.time_me
+def queue_v2(
+    refresh: bool,
+    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,
+    pool_match: Optional[str] = None,
+    page: Optional[int] = None,
+    limit: Optional[int] = None,
+    statuses: Optional[List[str]] = None,
+    fields: Optional[List[str]] = None,
 ) -> Tuple[List[Dict[str, Any]], int, Dict[str, int], int]:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
-    """Gets statuses of managed jobs.
+    """Gets statuses of managed jobs with filtering.
 
     Please refer to sky.cli.job_queue for documentation.
 
@@ -633,20 +816,23 @@ def queue(
         if page is not None:
             raise ValueError('Limit must be specified when page is specified')
 
-    handle = _maybe_restart_controller(refresh,
-                                       stopped_message='No in-progress '
-                                       'managed jobs.',
-                                       spinner_message='Checking '
-                                       'managed jobs')
+    with metrics_lib.time_it('jobs.queue.restart_controller', group='jobs'):
+        handle = _maybe_restart_controller(refresh,
+                                           stopped_message='No in-progress '
+                                           'managed jobs.',
+                                           spinner_message='Checking '
+                                           'managed jobs')
     backend = backend_utils.get_backend_from_handle(handle)
     assert isinstance(backend, backends.CloudVmRayBackend)
 
     user_hashes: Optional[List[Optional[str]]] = None
+    show_jobs_without_user_hash = False
     if not all_users:
         user_hashes = [common_utils.get_user_hash()]
         # For backwards compatibility, we show jobs that do not have a
         # user_hash. TODO(cooperc): Remove before 0.12.0.
         user_hashes.append(None)
+        show_jobs_without_user_hash = True
     elif user_match is not None:
         users = global_user_state.get_user_by_name_match(user_match)
         if not users:
@@ -654,70 +840,109 @@ def queue(
         user_hashes = [user.id for user in users]
 
     accessible_workspaces = list(workspaces_core.get_workspaces().keys())
-    code = managed_job_utils.ManagedJobCodeGen.get_job_table(
-        skip_finished, accessible_workspaces, job_ids, workspace_match,
-        name_match, pool_match, page, limit, user_hashes, statuses)
-    returncode, job_table_payload, stderr = backend.run_on_head(
-        handle,
-        code,
-        require_outputs=True,
-        stream_logs=False,
-        separate_stderr=True)
+
+    if handle.is_grpc_enabled_with_flag:
+        try:
+            request = managed_jobsv1_pb2.GetJobTableRequest(
+                skip_finished=skip_finished,
+                accessible_workspaces=(managed_jobsv1_pb2.Workspaces(
+                    workspaces=accessible_workspaces)),
+                job_ids=managed_jobsv1_pb2.JobIds(
+                    ids=job_ids) if job_ids is not None else None,
+                workspace_match=workspace_match,
+                name_match=name_match,
+                pool_match=pool_match,
+                page=page,
+                limit=limit,
+                # Remove None from user_hashes, as the gRPC server uses the
+                # show_jobs_without_user_hash flag instead.
+                user_hashes=managed_jobsv1_pb2.UserHashes(hashes=[
+                    user_hash for user_hash in user_hashes
+                    if user_hash is not None
+                ]) if user_hashes is not None else None,
+                statuses=managed_jobsv1_pb2.Statuses(
+                    statuses=statuses) if statuses is not None else None,
+                fields=managed_jobsv1_pb2.Fields(
+                    fields=fields) if fields is not None else None,
+                show_jobs_without_user_hash=show_jobs_without_user_hash,
+            )
+            response = backend_utils.invoke_skylet_with_retries(
+                lambda: cloud_vm_ray_backend.SkyletClient(
+                    handle.get_grpc_channel()).get_managed_job_table(request))
+            jobs = managed_job_utils.decode_managed_job_protos(response.jobs)
+            return jobs, response.total, dict(
+                response.status_counts), response.total_no_filter
+        except exceptions.SkyletMethodNotImplementedError:
+            pass
+
+    with metrics_lib.time_it('jobs.queue.generate_code', group='jobs'):
+        code = managed_job_utils.ManagedJobCodeGen.get_job_table(
+            skip_finished, accessible_workspaces, job_ids, workspace_match,
+            name_match, pool_match, page, limit, user_hashes, statuses, fields)
+    with metrics_lib.time_it('jobs.queue.run_on_head', group='jobs'):
+        returncode, job_table_payload, stderr = backend.run_on_head(
+            handle,
+            code,
+            require_outputs=True,
+            stream_logs=False,
+            separate_stderr=True)
 
     if returncode != 0:
         logger.error(job_table_payload + stderr)
         raise RuntimeError('Failed to fetch managed jobs with returncode: '
                            f'{returncode}.\n{job_table_payload + stderr}')
 
-    (jobs, total, result_type, total_no_filter, status_counts
-    ) = managed_job_utils.load_managed_job_queue(job_table_payload)
+    with metrics_lib.time_it('jobs.queue.load_job_queue', group='jobs'):
+        (jobs, total, result_type, total_no_filter, status_counts
+        ) = managed_job_utils.load_managed_job_queue(job_table_payload)
 
     if result_type == managed_job_utils.ManagedJobQueueResultType.DICT:
         return jobs, total, status_counts, total_no_filter
 
     # Backward compatibility for old jobs controller without filtering
     # TODO(hailong): remove this after 0.12.0
-    if not all_users:
+    with metrics_lib.time_it('jobs.queue.filter_and_process', group='jobs'):
+        if not all_users:
 
-        def user_hash_matches_or_missing(job: Dict[str, Any]) -> bool:
-            user_hash = job.get('user_hash', None)
-            if user_hash is None:
-                # For backwards compatibility, we show jobs that do not have a
-                # user_hash. TODO(cooperc): Remove before 0.12.0.
-                return True
-            return user_hash == common_utils.get_user_hash()
+            def user_hash_matches_or_missing(job: Dict[str, Any]) -> bool:
+                user_hash = job.get('user_hash', None)
+                if user_hash is None:
+                    # For backwards compatibility, we show jobs that do not have
+                    # a user_hash. TODO(cooperc): Remove before 0.12.0.
+                    return True
+                return user_hash == common_utils.get_user_hash()
 
-        jobs = list(filter(user_hash_matches_or_missing, jobs))
+            jobs = list(filter(user_hash_matches_or_missing, jobs))
 
-    jobs = list(
-        filter(
-            lambda job: job.get('workspace', skylet_constants.
-                                SKYPILOT_DEFAULT_WORKSPACE) in
-            accessible_workspaces, jobs))
-
-    if skip_finished:
-        # Filter out the finished jobs. If a multi-task job is partially
-        # finished, we will include all its tasks.
-        non_finished_tasks = list(
-            filter(lambda job: not job['status'].is_terminal(), jobs))
-        non_finished_job_ids = {job['job_id'] for job in non_finished_tasks}
         jobs = list(
-            filter(lambda job: job['job_id'] in non_finished_job_ids, jobs))
-
-    if job_ids:
-        jobs = [job for job in jobs if job['job_id'] in job_ids]
-
-    filtered_jobs, total, status_counts = managed_job_utils.filter_jobs(
-        jobs,
-        workspace_match,
-        name_match,
-        pool_match,
-        page=page,
-        limit=limit,
-        user_match=user_match,
-        enable_user_match=True,
-        statuses=statuses,
-    )
+            filter(
+                lambda job: job.get('workspace', skylet_constants.
+                                    SKYPILOT_DEFAULT_WORKSPACE) in
+                accessible_workspaces, jobs))
+
+        if skip_finished:
+            # Filter out the finished jobs. If a multi-task job is partially
+            # finished, we will include all its tasks.
+            non_finished_tasks = list(
+                filter(lambda job: not job['status'].is_terminal(), jobs))
+            non_finished_job_ids = {job['job_id'] for job in non_finished_tasks}
+            jobs = list(
+                filter(lambda job: job['job_id'] in non_finished_job_ids, jobs))
+
+        if job_ids:
+            jobs = [job for job in jobs if job['job_id'] in job_ids]
+
+        filtered_jobs, total, status_counts = managed_job_utils.filter_jobs(
+            jobs,
+            workspace_match,
+            name_match,
+            pool_match,
+            page=page,
+            limit=limit,
+            user_match=user_match,
+            enable_user_match=True,
+            statuses=statuses,
+        )
     return filtered_jobs, total, status_counts, total_no_filter
 
 
@@ -760,33 +985,60 @@ def cancel(name: Optional[str] = None,
                     'Can only specify one of JOB_IDS, name, pool, or all/'
                     f'all_users. Provided {" ".join(arguments)!r}.')
 
+        job_ids = None if (all_users or all) else job_ids
+
         backend = backend_utils.get_backend_from_handle(handle)
         assert isinstance(backend, backends.CloudVmRayBackend)
-        if all_users:
-            code = managed_job_utils.ManagedJobCodeGen.cancel_jobs_by_id(
-                None, all_users=True)
-        elif all:
-            code = managed_job_utils.ManagedJobCodeGen.cancel_jobs_by_id(None)
-        elif job_ids:
-            code = managed_job_utils.ManagedJobCodeGen.cancel_jobs_by_id(
-                job_ids)
-        elif name is not None:
-            code = managed_job_utils.ManagedJobCodeGen.cancel_job_by_name(name)
-        else:
-            assert pool is not None, (job_ids, name, pool, all)
-            code = managed_job_utils.ManagedJobCodeGen.cancel_jobs_by_pool(pool)
-        # The stderr is redirected to stdout
-        returncode, stdout, stderr = backend.run_on_head(handle,
-                                                         code,
-                                                         require_outputs=True,
-                                                         stream_logs=False)
-        try:
-            subprocess_utils.handle_returncode(returncode, code,
-                                               'Failed to cancel managed job',
-                                               stdout + stderr)
-        except exceptions.CommandError as e:
-            with ux_utils.print_exception_no_traceback():
-                raise RuntimeError(e.error_msg) from e
+
+        use_legacy = not handle.is_grpc_enabled_with_flag
+
+        if not use_legacy:
+            current_workspace = skypilot_config.get_active_workspace()
+            try:
+                request = managed_jobsv1_pb2.CancelJobsRequest(
+                    current_workspace=current_workspace)
+
+                if all_users or all or job_ids:
+                    request.all_users = all_users
+                    if all:
+                        request.user_hash = common_utils.get_user_hash()
+                    if job_ids is not None:
+                        request.job_ids.CopyFrom(
+                            managed_jobsv1_pb2.JobIds(ids=job_ids))
+                elif name is not None:
+                    request.job_name = name
+                else:
+                    assert pool is not None, (job_ids, name, pool, all)
+                    request.pool_name = pool
+
+                response = backend_utils.invoke_skylet_with_retries(
+                    lambda: cloud_vm_ray_backend.SkyletClient(
+                        handle.get_grpc_channel()).cancel_managed_jobs(request))
+                stdout = response.message
+            except exceptions.SkyletMethodNotImplementedError:
+                use_legacy = True
+
+        if use_legacy:
+            if all_users or all or job_ids:
+                code = managed_job_utils.ManagedJobCodeGen.cancel_jobs_by_id(
+                    job_ids, all_users=all_users)
+            elif name is not None:
+                code = managed_job_utils.ManagedJobCodeGen.cancel_job_by_name(
+                    name)
+            else:
+                assert pool is not None, (job_ids, name, pool, all)
+                code = managed_job_utils.ManagedJobCodeGen.cancel_jobs_by_pool(
+                    pool)
+            # The stderr is redirected to stdout
+            returncode, stdout, stderr = backend.run_on_head(
+                handle, code, require_outputs=True, stream_logs=False)
+            try:
+                subprocess_utils.handle_returncode(
+                    returncode, code, 'Failed to cancel managed job',
+                    stdout + stderr)
+            except exceptions.CommandError as e:
+                with ux_utils.print_exception_no_traceback():
+                    raise RuntimeError(e.error_msg) from e
 
         logger.info(stdout)
         if 'Multiple jobs found with name' in stdout:
@@ -901,9 +1153,10 @@ def pool_apply(
     task: 'sky.Task',
     pool_name: str,
     mode: serve_utils.UpdateMode = serve_utils.DEFAULT_UPDATE_MODE,
+    workers: Optional[int] = None,
 ) -> None:
     """Apply a config to a pool."""
-    return impl.apply(task, pool_name, mode, pool=True)
+    return impl.apply(task, workers, pool_name, mode, pool=True)
 
 
 @usage_lib.entrypoint