skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/skylet/job_lib.py

@@ -8,27 +8,32 @@ import json
 import os
 import pathlib
 import shlex
-import subprocess
+import signal
+import sqlite3
 import time
-import typing
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Sequence
 
 import colorama
 import filelock
 import psutil
 
+from sky import global_user_state
 from sky import sky_logging
 from sky.skylet import constants
 from sky.utils import common_utils
 from sky.utils import db_utils
 from sky.utils import log_utils
-
-if typing.TYPE_CHECKING:
-    from ray.dashboard.modules.job import pydantic_models as ray_pydantic
+from sky.utils import message_utils
+from sky.utils import subprocess_utils
 
 logger = sky_logging.init_logger(__name__)
 
+_LINUX_NEW_LINE = '\n'
 _JOB_STATUS_LOCK = '~/.sky/locks/.job_{}.lock'
+# JOB_CMD_IDENTIFIER is used for identifying the process retrieved
+# with pid is the same driver process to guard against the case where
+# the same pid is reused by a different process.
+JOB_CMD_IDENTIFIER = 'echo "SKYPILOT_JOB_ID <{}>"'
 
 
 def _get_lock_path(job_id: int) -> str:
@@ -48,6 +53,7 @@ class JobInfoLoc(enum.IntEnum):
     START_AT = 6
     END_AT = 7
     RESOURCES = 8
+    PID = 9
 
 
 _DB_PATH = os.path.expanduser('~/.sky/jobs.db')
@@ -55,6 +61,31 @@ os.makedirs(pathlib.Path(_DB_PATH).parents[0], exist_ok=True)
 
 
 def create_table(cursor, conn):
