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

sky/skylet/job_lib.py

@@ -23,20 +23,22 @@ from sky import global_user_state
 from sky import sky_logging
 from sky.adaptors import common as adaptors_common
 from sky.skylet import constants
+from sky.skylet import runtime_utils
 from sky.utils import common_utils
-from sky.utils import log_utils
 from sky.utils import message_utils
 from sky.utils import subprocess_utils
 from sky.utils.db import db_utils
 
 if typing.TYPE_CHECKING:
     import psutil
+
+    from sky.schemas.generated import jobsv1_pb2
 else:
     psutil = adaptors_common.LazyImport('psutil')
+    jobsv1_pb2 = adaptors_common.LazyImport('sky.schemas.generated.jobsv1_pb2')
 
 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
@@ -82,13 +84,9 @@ def create_table(cursor, conn):
             # 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.
+    # can be in two states:
     # 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).
+    #    INIT state, the pid will be set to 0.
     # >=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.
@@ -144,7 +142,7 @@ def init_db(func):
 
         with _db_init_lock:
             if _DB is None:
-                db_path = os.path.expanduser('~/.sky/jobs.db')
+                db_path = runtime_utils.get_runtime_dir_path('.sky/jobs.db')
                 os.makedirs(pathlib.Path(db_path).parents[0], exist_ok=True)
                 _DB = db_utils.SQLiteConn(db_path, create_table)
         return func(*args, **kwargs)
@@ -220,6 +218,45 @@ class JobStatus(enum.Enum):
         color = _JOB_STATUS_TO_COLOR[self]
         return f'{color}{self.value}{colorama.Style.RESET_ALL}'
 
+    @classmethod
+    def from_protobuf(
+            cls,
+            protobuf_value: 'jobsv1_pb2.JobStatus') -> Optional['JobStatus']:
+        """Convert protobuf JobStatus enum to Python enum value."""
+        protobuf_to_enum = {
+            jobsv1_pb2.JOB_STATUS_INIT: cls.INIT,
+            jobsv1_pb2.JOB_STATUS_PENDING: cls.PENDING,
+            jobsv1_pb2.JOB_STATUS_SETTING_UP: cls.SETTING_UP,
+            jobsv1_pb2.JOB_STATUS_RUNNING: cls.RUNNING,
+            jobsv1_pb2.JOB_STATUS_FAILED_DRIVER: cls.FAILED_DRIVER,
+            jobsv1_pb2.JOB_STATUS_SUCCEEDED: cls.SUCCEEDED,
+            jobsv1_pb2.JOB_STATUS_FAILED: cls.FAILED,
+            jobsv1_pb2.JOB_STATUS_FAILED_SETUP: cls.FAILED_SETUP,
+            jobsv1_pb2.JOB_STATUS_CANCELLED: cls.CANCELLED,
+            jobsv1_pb2.JOB_STATUS_UNSPECIFIED: None,
+        }
+        if protobuf_value not in protobuf_to_enum:
+            raise ValueError(
+                f'Unknown protobuf JobStatus value: {protobuf_value}')
+        return protobuf_to_enum[protobuf_value]
+
+    def to_protobuf(self) -> 'jobsv1_pb2.JobStatus':
+        """Convert this Python enum value to protobuf enum value."""
+        enum_to_protobuf = {
+            JobStatus.INIT: jobsv1_pb2.JOB_STATUS_INIT,
+            JobStatus.PENDING: jobsv1_pb2.JOB_STATUS_PENDING,
+            JobStatus.SETTING_UP: jobsv1_pb2.JOB_STATUS_SETTING_UP,
+            JobStatus.RUNNING: jobsv1_pb2.JOB_STATUS_RUNNING,
+            JobStatus.FAILED_DRIVER: jobsv1_pb2.JOB_STATUS_FAILED_DRIVER,
+            JobStatus.SUCCEEDED: jobsv1_pb2.JOB_STATUS_SUCCEEDED,
+            JobStatus.FAILED: jobsv1_pb2.JOB_STATUS_FAILED,
+            JobStatus.FAILED_SETUP: jobsv1_pb2.JOB_STATUS_FAILED_SETUP,
+            JobStatus.CANCELLED: jobsv1_pb2.JOB_STATUS_CANCELLED,
+        }
+        if self not in enum_to_protobuf:
+            raise ValueError(f'Unknown JobStatus value: {self}')
+        return enum_to_protobuf[self]
+
 
 # We have two steps for job submissions:
 # 1. Client reserve a job id from the job table by adding a INIT state job.
