[jobs] keep track of a managed job's owner/creator

Author: cg505

sky/cli.py

@@ -1379,12 +1379,14 @@ def exec(cluster: Optional[str], cluster_option: Optional[str],
 def _handle_jobs_queue_request(
         request_id: str,
         show_all: bool,
+        show_user: bool,
         limit_num_jobs_to_show: bool = False,
         is_called_by_user: bool = False) -> Tuple[Optional[int], str]:
     """Get the in-progress managed jobs.
 
     Args:
         show_all: Show all information of each job (e.g., region, price).
+        show_user: Show the user who submitted the job.
         limit_num_jobs_to_show: If True, limit the number of jobs to show to
             _NUM_MANAGED_JOBS_TO_SHOW_IN_STATUS, which is mainly used by
             `sky status`.
@@ -1452,6 +1454,7 @@ def _handle_jobs_queue_request(
                             if limit_num_jobs_to_show else None)
         msg = managed_jobs.format_job_table(managed_jobs_,
                                             show_all=show_all,
+                                            show_user=show_user,
                                             max_jobs=max_jobs_to_show)
     return num_in_progress_jobs, msg
 
@@ -1561,7 +1564,9 @@ def _status_kubernetes(show_all: bool):
         click.echo(f'\n{colorama.Fore.CYAN}{colorama.Style.BRIGHT}'
                    f'Managed jobs'
                    f'{colorama.Style.RESET_ALL}')
-        msg = managed_jobs.format_job_table(all_jobs, show_all=show_all)
+        msg = managed_jobs.format_job_table(all_jobs,
+                                            show_all=show_all,
+                                            show_user=False)
         click.echo(msg)
     if any(['sky-serve-controller' in c.cluster_name for c in all_clusters]):
         # TODO: Parse serve controllers and show services separately.
@@ -1779,7 +1784,8 @@ def status(verbose: bool, refresh: bool, ip: bool, endpoints: bool,
     show_managed_jobs = show_managed_jobs and not any([clusters, ip, endpoints])
     if show_managed_jobs:
         managed_jobs_queue_request_id = managed_jobs.queue(refresh=False,
-                                                           skip_finished=True)
+                                                           skip_finished=True,
+                                                           all_users=all_users)
     show_endpoints = endpoints or endpoint is not None
     show_single_endpoint = endpoint is not None
     show_services = show_services and not any([clusters, ip, endpoints])
@@ -1859,6 +1865,7 @@ def status(verbose: bool, refresh: bool, ip: bool, endpoints: bool,
                 num_in_progress_jobs, msg = _handle_jobs_queue_request(
                     managed_jobs_queue_request_id,
                     show_all=False,
+                    show_user=False,
                     limit_num_jobs_to_show=not all,
                     is_called_by_user=False)
             except KeyboardInterrupt:
@@ -2751,7 +2758,7 @@ def start(
 def down(
     clusters: List[str],
     all: bool,  # pylint: disable=redefined-builtin
-    all_users: bool,  # pylint: disable=redefined-builtin
+    all_users: bool,
     yes: bool,
     purge: bool,
     async_call: bool,
@@ -2812,7 +2819,9 @@ def _hint_or_raise_for_down_jobs_controller(controller_name: str,
     with rich_utils.client_status(
             '[bold cyan]Checking for in-progress managed jobs[/]'):
         try:
-            request_id = managed_jobs.queue(refresh=False, skip_finished=True)
+            request_id = managed_jobs.queue(refresh=False,
+                                            skip_finished=True,
+                                            all_users=True)
             managed_jobs_ = sdk.stream_and_get(request_id)
         except exceptions.ClusterNotUpError as e:
             if controller.value.connection_error_hint in str(e):
@@ -2836,7 +2845,9 @@ def _hint_or_raise_for_down_jobs_controller(controller_name: str,
            'jobs (output of `sky jobs queue`) will be lost.')
     click.echo(msg)
     if managed_jobs_:
-        job_table = managed_jobs.format_job_table(managed_jobs_, show_all=False)
+        job_table = managed_jobs.format_job_table(managed_jobs_,
+                                                  show_all=False,
+                                                  show_user=True)
         msg = controller.value.decline_down_for_dirty_controller_hint
         # Add prefix to each line to align with the bullet point.
         msg += '\n'.join(
@@ -3905,9 +3916,16 @@ def jobs_launch(
               is_flag=True,
               required=False,
               help='Show only pending/running jobs\' information.')
+@click.option('--all-users',
+              '-u',
+              default=False,
+              is_flag=True,
+              required=False,
+              help='Show jobs from all users.')
 @usage_lib.entrypoint
 # pylint: disable=redefined-builtin
-def jobs_queue(verbose: bool, refresh: bool, skip_finished: bool):
+def jobs_queue(verbose: bool, refresh: bool, skip_finished: bool,
+               all_users: bool):
     """Show statuses of managed jobs.
 
     Each managed jobs can have one of the following statuses:
@@ -3964,9 +3982,10 @@ def jobs_queue(verbose: bool, refresh: bool, skip_finished: bool):
     click.secho('Fetching managed job statuses...', fg='cyan')
     with rich_utils.client_status('[cyan]Checking managed jobs[/]'):
         managed_jobs_request_id = managed_jobs.queue(
-            refresh=refresh, skip_finished=skip_finished)
+            refresh=refresh, skip_finished=skip_finished, all_users=all_users)
         _, msg = _handle_jobs_queue_request(managed_jobs_request_id,
                                             show_all=verbose,
+                                            show_user=all_users,
                                             is_called_by_user=True)
     if not skip_finished:
         in_progress_only_hint = ''