+    # Enable WAL mode to avoid locking issues.
+    # See: issue #3863, #1441 and PR #1509
+    # https://github.com/microsoft/WSL/issues/2395
+    # TODO(romilb): We do not enable WAL for WSL because of known issue in WSL.
+    #  This may cause the database locked problem from WSL issue #1441.
+    if not common_utils.is_wsl():
+        try:
+            cursor.execute('PRAGMA journal_mode=WAL')
+        except sqlite3.OperationalError as e:
+            if 'database is locked' not in str(e):
+                raise
+            # If the database is locked, it is OK to continue, as the WAL mode
+            # is not critical and is likely to be enabled by other processes.
+
+    # Pid column is used for keeping track of the driver process of a job. It
+    # can be in three states:
+    # -1: The job was submitted with SkyPilot older than #4318, where we use
+    #     ray job submit to submit the job, i.e. no pid is recorded. This is for
+    #     backward compatibility and should be removed after 0.10.0.
+    # 0: The job driver process has never been started. When adding a job with
+    #    INIT state, the pid will be set to 0 (the default -1 value is just for
+    #    backward compatibility).
+    # >=0: The job has been started. The pid is the driver process's pid.
+    #      The driver can be actually running or finished.
+    # TODO(SKY-1213): username is actually user hash, should rename.
     cursor.execute("""\
         CREATE TABLE IF NOT EXISTS jobs (
         job_id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -63,7 +94,10 @@ def create_table(cursor, conn):
         submitted_at FLOAT,
         status TEXT,
         run_timestamp TEXT CANDIDATE KEY,
-        start_at FLOAT DEFAULT -1)""")
+        start_at FLOAT DEFAULT -1,
+        end_at FLOAT DEFAULT NULL,
+        resources TEXT DEFAULT NULL,
+        pid INTEGER DEFAULT -1)""")
 
     cursor.execute("""CREATE TABLE IF NOT EXISTS pending_jobs(
         job_id INTEGER,
@@ -74,7 +108,8 @@ def create_table(cursor, conn):
 
     db_utils.add_column_to_table(cursor, conn, 'jobs', 'end_at', 'FLOAT')
     db_utils.add_column_to_table(cursor, conn, 'jobs', 'resources', 'TEXT')
-
+    db_utils.add_column_to_table(cursor, conn, 'jobs', 'pid',
+                                 'INTEGER DEFAULT -1')
     conn.commit()
 
 
@@ -84,7 +119,7 @@ _CONN = _DB.conn
 
 
 class JobStatus(enum.Enum):
-    """Job status"""
+    """Job status enum."""
 
     # 3 in-flux states: each can transition to any state below it.
     # The `job_id` has been generated, but the generated ray program has
@@ -93,54 +128,75 @@ class JobStatus(enum.Enum):
     # In the 'jobs' table, the `submitted_at` column will be set to the current
     # time, when the job is firstly created (in the INIT state).
     INIT = 'INIT'
+    """The job has been submitted, but not started yet."""
     # The job is waiting for the required resources. (`ray job status`
     # shows RUNNING as the generated ray program has started, but blocked
     # by the placement constraints.)
     PENDING = 'PENDING'
-    # Running the user's setup script (only in effect if --detach-setup is
-    # set). Our update_job_status() can temporarily (for a short period) set
+    """The job is waiting for required resources."""
+    # Running the user's setup script.
+    # Our update_job_status() can temporarily (for a short period) set
     # the status to SETTING_UP, if the generated ray program has not set
     # the status to PENDING or RUNNING yet.
     SETTING_UP = 'SETTING_UP'
+    """The job is running the user's setup script."""
     # The job is running.
     # In the 'jobs' table, the `start_at` column will be set to the current
     # time, when the job is firstly transitioned to RUNNING.
     RUNNING = 'RUNNING'
+    """The job is running."""
+    # The job driver process failed. This happens when the job driver process
+    # finishes when the status in job table is still not set to terminal state.
+    # We should keep this state before the SUCCEEDED, as our job status update
+    # relies on the order of the statuses to keep the latest status.
+    FAILED_DRIVER = 'FAILED_DRIVER'
+    """The job driver process failed."""
     # 3 terminal states below: once reached, they do not transition.
     # The job finished successfully.
     SUCCEEDED = 'SUCCEEDED'
+    """The job finished successfully."""
     # The job fails due to the user code or a system restart.
     FAILED = 'FAILED'
-    # The job setup failed (only in effect if --detach-setup is set). It
-    # needs to be placed after the `FAILED` state, so that the status
-    # set by our generated ray program will not be overwritten by
-    # ray's job status (FAILED).
-    # This is for a better UX, so that the user can find out the reason
-    # of the failure quickly.
+    """The job fails due to the user code."""
+    # The job setup failed. It needs to be placed after the `FAILED` state,
+    # so that the status set by our generated ray program will not be
+    # overwritten by ray's job status (FAILED). This is for a better UX, so
+    # that the user can find out the reason of the failure quickly.
     FAILED_SETUP = 'FAILED_SETUP'
+    """The job setup failed."""
     # The job is cancelled by the user.
     CANCELLED = 'CANCELLED'
+    """The job is cancelled by the user."""
 
     @classmethod
     def nonterminal_statuses(cls) -> List['JobStatus']:
         return [cls.INIT, cls.SETTING_UP, cls.PENDING, cls.RUNNING]
 
-    def is_terminal(self):
+    def is_terminal(self) -> bool:
         return self not in self.nonterminal_statuses()
 
-    def __lt__(self, other):
+    @classmethod
+    def user_code_failure_states(cls) -> Sequence['JobStatus']:
+        return (cls.FAILED, cls.FAILED_SETUP)
+
+    def __lt__(self, other: 'JobStatus') -> bool:
         return list(JobStatus).index(self) < list(JobStatus).index(other)
 
-    def colored_str(self):
+    def colored_str(self) -> str:
         color = _JOB_STATUS_TO_COLOR[self]
         return f'{color}{self.value}{colorama.Style.RESET_ALL}'
 
 
-# Only update status of the jobs after this many seconds of job submission,
-# to avoid race condition with `ray job` to make sure it job has been
-# correctly updated.
+# We have two steps for job submissions:
+# 1. Client reserve a job id from the job table by adding a INIT state job.
+# 2. Client updates the job status to PENDING by actually submitting the job's
+#    command to the scheduler.
+# In normal cases, the two steps happens very close to each other through two
+# consecutive SSH connections.
+# We should update status for INIT job that has been staying in INIT state for
+# a while (60 seconds), which likely fails to reach step 2.
 # TODO(zhwu): This number should be tuned based on heuristics.
-_PENDING_SUBMIT_GRACE_PERIOD = 60
+_INIT_SUBMIT_GRACE_PERIOD = 60
 
 _PRE_RESOURCE_STATUSES = [JobStatus.PENDING]
 
@@ -163,17 +219,38 @@ class JobScheduler:
         _CURSOR.execute((f'UPDATE pending_jobs SET submit={int(time.time())} '
                          f'WHERE job_id={job_id!r}'))
         _CONN.commit()
-        subprocess.Popen(run_cmd, shell=True, stdout=subprocess.DEVNULL)
+        pid = subprocess_utils.launch_new_process_tree(run_cmd)
+        # TODO(zhwu): Backward compatibility, remove this check after 0.10.0.
+        # This is for the case where the job is submitted with SkyPilot older
+        # than #4318, using ray job submit.
+        if 'job submit' in run_cmd:
+            pid = -1
+        _CURSOR.execute((f'UPDATE jobs SET pid={pid} '
+                         f'WHERE job_id={job_id!r}'))
+        _CONN.commit()
 
     def schedule_step(self, force_update_jobs: bool = False) -> None:
-        jobs = self._get_jobs()
-        if len(jobs) > 0 or force_update_jobs:
+        if force_update_jobs:
             update_status()
+        pending_job_ids = self._get_pending_job_ids()
         # TODO(zhwu, mraheja): One optimization can be allowing more than one
         # job staying in the pending state after ray job submit, so that to be
         # faster to schedule a large amount of jobs.
-        for job_id, run_cmd, submit, created_time in jobs:
+        for job_id in pending_job_ids:
             with filelock.FileLock(_get_lock_path(job_id)):
+                pending_job = _get_pending_job(job_id)
+                if pending_job is None:
+                    # Pending job can be removed by another thread, due to the
+                    # job being scheduled already.
+                    continue
+                run_cmd = pending_job['run_cmd']
+                submit = pending_job['submit']
+                created_time = pending_job['created_time']
+                # We don't have to refresh the job status before checking, as
+                # the job status will only be stale in rare cases where ray job
+                # crashes; or the job stays in INIT state for a long time.
+                # In those cases, the periodic JobSchedulerEvent event will
+                # update the job status every 300 seconds.
                 status = get_status_no_lock(job_id)
                 if (status not in _PRE_RESOURCE_STATUSES or
                         created_time < psutil.boot_time()):
@@ -187,8 +264,8 @@ class JobScheduler:
                 self._run_job(job_id, run_cmd)
                 return
 
-    def _get_jobs(self) -> List[Tuple[int, str, int, int]]:
-        """Returns the metadata for jobs in the pending jobs table
+    def _get_pending_job_ids(self) -> List[int]:
+        """Returns the job ids in the pending jobs table
 
         The information contains job_id, run command, submit time,
         creation time.
@@ -199,9 +276,10 @@ class JobScheduler:
 class FIFOScheduler(JobScheduler):
     """First in first out job scheduler"""
 
-    def _get_jobs(self) -> List[Tuple[int, str, int, int]]:
-        return list(
-            _CURSOR.execute('SELECT * FROM pending_jobs ORDER BY job_id'))
+    def _get_pending_job_ids(self) -> List[int]:
+        rows = _CURSOR.execute(
+            'SELECT job_id FROM pending_jobs ORDER BY job_id').fetchall()
+        return [row[0] for row in rows]
 
 
 scheduler = FIFOScheduler()
@@ -211,59 +289,13 @@ _JOB_STATUS_TO_COLOR = {
     JobStatus.SETTING_UP: colorama.Fore.BLUE,
     JobStatus.PENDING: colorama.Fore.BLUE,
     JobStatus.RUNNING: colorama.Fore.GREEN,
+    JobStatus.FAILED_DRIVER: colorama.Fore.RED,
     JobStatus.SUCCEEDED: colorama.Fore.GREEN,
     JobStatus.FAILED: colorama.Fore.RED,
     JobStatus.FAILED_SETUP: colorama.Fore.RED,
     JobStatus.CANCELLED: colorama.Fore.YELLOW,
 }
 
-_RAY_TO_JOB_STATUS_MAP = {
-    # These are intentionally set this way, because:
-    # 1. when the ray status indicates the job is PENDING the generated
-    # python program has been `ray job submit` from the job queue
-    # and is now PENDING
-    # 2. when the ray status indicates the job is RUNNING the job can be in
-    # setup or resources may not be allocated yet, i.e. the job should be
-    # PENDING.
-    # For case 2, update_job_status() would compare this mapped PENDING to
-    # the status in our jobs DB and take the max. This is because the job's
-    # generated ray program is the only place that can determine a job has
-    # reserved resources and actually started running: it will set the
-    # status in the DB to SETTING_UP or RUNNING.
-    # If there is no setup specified in the task, as soon as it is started
-    # (ray's status becomes RUNNING), i.e. it will be very rare that the job
-    # will be set to SETTING_UP by the update_job_status, as our generated
-    # ray program will set the status to PENDING immediately.
-    'PENDING': JobStatus.PENDING,
-    'RUNNING': JobStatus.PENDING,
-    'SUCCEEDED': JobStatus.SUCCEEDED,
-    'FAILED': JobStatus.FAILED,
-    'STOPPED': JobStatus.CANCELLED,
-}
-
-
-def _create_ray_job_submission_client():
-    """Import the ray job submission client."""
-    try:
-        import ray  # pylint: disable=import-outside-toplevel
-    except ImportError:
-        logger.error('Failed to import ray')
-        raise
-    try:
-        # pylint: disable=import-outside-toplevel
-        from ray import job_submission
-    except ImportError:
-        logger.error(
-            f'Failed to import job_submission with ray=={ray.__version__}')
-        raise
-    port = get_job_submission_port()
-    return job_submission.JobSubmissionClient(
-        address=f'http://127.0.0.1:{port}')
-
-
-def make_ray_job_id(sky_job_id: int) -> str:
-    return f'{sky_job_id}-{getpass.getuser()}'
-
 
 def make_job_command_with_user_switching(username: str,
                                          command: str) -> List[str]:
@@ -275,9 +307,10 @@ def add_job(job_name: str, username: str, run_timestamp: str,
     """Atomically reserve the next available job id for the user."""
     job_submitted_at = time.time()
     # job_id will autoincrement with the null value
-    _CURSOR.execute('INSERT INTO jobs VALUES (null, ?, ?, ?, ?, ?, ?, null, ?)',
-                    (job_name, username, job_submitted_at, JobStatus.INIT.value,
-                     run_timestamp, None, resources_str))
+    _CURSOR.execute(
+        'INSERT INTO jobs VALUES (null, ?, ?, ?, ?, ?, ?, null, ?, 0)',
+        (job_name, username, job_submitted_at, JobStatus.INIT.value,
+         run_timestamp, None, resources_str))
     _CONN.commit()
     rows = _CURSOR.execute('SELECT job_id FROM jobs WHERE run_timestamp=(?)',
                            (run_timestamp,))
@@ -362,12 +395,12 @@ def get_statuses_payload(job_ids: List[Optional[int]]) -> str:
     statuses = {job_id: None for job_id in job_ids}
     for (job_id, status) in rows:
         statuses[job_id] = status
-    return common_utils.encode_payload(statuses)
+    return message_utils.encode_payload(statuses)
 
 
 def load_statuses_payload(
         statuses_payload: str) -> Dict[Optional[int], Optional[JobStatus]]:
-    original_statuses = common_utils.decode_payload(statuses_payload)
+    original_statuses = message_utils.decode_payload(statuses_payload)
     statuses = dict()
     for job_id, status in original_statuses.items():
         # json.dumps will convert all keys to strings. Integers will
@@ -405,8 +438,8 @@ def get_job_submitted_or_ended_timestamp_payload(job_id: int,
     rows = _CURSOR.execute(f'SELECT {field} FROM jobs WHERE job_id=(?)',
                            (job_id,))
     for (timestamp,) in rows:
-        return common_utils.encode_payload(timestamp)
-    return common_utils.encode_payload(None)
+        return message_utils.encode_payload(timestamp)
+    return message_utils.encode_payload(None)
 
 
 def get_ray_port():
@@ -452,35 +485,26 @@ def _get_records_from_rows(rows) -> List[Dict[str, Any]]:
             'start_at': row[JobInfoLoc.START_AT.value],
             'end_at': row[JobInfoLoc.END_AT.value],
             'resources': row[JobInfoLoc.RESOURCES.value],
+            'pid': row[JobInfoLoc.PID.value],
         })
     return records
 
 
 def _get_jobs(
-        username: Optional[str],
+        user_hash: Optional[str],
         status_list: Optional[List[JobStatus]] = None) -> List[Dict[str, Any]]:
     """Returns jobs with the given fields, sorted by job_id, descending."""
     if status_list is None:
         status_list = list(JobStatus)
-    status_str_list = [status.value for status in status_list]
-    if username is None:
-        rows = _CURSOR.execute(
-            f"""\
-            SELECT * FROM jobs
-            WHERE status IN ({','.join(['?'] * len(status_list))})
-            ORDER BY job_id DESC""",
-            (*status_str_list,),
-        )
-    else:
-        rows = _CURSOR.execute(
-            f"""\
-            SELECT * FROM jobs
-            WHERE status IN ({','.join(['?'] * len(status_list))})
-            AND username=(?)
-            ORDER BY job_id DESC""",
-            (*status_str_list, username),
-        )
-
+    status_str_list = [repr(status.value) for status in status_list]
+    filter_str = f'WHERE status IN ({",".join(status_str_list)})'
+    params = []
+    if user_hash is not None:
+        # We use the old username field for compatibility.
+        filter_str += ' AND username=(?)'
+        params.append(user_hash)
+    rows = _CURSOR.execute(
+        f'SELECT * FROM jobs {filter_str} ORDER BY job_id DESC', params)
     records = _get_records_from_rows(rows)
     return records
 
@@ -497,16 +521,35 @@ def _get_jobs_by_ids(job_ids: List[int]) -> List[Dict[str, Any]]:
     return records
 
 
-def _get_pending_jobs():
+def _get_pending_job(job_id: int) -> Optional[Dict[str, Any]]:
     rows = _CURSOR.execute(
-        'SELECT job_id, created_time, submit FROM pending_jobs')
-    rows = list(rows)
-    return {
-        job_id: {
+        'SELECT created_time, submit, run_cmd FROM pending_jobs '
+        f'WHERE job_id={job_id!r}')
+    for row in rows:
+        created_time, submit, run_cmd = row
+        return {
             'created_time': created_time,
-            'submit': submit
-        } for job_id, created_time, submit in rows
-    }
+            'submit': submit,
+            'run_cmd': run_cmd
+        }
+    return None
+
+
+def _is_job_driver_process_running(job_pid: int, job_id: int) -> bool:
+    """Check if the job driver process is running.
+
+    We check the cmdline to avoid the case where the same pid is reused by a
+    different process.
+    """
+    if job_pid <= 0:
+        return False
+    try:
+        job_process = psutil.Process(job_pid)
+        return job_process.is_running() and any(
+            JOB_CMD_IDENTIFIER.format(job_id) in line
+            for line in job_process.cmdline())
+    except psutil.NoSuchProcess:
+        return False
 
 
 def update_job_status(job_ids: List[int],
@@ -520,94 +563,131 @@ def update_job_status(job_ids: List[int],
     during job cancelling, we still need this to handle the staleness problem,
     caused by instance restarting and other corner cases (if any).
 
-    This function should only be run on the remote instance with ray==2.4.0.
+    This function should only be run on the remote instance with ray>=2.4.0.
     """
-    if len(job_ids) == 0:
+    echo = logger.info if not silent else logger.debug
+    if not job_ids:
         return []
 
-    # TODO: if too slow, directly query against redis.
-    ray_job_ids = [make_ray_job_id(job_id) for job_id in job_ids]
-
-    job_client = _create_ray_job_submission_client()
-
-    # In ray 2.4.0, job_client.list_jobs returns a list of JobDetails,
-    # which contains the job status (str) and submission_id (str).
-    job_detail_lists: List['ray_pydantic.JobDetails'] = job_client.list_jobs()
-
-    pending_jobs = _get_pending_jobs()
-    job_details = {}
-    ray_job_ids_set = set(ray_job_ids)
-    for job_detail in job_detail_lists:
-        if job_detail.submission_id in ray_job_ids_set:
-            job_details[job_detail.submission_id] = job_detail
-    job_statuses: List[Optional[JobStatus]] = [None] * len(ray_job_ids)
-    for i, ray_job_id in enumerate(ray_job_ids):
-        job_id = job_ids[i]
-        if ray_job_id in job_details:
-            ray_status = job_details[ray_job_id].status
-            job_statuses[i] = _RAY_TO_JOB_STATUS_MAP[ray_status]
-        if job_id in pending_jobs:
-            if pending_jobs[job_id]['created_time'] < psutil.boot_time():
-                logger.info(
-                    f'Job {job_id} is stale, setting to FAILED: '
-                    f'created_time={pending_jobs[job_id]["created_time"]}, '
-                    f'boot_time={psutil.boot_time()}')
-                # The job is stale as it is created before the instance
-                # is booted, e.g. the instance is rebooted.
-                job_statuses[i] = JobStatus.FAILED
-            # Gives a 60 second grace period between job being submit from
-            # the pending table until appearing in ray jobs.
-            if (pending_jobs[job_id]['submit'] > 0 and
-                    pending_jobs[job_id]['submit'] <
-                    time.time() - _PENDING_SUBMIT_GRACE_PERIOD):
-                # For jobs submitted outside of the grace period, we will
-                # consider the ray job status.
-                continue
-            else:
-                # Reset the job status to PENDING even though it may not appear
-                # in the ray jobs, so that it will not be considered as stale.
-                job_statuses[i] = JobStatus.PENDING
-
-    assert len(job_statuses) == len(job_ids), (job_statuses, job_ids)
-
     statuses = []
-    for job_id, status in zip(job_ids, job_statuses):
+    for job_id in job_ids:
         # Per-job status lock is required because between the job status
         # query and the job status update, the job status in the databse
         # can be modified by the generated ray program.
         with filelock.FileLock(_get_lock_path(job_id)):
-            original_status = get_status_no_lock(job_id)
+            status = None
+            job_record = _get_jobs_by_ids([job_id])[0]
+            original_status = job_record['status']
+            job_submitted_at = job_record['submitted_at']
+            job_pid = job_record['pid']
+
+            pid_query_time = time.time()
+            failed_driver_transition_message = None
+            if original_status == JobStatus.INIT:
+                if (job_submitted_at >= psutil.boot_time() and job_submitted_at
+                        >= pid_query_time - _INIT_SUBMIT_GRACE_PERIOD):
+                    # The job id is reserved, but the job is not submitted yet.
+                    # We should keep it in INIT.
+                    status = JobStatus.INIT
+                else:
+                    # We always immediately submit job after the job id is
+                    # allocated, i.e. INIT -> PENDING, if a job stays in INIT
+                    # for too long, it is likely the job submission process
+                    # was killed before the job is submitted. We should set it
+                    # to FAILED then. Note, if ray job indicates the job is
+                    # running, we will change status to PENDING below.
+                    failed_driver_transition_message = (
+                        f'INIT job {job_id} is stale, setting to FAILED_DRIVER')
+                    status = JobStatus.FAILED_DRIVER
+
+            # job_pid is 0 if the job is not submitted yet.
+            # job_pid is -1 if the job is submitted with SkyPilot older than
+            # #4318, using ray job submit. We skip the checking for those
+            # jobs.
+            if job_pid > 0:
+                if _is_job_driver_process_running(job_pid, job_id):
+                    status = JobStatus.PENDING
+                else:
+                    # By default, if the job driver process does not exist,
+                    # the actual SkyPilot job is one of the following:
+                    # 1. Still pending to be submitted.
+                    # 2. Submitted and finished.
+                    # 3. Driver failed without correctly setting the job
+                    #    status in the job table.
+                    # Although we set the status to FAILED_DRIVER, it can be
+                    # overridden to PENDING if the job is not submitted, or
+                    # any other terminal status if the job driver process
+                    # finished correctly.
+                    failed_driver_transition_message = (
+                        f'Job {job_id} driver process is not running, but '
+                        'the job state is not in terminal states, setting '
+                        'it to FAILED_DRIVER')
+                    status = JobStatus.FAILED_DRIVER
+            elif job_pid < 0:
+                # TODO(zhwu): Backward compatibility, remove after 0.9.0.
+                # We set the job status to PENDING instead of actually
+                # checking ray job status and let the status in job table
+                # take effect in the later max.
+                status = JobStatus.PENDING
+
+            pending_job = _get_pending_job(job_id)
+            if pending_job is not None:
+                if pending_job['created_time'] < psutil.boot_time():
+                    failed_driver_transition_message = (
+                        f'Job {job_id} is stale, setting to FAILED_DRIVER: '
+                        f'created_time={pending_job["created_time"]}, '
+                        f'boot_time={psutil.boot_time()}')
+                    # The job is stale as it is created before the instance
+                    # is booted, e.g. the instance is rebooted.
+                    status = JobStatus.FAILED_DRIVER
+                elif pending_job['submit'] <= 0:
+                    # The job is not submitted (submit <= 0), we set it to
+                    # PENDING.
+                    # For submitted jobs, the driver should have been started,
+                    # because the job_lib.JobScheduler.schedule_step() have
+                    # the submit field and driver process pid set in the same
+                    # job lock.
+                    # The job process check in the above section should
+                    # correctly figured out the status and we don't overwrite
+                    # it here. (Note: the FAILED_DRIVER status will be
+                    # overridden by the actual job terminal status in the table
+                    # if the job driver process finished correctly.)
+                    status = JobStatus.PENDING
+
             assert original_status is not None, (job_id, status)
             if status is None:
+                # The job is submitted but the job driver process pid is not
+                # set in the database. This is guarding against the case where
+                # the schedule_step() function is interrupted (e.g., VM stop)
+                # at the middle of starting a new process and setting the pid.
                 status = original_status
                 if (original_status is not None and
                         not original_status.is_terminal()):
-                    logger.info(f'Ray job status for job {job_id} is None, '
-                                'setting it to FAILED.')
-                    # The job may be stale, when the instance is restarted
-                    # (the ray redis is volatile). We need to reset the
-                    # status of the task to FAILED if its original status
-                    # is RUNNING or PENDING.
-                    status = JobStatus.FAILED
+                    echo(f'Job {job_id} status is None, setting it to '
+                         'FAILED_DRIVER.')
+                    # The job may be stale, when the instance is restarted. We
+                    # need to reset the job status to FAILED_DRIVER if its
+                    # original status is in nonterminal_statuses.
+                    echo(f'Job {job_id} is in a unknown state, setting it to '
+                         'FAILED_DRIVER')
+                    status = JobStatus.FAILED_DRIVER
                     _set_status_no_lock(job_id, status)
-                    if not silent:
-                        logger.info(f'Updated job {job_id} status to {status}')
             else:
                 # Taking max of the status is necessary because:
-                # 1. It avoids race condition, where the original status has
-                # already been set to later state by the job. We skip the
-                # update.
-                # 2. _RAY_TO_JOB_STATUS_MAP would map `ray job status`'s
-                # `RUNNING` to our JobStatus.SETTING_UP; if a job has already
-                # been set to JobStatus.PENDING or JobStatus.RUNNING by the
-                # generated ray program, `original_status` (job status from our
-                # DB) would already have that value. So we take the max here to
-                # keep it at later status.
+                # 1. The original status has already been set to later
+                #    terminal state by a finished job driver.
+                # 2. Job driver process check would map any running job process
+                #    to `PENDING`, so we need to take the max to keep it at
+                #    later status for jobs actually started in SETTING_UP or
+                #    RUNNING.
                 status = max(status, original_status)
+                assert status is not None, (job_id, status, original_status)
                 if status != original_status:  # Prevents redundant update.
                     _set_status_no_lock(job_id, status)
-                    if not silent:
-                        logger.info(f'Updated job {job_id} status to {status}')
+                    echo(f'Updated job {job_id} status to {status}')
+                    if (status == JobStatus.FAILED_DRIVER and
+                            failed_driver_transition_message is not None):
+                        echo(failed_driver_transition_message)
         statuses.append(status)
     return statuses
 
@@ -620,18 +700,14 @@ def fail_all_jobs_in_progress() -> None:
         f"""\
         UPDATE jobs SET status=(?)
         WHERE status IN ({','.join(['?'] * len(in_progress_status))})
-        """, (JobStatus.FAILED.value, *in_progress_status))
+        """, (JobStatus.FAILED_DRIVER.value, *in_progress_status))
     _CONN.commit()
 
 
 def update_status() -> None:
     # This will be called periodically by the skylet to update the status
     # of the jobs in the database, to avoid stale job status.
-    # NOTE: there might be a INIT job in the database set to FAILED by this
-    # function, as the ray job status does not exist due to the app
-    # not submitted yet. It will be then reset to PENDING / RUNNING when the
-    # app starts.
-    nonterminal_jobs = _get_jobs(username=None,
+    nonterminal_jobs = _get_jobs(user_hash=None,
                                  status_list=JobStatus.nonterminal_statuses())
     nonterminal_job_ids = [job['job_id'] for job in nonterminal_jobs]
 
@@ -661,13 +737,14 @@ def format_job_queue(jobs: List[Dict[str, Any]]):
         print(format_job_queue(jobs))
     """
     job_table = log_utils.create_table([
-        'ID', 'NAME', 'SUBMITTED', 'STARTED', 'DURATION', 'RESOURCES', 'STATUS',
-        'LOG'
+        'ID', 'NAME', 'USER', 'SUBMITTED', 'STARTED', 'DURATION', 'RESOURCES',
+        'STATUS', 'LOG'
     ])
     for job in jobs:
         job_table.add_row([
             job['job_id'],
             job['job_name'],
+            job['username'],
             log_utils.readable_time_duration(job['submitted_at']),
             log_utils.readable_time_duration(job['start_at']),
             log_utils.readable_time_duration(job['start_at'],
@@ -680,11 +757,11 @@ def format_job_queue(jobs: List[Dict[str, Any]]):
     return job_table
 
 
-def dump_job_queue(username: Optional[str], all_jobs: bool) -> str:
+def dump_job_queue(user_hash: Optional[str], all_jobs: bool) -> str:
     """Get the job queue in encoded json format.
 
     Args:
-        username: The username to show jobs for. Show all the users if None.
+        user_hash: The user hash to show jobs for. Show all the users if None.
         all_jobs: Whether to show all jobs, not just the pending/running ones.
     """
     status_list: Optional[List[JobStatus]] = [
@@ -693,12 +770,12 @@ def dump_job_queue(username: Optional[str], all_jobs: bool) -> str:
     if all_jobs:
         status_list = None
 
-    jobs = _get_jobs(username, status_list=status_list)
+    jobs = _get_jobs(user_hash, status_list=status_list)
     for job in jobs:
         job['status'] = job['status'].value
         job['log_path'] = os.path.join(constants.SKY_LOGS_DIRECTORY,
                                        job.pop('run_timestamp'))
-    return common_utils.encode_payload(jobs)
+    return message_utils.encode_payload(jobs)
 
 
 def load_job_queue(payload: str) -> List[Dict[str, Any]]:
@@ -707,68 +784,118 @@ def load_job_queue(payload: str) -> List[Dict[str, Any]]:
     Args:
         payload: The encoded payload string to load.
     """
-    jobs = common_utils.decode_payload(payload)
+    jobs = message_utils.decode_payload(payload)
     for job in jobs:
         job['status'] = JobStatus(job['status'])
+        job['user_hash'] = job['username']
+        job['username'] = global_user_state.get_user(job['user_hash']).name
     return jobs
 
 
+# TODO(zhwu): Backward compatibility for jobs submitted before #4318, remove
+# after 0.10.0.
+def _create_ray_job_submission_client():
+    """Import the ray job submission client."""
+    try:
+        import ray  # pylint: disable=import-outside-toplevel
+    except ImportError:
+        logger.error('Failed to import ray')
+        raise
+    try:
+        # pylint: disable=import-outside-toplevel
+        from ray import job_submission
+    except ImportError:
+        logger.error(
+            f'Failed to import job_submission with ray=={ray.__version__}')
+        raise
+    port = get_job_submission_port()
+    return job_submission.JobSubmissionClient(
+        address=f'http://127.0.0.1:{port}')
+
+
+def _make_ray_job_id(sky_job_id: int) -> str:
+    return f'{sky_job_id}-{getpass.getuser()}'
+
+
 def cancel_jobs_encoded_results(jobs: Optional[List[int]],
-                                cancel_all: bool = False) -> str:
+                                cancel_all: bool = False,
+                                user_hash: Optional[str] = None) -> str:
     """Cancel jobs.
 
     Args:
-        jobs: Job IDs to cancel. (See `cancel_all` for special semantics.)
-        cancel_all: Whether to cancel all jobs. If True, asserts `jobs` is
-            set to None. If False and `jobs` is None, cancel the latest
-            running job.
+        jobs: Job IDs to cancel.
+        cancel_all: Whether to cancel all jobs.
+        user_hash: If specified, cancels the jobs for the specified user only.
+            Otherwise, applies to all users.
 
     Returns:
         Encoded job IDs that are actually cancelled. Caller should use
-        common_utils.decode_payload() to parse.
+        message_utils.decode_payload() to parse.
     """
-    if cancel_all:
-        # Cancel all in-progress jobs.
-        assert jobs is None, ('If cancel_all=True, usage is to set jobs=None')
-        job_records = _get_jobs(
-            None, [JobStatus.PENDING, JobStatus.SETTING_UP, JobStatus.RUNNING])
-    else:
-        if jobs is None:
-            # Cancel the latest (largest job ID) running job.
-            job_records = _get_jobs(None, [JobStatus.RUNNING])[:1]
-        else:
-            # Cancel jobs with specified IDs.
-            job_records = _get_jobs_by_ids(jobs)
-
-    # TODO(zhwu): `job_client.stop_job` will wait for the jobs to be killed, but
-    # when the memory is not enough, this will keep waiting.
-    job_client = _create_ray_job_submission_client()
-    cancelled_ids = []
+    job_records = []
+    all_status = [JobStatus.PENDING, JobStatus.SETTING_UP, JobStatus.RUNNING]
+    if jobs is None and not cancel_all:
+        # Cancel the latest (largest job ID) running job from current user.
+        job_records = _get_jobs(user_hash, [JobStatus.RUNNING])[:1]
+    elif cancel_all:
+        job_records = _get_jobs(user_hash, all_status)
+    if jobs is not None:
+        job_records.extend(_get_jobs_by_ids(jobs))
 
+    cancelled_ids = []
     # Sequentially cancel the jobs to avoid the resource number bug caused by
     # ray cluster (tracked in #1262).
-    for job in job_records:
-        job_id = make_ray_job_id(job['job_id'])
+    for job_record in job_records:
+        job_id = job_record['job_id']
         # Job is locked to ensure that pending queue does not start it while
         # it is being cancelled
-        with filelock.FileLock(_get_lock_path(job['job_id'])):
-            try:
-                job_client.stop_job(job_id)
-            except RuntimeError as e:
-                # If the request to the job server fails, we should not
-                # set the job to CANCELLED.
-                if 'does not exist' not in str(e):
-                    logger.warning(str(e))
-                    continue
-
-            if job['status'] in [
+        with filelock.FileLock(_get_lock_path(job_id)):
+            job = _get_jobs_by_ids([job_id])[0]
+            if _is_job_driver_process_running(job['pid'], job_id):
+                # Not use process.terminate() as that will only terminate the
+                # process shell process, not the ray driver process
+                # under the shell.
+                #
+                # We don't kill all the children of the process, like
+                # subprocess_utils.kill_process_daemon() does, but just the
+                # process group here, because the underlying job driver can
+                # start other jobs with `schedule_step`, causing the other job
+                # driver processes to be children of the current job driver
+                # process.
+                #
+                # Killing the process group is enough as the underlying job
+                # should be able to clean itself up correctly by ray driver.
+                #
+                # The process group pid should be the same as the job pid as we
+                # use start_new_session=True, but we use os.getpgid() to be
+                # extra cautious.
+                job_pgid = os.getpgid(job['pid'])
+                os.killpg(job_pgid, signal.SIGTERM)
+                # We don't have to start a daemon to forcefully kill the process
+                # as our job driver process will clean up the underlying
+                # child processes.
+            elif job['pid'] < 0:
+                try:
+                    # TODO(zhwu): Backward compatibility, remove after 0.9.0.
+                    # The job was submitted with ray job submit before #4318.
+                    job_client = _create_ray_job_submission_client()
+                    job_client.stop_job(_make_ray_job_id(job['job_id']))
+                except RuntimeError as e:
+                    # If the request to the job server fails, we should not
+                    # set the job to CANCELLED.
+                    if 'does not exist' not in str(e):
+                        logger.warning(str(e))
+                        continue
+            # Get the job status again to avoid race condition.
+            job_status = get_status_no_lock(job['job_id'])
+            if job_status in [
                     JobStatus.PENDING, JobStatus.SETTING_UP, JobStatus.RUNNING
             ]:
                 _set_status_no_lock(job['job_id'], JobStatus.CANCELLED)
                 cancelled_ids.append(job['job_id'])
 
         scheduler.schedule_step()
-    return common_utils.encode_payload(cancelled_ids)
+    return message_utils.encode_payload(cancelled_ids)
 
 
 def get_run_timestamp(job_id: Optional[int]) -> Optional[str]:
@@ -797,7 +924,7 @@ def run_timestamp_with_globbing_payload(job_ids: List[Optional[str]]) -> str:
         job_id = row[JobInfoLoc.JOB_ID.value]
         run_timestamp = row[JobInfoLoc.RUN_TIMESTAMP.value]
         run_timestamps[str(job_id)] = run_timestamp
-    return common_utils.encode_payload(run_timestamps)
+    return message_utils.encode_payload(run_timestamps)
 
 
 class JobLibCodeGen:
@@ -812,14 +939,6 @@ class JobLibCodeGen:
         'import os',
         'import getpass',
         'from sky.skylet import job_lib, log_lib, constants',
-        # Backward compatibility for old skylet lib version on the remote
-        # machine. The `job_owner` argument was removed in #3037, and if the
-        # remote machine has an old SkyPilot version before that, we need to
-        # pass the `job_owner` argument to the job_lib functions.
-        # TODO(zhwu): Remove this in 0.7.0 release.
-        'job_owner_kwargs = {} '
-        'if getattr(constants, "SKYLET_LIB_VERSION", 0) >= 1 '
-        'else {"job_owner": getpass.getuser()}',
     ]
 
     @classmethod
@@ -828,10 +947,17 @@ class JobLibCodeGen:
         if job_name is None:
             job_name = '-'
         code = [
-            'job_id = job_lib.add_job('
-            f'{job_name!r}, '
-            f'{username!r}, '
-            f'{run_timestamp!r}, '
+            # We disallow job submission when SKYLET_VERSION is older than 9, as
+            # it was using ray job submit before #4318, and switched to raw
+            # process. Using the old skylet version will cause the job status
+            # to be stuck in PENDING state or transition to FAILED_DRIVER state.
+            '\nif int(constants.SKYLET_VERSION) < 9: '
+            'raise RuntimeError("SkyPilot runtime is too old, which does not '
+            'support submitting jobs.")',
+            '\njob_id = job_lib.add_job('
+            f'{job_name!r},'
+            f'{username!r},'
+            f'{run_timestamp!r},'
             f'{resources_str!r})',
             'print("Job ID: " + str(job_id), flush=True)',
         ]
@@ -839,35 +965,49 @@ class JobLibCodeGen:
 
     @classmethod
     def queue_job(cls, job_id: int, cmd: str) -> str:
-        code = ['job_lib.scheduler.queue('
-                f'{job_id!r},'
-                f'{cmd!r})']
+        code = [
+            'job_lib.scheduler.queue('
+            f'{job_id!r},'
+            f'{cmd!r})',
+        ]
         return cls._build(code)
 
     @classmethod
     def update_status(cls) -> str:
-        code = ['job_lib.update_status(**job_owner_kwargs)']
+        code = ['job_lib.update_status()']
         return cls._build(code)
 
     @classmethod
-    def get_job_queue(cls, username: Optional[str], all_jobs: bool) -> str:
+    def get_job_queue(cls, user_hash: Optional[str], all_jobs: bool) -> str:
+        # TODO(SKY-1214): combine get_job_queue with get_job_statuses.
         code = [
             'job_queue = job_lib.dump_job_queue('
-            f'{username!r}, {all_jobs})', 'print(job_queue, flush=True)'
+            f'{user_hash!r}, {all_jobs})',
+            'print(job_queue, flush=True)',
         ]
         return cls._build(code)
 
     @classmethod
     def cancel_jobs(cls,
                     job_ids: Optional[List[int]],
-                    cancel_all: bool = False) -> str:
+                    cancel_all: bool = False,
+                    user_hash: Optional[str] = None) -> str:
         """See job_lib.cancel_jobs()."""
         code = [
             (f'cancelled = job_lib.cancel_jobs_encoded_results('
-             f' {job_ids!r}, {cancel_all}, **job_owner_kwargs)'),
+             f'jobs={job_ids!r}, cancel_all={cancel_all}, '
+             f'user_hash={user_hash!r})'),
             # Print cancelled IDs. Caller should parse by decoding.
             'print(cancelled, flush=True)',
         ]
