add Graphband.has_more_jobs (#38)

* add simple test case

* add additional test case

* test killed (apparently)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* lint

* test killed retries

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: PythonFZ

laufband/graphband.py

@@ -161,6 +161,7 @@ def __init__(
         # here we keep track of failed job data to be retried later
         self._failed_job_cache = {}
         self._iterator = None
+        self._iterator_completed = False
         self._heartbeat_timeout = heartbeat_timeout
         self._heartbeat_interval = heartbeat_interval
         self._labels = frozenset(labels or [])
@@ -257,6 +258,114 @@ def db(self) -> str:
         """Return the database URL."""
         return self._db
 
+    @property
+    def has_more_jobs(self) -> bool:
+        """Check if there are any more jobs that this worker could process.
+
+        Returns
+        -------
+        bool
+            True if there are jobs that this worker could potentially process,
+            False if all jobs are either completed or permanently failed/killed.
+
+        Notes
+        -----
+        Simple logic: Check if there are processable jobs in the failed cache
+        or incomplete tasks in the database that match our labels and haven't
+        exceeded retry limits.
+
+        Jobs with label mismatches are ignored since this worker would never
+        pick them up.
+
+        Examples
+        --------
+        >>> pbar = Graphband(tasks, labels={'worker-a'})
+        >>> list(pbar)  # Process all available jobs
+        >>> pbar.has_more_jobs  # Should be False if no more jobs for this worker
+        False
+        """
+        if self.disabled:
+            return False
+
+        # Simple logic - check failed cache and database
+        retryable_failed_jobs = 0
+        incomplete_jobs = 0
+
+        with self.lock:
+            with Session(self._engine) as session:
+                # Check failed job cache
+                for task in self._failed_job_cache.values():
+                    if not task.requirements.issubset(self.labels):
+                        continue
+
+                    task_entry = session.get(TaskEntry, task.id)
+                    if task_entry is None:
+                        retryable_failed_jobs += 1  # New task, can be processed
+                    elif (
+                        task_entry.failed_retries < self._max_failed_retries
+                        and task_entry.killed_retries < self._max_killed_retries
+                    ):
+                        retryable_failed_jobs += 1
+
+                # Check database for incomplete tasks
+                workflow = (
+                    session.query(WorkflowEntry)
+                    .filter(WorkflowEntry.id == "main")
+                    .first()
+                )
+                if workflow is not None:
+                    for task_entry in workflow.tasks:
+                        # Skip if labels don't match
+                        if not set(task_entry.requirements).issubset(self.labels):
+                            continue
+
+                        # Count incomplete tasks that can be retried
+                        if (
+                            not task_entry.completed
+                            and task_entry.failed_retries < self._max_failed_retries
+                            and task_entry.killed_retries < self._max_killed_retries
+                        ):
+                            incomplete_jobs += 1
+
+                    # If no incomplete jobs found but iterator hasn't completed,
+                    # check if we have any tasks matching our labels
+                    if incomplete_jobs == 0 and not self._iterator_completed:
+                        # Check if all tasks matching our labels are complete
+                        total_matching_tasks = len(
+                            [
+                                t
+                                for t in workflow.tasks
+                                if set(t.requirements).issubset(self.labels)
+                            ]
+                        )
+                        completed_matching_tasks = len(
+                            [
+                                t
+                                for t in workflow.tasks
+                                if set(t.requirements).issubset(self.labels)
+                                and t.completed
+                            ]
+                        )
+
+                        # If we have tasks in DB and all matching ones are complete,
+                        # and no retryable failed jobs, then no more jobs
+                        if (
+                            total_matching_tasks > 0
+                            and completed_matching_tasks == total_matching_tasks
+                            and retryable_failed_jobs == 0
+                        ):
+                            return False
+
+                        # Otherwise, there might be more tasks to process
+                        return True
+                else:
+                    # No workflow exists yet - if we haven't completed iteration,
+                    # assume there are jobs from the original graph
+                    if not self._iterator_completed:
+                        return True
+
+        return (retryable_failed_jobs + incomplete_jobs) > 0
+
     def __iter__(self) -> Iterator[Task[TaskTypeVar]]:
         """The generator that handles the iteration logic."""
 
@@ -278,6 +387,7 @@ def __iter__(self) -> Iterator[Task[TaskTypeVar]]:
             yield from self._iterator
             return
 
+        completed_naturally = True
         for task in self._iterator:
             if not task.requirements.issubset(self.labels):
                 log.debug(
@@ -389,6 +499,7 @@ def __iter__(self) -> Iterator[Task[TaskTypeVar]]:
 
                         worker.status = WorkerStatus.IDLE
                         session.commit()
+                    completed_naturally = False
                     break
             with self.lock:
                 with Session(self._engine) as session:
@@ -421,6 +532,9 @@ def __iter__(self) -> Iterator[Task[TaskTypeVar]]:
                     worker.status = WorkerStatus.IDLE
                     session.commit()
             if self._close_trigger:
+                completed_naturally = False
                 break
+        if completed_naturally:
+            self._iterator_completed = True
         self._thread_event.set()
         self._heartbeat_thread.join()

tests/test_graphband.py

@@ -525,3 +525,210 @@ def test_multi_dependency_graph_task(tmp_path):
             assert sorted(
                 [e.id for e in entries[task_id].current_status.dependencies]
             ) == sorted(deps)
+
+
+def test_has_more_jobs(tmp_path):
+    worker = Graphband(
+        sequential_task(),
+        db=f"sqlite:///{tmp_path}/graphband.sqlite",
+        lock=Lock(f"{tmp_path}/graphband.lock"),
+    )
+    assert worker.has_more_jobs is True
+    for item in worker:
+        if item.id == "task_5":
+            break
+    # there are still 4 remaining tasks
+    assert worker.has_more_jobs is True
+
+    assert len(list(worker)) == 4
+    # won't pick up the failed job
+    assert worker.has_more_jobs is False
+
+    w2 = Graphband(
+        sequential_task(),
+        db=f"sqlite:///{tmp_path}/graphband.sqlite",
+        lock=Lock(f"{tmp_path}/graphband.lock"),
+        max_failed_retries=2,
+        identifier="retry-worker",
+    )
+    # will pick up the failed job
+    assert w2.has_more_jobs is True
+    assert len(list(w2)) == 1
+    assert w2.has_more_jobs is False
+
+    w3 = Graphband(
+        sequential_task(),
+        db=f"sqlite:///{tmp_path}/graphband.sqlite",
+        lock=Lock(f"{tmp_path}/graphband.lock"),
+        max_failed_retries=2,
+        identifier="retry-worker-2",
+    )
+    assert w3.has_more_jobs is False  # all jobs are no completed successfully
+
+
+def blocked_dependency_graph_task():
+    """Create a graph where task b blocks task c due to label requirements.
+
+    Graph structure:
+    a --> b --> c
+    Where b requires 'special-worker' label, but a and c require 'main' label.
+    """
+    digraph = nx.DiGraph()
+    edges = [
+        ("a", "b"),
+        ("b", "c"),
+    ]
+    digraph.add_edges_from(edges)
+    digraph.nodes["b"]["requirements"] = {"special-worker"}
+    for node in nx.topological_sort(digraph):
+        yield Task(
+            id=node,
+            data=node,
+            dependencies=set(digraph.predecessors(node)),
+            requirements=digraph.nodes[node].get("requirements", {"main"}),
+        )
+
+
+def test_has_more_jobs_with_blocked_dependencies(tmp_path):
+    """Test has_more_jobs when dependencies are blocked by label mismatches.
+
+    This test verifies the scenario where:
+    - a --> b --> c dependency chain
+    - b needs a different label than a, c
+    - has_more_jobs for the a/c worker should be true until c is completed
+    - but the worker won't be able to pick up c because b isn't completed
+    """
+    # Worker that can process 'main' tasks (a and c) but not 'special-worker' tasks (b)
+    main_worker = Graphband(
+        blocked_dependency_graph_task(),
+        db=f"sqlite:///{tmp_path}/graphband.sqlite",
+        lock=Lock(f"{tmp_path}/graphband.lock"),
+        labels={"main"},
+        identifier="main-worker",
+    )
+
+    # Initially should have jobs available
+    assert main_worker.has_more_jobs is True
+
+    # Process available tasks - should only get task 'a'
+    items = list(main_worker)
+    assert len(items) == 1
+    assert items[0].id == "a"
+
+    # After processing 'a', should still have more jobs (task 'c' exists but is blocked)
+    # This is the key behavior: has_more_jobs returns True even though this worker
+    # cannot make progress because 'c' depends on 'b' which requires different labels
+    assert main_worker.has_more_jobs is True
+
+    # Trying to iterate again should yield nothing since 'c' is blocked by 'b'
+    items_second = list(main_worker)
+    assert len(items_second) == 0
+
+    # Should still report more jobs available (task 'c' is incomplete)
+    assert main_worker.has_more_jobs is True
+
+    # Now create a worker that can process the 'special-worker' task 'b'
+    special_worker = Graphband(
+        blocked_dependency_graph_task(),
+        db=f"sqlite:///{tmp_path}/graphband.sqlite",
+        lock=Lock(f"{tmp_path}/graphband.lock"),
+        labels={"special-worker"},
+        identifier="special-worker",
+    )
+
+    # Special worker should have jobs (task 'b')
+    assert special_worker.has_more_jobs is True
+
+    # Process task 'b'
+    items_special = list(special_worker)
+    assert len(items_special) == 1
+    assert items_special[0].id == "b"
+
+    # After 'b' is complete, special worker should have no more jobs
+    assert special_worker.has_more_jobs is False
+
+    # Now main worker should be able to process task 'c'
+    assert main_worker.has_more_jobs is True
+
+    items_final = list(main_worker)
+    assert len(items_final) == 1
+    assert items_final[0].id == "c"
+
+    # Finally, no more jobs for either worker
+    assert main_worker.has_more_jobs is False
+    assert special_worker.has_more_jobs is False
+
+
+def test_has_more_jobs_with_killed_workers(tmp_path):
+    """Test has_more_jobs behavior when workers are killed
+    and tasks exceed retry limits."""
+    # Test case where killed tasks cannot be retried (max_killed_retries=0)
+    lock_path = f"{tmp_path}/graphband.lock"
+    db = f"sqlite:///{tmp_path}/graphband.sqlite"
+    file = f"{tmp_path}/output.txt"
+
+    # Start a worker that will be killed with no retry allowance
+    proc = multiprocessing.Process(
+        target=task_worker,
+        args=(sequential_task, lock_path, db, file, 3),
+        kwargs={
+            "heartbeat_timeout": 1,
+            "heartbeat_interval": 0.5,
+            "max_killed_retries": 0,  # No retries allowed for killed tasks
+            "identifier": "killed-worker",
+        },
+    )
+    proc.start()
+    time.sleep(2)  # Let the worker start one task
+    proc.kill()
+    proc.join()
+
+    time.sleep(2)
+    # need to start another worker to mark the job as killed in the db
+    _ = Graphband(
+        sequential_task(),
+        db=db,
+        lock=Lock(lock_path),
+        heartbeat_timeout=2,
+        heartbeat_interval=1,
+        identifier="update-worker",
+    )
+    time.sleep(2)
+
+    # Verify the killed task is permanently blocked
+    engine = create_engine(db)
+    with Session(engine) as session:
+        tasks = session.query(TaskEntry).all()
+        # Verify we have exactly one killed task and 9 completed tasks
+        killed_tasks = [
+            t for t in tasks if t.current_status.status == TaskStatusEnum.KILLED
+        ]
+        assert len(killed_tasks) == 1
+
+    no_retries_worker = Graphband(
+        sequential_task(),
+        db=db,
+        lock=Lock(lock_path),
+        heartbeat_timeout=2,
+        heartbeat_interval=1,
+        max_killed_retries=0,  # No retries allowed
+        identifier="no-retries-worker",
+    )
+
+    retries_worker = Graphband(
+        sequential_task(),
+        db=db,
+        lock=Lock(lock_path),
+        heartbeat_timeout=2,
+        heartbeat_interval=1,
+        max_killed_retries=2,  # Allow retries
+        identifier="retries-worker",
+    )
+
+    assert no_retries_worker.has_more_jobs is True
+    assert retries_worker.has_more_jobs is True
+    assert len(list(no_retries_worker)) == 9
+    assert no_retries_worker.has_more_jobs is False
+    assert retries_worker.has_more_jobs is True
+    assert len(list(retries_worker)) == 1
+    assert retries_worker.has_more_jobs is False