@@ -261,11 +298,7 @@ class JobScheduler:
              f'WHERE job_id={job_id!r}'))
         _DB.conn.commit()
         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
+
         _DB.cursor.execute((f'UPDATE jobs SET pid={pid} '
                             f'WHERE job_id={job_id!r}'))
         _DB.conn.commit()
@@ -475,6 +508,11 @@ def get_status(job_id: int) -> Optional[JobStatus]:
 
 @init_db
 def get_statuses_payload(job_ids: List[Optional[int]]) -> str:
+    return message_utils.encode_payload(get_statuses(job_ids))
+
+
+@init_db
+def get_statuses(job_ids: List[int]) -> Dict[int, Optional[str]]:
     assert _DB is not None
     # Per-job lock is not required here, since the staled job status will not
     # affect the caller.
@@ -482,10 +520,51 @@ def get_statuses_payload(job_ids: List[Optional[int]]) -> str:
     rows = _DB.cursor.execute(
         f'SELECT job_id, status FROM jobs WHERE job_id IN ({query_str})',
         job_ids)
-    statuses = {job_id: None for job_id in job_ids}
+    statuses: Dict[int, Optional[str]] = {job_id: None for job_id in job_ids}
     for (job_id, status) in rows:
         statuses[job_id] = status
-    return message_utils.encode_payload(statuses)
+    return statuses
+
+
+@init_db
+def get_jobs_info(user_hash: Optional[str] = None,
+                  all_jobs: bool = False) -> List['jobsv1_pb2.JobInfo']:
+    """Get detailed job information.
+
+    Similar to dump_job_queue but returns structured protobuf objects instead
+    of encoded strings.
+
+    Args:
+        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.
+    """
+    assert _DB is not None
+
+    status_list: Optional[List[JobStatus]] = [
+        JobStatus.SETTING_UP, JobStatus.PENDING, JobStatus.RUNNING
+    ]
+    if all_jobs:
+        status_list = None
+
+    jobs = _get_jobs(user_hash, status_list=status_list)
+    jobs_info = []
+    for job in jobs:
+        jobs_info.append(
+            jobsv1_pb2.JobInfo(job_id=job['job_id'],
+                               job_name=job['job_name'],
+                               username=job['username'],
+                               submitted_at=job['submitted_at'],
+                               status=job['status'].to_protobuf(),
+                               run_timestamp=job['run_timestamp'],
+                               start_at=job['start_at'],
+                               end_at=job['end_at'],
+                               resources=job['resources'],
+                               pid=job['pid'],
+                               log_path=os.path.join(
+                                   constants.SKY_LOGS_DIRECTORY,
+                                   job['run_timestamp']),
+                               metadata=json.dumps(job['metadata'])))
+    return jobs_info
 
 
 def load_statuses_payload(
@@ -524,16 +603,27 @@ def get_job_submitted_or_ended_timestamp_payload(job_id: int,
     PENDING state.
 
     The normal job duration will use `start_at` instead of `submitted_at` (in
-    `format_job_queue()`), because the job may stay in PENDING if the cluster is
-    busy.
+    `table_utils.format_job_queue()`), because the job may stay in PENDING if
+    the cluster is busy.
+    """
+    return message_utils.encode_payload(
+        get_job_submitted_or_ended_timestamp(job_id, get_ended_time))
+
+
+@init_db
+def get_job_submitted_or_ended_timestamp(
+        job_id: int, get_ended_time: bool) -> Optional[float]:
+    """Get the job submitted timestamp.
+
+    Returns the raw timestamp or None if job doesn't exist.
     """
     assert _DB is not None
     field = 'end_at' if get_ended_time else 'submitted_at'
     rows = _DB.cursor.execute(f'SELECT {field} FROM jobs WHERE job_id=(?)',
                               (job_id,))
     for (timestamp,) in rows:
-        return message_utils.encode_payload(timestamp)
-    return message_utils.encode_payload(None)
+        return timestamp
+    return None
 
 
 def get_ray_port():
@@ -542,7 +632,8 @@ def get_ray_port():
     If the port file does not exist, the cluster was launched before #1790,
     return the default port.
     """
-    port_path = os.path.expanduser(constants.SKY_REMOTE_RAY_PORT_FILE)
+    port_path = runtime_utils.get_runtime_dir_path(
+        constants.SKY_REMOTE_RAY_PORT_FILE)
     if not os.path.exists(port_path):
         return 6379
     port = json.load(open(port_path, 'r', encoding='utf-8'))['ray_port']
@@ -555,7 +646,8 @@ def get_job_submission_port():
     If the port file does not exist, the cluster was launched before #1790,
     return the default port.
     """
-    port_path = os.path.expanduser(constants.SKY_REMOTE_RAY_PORT_FILE)
+    port_path = runtime_utils.get_runtime_dir_path(
+        constants.SKY_REMOTE_RAY_PORT_FILE)
     if not os.path.exists(port_path):
         return 8265
     port = json.load(open(port_path, 'r',
@@ -673,7 +765,7 @@ def update_job_status(job_ids: List[int],
     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
+        # query and the job status update, the job status in the database
         # can be modified by the generated ray program.
         with filelock.FileLock(_get_lock_path(job_id)):
             status = None
@@ -724,12 +816,6 @@ def update_job_status(job_ids: List[int],
                         '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.10.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:
@@ -842,35 +928,6 @@ def is_cluster_idle() -> bool:
     assert False, 'Should not reach here'
 
 
-def format_job_queue(jobs: List[Dict[str, Any]]):
-    """Format the job queue for display.
-
-    Usage:
-        jobs = get_job_queue()
-        print(format_job_queue(jobs))
-    """
-    job_table = log_utils.create_table([
-        'ID', 'NAME', 'USER', 'SUBMITTED', 'STARTED', 'DURATION', 'RESOURCES',
-        'STATUS', 'LOG', 'GIT COMMIT'
-    ])
-    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'],
-                                             job['end_at'],
-                                             absolute=True),
-            job['resources'],
-            job['status'].colored_str(),
-            job['log_path'],
-            job.get('metadata', {}).get('git_commit', '-'),
-        ])
-    return job_table
-
-
 def dump_job_queue(user_hash: Optional[str], all_jobs: bool) -> str:
     """Get the job queue in encoded json format.
 
@@ -907,27 +964,6 @@ 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()}'
 