@@ -3989,16 +4008,23 @@ def jobs_queue(verbose: bool, refresh: bool, skip_finished: bool):
               is_flag=True,
               default=False,
               required=False,
-              help='Cancel all managed jobs.')
+              help='Cancel all managed jobs for the current user.')
 @click.option('--yes',
               '-y',
               is_flag=True,
               default=False,
               required=False,
               help='Skip confirmation prompt.')
+@click.option('--all-users',
+              '-u',
+              is_flag=True,
+              default=False,
+              required=False,
+              help='Cancel all managed jobs from all users.')
 @usage_lib.entrypoint
 # pylint: disable=redefined-builtin
-def jobs_cancel(name: Optional[str], job_ids: Tuple[int], all: bool, yes: bool):
+def jobs_cancel(name: Optional[str], job_ids: Tuple[int], all: bool, yes: bool,
+                all_users: bool):
     """Cancel managed jobs.
 
     You can provide either a job name or a list of job IDs to be cancelled.
@@ -4015,25 +4041,33 @@ def jobs_cancel(name: Optional[str], job_ids: Tuple[int], all: bool, yes: bool):
       $ sky jobs cancel 1 2 3
     """
     job_id_str = ','.join(map(str, job_ids))
-    if sum([bool(job_ids), name is not None, all]) != 1:
-        argument_str = f'--job-ids {job_id_str}' if job_ids else ''
-        argument_str += f' --name {name}' if name is not None else ''
-        argument_str += ' --all' if all else ''
+    if sum([bool(job_ids), name is not None, all, all_users]) != 1:
+        arguments = []
+        arguments += [f'--job-ids {job_id_str}'] if job_ids else []
+        arguments += [f'--name {name}'] if name is not None else []
+        arguments += ['--all'] if all else []
+        arguments += ['--all-users'] if all_users else []
         raise click.UsageError(
             'Can only specify one of JOB_IDS or --name or --all. '
-            f'Provided {argument_str!r}.')
+            f'Provided {" ".join(arguments)!r}.')
 
     if not yes:
         job_identity_str = (f'managed jobs with IDs {job_id_str}'
                             if job_ids else repr(name))
         if all:
             job_identity_str = 'all managed jobs'
+        if all_users:
+            job_identity_str = 'all managed jobs FOR ALL USERS'
         click.confirm(f'Cancelling {job_identity_str}. Proceed?',
                       default=True,
                       abort=True,
                       show_default=True)
 
-    sdk.stream_and_get(managed_jobs.cancel(job_ids=job_ids, name=name, all=all))
+    sdk.stream_and_get(
+        managed_jobs.cancel(job_ids=job_ids,
+                            name=name,
+                            all=all,
+                            all_users=all_users))
 
 
 @jobs.command('logs', cls=_DocumentedCodeCommand)

