skypilot-nightly 1.0.0.dev20241114__py3-none-any.whl → 1.0.0.dev20241116__py3-none-any.whl

sky/skylet/job_lib.py

@@ -8,6 +8,7 @@ import json
 import os
 import pathlib
 import shlex
+import signal
 import sqlite3
 import subprocess
 import time
@@ -27,6 +28,10 @@ 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:
@@ -46,6 +51,7 @@ class JobInfoLoc(enum.IntEnum):
     START_AT = 6
     END_AT = 7
     RESOURCES = 8
+    PID = 9
 
 
 _DB_PATH = os.path.expanduser('~/.sky/jobs.db')
@@ -67,6 +73,16 @@ def create_table(cursor, conn):
             # 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.
     cursor.execute("""\
         CREATE TABLE IF NOT EXISTS jobs (
         job_id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -75,7 +91,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,
@@ -86,7 +105,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()
 
 
@@ -118,6 +138,11 @@ class JobStatus(enum.Enum):
     # 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 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'
     # 3 terminal states below: once reached, they do not transition.
     # The job finished successfully.
     SUCCEEDED = 'SUCCEEDED'
@@ -148,11 +173,16 @@ class JobStatus(enum.Enum):
         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]
 
@@ -175,7 +205,39 @@ 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)
+        # Use nohup to ensure the job driver process is a separate process tree,
+        # instead of being a child of the current process. This is important to
+        # avoid a chain of driver processes (job driver can call schedule_step()
+        # to submit new jobs, and the new job can also call schedule_step()
+        # recursively).
+        #
+        # echo $! will output the PID of the last background process started
+        # in the current shell, so we can retrieve it and record in the DB.
+        #
+        # TODO(zhwu): A more elegant solution is to use another daemon process
+        # to be in charge of starting these driver processes, instead of
+        # starting them in the current process.
+        wrapped_cmd = (f'nohup bash -c {shlex.quote(run_cmd)} '
+                       '</dev/null >/dev/null 2>&1 & echo $!')
+        proc = subprocess.run(wrapped_cmd,
+                              stdout=subprocess.PIPE,
+                              stderr=subprocess.PIPE,
+                              stdin=subprocess.DEVNULL,
+                              start_new_session=True,
+                              check=True,
+                              shell=True,
+                              text=True)
+        # Get the PID of the detached process
+        pid = int(proc.stdout.strip())
+
+        # 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:
         if force_update_jobs:
@@ -237,59 +299,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]:
@@ -301,9 +317,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,))
@@ -478,6 +495,7 @@ 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
 
@@ -537,6 +555,23 @@ def _get_pending_job(job_id: int) -> Optional[Dict[str, Any]]:
     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],
                       silent: bool = False) -> List[JobStatus]:
     """Updates and returns the job statuses matching our `JobStatus` semantics.
@@ -554,11 +589,8 @@ def update_job_status(job_ids: List[int],
     if len(job_ids) == 0:
         return []
 
-    ray_job_ids = [make_ray_job_id(job_id) for job_id in job_ids]
-    job_client = _create_ray_job_submission_client()
-
     statuses = []
-    for job_id, ray_job_id in zip(job_ids, ray_job_ids):
+    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.
@@ -567,11 +599,13 @@ def update_job_status(job_ids: List[int],
             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']
 
-            ray_job_query_time = time.time()
+            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
-                        >= ray_job_query_time - _PENDING_SUBMIT_GRACE_PERIOD):
+                        >= 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
@@ -582,75 +616,98 @@ def update_job_status(job_ids: List[int],
                     # 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.
-                    echo(f'INIT job {job_id} is stale, setting to FAILED')
-                    status = JobStatus.FAILED
-
-            try:
-                # Querying status within the lock is safer than querying
-                # outside, as it avoids the race condition when job table is
-                # updated after the ray job status query.
-                # Also, getting per-job status is faster than querying all jobs,
-                # when there are significant number of finished jobs.
-                # Reference: getting 124 finished jobs takes 0.038s, while
-                # querying a single job takes 0.006s, 10 jobs takes 0.066s.
-                # TODO: if too slow, directly query against redis.
-                ray_job_status = job_client.get_job_status(ray_job_id)
-                status = _RAY_TO_JOB_STATUS_MAP[ray_job_status.value]
-            except RuntimeError:
-                # Job not found.
-                pass
+                    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():
-                    echo(f'Job {job_id} is stale, setting to FAILED: '
-                         f'created_time={pending_job["created_time"]}, '
-                         f'boot_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
-                # Gives a 60 second grace period between job being submit from
-                # the pending table until appearing in ray jobs. For jobs
-                # submitted outside of the grace period, we will consider the
-                # ray job status.
-
-                if not (pending_job['submit'] > 0 and pending_job['submit'] <
-                        ray_job_query_time - _PENDING_SUBMIT_GRACE_PERIOD):
-                    # 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.
+                    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()):
-                    echo(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)
-                    echo(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)
                     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
 
@@ -663,17 +720,13 @@ 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,
                                  status_list=JobStatus.nonterminal_statuses())
     nonterminal_job_ids = [job['job_id'] for job in nonterminal_jobs]
@@ -756,6 +809,31 @@ def load_job_queue(payload: str) -> List[Dict[str, Any]]:
     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 jobs.
@@ -783,27 +861,51 @@ def cancel_jobs_encoded_results(jobs: Optional[List[int]],
             # 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 = []
 
     # 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
-
+        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 [
@@ -865,10 +967,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)',
         ]
@@ -876,9 +985,11 @@ 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

sky/skylet/log_lib.py

@@ -183,40 +183,7 @@ def run_with_log(
                           shell=shell,
                           **kwargs) as proc:
         try:
-            # The proc can be defunct if the python program is killed. Here we
-            # open a new subprocess to gracefully kill the proc, SIGTERM
-            # and then SIGKILL the process group.
-            # Adapted from ray/dashboard/modules/job/job_manager.py#L154
-            parent_pid = os.getpid()
-            daemon_script = os.path.join(
-                os.path.dirname(os.path.abspath(job_lib.__file__)),
-                'subprocess_daemon.py')
-            python_path = subprocess.check_output(
-                constants.SKY_GET_PYTHON_PATH_CMD,
-                shell=True,
-                stderr=subprocess.DEVNULL,
-                encoding='utf-8').strip()
-            daemon_cmd = [
-                python_path,
-                daemon_script,
-                '--parent-pid',
-                str(parent_pid),
-                '--proc-pid',
-                str(proc.pid),
-            ]
-
-            # We do not need to set `start_new_session=True` here, as the
-            # daemon script will detach itself from the parent process with
-            # fork to avoid being killed by ray job. See the reason we
-            # daemonize the process in `sky/skylet/subprocess_daemon.py`.
-            subprocess.Popen(
-                daemon_cmd,
-                # Suppress output
-                stdout=subprocess.DEVNULL,
-                stderr=subprocess.DEVNULL,
-                # Disable input
-                stdin=subprocess.DEVNULL,
-            )
+            subprocess_utils.kill_process_daemon(proc.pid)
             stdout = ''
             stderr = ''