@@ -947,6 +983,13 @@ def cancel_jobs_encoded_results(jobs: Optional[List[int]],
         Encoded job IDs that are actually cancelled. Caller should use
         message_utils.decode_payload() to parse.
     """
+    return message_utils.encode_payload(cancel_jobs(jobs, cancel_all,
+                                                    user_hash))
+
+
+def cancel_jobs(jobs: Optional[List[int]],
+                cancel_all: bool = False,
+                user_hash: Optional[str] = None) -> List[int]:
     job_records = []
     all_status = [JobStatus.PENDING, JobStatus.SETTING_UP, JobStatus.RUNNING]
     if jobs is None and not cancel_all:
@@ -989,18 +1032,6 @@ def cancel_jobs_encoded_results(jobs: Optional[List[int]],
                 # 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.10.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 [
@@ -1010,7 +1041,7 @@ def cancel_jobs_encoded_results(jobs: Optional[List[int]],
                 cancelled_ids.append(job['job_id'])
 
         scheduler.schedule_step()
-    return message_utils.encode_payload(cancelled_ids)
+    return cancelled_ids
 
 
 @init_db
@@ -1030,6 +1061,17 @@ def get_run_timestamp(job_id: Optional[int]) -> Optional[str]:
 
 @init_db
 def get_log_dir_for_jobs(job_ids: List[Optional[str]]) -> str:
+    """Returns the relative paths to the log files for jobs with globbing,
+    encoded."""
+    job_to_dir = get_job_log_dirs(job_ids)
+    job_to_dir_str: Dict[str, str] = {}
+    for job_id, log_dir in job_to_dir.items():
+        job_to_dir_str[str(job_id)] = log_dir
+    return message_utils.encode_payload(job_to_dir_str)
+
+
+@init_db
+def get_job_log_dirs(job_ids: List[int]) -> Dict[int, str]:
     """Returns the relative paths to the log files for jobs with globbing."""
     assert _DB is not None
     query_str = ' OR '.join(['job_id GLOB (?)'] * len(job_ids))
@@ -1038,16 +1080,16 @@ def get_log_dir_for_jobs(job_ids: List[Optional[str]]) -> str:
             SELECT * FROM jobs
             WHERE {query_str}""", job_ids)
     rows = _DB.cursor.fetchall()