sky/jobs/client/sdk.py

@@ -85,14 +85,16 @@ def launch(
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 def queue(refresh: bool,
-          skip_finished: bool = False) -> server_common.RequestId:
+          skip_finished: bool = False,
+          all_users: bool = False) -> server_common.RequestId:
     """Gets statuses of managed jobs.
 
     Please refer to sky.cli.job_queue for documentation.
 
     Args:
         refresh: Whether to restart the jobs controller if it is stopped.
         skip_finished: Whether to skip finished jobs.
+        all_users: Whether to show all users' jobs.
 
     Returns:
         The request ID of the queue request.
@@ -126,6 +128,7 @@ def queue(refresh: bool,
     body = payloads.JobsQueueBody(
         refresh=refresh,
         skip_finished=skip_finished,
+        all_users=all_users,
     )
     response = requests.post(
         f'{server_common.get_server_url()}/jobs/queue',
@@ -138,9 +141,10 @@ def queue(refresh: bool,
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
 def cancel(
-        name: Optional[str] = None,
-        job_ids: Optional[List[int]] = None,
-        all: bool = False,  # pylint: disable=redefined-builtin
+    name: Optional[str] = None,
+    job_ids: Optional[List[int]] = None,
+    all: bool = False,  # pylint: disable=redefined-builtin
+    all_users: bool = False,
 ) -> server_common.RequestId:
     """Cancels managed jobs.
 
@@ -150,6 +154,7 @@ def cancel(
         name: Name of the managed job to cancel.
         job_ids: IDs of the managed jobs to cancel.
         all: Whether to cancel all managed jobs.
+        all_users: Whether to cancel all managed jobs from all users.
 
     Returns:
         The request ID of the cancel request.
@@ -162,6 +167,7 @@ def cancel(
         name=name,
         job_ids=job_ids,
         all=all,
+        all_users=all_users,
     )
     response = requests.post(
         f'{server_common.get_server_url()}/jobs/cancel',

sky/jobs/dashboard/dashboard.py

@@ -16,6 +16,7 @@
 import yaml
 
 from sky import jobs as managed_jobs
+from sky.client import sdk
 from sky.jobs import constants as managed_job_constants
 from sky.utils import common_utils
 from sky.utils import controller_utils
@@ -134,14 +135,16 @@ def _extract_launch_history(log_content: str) -> str:
 def home():
     if not _is_running_on_jobs_controller():
         # Experimental: run on laptop (refresh is very slow).
-        all_managed_jobs = managed_jobs.queue(refresh=True, skip_finished=False)
+        request_id = managed_jobs.queue(refresh=True, skip_finished=False)
+        all_managed_jobs = sdk.get(request_id)
     else:
         job_table = managed_jobs.dump_managed_job_queue()
         all_managed_jobs = managed_jobs.load_managed_job_queue(job_table)
 
     timestamp = datetime.datetime.now(datetime.timezone.utc)
     rows = managed_jobs.format_job_table(all_managed_jobs,
                                          show_all=True,
+                                         show_user=False,
                                          return_rows=True)
 
     status_counts = collections.defaultdict(int)

sky/jobs/scheduler.py

@@ -49,6 +49,7 @@
 from sky.jobs import constants as managed_job_constants
 from sky.jobs import state
 from sky.skylet import constants
+from sky.utils import common_utils
 from sky.utils import subprocess_utils
 
 logger = sky_logging.init_logger('sky.jobs.controller')
@@ -190,9 +191,12 @@ def submit_job(job_id: int, dag_yaml_path: str, env_file_path: str) -> None:
     PENDING. It will tell the scheduler to try and start the job controller, if
     there are resources available. It may block to acquire the lock, so it
     should not be on the critical path for `sky jobs launch -d`.
+
+    The user hash should be set (e.g. via SKYPILOT_USER_ID) before calling this.
     """
     with filelock.FileLock(_get_lock_path()):
-        state.scheduler_set_waiting(job_id, dag_yaml_path, env_file_path)
+        state.scheduler_set_waiting(job_id, dag_yaml_path, env_file_path,
+                                    common_utils.get_user_hash())
     maybe_schedule_next_jobs()
 
 

sky/jobs/server/core.py

@@ -346,7 +346,9 @@ def _maybe_restart_controller(
 
 
 @usage_lib.entrypoint
-def queue(refresh: bool, skip_finished: bool = False) -> List[Dict[str, Any]]:
+def queue(refresh: bool,
+          skip_finished: bool = False,
+          all_users: bool = False) -> List[Dict[str, Any]]:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
     """Gets statuses of managed jobs.
 
@@ -394,6 +396,18 @@ def queue(refresh: bool, skip_finished: bool = False) -> List[Dict[str, Any]]:
                            f'{returncode}')
 
     jobs = managed_job_utils.load_managed_job_queue(job_table_payload)
+
+    if not all_users:
+        # For backwards compatibility, we show jobs that have user_hash None.
+        def user_hash_matches_or_missing(job: Dict[str, Any]) -> bool:
+            if 'user_hash' not in job:
+                return True
+            if job['user_hash'] is None:
+                return True
+            return job['user_hash'] == common_utils.get_user_hash()
+
+        jobs = list(filter(user_hash_matches_or_missing, jobs))
+
     if skip_finished:
         # Filter out the finished jobs. If a multi-task job is partially
         # finished, we will include all its tasks.
@@ -402,14 +416,16 @@ def queue(refresh: bool, skip_finished: bool = False) -> List[Dict[str, Any]]:
         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))
