Repository: skypilot-org/skypilot
Pull request #4310 (2024-11-09 04:33:27)

[Core] Avoid job scheduling race condition
<!-- Describe the changes in this PR -->

Another race condition in job scheduling besides #4264 ...

The pending jobs should be queried during the pending loop, otherwise, a same job can be submitted twice to ray job in the following condition.
1. Two threads calling `schedule_step()` and both get the list of pending jobs
2. The first thread submitted the first pending job
3. The second thread look at the stale pending job list and found the first job not submitted yet, so it will submit the job again.

23 jobs in parallel now, a bit more than #4264
```
290  sky-cmd  8 mins ago      -               -         1x[CPU:1+]  PENDING    ~/sky_logs/sky-2024-11-09-04-39-25-046022  
289  sky-cmd  8 mins ago      a few secs ago  2s        1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-23-183791  
288  sky-cmd  8 mins ago      a few secs ago  5s        1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-22-702691  
287  sky-cmd  8 mins ago      a few secs ago  7s        1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-22-167590  
286  sky-cmd  8 mins ago      a few secs ago  10s       1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-21-982999  
285  sky-cmd  8 mins ago      12 secs ago     12s       1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-21-824679  
284  sky-cmd  8 mins ago      15 secs ago     15s       1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-20-840780  
283  sky-cmd  8 mins ago      18 secs ago     18s       1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-19-944765  
282  sky-cmd  8 mins ago      20 secs ago     20s       1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-18-955162  
281  sky-cmd  8 mins ago      23 secs ago     23s       1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-17-037321  
280  sky-cmd  8 mins ago      25 secs ago     25s       1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-16-274941  
279  sky-cmd  8 mins ago      28 secs ago     28s       1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-15-639855  
278  sky-cmd  8 mins ago      30 secs ago     30s       1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-15-495491  
277  sky-cmd  8 mins ago      33 secs ago     33s       1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-15-271224  
276  sky-cmd  8 mins ago      36 secs ago     36s       1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-14-332769  
275  sky-cmd  8 mins ago      38 secs ago     38s       1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-13-681088  
274  sky-cmd  8 mins ago      41 secs ago     41s       1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-13-161247  
273  sky-cmd  8 mins ago      44 secs ago     44s       1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-11-060571  
272  sky-cmd  8 mins ago      46 secs ago     46s       1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-10-313379  
271  sky-cmd  8 mins ago      49 secs ago     49s       1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-09-829889  
270  sky-cmd  8 mins ago      52 secs ago     52s       1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-08-895334  
269  sky-cmd  8 mins ago      54 secs ago     54s       1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-08-861447  
268  sky-cmd  8 mins ago      57 secs ago     57s       1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-08-065965  
267  sky-cmd  8 mins ago      59 secs ago     59s       1x[CPU:1+]  RUNNING    ~/sky_logs/sky-2024-11-09-04-39-07-076732  
266  sky-cmd  8 mins ago      1 min ago       1m        1x[CPU:1+]  SUCCEEDED  ~/sky_logs/sky-2024-11-09-04-39-06-624532
```

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [ ] Code formatting: `bash format.sh`
- [x] Any manual or new tests for this PR (please specify below)
- [ ] All smoke tests: `pytest tests/test_smoke.py` 
- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
- [ ] Backward compatibility tests: `conda deactivate; bash -i tests/backward_compatibility_tests.sh`



Patch:
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index 68d9bfecf..0198ca45d 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -3711,7 +3711,8 @@ class CloudVmRayBackend(backends.Backend['CloudVmRayResourceHandle']):
                   handle: CloudVmRayResourceHandle,
                   job_id: Optional[int],
                   managed_job_id: Optional[int] = None,
-                  follow: bool = True) -> int:
+                  follow: bool = True,
+                  tail: int = 0) -> int:
         """Tail the logs of a job.
 
         Args:
@@ -3719,10 +3720,13 @@ class CloudVmRayBackend(backends.Backend['CloudVmRayResourceHandle']):
             job_id: The job ID to tail the logs of.
             managed_job_id: The managed job ID for display purpose only.
             follow: Whether to follow the logs.
+            tail: The number of lines to display from the end of the
+                log file. If 0, print all lines.
         """
         code = job_lib.JobLibCodeGen.tail_logs(job_id,
                                                managed_job_id=managed_job_id,
-                                               follow=follow)
+                                               follow=follow,
+                                               tail=tail)
         if job_id is None and managed_job_id is None:
             logger.info(
                 'Job ID not provided. Streaming the logs of the latest job.')