-    job_to_dir = {}
+    job_to_dir: Dict[int, str] = {}
     for row in rows:
         job_id = row[JobInfoLoc.JOB_ID.value]
         if row[JobInfoLoc.LOG_PATH.value]:
-            job_to_dir[str(job_id)] = row[JobInfoLoc.LOG_PATH.value]
+            job_to_dir[job_id] = row[JobInfoLoc.LOG_PATH.value]
         else:
             run_timestamp = row[JobInfoLoc.RUN_TIMESTAMP.value]
-            job_to_dir[str(job_id)] = os.path.join(constants.SKY_LOGS_DIRECTORY,
-                                                   run_timestamp)
-    return message_utils.encode_payload(job_to_dir)
+            job_to_dir[job_id] = os.path.join(constants.SKY_LOGS_DIRECTORY,
+                                              run_timestamp)
+    return job_to_dir
 
 
 class JobLibCodeGen:
@@ -1176,15 +1218,10 @@ class JobLibCodeGen:
              f'  log_dir = None if run_timestamp is None else os.path.join({constants.SKY_LOGS_DIRECTORY!r}, run_timestamp)'
             ),
             # Add a newline to leave the if indent block above.
-            f'\ntail_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)',
+            f'\nlog_lib.tail_logs(job_id=job_id, log_dir=log_dir, managed_job_id={managed_job_id!r}, follow={follow}, tail={tail})',
             # After tailing, check the job status and exit with appropriate code
             'job_status = job_lib.get_status(job_id)',
-            # Backward compatibility for returning exit code: Skylet versions 2
-            # and older did not have JobExitCode, so we use 0 for those versions
-            # TODO: Remove this special handling after 0.10.0.
-            'exit_code = exceptions.JobExitCode.from_job_status(job_status) if getattr(constants, "SKYLET_LIB_VERSION", 1) > 2 else 0',
+            'exit_code = exceptions.JobExitCode.from_job_status(job_status)',
             # Fix for dashboard: When follow=False and job is still running (NOT_FINISHED=101),
             # exit with success (0) since fetching current logs is a successful operation.
             # This prevents shell wrappers from printing "command terminated with exit code 101".

sky/skylet/log_lib.py

@@ -8,11 +8,13 @@ import functools
 import io
 import multiprocessing.pool
 import os
+import queue as queue_lib
 import shlex
 import subprocess
 import sys
 import tempfile
 import textwrap
+import threading
 import time
 from typing import (Deque, Dict, Iterable, Iterator, List, Optional, TextIO,
                     Tuple, Union)
@@ -39,6 +41,11 @@ logger = sky_logging.init_logger(__name__)
 
 LOG_FILE_START_STREAMING_AT = 'Waiting for task resources on '
 
+# 16-64KiB seems to be the sweet spot:
+# https://github.com/grpc/grpc.github.io/issues/371
+# TODO(kevin): Benchmark this ourselves and verify.
+DEFAULT_LOG_CHUNK_SIZE = 16 * 1024  # 16KiB
+
 
 class _ProcessingArgs:
     """Arguments for processing logs."""
@@ -213,7 +220,14 @@ def run_with_log(
                           stdin=stdin,
                           **kwargs) as proc:
         try:
-            subprocess_utils.kill_process_daemon(proc.pid)
+            if ctx is not None:
+                # When runs in coroutine, use kill_pg if available to avoid
+                # the overhead of refreshing the process tree in the daemon.
+                subprocess_utils.kill_process_daemon(proc.pid, use_kill_pg=True)
+            else:
+                # For backward compatibility, do not specify use_kill_pg by
+                # default.
+                subprocess_utils.kill_process_daemon(proc.pid)
             stdout = ''
             stderr = ''
             stdout_stream_handler = None