+        # TODO(zhwu): Backward compatibility, remove after 0.9.0.
+        if user_hash is None:
+            code = [
+                (f'cancelled = job_lib.cancel_jobs_encoded_results('
+                 f' {job_ids!r}, {cancel_all})'),
+                # Print cancelled IDs. Caller should parse by decoding.
+                'print(cancelled, flush=True)',
+            ]
         return cls._build(code)
 
     @classmethod
@@ -880,14 +1020,19 @@ class JobLibCodeGen:
     def tail_logs(cls,
                   job_id: Optional[int],
                   managed_job_id: Optional[int],
-                  follow: bool = True) -> str:
+                  follow: bool = True,
+                  tail: int = 0) -> str:
         # pylint: disable=line-too-long
+
         code = [
+            # We use != instead of is not because 1 is not None will print a warning:
+            # <stdin>:1: SyntaxWarning: "is not" with a literal. Did you mean "!="?
             f'job_id = {job_id} if {job_id} != None else job_lib.get_latest_job_id()',
             'run_timestamp = job_lib.get_run_timestamp(job_id)',
             f'log_dir = None if run_timestamp is None else os.path.join({constants.SKY_LOGS_DIRECTORY!r}, run_timestamp)',
-            f'log_lib.tail_logs(job_id=job_id, log_dir=log_dir, '
-            f'managed_job_id={managed_job_id!r}, follow={follow}, **job_owner_kwargs)',
+            f'tail_log_kwargs = {{"job_id": job_id, "log_dir": log_dir, "managed_job_id": {managed_job_id!r}, "follow": {follow}}}',
+            f'{_LINUX_NEW_LINE}if getattr(constants, "SKYLET_LIB_VERSION", 1) > 1: tail_log_kwargs["tail"] = {tail}',
+            f'{_LINUX_NEW_LINE}log_lib.tail_logs(**tail_log_kwargs)',
         ]
         return cls._build(code)