+
     return jobs
 
 
 @usage_lib.entrypoint
 # pylint: disable=redefined-builtin
 def cancel(name: Optional[str] = None,
            job_ids: Optional[List[int]] = None,
-           all: bool = False) -> None:
+           all: bool = False,
+           all_users: bool = False) -> None:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
     """Cancels managed jobs.
 
@@ -425,18 +441,23 @@ def cancel(name: Optional[str] = None,
         stopped_message='All managed jobs should have finished.')
 
     job_id_str = ','.join(map(str, job_ids))
-    if sum([bool(job_ids), name is not None, all]) != 1:
-        argument_str = f'job_ids={job_id_str}' if job_ids else ''
-        argument_str += f' name={name}' if name is not None else ''
-        argument_str += ' all' if all else ''
+    if sum([bool(job_ids), name is not None, all, all_users]) != 1:
+        arguments = []
+        arguments += [f'job_ids={job_id_str}'] if job_ids else []
+        arguments += [f'name={name}'] if name is not None else []
+        arguments += ['all'] if all else []
+        arguments += ['all_users'] if all_users else []
         with ux_utils.print_exception_no_traceback():
             raise ValueError('Can only specify one of JOB_IDS or name or all. '
-                             f'Provided {argument_str!r}.')
+                             f'Provided {" ".join(arguments)!r}.')
 
     backend = backend_utils.get_backend_from_handle(handle)
     assert isinstance(backend, backends.CloudVmRayBackend)
     if all:
         code = managed_job_utils.ManagedJobCodeGen.cancel_jobs_by_id(None)
+    elif all_users:
+        code = managed_job_utils.ManagedJobCodeGen.cancel_jobs_by_id(
+            None, all_users=True)
     elif job_ids:
         code = managed_job_utils.ManagedJobCodeGen.cancel_jobs_by_id(job_ids)
     else:

sky/jobs/server/server.py

@@ -109,9 +109,18 @@ async def download_logs(
 @router.get('/dashboard')
 async def dashboard(request: fastapi.Request,
                     user_hash: str) -> fastapi.Response:
+    # TODO(cooperc): Support showing only jobs for a specific user.
+
+    # FIX(zhwu/cooperc/eric): Fix log downloading (assumes global
+    # /download_log/xx route)
+
     # Note: before #4717, each user had their own controller, and thus their own
     # dashboard. Now, all users share the same controller, so this isn't really
     # necessary. TODO(cooperc): clean up.
+
+    # TODO: Put this in an executor to avoid blocking the main server thread.
+    # It can take a long time if it needs to check the controller status.
+
     # Find the port for the dashboard of the user
     os.environ[constants.USER_ID_ENV_VAR] = user_hash
     server_common.reload_for_new_request(client_entrypoint=None,