@@ -264,7 +278,6 @@ def run_with_log(
                 stdout, stderr = context_utils.pipe_and_wait_process(
                     ctx,
                     proc,
-                    cancel_callback=subprocess_utils.kill_children_processes,
                     stdout_stream_handler=stdout_stream_handler,
                     stderr_stream_handler=stderr_stream_handler)
             elif process_stream:
@@ -354,6 +367,17 @@ def run_bash_command_with_log(bash_command: str,
                             shell=True)
 
 
+def run_bash_command_with_log_and_return_pid(
+        bash_command: str,
+        log_path: str,
+        env_vars: Optional[Dict[str, str]] = None,
+        stream_logs: bool = False,
+        with_ray: bool = False):
+    return_code = run_bash_command_with_log(bash_command, log_path, env_vars,
+                                            stream_logs, with_ray)
+    return {'return_code': return_code, 'pid': os.getpid()}
+
+
 def _follow_job_logs(file,
                      job_id: int,
                      start_streaming: bool,
@@ -395,9 +419,9 @@ def _follow_job_logs(file,
                     wait_last_logs = False
                     continue
                 status_str = status.value if status is not None else 'None'
-                print(ux_utils.finishing_message(
-                    f'Job finished (status: {status_str}).'),
-                      flush=True)
+                finish = ux_utils.finishing_message(
+                    f'Job finished (status: {status_str}).')
+                yield finish + '\n'
                 return
 
             time.sleep(SKY_LOG_TAILING_GAP_SECONDS)
@@ -552,3 +576,207 @@ def tail_logs(job_id: Optional[int],
         except FileNotFoundError:
             print(f'{colorama.Fore.RED}ERROR: Logs for job {job_id} (status:'
                   f' {status.value}) does not exist.{colorama.Style.RESET_ALL}')
+
+
+def tail_logs_iter(job_id: Optional[int],
+                   log_dir: Optional[str],
+                   managed_job_id: Optional[int] = None,
+                   follow: bool = True,
+                   tail: int = 0) -> Iterator[str]:
+    """Tail the logs of a job. This is mostly the same as tail_logs, but
+    returns an iterator instead of printing to stdout/stderr."""
+    if job_id is None:
+        # This only happens when job_lib.get_latest_job_id() returns None,
+        # which means no job has been submitted to this cluster. See
+        # sky.skylet.job_lib.JobLibCodeGen.tail_logs for more details.
+        logger.info('Skip streaming logs as no job has been submitted.')
+        return
+    job_str = f'job {job_id}'
+    if managed_job_id is not None:
+        job_str = f'managed job {managed_job_id}'
+    if log_dir is None:
+        msg = f'{job_str.capitalize()} not found (see `sky queue`).'
+        yield msg + '\n'
+        return
+    logger.debug(f'Tailing logs for job, real job_id {job_id}, managed_job_id '
+                 f'{managed_job_id}.')
+    log_path = os.path.join(log_dir, 'run.log')
+    log_path = os.path.expanduser(log_path)
+
+    status = job_lib.update_job_status([job_id], silent=True)[0]
+
+    # Wait for the log to be written. This is needed due to the `ray submit`
+    # will take some time to start the job and write the log.
+    retry_cnt = 0
+    while status is not None and not status.is_terminal():
+        retry_cnt += 1
+        if os.path.exists(log_path) and status != job_lib.JobStatus.INIT:
+            break
+        if retry_cnt >= SKY_LOG_WAITING_MAX_RETRY:
+            err = (f'{colorama.Fore.RED}ERROR: Logs for '
+                   f'{job_str} (status: {status.value}) does not exist '
+                   f'after retrying {retry_cnt} times.'
+                   f'{colorama.Style.RESET_ALL}')
+            yield err + '\n'
+            return
+        waiting = (f'INFO: Waiting {SKY_LOG_WAITING_GAP_SECONDS}s for the logs '
+                   'to be written...')
+        yield waiting + '\n'
+        time.sleep(SKY_LOG_WAITING_GAP_SECONDS)
+        status = job_lib.update_job_status([job_id], silent=True)[0]
+
+    start_stream_at = LOG_FILE_START_STREAMING_AT
+    # Explicitly declare the type to avoid mypy warning.
+    lines: Iterable[str] = []
+    if follow and status in [
+            job_lib.JobStatus.SETTING_UP,
+            job_lib.JobStatus.PENDING,
+            job_lib.JobStatus.RUNNING,
+    ]:
+        # Not using `ray job logs` because it will put progress bar in
+        # multiple lines.
+        with open(log_path, 'r', newline='', encoding='utf-8') as log_file:
+            # Using `_follow` instead of `tail -f` to streaming the whole
+            # log and creating a new process for tail.
+            start_streaming = False
+            if tail > 0:
+                head_lines_of_log_file = _peek_head_lines(log_file)
+                lines = collections.deque(log_file, maxlen=tail)
+                start_streaming = _should_stream_the_whole_tail_lines(
+                    head_lines_of_log_file, lines, start_stream_at)
+                for line in lines:
+                    if start_stream_at in line:
+                        start_streaming = True
+                    if start_streaming:
+                        yield line
+            # Now, the cursor is at the end of the last lines
+            # if tail > 0
+            for line in _follow_job_logs(log_file,
+                                         job_id=job_id,
+                                         start_streaming=start_streaming,
+                                         start_streaming_at=start_stream_at):
+                yield line
+    else:
+        try:
+            start_streaming = False
+            with open(log_path, 'r', encoding='utf-8') as log_file:
+                if tail > 0:
+                    # If tail > 0, we need to read the last n lines.
+                    # We use double ended queue to rotate the last n lines.
+                    head_lines_of_log_file = _peek_head_lines(log_file)
+                    lines = collections.deque(log_file, maxlen=tail)
+                    start_streaming = _should_stream_the_whole_tail_lines(
+                        head_lines_of_log_file, lines, start_stream_at)
+                else:
+                    lines = log_file
+                for line in lines:
+                    if start_stream_at in line:
+                        start_streaming = True
+                    if start_streaming:
+                        yield line
+                status_str = status.value if status is not None else 'None'
+                # Only show "Job finished" for actually terminal states
+                if status is not None and status.is_terminal():
+                    finish = ux_utils.finishing_message(
+                        f'Job finished (status: {status_str}).')
+                    yield finish + '\n'
+            return
+        except FileNotFoundError:
+            err = (
+                f'{colorama.Fore.RED}ERROR: Logs for job {job_id} (status:'
+                f' {status.value}) does not exist.{colorama.Style.RESET_ALL}')
+            yield err + '\n'
+
+
+class LogBuffer:
+    """In-memory buffer for chunking log lines for streaming."""
+
+    def __init__(self, max_chars: int = DEFAULT_LOG_CHUNK_SIZE):
+        """Initialize the log buffer.
+
+        Args:
+            max_chars: Maximum buffer size (in characters, not bytes) before
+                       flushing. The actual amount of bytes (UTF-8 encoding)
+                       could be more than this, depending on the characters,
+                       i.e. ASCII characters take 1 byte, while others
+                       may take 2-4 bytes. But this is fine as our default
+                       chunk size is well below the default value of
+                       grpc.max_receive_message_length which is 4MB.
+        """
+        self.max_chars = max_chars
+        self._buffer = io.StringIO()
+
+    def _should_flush(self) -> bool:
+        return self._buffer.tell() >= self.max_chars
+
+    def flush(self) -> str:
+        """Get the current buffered content and clear the buffer.
+
+        Returns:
+            The buffered log lines as a single string
+        """
+        if not self._buffer.tell():
+            return ''
+        chunk = self._buffer.getvalue()
+        self._buffer.truncate(0)
+        self._buffer.seek(0)
+        return chunk
+
+    def write(self, line: str) -> bool:
+        """Add a line to the buffer.
+
+        Args:
+            line: The log line to add
+
+        Returns:
+            True if buffer should be flushed after adding the line
+        """
+        self._buffer.write(line)
+        return self._should_flush()
+
+    def close(self):
+        self._buffer.close()
+
+
+def buffered_iter_with_timeout(buffer: LogBuffer, iterable: Iterable[str],
+                               timeout: float) -> Iterable[str]:
+    """Iterates over an iterable, writing each item to a buffer,
+        and flushing the buffer when it is full or no item is
+        yielded within the timeout duration."""
+    # TODO(kevin): Simplify this using asyncio.timeout, once we move
+    # the skylet event loop and gRPC server to asyncio.
+    # https://docs.python.org/3/library/asyncio-task.html#timeouts
+
+    queue: queue_lib.Queue = queue_lib.Queue()
+    sentinel = object()
+
+    def producer():
+        try:
+            for item in iterable:
+                queue.put(item)
+        finally:
+            queue.put(sentinel)
+
+    thread = threading.Thread(target=producer, daemon=True)
+    thread.start()
+
+    while True:
+        try:
+            item = queue.get(timeout=timeout)
+        except queue_lib.Empty:
+            out = buffer.flush()
+            if out:
+                yield out
+            continue
+
+        if item is sentinel:
+            thread.join()
+            out = buffer.flush()
+            if out:
+                yield out
+            return
+
+        if buffer.write(item):
+            out = buffer.flush()
+            if out:
+                yield out