diff --git a/sky/cli.py b/sky/cli.py
index 29e3b2e51..c74b2f3ad 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -2022,6 +2022,12 @@ def queue(clusters: List[str], skip_finished: bool, all_users: bool):
     help=('Follow the logs of a job. '
           'If --no-follow is specified, print the log so far and exit. '
           '[default: --follow]'))
+@click.option(
+    '--tail',
+    default=0,
+    type=int,
+    help=('The number of lines to display from the end of the log file. '
+          'Default is 0, which means print all lines.'))
 @click.argument('cluster',
                 required=True,
                 type=str,
@@ -2035,6 +2041,7 @@ def logs(
     sync_down: bool,
     status: bool,  # pylint: disable=redefined-outer-name
     follow: bool,
+    tail: int,
 ):
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
     """Tail the log of a job.
@@ -2101,7 +2108,7 @@ def logs(
                 click.secho(f'Job {id_str}not found', fg='red')
             sys.exit(1)
 
-    core.tail_logs(cluster, job_id, follow)
+    core.tail_logs(cluster, job_id, follow, tail)
 
 
 @cli.command()
diff --git a/sky/core.py b/sky/core.py
index 496b8b8ad..4bb12f4a2 100644
--- a/sky/core.py
+++ b/sky/core.py
@@ -742,7 +742,8 @@ def cancel(
 @usage_lib.entrypoint
 def tail_logs(cluster_name: str,
               job_id: Optional[int],
-              follow: bool = True) -> None:
+              follow: bool = True,
+              tail: int = 0) -> None:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
     """Tail the logs of a job.
 
@@ -775,7 +776,7 @@ def tail_logs(cluster_name: str,
         f'{colorama.Style.RESET_ALL}')
 
     usage_lib.record_cluster_name_for_current_operation(cluster_name)
-    backend.tail_logs(handle, job_id, follow=follow)
+    backend.tail_logs(handle, job_id, follow=follow, tail=tail)
 
 
 @usage_lib.entrypoint
diff --git a/sky/dag.py b/sky/dag.py
index dc580d0cf..19df51070 100644
--- a/sky/dag.py
+++ b/sky/dag.py
@@ -56,21 +56,25 @@ class Dag:
         return self.graph
 
     def is_chain(self) -> bool:
-        # NOTE: this method assumes that the graph has no cycle.
-        is_chain = True
-        visited_zero_out_degree = False
-        for node in self.graph.nodes:
-            out_degree = self.graph.out_degree(node)
-            if out_degree > 1:
-                is_chain = False
-                break
-            elif out_degree == 0:
-                if visited_zero_out_degree:
-                    is_chain = False
-                    break
-                else:
-                    visited_zero_out_degree = True
-        return is_chain
+        """Check if the DAG is a linear chain of tasks."""
+
+        nodes = list(self.graph.nodes)
+
+        if len(nodes) == 0:
+            return True
+
+        in_degrees = [self.graph.in_degree(node) for node in nodes]
+        out_degrees = [self.graph.out_degree(node) for node in nodes]
+
+        # Check out-degrees: all <= 1 and exactly one node has out_degree == 0
+        out_degree_condition = (all(degree <= 1 for degree in out_degrees) and
+                                sum(degree == 0 for degree in out_degrees) == 1)
+
+        # Check in-degrees: all <= 1 and exactly one node has in_degree == 0
+        in_degree_condition = (all(degree <= 1 for degree in in_degrees) and
+                               sum(degree == 0 for degree in in_degrees) == 1)
+
+        return out_degree_condition and in_degree_condition
 
 
 class _DagContext(threading.local):
diff --git a/sky/skylet/constants.py b/sky/skylet/constants.py
index a9b8013ca..91476cf8f 100644
--- a/sky/skylet/constants.py
+++ b/sky/skylet/constants.py
@@ -79,7 +79,7 @@ SKYLET_VERSION = '8'
 # The version of the lib files that skylet/jobs use. Whenever there is an API
 # change for the job_lib or log_lib, we need to bump this version, so that the
 # user can be notified to update their SkyPilot version on the remote cluster.
-SKYLET_LIB_VERSION = 1
+SKYLET_LIB_VERSION = 2
 SKYLET_VERSION_FILE = '~/.sky/skylet_version'
 
 # `sky jobs dashboard`-related
diff --git a/sky/skylet/job_lib.py b/sky/skylet/job_lib.py
index fba015618..ee7aee85f 100644
--- a/sky/skylet/job_lib.py
+++ b/sky/skylet/job_lib.py
@@ -11,8 +11,7 @@ import shlex
 import sqlite3
 import subprocess
 import time
-import typing
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional
 
 import colorama
 import filelock
@@ -24,11 +23,9 @@ 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
-
 logger = sky_logging.init_logger(__name__)
 
+_LINUX_NEW_LINE = '\n'
 _JOB_STATUS_LOCK = '~/.sky/locks/.job_{}.lock'
 
 
@@ -183,12 +180,20 @@ class JobScheduler:
     def schedule_step(self, force_update_jobs: bool = False) -> None:
         if force_update_jobs:
             update_status()
-        pending_jobs = self._get_pending_jobs()
+        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 pending_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.
@@ -207,8 +212,8 @@ class JobScheduler:
                 self._run_job(job_id, run_cmd)
                 return
 
-    def _get_pending_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.
@@ -219,9 +224,10 @@ class JobScheduler:
 class FIFOScheduler(JobScheduler):
     """First in first out job scheduler"""
 
-    def _get_pending_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()
@@ -518,11 +524,16 @@ def _get_jobs_by_ids(job_ids: List[int]) -> List[Dict[str, Any]]:
 
 
 def _get_pending_job(job_id: int) -> Optional[Dict[str, Any]]:
-    rows = _CURSOR.execute('SELECT created_time, submit FROM pending_jobs '
-                           f'WHERE job_id={job_id!r}')
+    rows = _CURSOR.execute(
+        'SELECT created_time, submit, run_cmd FROM pending_jobs '
+        f'WHERE job_id={job_id!r}')
     for row in rows:
-        created_time, submit = row
-        return {'created_time': created_time, 'submit': submit}
+        created_time, submit, run_cmd = row
+        return {
+            'created_time': created_time,
+            'submit': submit,
+            'run_cmd': run_cmd
+        }
     return None
 
 
@@ -602,6 +613,7 @@ def update_job_status(job_ids: List[int],
                 # 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
@@ -792,7 +804,9 @@ def cancel_jobs_encoded_results(jobs: Optional[List[int]],
                     logger.warning(str(e))
                     continue
 
-            if job['status'] in [
+            # 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)
@@ -903,14 +917,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})',
+            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)
 
diff --git a/sky/skylet/log_lib.py b/sky/skylet/log_lib.py
index eb6444007..391fa8c4f 100644
--- a/sky/skylet/log_lib.py
+++ b/sky/skylet/log_lib.py
@@ -2,6 +2,7 @@
 
 This is a remote utility module that provides logging functionality.
 """
+import collections
 import copy
 import io
 import multiprocessing.pool
@@ -12,7 +13,8 @@ import sys
 import tempfile
 import textwrap
 import time
-from typing import Dict, Iterator, List, Optional, Tuple, Union
+from typing import (Deque, Dict, Iterable, Iterator, List, Optional, TextIO,
+                    Tuple, Union)
 
 import colorama
 
@@ -26,6 +28,9 @@ from sky.utils import ux_utils
 _SKY_LOG_WAITING_GAP_SECONDS = 1
 _SKY_LOG_WAITING_MAX_RETRY = 5
 _SKY_LOG_TAILING_GAP_SECONDS = 0.2
+# Peek the head of the lines to check if we need to start
+# streaming when tail > 0.
+PEEK_HEAD_LINES_FOR_START_STREAM = 20
 
 logger = sky_logging.init_logger(__name__)
 
@@ -330,6 +335,7 @@ def run_bash_command_with_log(bash_command: str,
 
 def _follow_job_logs(file,
                      job_id: int,
+                     start_streaming: bool,
                      start_streaming_at: str = '') -> Iterator[str]:
     """Yield each line from a file as they are written.
 
@@ -338,7 +344,6 @@ def _follow_job_logs(file,
     # No need to lock the status here, as the while loop can handle
     # the older status.
     status = job_lib.get_status_no_lock(job_id)
-    start_streaming = False
     wait_last_logs = True
     while True:
         tmp = file.readline()
@@ -378,10 +383,45 @@ def _follow_job_logs(file,
             status = job_lib.get_status_no_lock(job_id)
 
 
+def _peek_head_lines(log_file: TextIO) -> List[str]:
+    """Peek the head of the file."""
+    lines = [
+        log_file.readline() for _ in range(PEEK_HEAD_LINES_FOR_START_STREAM)
+    ]
+    # Reset the file pointer to the beginning
+    log_file.seek(0, os.SEEK_SET)
+    return [line for line in lines if line]
+
+
+def _should_stream_the_whole_tail_lines(head_lines_of_log_file: List[str],
+                                        tail_lines: Deque[str],
+                                        start_stream_at: str) -> bool:
+    """Check if the entire tail lines should be streamed."""
+    # See comment:
+    # https://github.com/skypilot-org/skypilot/pull/4241#discussion_r1833611567
+    # for more details.
+    # Case 1: If start_stream_at is found at the head of the tail lines,
+    # we should not stream the whole tail lines.
+    for index, line in enumerate(tail_lines):
+        if index >= PEEK_HEAD_LINES_FOR_START_STREAM:
+            break
+        if start_stream_at in line:
+            return False
+    # Case 2: If start_stream_at is found at the head of log file, but not at
+    # the tail lines, we need to stream the whole tail lines.
+    for line in head_lines_of_log_file:
+        if start_stream_at in line:
+            return True
+    # Case 3: If start_stream_at is not at the head, and not found at the tail
+    # lines, we should not stream the whole tail lines.
+    return False
+
+
 def tail_logs(job_id: Optional[int],
               log_dir: Optional[str],
               managed_job_id: Optional[int] = None,
-              follow: bool = True) -> None:
+              follow: bool = True,
+              tail: int = 0) -> None:
     """Tail the logs of a job.
 
     Args:
@@ -390,6 +430,8 @@ def tail_logs(job_id: Optional[int],
         managed_job_id: The managed job id (for logging info only to avoid
             confusion).
         follow: Whether to follow the logs or print the logs so far and exit.
+        tail: The number of lines to display from the end of the log file,
+            if 0, print all lines.
     """
     if job_id is None:
         # This only happens when job_lib.get_latest_job_id() returns None,
@@ -430,6 +472,8 @@ def tail_logs(job_id: Optional[int],
         status = job_lib.update_job_status([job_id], silent=True)[0]
 
     start_stream_at = 'Waiting for task resources on '
+    # 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,
@@ -440,18 +484,43 @@ def tail_logs(job_id: Optional[int],
         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:
+                        print(line, end='')
+                # Flush the last n lines
+                print(end='', flush=True)
+            # 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):
                 print(line, end='', flush=True)
     else:
         try:
-            start_stream = False
-            with open(log_path, 'r', encoding='utf-8') as f:
-                for line in f.readlines():
+            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_stream = True
-                    if start_stream:
+                        start_streaming = True
+                    if start_streaming:
                         print(line, end='', flush=True)
         except FileNotFoundError:
             print(f'{colorama.Fore.RED}ERROR: Logs for job {job_id} (status:'


Test patch:
diff --git a/tests/unit_tests/test_dag.py b/tests/unit_tests/test_dag.py
new file mode 100644
index 000000000..9bc626b00
--- /dev/null
+++ b/tests/unit_tests/test_dag.py
@@ -0,0 +1,148 @@
+"""Unit tests for sky.dag."""
+import pytest
+
+import sky
+
+
+@pytest.fixture
+def empty_dag():
+    """Fixture for an empty DAG."""
+    with sky.Dag() as dag:
+        yield dag
+
+
+@pytest.fixture
+def single_task_dag():
+    """Fixture for a DAG with single task."""
+    with sky.Dag() as dag:
+        sky.Task()
+        yield dag
+
+
+@pytest.fixture
+def linear_three_task_dag():
+    """Fixture for a DAG with three tasks in a linear chain."""
+    with sky.Dag() as dag:
+        task1 = sky.Task()
+        task2 = sky.Task()
+        task3 = sky.Task()
+        dag.add_edge(task1, task2)
+        dag.add_edge(task2, task3)
+        yield dag
+
+
+@pytest.fixture
+def branching_dag():
+    """Fixture for a DAG with one task branching into two tasks."""
+    with sky.Dag() as dag:
+        dag._tasks = [sky.Task() for _ in range(3)]
+        dag.add_edge(dag.tasks[0], dag.tasks[1])
+        dag.add_edge(dag.tasks[0], dag.tasks[2])
+        yield dag
+
+
+@pytest.fixture
+def merging_dag():
+    """Fixture for a DAG with two tasks merging into one task."""
+    with sky.Dag() as dag:
+        dag._tasks = [sky.Task() for _ in range(3)]
+        dag.add_edge(dag.tasks[0], dag.tasks[2])
+        dag.add_edge(dag.tasks[1], dag.tasks[2])
+        yield dag
+
+
+@pytest.fixture
+def multi_parent_dag():
+    """Fixture for a DAG with two tasks merging into one task."""
+    with sky.Dag() as dag:
+        dag._tasks = [sky.Task() for _ in range(3)]
+        dag.add_edge(dag.tasks[0], dag.tasks[2])
+        dag.add_edge(dag.tasks[1], dag.tasks[2])
+        yield dag
+
+
+@pytest.fixture
+def diamond_dag():
+    """Fixture for a diamond-shaped DAG: A -> (B,C) -> D."""
+    with sky.Dag() as dag:
+        dag._tasks = [sky.Task() for _ in range(4)]
+        dag.add_edge(dag.tasks[0], dag.tasks[1])
+        dag.add_edge(dag.tasks[0], dag.tasks[2])
+        dag.add_edge(dag.tasks[1], dag.tasks[3])
+        dag.add_edge(dag.tasks[2], dag.tasks[3])
+        yield dag
+
+
+@pytest.fixture
+def cyclic_dag():
+    """Fixture for a DAG with a cycle (invalid DAG)."""
+    with sky.Dag() as dag:
+        task1 = sky.Task()
+        task2 = sky.Task()
+        dag.add_edge(task1, task2)
+        dag.add_edge(task2, task1)  # Create cycle
+        yield dag
+
+
+@pytest.mark.parametrize('dag_fixture,description', [
+    ('empty_dag', 'Empty DAG'),
+    ('single_task_dag', 'Single task DAG'),
+    ('linear_three_task_dag', 'Linear chain of three tasks'),
+])
+def test_is_chain_true_cases(request, dag_fixture, description):
+    """Test cases where is_chain() should return True."""
+    dag = request.getfixturevalue(dag_fixture)
+    assert dag.is_chain(), f"Failed for case: {description}"
+
+
+@pytest.mark.parametrize('dag_fixture,description', [
+    ('branching_dag', 'Branching DAG'),
+    ('merging_dag', 'Merging DAG'),
+    ('diamond_dag', 'Diamond DAG'),
+])
+def test_is_chain_false_cases(request, dag_fixture, description):
+    """Test cases where is_chain() should return False."""
+    dag = request.getfixturevalue(dag_fixture)
+    assert not dag.is_chain(), f"Failed for case: {description}"
+
+
+@pytest.mark.parametrize('dag_fixture,description,expected_old,expected_new', [
+    ('linear_three_task_dag', 'Linear chain of three tasks', True, True),
+    ('multi_parent_dag', 'DAG with two tasks merging into one task', True,
+     False),
+])
+def test_is_chain_regression(request, dag_fixture, description, expected_old,
+                             expected_new):
+    """Regression test comparing new implementation with old behavior."""
+    dag = request.getfixturevalue(dag_fixture)
+
+    def old_is_chain(dag):
+        # Old implementation
+        is_chain = True
+        visited_zero_out_degree = False
+        for node in dag.graph.nodes:
+            out_degree = dag.graph.out_degree(node)
+            if out_degree > 1:
+                is_chain = False
+                break
+            elif out_degree == 0:
+                if visited_zero_out_degree:
+                    is_chain = False
+                    break
+                else:
+                    visited_zero_out_degree = True
+        return is_chain
+
+    assert dag.is_chain() == expected_new, f"Failed for case: {description}"
+    assert old_is_chain(dag) == expected_old, f"Failed for case: {description}"
+
+
+# TODO(andy): Currently cyclic DAGs are not detected and is_chain() simply
+# returns False. Once we implement cycle detection that raises an error,
+# update this test to use pytest.raises.
+@pytest.mark.xfail(reason="Cycle detection not implemented yet")
+def test_is_chain_with_cycle(cyclic_dag):
+    """Test is_chain() with cyclic graph.
+    """
+    with pytest.raises(ValueError):
+        cyclic_dag.is_chain()