From 4bca92285dccab9c0f7dd7e88ecdfe36804c91c3 Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Thu, 29 Oct 2026 17:32:21 +0530 Subject: [PATCH 1/8] UN-4318 [FEAT] List, count and abort stuck executions - GET execution/?stuck=true and execution/stuck-count/ for PENDING/EXECUTING runs idle past STUCK_EXECUTION_THRESHOLD_SECONDS (default 9000). - POST execution/abort/ stops runs the caller owns (or any run, for org admins and service accounts), with a per-id outcome and an audit event. - Workers check the run's status before each file and stop once it is STOPPED; the file in progress finishes. Co-Authored-By: Claude Opus 5.5 --- backend/backend/settings/base.py | 4 + backend/sample.env | 2 + backend/utils/organization_utils.py | 5 +- backend/workflow_manager/execution/README.md | 160 +++++++++++++ backend/workflow_manager/execution/abort.py | 163 +++++++++++++ backend/workflow_manager/execution/entity.py | 40 ++++ backend/workflow_manager/execution/enum.py | 10 +- backend/workflow_manager/execution/filter.py | 55 ++--- .../execution/serializer/stuck_abort.py | 50 ++++ backend/workflow_manager/execution/stuck.py | 65 +++++ .../execution/tests/stuck_abort_base.py | 116 +++++++++ .../execution/tests/test_abort_executions.py | 224 ++++++++++++++++++ .../tests/test_execution_contract.py | 172 ++++++++++++++ .../execution/tests/test_stuck_executions.py | 132 +++++++++++ backend/workflow_manager/execution/urls.py | 12 +- .../execution/views/__init__.py | 1 + .../execution/views/execution.py | 31 ++- .../execution/views/stuck_abort.py | 71 ++++++ backend/workflow_manager/internal_views.py | 22 ++ workers/file_processing/tasks.py | 66 ++++++ workers/shared/api/internal_client.py | 8 + workers/shared/clients/execution_client.py | 13 + workers/tests/test_abort_between_files.py | 216 +++++++++++++++++ 23 files changed, 1600 insertions(+), 38 deletions(-) create mode 100644 backend/workflow_manager/execution/README.md create mode 100644 backend/workflow_manager/execution/abort.py create mode 100644 backend/workflow_manager/execution/entity.py create mode 100644 backend/workflow_manager/execution/serializer/stuck_abort.py create mode 100644 backend/workflow_manager/execution/stuck.py create mode 100644 backend/workflow_manager/execution/tests/stuck_abort_base.py create mode 100644 backend/workflow_manager/execution/tests/test_abort_executions.py create mode 100644 backend/workflow_manager/execution/tests/test_execution_contract.py create mode 100644 backend/workflow_manager/execution/tests/test_stuck_executions.py create mode 100644 backend/workflow_manager/execution/views/stuck_abort.py create mode 100644 workers/tests/test_abort_between_files.py diff --git a/backend/backend/settings/base.py b/backend/backend/settings/base.py index 687b5182f8..54c2946d6b 100644 --- a/backend/backend/settings/base.py +++ b/backend/backend/settings/base.py @@ -407,6 +407,10 @@ def validate_agent_kv_path_prefix(raw, path_prefix, deployment_prefix): INSTANT_WF_POLLING_TIMEOUT = int( os.environ.get("INSTANT_WF_POLLING_TIMEOUT", "300") ) # 5 minutes +# A PENDING/EXECUTING run with no activity for this long is listed as stuck. +STUCK_EXECUTION_THRESHOLD_SECONDS = int( + os.environ.get("STUCK_EXECUTION_THRESHOLD_SECONDS", "9000") +) # ETL Pipeline minimum schedule interval (in seconds) # Default: 1800 seconds (30 minutes) diff --git a/backend/sample.env b/backend/sample.env index 33d4f535b0..ba8ea943c8 100644 --- a/backend/sample.env +++ b/backend/sample.env @@ -288,6 +288,8 @@ EXECUTION_RESULT_TTL_SECONDS=86400 EXECUTION_CACHE_TTL_SECONDS=86400 # Instant workflow polling timeout in seconds (5 minutes) INSTANT_WF_POLLING_TIMEOUT=300 +# A PENDING/EXECUTING run with no activity for this long is listed as stuck (2.5 hours) +STUCK_EXECUTION_THRESHOLD_SECONDS=9000 # Maximum number of batches (i.e., parallel tasks) created for a single workflow execution (1 file at a time) MAX_PARALLEL_FILE_BATCHES=1 diff --git a/backend/utils/organization_utils.py b/backend/utils/organization_utils.py index b81252899d..f0cb742f4d 100644 --- a/backend/utils/organization_utils.py +++ b/backend/utils/organization_utils.py @@ -6,6 +6,7 @@ from typing import Any from account_v2.models import Organization +from django.db.models import QuerySet from django.shortcuts import get_object_or_404 logger = logging.getLogger(__name__) @@ -71,7 +72,9 @@ def get_organization_context(organization: Organization) -> dict[str, Any]: } -def filter_queryset_by_organization(queryset, request, organization_field="organization"): +def filter_queryset_by_organization( + queryset: QuerySet, request: Any, organization_field: str = "organization" +) -> QuerySet: """Filter a Django queryset by the request's organization context. Fails closed. For every caller, this function is the only tenant boundary diff --git a/backend/workflow_manager/execution/README.md b/backend/workflow_manager/execution/README.md new file mode 100644 index 0000000000..2515c59285 --- /dev/null +++ b/backend/workflow_manager/execution/README.md @@ -0,0 +1,160 @@ +# Executions: stuck runs and abort + +Finding executions that stopped making progress, and stopping running ones +(PB-T08, UN-4318). This README and the `@extend_schema` annotations in +`views/stuck_abort.py` and `views/execution.py` are the published contract. +Changes to anything below are contract changes. + +- [Endpoints](#endpoints) +- [What counts as stuck](#what-counts-as-stuck) +- [Listing stuck executions](#listing-stuck-executions) +- [Counting stuck executions](#counting-stuck-executions) +- [Aborting executions](#aborting-executions) +- [Deploying](#deploying) +- [Limits](#limits) + +## Endpoints + +Mounted under the organisation path: `/api/v1/unstract/{org_id}/execution/`. + +| Method | Path | Returns | +|---|---|---| +| GET | `execution/?stuck=true` | The existing execution list, narrowed to stuck runs | +| GET | `execution/stuck-count/` | Stuck runs per entity, and the threshold | +| POST | `execution/abort/` | One outcome per requested id | + +Everything here is additive. The existing list, detail and log endpoints keep +their paths, parameters, response fields and permissions. Without `stuck`, or +with `stuck=false`, the list response is the same as before. + +## What counts as stuck + +An execution is stuck when its status is `PENDING` or `EXECUTING` and its last +activity is older than the threshold. Last activity is the later of the +execution row's `modified_at` and the newest `modified_at` of its file rows. + +The threshold is `STUCK_EXECUTION_THRESHOLD_SECONDS` (env, default **9000**, +2.5 hours). It is returned as `threshold_seconds` by `stuck-count/`. + +## Listing stuck executions + +`GET execution/?stuck=true` takes the list's existing filters, ordering and +pagination. Each row has the usual fields plus: + +| Field | Type | Meaning | +|---|---|---| +| `stuck_since` | ISO datetime | The time of last activity | + +`stuck` accepts `true`/`1` and `false`/`0`, case-insensitive. Any other value +is ignored, as if `stuck` were absent. Visibility is the same as the list's. + +## Counting stuck executions + +`GET execution/stuck-count/` returns 200: + +```json +{ + "threshold_seconds": 9000, + "total": 3, + "counts": {"API": 1, "ETL": 2, "TASK": 0, "WF": 0} +} +``` + +- `counts` uses the same keys as the list's `execution_entity` filter: `API` + deployments, `ETL` and `TASK` pipelines, and `WF` for workflow runs without a + pipeline. +- `total` counts every stuck run the caller can see. It can exceed the sum of + `counts` when a run's pipeline was deleted or is of another type. +- Visibility is the same as the list's. + +## Aborting executions + +`POST execution/abort/` with: + +```json +{"ids": ["3f0d6c1e-5a0b-4a39-9a7e-0d6f3c2b1a90", "..."]} +``` + +`ids` holds 1 to 100 execution ids; duplicates are answered once. It returns +200, with results in request order: + +```json +{ + "results": [ + {"id": "3f0d6c1e-5a0b-4a39-9a7e-0d6f3c2b1a90", "outcome": "aborted", "status": "STOPPED"}, + {"id": "57a4f9a2-2a4e-4a3c-8f0e-7c1d2b3a4f5e", "outcome": "already_finished", "status": "COMPLETED"}, + {"id": "9b1c2d3e-4f5a-4b6c-8d7e-0f1a2b3c4d5e", "outcome": "not_permitted", "status": null} + ], + "aborted": 1, + "already_finished": 1, + "not_permitted": 1 +} +``` + +| Outcome | When | `status` | +|---|---|---| +| `aborted` | The run was `PENDING` or `EXECUTING` and is now `STOPPED` | `STOPPED` | +| `already_finished` | The run was already `COMPLETED`, `ERROR` or `STOPPED`; nothing changed | its status | +| `not_permitted` | The caller may not abort it, or the id does not exist | `null` | + +Unknown ids and ids the caller cannot abort get the same answer, so the +endpoint does not reveal which ids exist. + +**Who may abort.** The owner of what started the run: an OWNER of its API +deployment or pipeline, or of its workflow when the run has no pipeline. Org +admins and service accounts (platform API keys) may abort any run in the +organisation. Shared users who are not owners get `not_permitted`. + +**What abort does,** in one transaction per id: + +1. Sets the execution to `STOPPED` with the error `Execution aborted`, which + also frees an API deployment's rate-limit slot and sends the existing + "was stopped" inbox notice. +2. Sets its `PENDING` and `EXECUTING` file rows to `STOPPED`. Finished file + rows, file history, destination writes and usage are kept. +3. Records an `execution.aborted` event in the platform audit log, with + target `execution`, the workflow name, the summary `was ` + and `previous_status`, `pipeline_id` and `entity` in its metadata. + +The pipeline's last-run status is not changed. + +**How the work stops.** Queued batches of a stopped run are skipped when a +worker picks them up. A running batch checks the status before each file and +stops once it is `STOPPED`: its remaining files are marked `STOPPED` and +counted as failed. The file +being processed when the abort lands finishes, and its own result replaces +`STOPPED` on its row. Later status writes from workers cannot undo the abort on +runs using the PG queue. + +### Errors + +Errors use the project's standard envelope (drf-standardized-errors). + +| Status | When | +|---|---| +| 400 | `ids` missing, empty, longer than 100, or holding a value that is not a UUID | +| 401 / 403 | Not authenticated | + +An unexpected failure part-way returns 500. Each id is aborted in its own +transaction, so the aborts made before it are kept (unless `ATOMIC_REQUESTS` is +on, which rolls the whole request back). Retrying is safe: ids already aborted +answer `already_finished`. + +## Deploying + +Deploy the backend before the workers. Workers read the status from the +internal `GET /internal/v1/workflow-execution//status/` endpoint. An older +backend answers 404, and the worker then processes the file anyway (fail +open), as it did before this change. + +## Limits + +- The file being processed when the abort lands is not interrupted. A worker + that is alive but hung keeps its claim until the reaper cleans it up. +- Runs from before the PG queue (`queue_message_id` empty) are not protected + from a late worker status write. +- If the batch that was running was the last one, the aggregating callback + still runs. It can update the pipeline's last-run status and fire webhooks, + while the execution stays `STOPPED`. +- A running batch stops between files only for `STOPPED`. A run that becomes + `ERROR` or `COMPLETED` mid-batch keeps processing its files, as before. diff --git a/backend/workflow_manager/execution/abort.py b/backend/workflow_manager/execution/abort.py new file mode 100644 index 0000000000..c8dcf12d7d --- /dev/null +++ b/backend/workflow_manager/execution/abort.py @@ -0,0 +1,163 @@ +"""Abort running executions: stop the run, keep processed files, audit it.""" + +import logging +from dataclasses import dataclass +from typing import Any +from uuid import UUID + +from api_v2.models import APIDeployment +from django.db import transaction +from django.db.models import Q, QuerySet +from permissions.roles import ResourceRole +from pipeline_v2.models import Pipeline +from platform_audit.actions import Action, TargetType +from platform_audit.recorder import record_platform_event +from tenant_account_v2.organization_member_service import OrganizationMemberService +from tenant_account_v2.sharing_helpers import resources_visible_via_memberships + +from workflow_manager.execution.entity import entity_annotation +from workflow_manager.execution.enum import AbortOutcome +from workflow_manager.file_execution.models import WorkflowFileExecution +from workflow_manager.workflow_v2.enums import ExecutionStatus +from workflow_manager.workflow_v2.models import Workflow, WorkflowExecution + +logger = logging.getLogger(__name__) + +ABORT_ERROR_MESSAGE = "Execution aborted" +MAX_ABORT_IDS = 100 + + +@dataclass(frozen=True) +class AbortResult: + """Outcome for one id. ``status`` is ``None`` when not permitted.""" + + id: UUID + outcome: AbortOutcome + status: str | None + + +@dataclass(frozen=True) +class AbortReport: + """Results in request order, with a count per outcome.""" + + results: list[AbortResult] + + def _count(self, outcome: AbortOutcome) -> int: + return sum(1 for result in self.results if result.outcome == outcome) + + @property + def aborted(self) -> int: + return self._count(AbortOutcome.ABORTED) + + @property + def already_finished(self) -> int: + return self._count(AbortOutcome.ALREADY_FINISHED) + + @property + def not_permitted(self) -> int: + return self._count(AbortOutcome.NOT_PERMITTED) + + +def abort_executions(request: Any, ids: list[UUID]) -> AbortReport: + """Abort each id the caller may abort; one transaction per id.""" + unique_ids = list(dict.fromkeys(ids)) + permitted = _abortable_entities(request.user, unique_ids) + results = [ + _abort_one(request, execution_id, permitted[execution_id]) + if execution_id in permitted + else AbortResult(execution_id, AbortOutcome.NOT_PERMITTED, None) + for execution_id in unique_ids + ] + return AbortReport(results) + + +def _abortable_entities(user: Any, ids: list[UUID]) -> dict[UUID, str | None]: + """Map each id ``user`` may abort to its entity, in one query. + + Owners of the run's pipeline or API deployment (or of its workflow, for a + run without one), org admins and service accounts may abort. + """ + rows = ( + _abortable(user) + .filter(pk__in=ids) + .annotate(entity=entity_annotation()) + .values_list("pk", "entity") + ) + return dict(rows) + + +def _abortable(user: Any) -> QuerySet: + visible = WorkflowExecution.objects.for_user(user) + if getattr(user, "is_service_account", False): + return visible + if OrganizationMemberService.is_user_organization_admin(user): + return visible + return visible.filter(_owner_condition(user)) + + +def _owner_condition(user: Any) -> Q: + def owned(model: Any) -> QuerySet: + return resources_visible_via_memberships( + model, user, role=str(ResourceRole.OWNER) + ) + + return ( + Q(pipeline_id__isnull=True, workflow_id__in=owned(Workflow)) + | Q(pipeline_id__in=owned(APIDeployment)) + | Q(pipeline_id__in=owned(Pipeline)) + ) + + +def _abort_one(request: Any, execution_id: UUID, entity: str | None) -> AbortResult: + """Stop one execution under its row lock, then record the audit event.""" + with transaction.atomic(): + execution = ( + WorkflowExecution.objects.select_for_update(no_key=True, of=("self",)) + .select_related("workflow") + .filter(pk=execution_id) + .first() + ) + if execution is None: + return AbortResult(execution_id, AbortOutcome.NOT_PERMITTED, None) + if execution.status in ExecutionStatus.terminal_values(): + return AbortResult( + execution_id, AbortOutcome.ALREADY_FINISHED, execution.status + ) + previous_status = execution.status + execution.update_execution( + status=ExecutionStatus.STOPPED, error=ABORT_ERROR_MESSAGE, row_locked=True + ) + _stop_open_files(execution) + _record_abort(request, execution, previous_status, entity) + return AbortResult(execution_id, AbortOutcome.ABORTED, execution.status) + + +def _stop_open_files(execution: WorkflowExecution) -> None: + """Mark file rows not yet terminal STOPPED; finished rows keep their status.""" + stopped = ( + WorkflowFileExecution.objects.filter(workflow_execution=execution) + .exclude(status__in=ExecutionStatus.terminal_values()) + .update(status=ExecutionStatus.STOPPED.value, execution_error=ABORT_ERROR_MESSAGE) + ) + logger.info( + "Aborted execution %s; stopped %s open file execution(s)", execution.id, stopped + ) + + +def _record_abort( + request: Any, execution: WorkflowExecution, previous_status: str, entity: str | None +) -> None: + pipeline_id = str(execution.pipeline_id) if execution.pipeline_id else None + record_platform_event( + action=Action.EXECUTION_ABORTED, + target_type=TargetType.EXECUTION, + target_id=execution.id, + target_name=execution.workflow_name or "", + request=request, + summary=f"was {previous_status}", + metadata={ + "previous_status": previous_status, + "pipeline_id": pipeline_id, + "entity": entity, + }, + ) diff --git a/backend/workflow_manager/execution/entity.py b/backend/workflow_manager/execution/entity.py new file mode 100644 index 0000000000..0afdcc62c4 --- /dev/null +++ b/backend/workflow_manager/execution/entity.py @@ -0,0 +1,40 @@ +"""Which kind of source started a run: API deployment, ETL/TASK pipeline or workflow.""" + +from api_v2.models import APIDeployment +from django.db.models import Case, CharField, Exists, OuterRef, Q, Value, When +from pipeline_v2.models import Pipeline + +from workflow_manager.execution.enum import ExecutionEntity + +_PIPELINE_TYPES = { + ExecutionEntity.ETL: Pipeline.PipelineType.ETL, + ExecutionEntity.TASK: Pipeline.PipelineType.TASK, +} + + +def entity_condition(entity: ExecutionEntity) -> Exists | Q: + """Condition matching executions started by ``entity``. + + EXISTS subqueries keep the check in SQL instead of fetching all ids. + """ + if entity is ExecutionEntity.API: + return Exists(APIDeployment.objects.filter(id=OuterRef("pipeline_id"))) + if entity is ExecutionEntity.WORKFLOW: + return Q(pipeline_id__isnull=True) + return Exists( + Pipeline.objects.filter( + id=OuterRef("pipeline_id"), pipeline_type=_PIPELINE_TYPES[entity] + ) + ) + + +def entity_annotation() -> Case: + """Annotation naming each execution's entity; ``None`` when it matches none.""" + return Case( + *( + When(entity_condition(entity), then=Value(entity.value)) + for entity in ExecutionEntity + ), + default=Value(None), + output_field=CharField(), + ) diff --git a/backend/workflow_manager/execution/enum.py b/backend/workflow_manager/execution/enum.py index 27900508c0..ba6f8fa1a7 100644 --- a/backend/workflow_manager/execution/enum.py +++ b/backend/workflow_manager/execution/enum.py @@ -1,4 +1,4 @@ -from enum import Enum +from enum import Enum, StrEnum class ExecutionEntity(Enum): @@ -6,3 +6,11 @@ class ExecutionEntity(Enum): API = "API" TASK = "TASK" WORKFLOW = "WF" + + +class AbortOutcome(StrEnum): + """Per-id answer of the abort endpoint.""" + + ABORTED = "aborted" + ALREADY_FINISHED = "already_finished" + NOT_PERMITTED = "not_permitted" diff --git a/backend/workflow_manager/execution/filter.py b/backend/workflow_manager/execution/filter.py index d62a96432a..9dbae68f88 100644 --- a/backend/workflow_manager/execution/filter.py +++ b/backend/workflow_manager/execution/filter.py @@ -1,13 +1,23 @@ -from api_v2.models import APIDeployment -from django.db.models import Exists, OuterRef, QuerySet +from typing import Any + +from django.db.models import QuerySet from django_filters import rest_framework as filters -from pipeline_v2.models import Pipeline +from django_filters.widgets import BooleanWidget from utils.date import DateRangePresets, DateTimeProcessor +from workflow_manager.execution.entity import entity_condition from workflow_manager.execution.enum import ExecutionEntity +from workflow_manager.execution.stuck import stuck_executions from workflow_manager.workflow_v2.enums import ExecutionStatus from workflow_manager.workflow_v2.models import WorkflowExecution +STUCK_PARAM = "stuck" + + +def stuck_requested(query_params: Any) -> bool: + """True when ``stuck`` parses as true, by the same rule as the filter.""" + return BooleanWidget().value_from_datadict(query_params, None, STUCK_PARAM) is True + class ExecutionFilter(filters.FilterSet): id = filters.UUIDFilter() @@ -29,6 +39,8 @@ class ExecutionFilter(filters.FilterSet): choices=DateRangePresets.choices(), method="filter_by_date_range", ) + # Documented on ExecutionViewSet's schema. + stuck = filters.BooleanFilter(method="filter_stuck") class Meta: model = WorkflowExecution @@ -40,38 +52,13 @@ def filter_execution_entity( """Filter executions by entity type efficiently. The queryset is already filtered by user permissions via get_queryset(), - so we use EXISTS subqueries to check entity type without fetching all IDs. - This is more efficient than values_list() for large datasets. + so EXISTS subqueries check the entity type without fetching all IDs. """ - if value == ExecutionEntity.API.value: - # Filter for API deployments using EXISTS for efficiency - return queryset.filter( - Exists(APIDeployment.objects.filter(id=OuterRef("pipeline_id"))) - ) - elif value == ExecutionEntity.ETL.value: - # Filter for ETL pipelines using EXISTS - return queryset.filter( - Exists( - Pipeline.objects.filter( - id=OuterRef("pipeline_id"), - pipeline_type=Pipeline.PipelineType.ETL, - ) - ) - ) - elif value == ExecutionEntity.TASK.value: - # Filter for TASK pipelines using EXISTS - return queryset.filter( - Exists( - Pipeline.objects.filter( - id=OuterRef("pipeline_id"), - pipeline_type=Pipeline.PipelineType.TASK, - ) - ) - ) - elif value == ExecutionEntity.WORKFLOW.value: - # Filter for workflow-level executions (no pipeline) - return queryset.filter(pipeline_id__isnull=True) - return queryset + return queryset.filter(entity_condition(ExecutionEntity(value))) + + def filter_stuck(self, queryset: QuerySet, name: str, value: bool) -> QuerySet: + """Narrow to stuck executions when true; false leaves the queryset as is.""" + return stuck_executions(queryset) if value else queryset def filter_by_date_range(self, queryset: QuerySet, name: str, value: str) -> QuerySet: """Filters Usages based on the provided date range.""" diff --git a/backend/workflow_manager/execution/serializer/stuck_abort.py b/backend/workflow_manager/execution/serializer/stuck_abort.py new file mode 100644 index 0000000000..c75cc25802 --- /dev/null +++ b/backend/workflow_manager/execution/serializer/stuck_abort.py @@ -0,0 +1,50 @@ +from rest_framework import serializers + +from workflow_manager.execution.abort import MAX_ABORT_IDS +from workflow_manager.execution.enum import AbortOutcome +from workflow_manager.execution.serializer.execution import ExecutionSerializer + + +class StuckExecutionSerializer(ExecutionSerializer): + """An execution row plus ``stuck_since``, served for ``?stuck=true``.""" + + stuck_since = serializers.DateTimeField( + read_only=True, help_text="Time of last activity on the run or its files." + ) + + +class StuckExecutionCountsSerializer(serializers.Serializer): + API = serializers.IntegerField() + ETL = serializers.IntegerField() + TASK = serializers.IntegerField() + WF = serializers.IntegerField() + + +class StuckExecutionCountSerializer(serializers.Serializer): + threshold_seconds = serializers.IntegerField() + total = serializers.IntegerField() + counts = StuckExecutionCountsSerializer() + + +class ExecutionAbortRequestSerializer(serializers.Serializer): + ids = serializers.ListField( + child=serializers.UUIDField(), + min_length=1, + max_length=MAX_ABORT_IDS, + help_text=f"1 to {MAX_ABORT_IDS} execution ids; duplicates are removed.", + ) + + +class ExecutionAbortResultSerializer(serializers.Serializer): + id = serializers.UUIDField() + outcome = serializers.ChoiceField(choices=[o.value for o in AbortOutcome]) + status = serializers.CharField( + allow_null=True, help_text="The run's status; null when not_permitted." + ) + + +class ExecutionAbortResponseSerializer(serializers.Serializer): + results = ExecutionAbortResultSerializer(many=True) + aborted = serializers.IntegerField() + already_finished = serializers.IntegerField() + not_permitted = serializers.IntegerField() diff --git a/backend/workflow_manager/execution/stuck.py b/backend/workflow_manager/execution/stuck.py new file mode 100644 index 0000000000..292d9f9a46 --- /dev/null +++ b/backend/workflow_manager/execution/stuck.py @@ -0,0 +1,65 @@ +"""Stuck executions: PENDING/EXECUTING runs with no activity past a threshold.""" + +from dataclasses import dataclass +from datetime import timedelta + +from django.conf import settings +from django.db.models import Count, OuterRef, QuerySet, Subquery +from django.db.models.functions import Greatest +from django.utils import timezone + +from workflow_manager.execution.entity import entity_condition +from workflow_manager.execution.enum import ExecutionEntity +from workflow_manager.file_execution.models import WorkflowFileExecution +from workflow_manager.workflow_v2.enums import ExecutionStatus +from workflow_manager.workflow_v2.models import WorkflowExecution + +ACTIVE_STATUSES = (ExecutionStatus.PENDING.value, ExecutionStatus.EXECUTING.value) + + +@dataclass(frozen=True) +class StuckCount: + """Stuck executions visible to a caller, in total and per entity.""" + + threshold_seconds: int + total: int + counts: dict[str, int] + + +def stuck_threshold_seconds() -> int: + return int(settings.STUCK_EXECUTION_THRESHOLD_SECONDS) + + +def stuck_executions(queryset: QuerySet) -> QuerySet: + """Narrow ``queryset`` to stuck runs, annotated with ``stuck_since``. + + ``stuck_since`` is the last activity: the later of the run's own + ``modified_at`` and its newest file row's ``modified_at``. + """ + cutoff = timezone.now() - timedelta(seconds=stuck_threshold_seconds()) + last_file_at = ( + WorkflowFileExecution.objects.filter(workflow_execution=OuterRef("pk")) + .order_by("-modified_at") + .values("modified_at")[:1] + ) + # Postgres GREATEST skips NULL, so a run with no files uses its own time. + return ( + queryset.filter(status__in=ACTIVE_STATUSES, modified_at__lt=cutoff) + .annotate(stuck_since=Greatest("modified_at", Subquery(last_file_at))) + .filter(stuck_since__lt=cutoff) + ) + + +def count_stuck(queryset: QuerySet) -> StuckCount: + """Count the stuck runs in ``queryset``, in one query.""" + per_entity = { + entity.value: Count("pk", filter=entity_condition(entity)) + for entity in ExecutionEntity + } + counts = WorkflowExecution.objects.filter( + pk__in=stuck_executions(queryset).values("pk") + ).aggregate(total=Count("pk"), **per_entity) + total = counts.pop("total") + return StuckCount( + threshold_seconds=stuck_threshold_seconds(), total=total, counts=counts + ) diff --git a/backend/workflow_manager/execution/tests/stuck_abort_base.py b/backend/workflow_manager/execution/tests/stuck_abort_base.py new file mode 100644 index 0000000000..72d2107f03 --- /dev/null +++ b/backend/workflow_manager/execution/tests/stuck_abort_base.py @@ -0,0 +1,116 @@ +"""Fixtures for the stuck-execution and abort tests (UN-4318). + +``self.owner`` owns the workflow, deployments and pipelines built here; +``self.admin`` is the only org admin; ``self.outsider`` has no share. +""" + +import secrets +from datetime import timedelta +from typing import Any +from unittest.mock import patch + +from api_v2.models import APIDeployment +from django.utils import timezone +from permissions.roles import ResourceRole +from pipeline_v2.models import Pipeline +from rest_framework.test import APIRequestFactory, force_authenticate +from tenant_account_v2.tests import GroupSharingTestBase + +from workflow_manager.file_execution.models import WorkflowFileExecution +from workflow_manager.workflow_v2.enums import ExecutionStatus +from workflow_manager.workflow_v2.models.execution import WorkflowExecution + +ADMIN_PREDICATE = ( + "tenant_account_v2.organization_member_service." + "OrganizationMemberService.is_user_organization_admin" +) +# Past the default 9000 s threshold, and well inside it. +OLD = timedelta(hours=3) +RECENT = timedelta(minutes=1) + + +class StuckAbortTestBase(GroupSharingTestBase): + def setUp(self) -> None: + super().setUp() + patcher = patch( + ADMIN_PREDICATE, + side_effect=lambda user: getattr(user, "email", None) == self.admin.email, + ) + patcher.start() + self.addCleanup(patcher.stop) + self.factory = APIRequestFactory() + + def deployment(self) -> APIDeployment: + deployment: APIDeployment = APIDeployment.objects.create( + api_name=f"api-{secrets.token_hex(4)}", + workflow=self.workflow, + organization=self.org, + created_by=self.owner, + ) + deployment.memberships.create(user=self.owner, role=ResourceRole.OWNER) + return deployment + + def pipeline(self, pipeline_type: Any = Pipeline.PipelineType.ETL) -> Pipeline: + pipeline: Pipeline = Pipeline.objects.create( + pipeline_name=f"pipe-{secrets.token_hex(4)}", + workflow=self.workflow, + organization=self.org, + created_by=self.owner, + pipeline_type=pipeline_type, + ) + pipeline.memberships.create(user=self.owner, role=ResourceRole.OWNER) + return pipeline + + def execution( + self, + resource: Any = None, + status: ExecutionStatus = ExecutionStatus.EXECUTING, + idle: timedelta = OLD, + ) -> WorkflowExecution: + """A PG-transport run whose last own write was ``idle`` ago.""" + execution: WorkflowExecution = WorkflowExecution.objects.create( + workflow=self.workflow, + pipeline_id=resource.id if resource else None, + execution_mode=WorkflowExecution.Mode.QUEUE, + execution_method=WorkflowExecution.Method.DIRECT, + execution_type=WorkflowExecution.Type.COMPLETE, + status=status.value, + queue_message_id=1, + ) + self.backdate(WorkflowExecution, execution.pk, idle) + execution.refresh_from_db() + return execution + + def file( + self, execution: WorkflowExecution, status: ExecutionStatus, idle: timedelta + ) -> WorkflowFileExecution: + row: WorkflowFileExecution = WorkflowFileExecution.objects.create( + workflow_execution=execution, file_name="f.pdf", status=status.value + ) + self.backdate(WorkflowFileExecution, row.pk, idle) + row.refresh_from_db() + return row + + @staticmethod + def backdate(model: Any, pk: Any, idle: timedelta) -> None: + # update() skips auto_now, so modified_at keeps the backdated value. + model.objects.filter(pk=pk).update(modified_at=timezone.now() - idle) + + def call( + self, + view: Any, + user: Any, + method: str = "get", + data: Any = None, + **url_kwargs: Any, + ) -> Any: + """Run ``view`` as ``user`` and every on-commit callback it registers.""" + if method == "get": + request = self.factory.get("/", data or {}) + else: + request = self.factory.post("/", data, format="json") + force_authenticate(request, user=user) + with self.captureOnCommitCallbacks(execute=True): + response = view(request, **url_kwargs) + response.render() + return response diff --git a/backend/workflow_manager/execution/tests/test_abort_executions.py b/backend/workflow_manager/execution/tests/test_abort_executions.py new file mode 100644 index 0000000000..35b1cff59f --- /dev/null +++ b/backend/workflow_manager/execution/tests/test_abort_executions.py @@ -0,0 +1,224 @@ +"""``POST execution/abort/`` (UN-4318): owner-or-admin, per-id outcomes, audit.""" + +import uuid +from typing import Any +from unittest.mock import MagicMock, patch + +from permissions.roles import ResourceRole +from pipeline_v2.models import Pipeline +from platform_audit.actions import Action, TargetType +from platform_audit.models import PlatformAuditEvent +from tenant_account_v2.tests import _add_viewers, _make_user + +from workflow_manager.execution.abort import ABORT_ERROR_MESSAGE, MAX_ABORT_IDS +from workflow_manager.execution.tests.stuck_abort_base import RECENT, StuckAbortTestBase +from workflow_manager.execution.views import ExecutionAbortView +from workflow_manager.internal_views import WorkflowExecutionInternalViewSet +from workflow_manager.workflow_v2.enums import ExecutionStatus +from workflow_manager.workflow_v2.models.execution import WorkflowExecution + +_abort = ExecutionAbortView.as_view() + + +class AbortExecutionTests(StuckAbortTestBase): + def _abort(self, user: Any, *ids: Any, raw: dict[str, Any] | None = None) -> Any: + body = raw if raw is not None else {"ids": [str(i) for i in ids]} + return self.call(_abort, user, method="post", data=body) + + def _outcomes(self, response: Any) -> dict[str, tuple[str, str | None]]: + self.assertEqual(response.status_code, 200, response.data) + return { + row["id"]: (row["outcome"], row["status"]) + for row in response.data["results"] + } + + def _status(self, execution: WorkflowExecution) -> str: + execution.refresh_from_db() + status: str = execution.status + return status + + def test_abort_stops_the_run_and_its_open_files_only(self) -> None: + execution = self.execution(self.deployment(), idle=RECENT) + done = self.file(execution, ExecutionStatus.COMPLETED, idle=RECENT) + running = self.file(execution, ExecutionStatus.EXECUTING, idle=RECENT) + queued = self.file(execution, ExecutionStatus.PENDING, idle=RECENT) + + response = self._abort(self.owner, execution.id) + + self.assertEqual( + self._outcomes(response), {str(execution.id): ("aborted", "STOPPED")} + ) + self.assertEqual( + (response.data["aborted"], response.data["already_finished"]), (1, 0) + ) + execution.refresh_from_db() + self.assertEqual(execution.status, ExecutionStatus.STOPPED.value) + self.assertEqual(execution.error_message, ABORT_ERROR_MESSAGE) + for row, expected in ( + (done, ExecutionStatus.COMPLETED), + (running, ExecutionStatus.STOPPED), + (queued, ExecutionStatus.STOPPED), + ): + row.refresh_from_db() + self.assertEqual(row.status, expected.value) + done.refresh_from_db() + self.assertIsNone(done.execution_error) + + def test_abort_is_audited(self) -> None: + deployment = self.deployment() + execution = self.execution(deployment) + + self._abort(self.owner, execution.id) + + event = PlatformAuditEvent.objects.get(action=Action.EXECUTION_ABORTED) + self.assertEqual(event.organization, self.org) + self.assertEqual(event.actor, self.owner) + self.assertEqual(event.target_type, TargetType.EXECUTION) + self.assertEqual(event.target_id, str(execution.id)) + self.assertEqual(event.target_name, self.workflow.workflow_name) + self.assertEqual(event.summary, "was EXECUTING") + self.assertEqual( + event.metadata, + { + "previous_status": "EXECUTING", + "pipeline_id": str(deployment.id), + "entity": "API", + }, + ) + + def test_finished_run_is_already_finished_and_not_audited(self) -> None: + execution = self.execution(status=ExecutionStatus.COMPLETED) + response = self._abort(self.owner, execution.id) + self.assertEqual( + self._outcomes(response), + {str(execution.id): ("already_finished", "COMPLETED")}, + ) + self.assertEqual(self._status(execution), ExecutionStatus.COMPLETED.value) + self.assertFalse(PlatformAuditEvent.objects.exists()) + + def test_mixed_batch_answers_each_id_in_one_200(self) -> None: + mine = self.execution(self.pipeline()) + finished = self.execution(status=ExecutionStatus.ERROR) + other = _make_user("other@example.com") + foreign_workflow_run = self._run_owned_by(other) + unknown = uuid.uuid4() + + response = self._abort( + self.owner, mine.id, finished.id, foreign_workflow_run.id, unknown + ) + + self.assertEqual( + self._outcomes(response), + { + str(mine.id): ("aborted", "STOPPED"), + str(finished.id): ("already_finished", "ERROR"), + str(foreign_workflow_run.id): ("not_permitted", None), + str(unknown): ("not_permitted", None), + }, + ) + self.assertEqual( + [r["id"] for r in response.data["results"]], + [str(mine.id), str(finished.id), str(foreign_workflow_run.id), str(unknown)], + ) + self.assertEqual( + ( + response.data["aborted"], + response.data["already_finished"], + response.data["not_permitted"], + ), + (1, 1, 2), + ) + self.assertEqual( + self._status(foreign_workflow_run), ExecutionStatus.EXECUTING.value + ) + + def _run_owned_by(self, user: Any) -> WorkflowExecution: + """A run of a deployment ``user`` owns, which ``self.owner`` only sees.""" + deployment = self.deployment() + deployment.memberships.filter(user=self.owner).delete() + deployment.memberships.create(user=user, role=ResourceRole.OWNER) + _add_viewers(deployment, self.owner) + return self.execution(deployment) + + def test_shared_viewer_is_not_permitted(self) -> None: + deployment = self.deployment() + _add_viewers(deployment, self.outsider) + execution = self.execution(deployment) + # Control: the viewer can see the run. + self.assertTrue( + WorkflowExecution.objects.for_user(self.outsider) + .filter(pk=execution.pk) + .exists() + ) + + response = self._abort(self.outsider, execution.id) + + self.assertEqual( + self._outcomes(response), {str(execution.id): ("not_permitted", None)} + ) + self.assertEqual(self._status(execution), ExecutionStatus.EXECUTING.value) + + def test_workflow_owner_cannot_abort_a_pipeline_run_they_do_not_own(self) -> None: + pipeline = self.pipeline() + pipeline.memberships.filter(user=self.owner).delete() + Pipeline.objects.filter(pk=pipeline.pk).update(shared_to_org=True) + execution = self.execution(pipeline) + outcome = self._outcomes(self._abort(self.owner, execution.id)) + self.assertEqual(outcome[str(execution.id)][0], "not_permitted") + + def test_co_owner_of_the_workflow_aborts_a_workflow_run(self) -> None: + self.workflow.memberships.create(user=self.member, role=ResourceRole.OWNER) + execution = self.execution(status=ExecutionStatus.PENDING) + outcome = self._outcomes(self._abort(self.member, execution.id)) + self.assertEqual(outcome[str(execution.id)], ("aborted", "STOPPED")) + + def test_org_admin_can_abort_any_run(self) -> None: + execution = self.execution(self.deployment()) + outcome = self._outcomes(self._abort(self.admin, execution.id)) + self.assertEqual(outcome[str(execution.id)], ("aborted", "STOPPED")) + event = PlatformAuditEvent.objects.get(action=Action.EXECUTION_ABORTED) + self.assertEqual(event.actor, self.admin) + + def test_service_account_can_abort(self) -> None: + service = _make_user("svc@example.com", is_service_account=True) + execution = self.execution(self.pipeline()) + outcome = self._outcomes(self._abort(service, execution.id)) + self.assertEqual(outcome[str(execution.id)], ("aborted", "STOPPED")) + + def test_duplicate_ids_are_answered_once(self) -> None: + execution = self.execution() + response = self._abort(self.owner, execution.id, execution.id) + self.assertEqual(len(response.data["results"]), 1) + self.assertEqual(PlatformAuditEvent.objects.count(), 1) + + def test_invalid_bodies_are_400(self) -> None: + too_many = [str(uuid.uuid4()) for _ in range(MAX_ABORT_IDS + 1)] + bodies: list[dict[str, Any]] = [ + {}, + {"ids": []}, + {"ids": ["not-a-uuid"]}, + {"ids": too_many}, + ] + for body in bodies: + with self.subTest(body=body): + response = self._abort(self.owner, raw=body) + self.assertEqual(response.status_code, 400) + + def test_requires_authentication(self) -> None: + request = self.factory.post("/", {"ids": [str(uuid.uuid4())]}, format="json") + response = _abort(request) + self.assertIn(response.status_code, (401, 403)) + + def test_later_completion_cannot_undo_the_abort(self) -> None: + execution = self.execution(self.pipeline()) + self._abort(self.owner, execution.id) + + view = WorkflowExecutionInternalViewSet() + request = MagicMock() + request.data = {"status": ExecutionStatus.COMPLETED.value} + fresh = WorkflowExecution.objects.get(pk=execution.pk) + with patch.object(view, "get_object", return_value=fresh): + out = view.update_status(request, id=str(execution.id)) + + self.assertEqual(out.data.get("reason"), "already_final") + self.assertEqual(self._status(execution), ExecutionStatus.STOPPED.value) diff --git a/backend/workflow_manager/execution/tests/test_execution_contract.py b/backend/workflow_manager/execution/tests/test_execution_contract.py new file mode 100644 index 0000000000..d6baccdc76 --- /dev/null +++ b/backend/workflow_manager/execution/tests/test_execution_contract.py @@ -0,0 +1,172 @@ +"""The stuck/abort additions leave the existing execution endpoints unchanged, +and their published contract generates (UN-4318).""" + +import uuid +from typing import Any + +from django.test import SimpleTestCase +from django.urls import resolve, reverse +from drf_spectacular.generators import SchemaGenerator +from pipeline_v2.models import Pipeline +from rest_framework.test import APIRequestFactory + +from workflow_manager.execution import urls +from workflow_manager.execution.tests.stuck_abort_base import StuckAbortTestBase +from workflow_manager.execution.views import ExecutionViewSet +from workflow_manager.internal_views import WorkflowExecutionInternalViewSet +from workflow_manager.workflow_v2.enums import ExecutionStatus + +_URLCONF = "workflow_manager.execution.urls" +_list = ExecutionViewSet.as_view({"get": "list"}) +_detail = ExecutionViewSet.as_view({"get": "retrieve"}) +_status = WorkflowExecutionInternalViewSet.as_view({"get": "execution_status"}) + +# The default row, pinned: a new key here is a contract change. +_EXECUTION_KEYS = { + "aggregated_total_pages_processed", + "attempts", + "created_at", + "dispatched_at", + "error_message", + "execution_method", + "execution_mode", + "execution_time", + "failed_files", + "id", + "modified_at", + "pipeline_id", + "pipeline_name", + "queue_message_id", + "result_acknowledged", + "status", + "successful_files", + "tags", + "total_files", + "workflow", + "workflow_name", +} + + +class ExistingResponsesUnchangedTests(StuckAbortTestBase): + def test_list_and_detail_keep_their_keys(self) -> None: + execution = self.execution() + listed = self.call(_list, self.owner) + self.assertEqual(set(listed.data["results"][0]), _EXECUTION_KEYS) + detail = self.call(_detail, self.owner, pk=str(execution.id)) + self.assertEqual(set(detail.data), _EXECUTION_KEYS) + + def test_stuck_false_is_the_same_response_as_no_parameter(self) -> None: + self.execution() + self.execution(status=ExecutionStatus.COMPLETED) + plain = self.call(_list, self.owner) + for value in ("false", "0", ""): + with self.subTest(stuck=value): + filtered = self.call(_list, self.owner, data={"stuck": value}) + self.assertEqual(filtered.content, plain.content) + + def test_entity_filter_still_selects_each_entity(self) -> None: + runs = { + "API": self.execution(self.deployment()), + "ETL": self.execution(self.pipeline(Pipeline.PipelineType.ETL)), + "TASK": self.execution(self.pipeline(Pipeline.PipelineType.TASK)), + "WF": self.execution(), + } + for entity, run in runs.items(): + with self.subTest(entity=entity): + response = self.call( + _list, self.owner, data={"execution_entity": entity} + ) + ids = [row["id"] for row in response.data["results"]] + self.assertEqual(ids, [str(run.id)]) + + def test_stuck_rows_add_only_stuck_since(self) -> None: + self.execution() + listed = self.call(_list, self.owner, data={"stuck": "true"}) + self.assertEqual( + set(listed.data["results"][0]), _EXECUTION_KEYS | {"stuck_since"} + ) + + +class InternalStatusActionTests(StuckAbortTestBase): + def _get(self, execution_id: Any, organization_id: str | None) -> Any: + request = APIRequestFactory().get("/") + request.organization_id = organization_id + return _status(request, id=str(execution_id)) + + def test_returns_the_status(self) -> None: + execution = self.execution(status=ExecutionStatus.STOPPED) + response = self._get(execution.id, self.org.organization_id) + self.assertEqual(response.status_code, 200) + self.assertEqual( + response.data, {"id": str(execution.id), "status": "STOPPED"} + ) + + def test_other_org_unknown_and_malformed_ids_are_404(self) -> None: + execution = self.execution() + cases = ( + (execution.id, "some-other-org"), + (execution.id, None), + (uuid.uuid4(), self.org.organization_id), + ("not-a-uuid", self.org.organization_id), + ) + for execution_id, organization_id in cases: + with self.subTest(execution_id=execution_id, org=organization_id): + response = self._get(execution_id, organization_id) + self.assertEqual(response.status_code, 404) + + +class RoutesUnchangedTests(SimpleTestCase): + def test_existing_names_reverse_to_the_same_paths(self) -> None: + pk = uuid.uuid4() + expected = { + ("execution-list", None): "/", + ("execution-detail", pk): f"/{pk}/", + ("execution-log", pk): f"/{pk}/logs/", + ("execution-log-export", pk): f"/{pk}/logs/export/", + ("execution-stuck-count", None): "/stuck-count/", + ("execution-abort", None): "/abort/", + } + for (name, key), path in expected.items(): + kwargs = {"pk": key} if key else None + with self.subTest(name=name): + self.assertEqual(reverse(name, urlconf=_URLCONF, kwargs=kwargs), path) + + def test_new_paths_do_not_shadow_detail(self) -> None: + self.assertEqual( + resolve("/stuck-count/", urlconf=_URLCONF).url_name, + "execution-stuck-count", + ) + self.assertEqual( + resolve(f"/{uuid.uuid4()}/", urlconf=_URLCONF).url_name, + "execution-detail", + ) + + +class SchemaContractTests(SimpleTestCase): + """The ``@extend_schema`` annotations and README.md are the contract.""" + + def test_schema_documents_the_new_operations(self) -> None: + schema = SchemaGenerator(patterns=urls.urlpatterns).get_schema( + request=None, public=True + ) + paths = schema["paths"] + count = paths["/stuck-count/"]["get"] + self.assertEqual(count["operationId"], "execution_stuck_count") + self.assertTrue({"200", "403"} <= set(count["responses"])) + + abort = paths["/abort/"]["post"] + self.assertEqual(abort["operationId"], "execution_abort") + self.assertTrue({"200", "400", "403"} <= set(abort["responses"])) + components = schema["components"]["schemas"] + self.assertIn("ids", components["ExecutionAbortRequest"]["properties"]) + self.assertEqual( + set(components["ExecutionAbortResponse"]["properties"]), + {"results", "aborted", "already_finished", "not_permitted"}, + ) + self.assertEqual( + set(components["StuckExecutionCount"]["properties"]), + {"threshold_seconds", "total", "counts"}, + ) + + list_params = {p["name"] for p in paths["/"]["get"]["parameters"]} + self.assertIn("stuck", list_params) diff --git a/backend/workflow_manager/execution/tests/test_stuck_executions.py b/backend/workflow_manager/execution/tests/test_stuck_executions.py new file mode 100644 index 0000000000..3ddfadc0fd --- /dev/null +++ b/backend/workflow_manager/execution/tests/test_stuck_executions.py @@ -0,0 +1,132 @@ +"""Stuck executions: the ``stuck=true`` list filter and ``stuck-count/`` (UN-4318).""" + +from datetime import timedelta +from typing import Any + +from django.test import override_settings +from django.utils.dateparse import parse_datetime +from pipeline_v2.models import Pipeline + +from workflow_manager.execution.stuck import count_stuck +from workflow_manager.execution.tests.stuck_abort_base import ( + OLD, + RECENT, + StuckAbortTestBase, +) +from workflow_manager.execution.views import ExecutionViewSet, StuckExecutionCountView +from workflow_manager.workflow_v2.enums import ExecutionStatus +from workflow_manager.workflow_v2.models.execution import WorkflowExecution + +_list = ExecutionViewSet.as_view({"get": "list"}) +_count = StuckExecutionCountView.as_view() + + +class StuckExecutionTests(StuckAbortTestBase): + def _stuck_rows(self, user: Any) -> dict[str, Any]: + response = self.call(_list, user, data={"stuck": "true"}) + self.assertEqual(response.status_code, 200) + return {row["id"]: row for row in response.data["results"]} + + def _count(self, user: Any) -> dict[str, Any]: + response = self.call(_count, user) + self.assertEqual(response.status_code, 200) + data: dict[str, Any] = response.data + return data + + def test_stuck_run_is_listed_with_stuck_since_and_counted(self) -> None: + execution = self.execution() + latest = self.file( + execution, ExecutionStatus.EXECUTING, idle=OLD - timedelta(minutes=10) + ) + self.file(execution, ExecutionStatus.COMPLETED, idle=OLD) + + rows = self._stuck_rows(self.owner) + + self.assertEqual(list(rows), [str(execution.id)]) + stuck_since = parse_datetime(rows[str(execution.id)]["stuck_since"]) + self.assertEqual(stuck_since, latest.modified_at) + self.assertEqual( + self._count(self.owner), + { + "threshold_seconds": 9000, + "total": 1, + "counts": {"API": 0, "ETL": 0, "TASK": 0, "WF": 1}, + }, + ) + + def test_run_without_files_is_stuck_since_its_own_last_write(self) -> None: + execution = self.execution(status=ExecutionStatus.PENDING) + rows = self._stuck_rows(self.owner) + self.assertEqual( + parse_datetime(rows[str(execution.id)]["stuck_since"]), + execution.modified_at, + ) + + def test_recent_file_activity_is_not_stuck(self) -> None: + execution = self.execution() + self.file(execution, ExecutionStatus.EXECUTING, idle=RECENT) + self.assertEqual(self._stuck_rows(self.owner), {}) + self.assertEqual(self._count(self.owner)["total"], 0) + + def test_recently_updated_run_is_not_stuck(self) -> None: + self.execution(idle=RECENT) + self.assertEqual(self._stuck_rows(self.owner), {}) + + def test_finished_runs_are_not_stuck(self) -> None: + for status in ( + ExecutionStatus.COMPLETED, + ExecutionStatus.ERROR, + ExecutionStatus.STOPPED, + ): + self.execution(status=status) + self.assertEqual(self._stuck_rows(self.owner), {}) + self.assertEqual(self._count(self.owner)["total"], 0) + + def test_runs_the_caller_cannot_see_are_excluded(self) -> None: + self.execution(self.deployment()) + self.assertEqual(self._stuck_rows(self.outsider), {}) + self.assertEqual(self._count(self.outsider)["total"], 0) + # Control: the owner does see it. + self.assertEqual(self._count(self.owner)["total"], 1) + + def test_counts_split_by_entity(self) -> None: + self.execution(self.deployment()) + self.execution(self.pipeline(Pipeline.PipelineType.ETL)) + self.execution(self.pipeline(Pipeline.PipelineType.ETL)) + self.execution(self.pipeline(Pipeline.PipelineType.TASK)) + self.execution() + + report = self._count(self.owner) + + self.assertEqual(report["total"], 5) + self.assertEqual(report["counts"], {"API": 1, "ETL": 2, "TASK": 1, "WF": 1}) + + def test_count_is_one_query(self) -> None: + self.execution(self.deployment()) + self.execution() + with self.assertNumQueries(1): + report = count_stuck(WorkflowExecution.objects.all()) + self.assertEqual(report.total, 2) + self.assertEqual(report.counts["API"], 1) + + def test_stuck_combines_with_the_entity_filter(self) -> None: + api_run = self.execution(self.deployment()) + self.execution() + response = self.call( + _list, self.owner, data={"stuck": "true", "execution_entity": "API"} + ) + ids = [row["id"] for row in response.data["results"]] + self.assertEqual(ids, [str(api_run.id)]) + + @override_settings(STUCK_EXECUTION_THRESHOLD_SECONDS=30) + def test_threshold_comes_from_settings(self) -> None: + execution = self.execution(idle=RECENT) + self.assertIn(str(execution.id), self._stuck_rows(self.owner)) + report = self._count(self.owner) + self.assertEqual(report["threshold_seconds"], 30) + self.assertEqual(report["total"], 1) + + def test_count_requires_authentication(self) -> None: + request = self.factory.get("/") + response = _count(request) + self.assertIn(response.status_code, (401, 403)) diff --git a/backend/workflow_manager/execution/urls.py b/backend/workflow_manager/execution/urls.py index c5bf5cc307..1684fded05 100644 --- a/backend/workflow_manager/execution/urls.py +++ b/backend/workflow_manager/execution/urls.py @@ -1,7 +1,11 @@ from django.urls import path from rest_framework.urlpatterns import format_suffix_patterns -from workflow_manager.execution.views import ExecutionViewSet +from workflow_manager.execution.views import ( + ExecutionAbortView, + ExecutionViewSet, + StuckExecutionCountView, +) from workflow_manager.workflow_v2.execution_log_view import ( WorkflowExecutionLogViewSet as ExecutionLogViewSet, ) @@ -25,5 +29,11 @@ execution_log_export, name="execution-log-export", ), + path( + "stuck-count/", + StuckExecutionCountView.as_view(), + name="execution-stuck-count", + ), + path("abort/", ExecutionAbortView.as_view(), name="execution-abort"), ] ) diff --git a/backend/workflow_manager/execution/views/__init__.py b/backend/workflow_manager/execution/views/__init__.py index 90a34404ac..e59f7e7297 100644 --- a/backend/workflow_manager/execution/views/__init__.py +++ b/backend/workflow_manager/execution/views/__init__.py @@ -1 +1,2 @@ from .execution import ExecutionViewSet # noqa: F401 +from .stuck_abort import ExecutionAbortView, StuckExecutionCountView # noqa: F401 diff --git a/backend/workflow_manager/execution/views/execution.py b/backend/workflow_manager/execution/views/execution.py index a80a1e354f..369ceeb590 100644 --- a/backend/workflow_manager/execution/views/execution.py +++ b/backend/workflow_manager/execution/views/execution.py @@ -1,13 +1,20 @@ import logging from django_filters.rest_framework import DjangoFilterBackend +from drf_spectacular.types import OpenApiTypes +from drf_spectacular.utils import OpenApiParameter, extend_schema, extend_schema_view from rest_framework import viewsets from rest_framework.permissions import BasePermission, IsAuthenticated from utils.filters.ordering_filter import DeterministicOrderingFilter from utils.pagination import CustomPagination -from workflow_manager.execution.filter import ExecutionFilter +from workflow_manager.execution.filter import ( + STUCK_PARAM, + ExecutionFilter, + stuck_requested, +) from workflow_manager.execution.serializer import ExecutionSerializer +from workflow_manager.execution.serializer.stuck_abort import StuckExecutionSerializer from workflow_manager.workflow_v2.models import WorkflowExecution logger = logging.getLogger(__name__) @@ -30,6 +37,21 @@ def has_object_permission(self, request, view, obj): return WorkflowExecution.objects.for_user(request.user).filter(id=obj.id).exists() +@extend_schema_view( + list=extend_schema( + parameters=[ + OpenApiParameter( + STUCK_PARAM, + OpenApiTypes.BOOL, + description=( + "true: only PENDING/EXECUTING runs with no activity for longer " + "than the stuck threshold; each row then also carries " + "stuck_since. See execution/README.md." + ), + ) + ] + ) +) class ExecutionViewSet(viewsets.ReadOnlyModelViewSet): permission_classes = [IsAuthenticated, UserWorkflowExecutionPermission] serializer_class = ExecutionSerializer @@ -42,3 +64,10 @@ class ExecutionViewSet(viewsets.ReadOnlyModelViewSet): def get_queryset(self): # Use the custom manager method to filter executions for the current user return WorkflowExecution.objects.for_user(self.request.user) + + def get_serializer_class(self) -> type[ExecutionSerializer]: + # Only ?stuck=true rows carry the stuck_since annotation. + request = getattr(self, "request", None) + if request is not None and stuck_requested(request.query_params): + return StuckExecutionSerializer + return ExecutionSerializer diff --git a/backend/workflow_manager/execution/views/stuck_abort.py b/backend/workflow_manager/execution/views/stuck_abort.py new file mode 100644 index 0000000000..d8ac769ae1 --- /dev/null +++ b/backend/workflow_manager/execution/views/stuck_abort.py @@ -0,0 +1,71 @@ +"""Stuck-execution count and abort. See ../README.md for the contract.""" + +from drf_spectacular.utils import OpenApiExample, extend_schema +from drf_standardized_errors.openapi_serializers import ( + ErrorResponse403Serializer, + ValidationErrorResponseSerializer, +) +from rest_framework.permissions import IsAuthenticated +from rest_framework.request import Request +from rest_framework.response import Response +from rest_framework.views import APIView + +from workflow_manager.execution.abort import abort_executions +from workflow_manager.execution.serializer.stuck_abort import ( + ExecutionAbortRequestSerializer, + ExecutionAbortResponseSerializer, + StuckExecutionCountSerializer, +) +from workflow_manager.execution.stuck import count_stuck +from workflow_manager.workflow_v2.models import WorkflowExecution + +_TAGS = ["execution"] + + +class StuckExecutionCountView(APIView): + permission_classes = [IsAuthenticated] + + @extend_schema( + operation_id="execution_stuck_count", + tags=_TAGS, + description="Stuck executions visible to the caller, per entity.", + responses={200: StuckExecutionCountSerializer, 403: ErrorResponse403Serializer}, + examples=[ + OpenApiExample( + "Counts", + value={ + "threshold_seconds": 9000, + "total": 3, + "counts": {"API": 1, "ETL": 2, "TASK": 0, "WF": 0}, + }, + response_only=True, + ) + ], + ) + def get(self, request: Request) -> Response: + report = count_stuck(WorkflowExecution.objects.for_user(request.user)) + return Response(StuckExecutionCountSerializer(report).data) + + +class ExecutionAbortView(APIView): + permission_classes = [IsAuthenticated] + + @extend_schema( + operation_id="execution_abort", + tags=_TAGS, + description=( + "Stop running executions. Only the owner or an org admin may abort; " + "every id gets its own outcome in a 200." + ), + request=ExecutionAbortRequestSerializer, + responses={ + 200: ExecutionAbortResponseSerializer, + 400: ValidationErrorResponseSerializer, + 403: ErrorResponse403Serializer, + }, + ) + def post(self, request: Request) -> Response: + body = ExecutionAbortRequestSerializer(data=request.data) + body.is_valid(raise_exception=True) + report = abort_executions(request, body.validated_data["ids"]) + return Response(ExecutionAbortResponseSerializer(report).data) diff --git a/backend/workflow_manager/internal_views.py b/backend/workflow_manager/internal_views.py index bbb7259d69..eba309bf5e 100644 --- a/backend/workflow_manager/internal_views.py +++ b/backend/workflow_manager/internal_views.py @@ -10,6 +10,7 @@ from django.shortcuts import get_object_or_404 from rest_framework import status, viewsets from rest_framework.decorators import action +from rest_framework.request import Request from rest_framework.response import Response from rest_framework.views import APIView from tool_instance_v2.models import ToolInstance @@ -515,6 +516,27 @@ def _get_organization_context(self, execution: WorkflowExecution) -> dict: "settings": {}, } + # Named execution_status because `status` is the imported DRF module. + @action(detail=True, methods=["get"], url_path="status") + def execution_status(self, request: Request, id: str | None = None) -> Response: + """Lightweight status read; workers poll it between files.""" + try: + execution_id = uuid.UUID(str(id)) + except ValueError: + execution_id = None + executions = filter_queryset_by_organization( + WorkflowExecution.objects.filter(id=execution_id), + request, + "workflow__organization", + ) + execution_status = executions.values_list("status", flat=True).first() + if execution_status is None: + return Response( + {"error": "WorkflowExecution not found"}, + status=status.HTTP_404_NOT_FOUND, + ) + return Response({"id": str(execution_id), "status": execution_status}) + @action(detail=True, methods=["post"]) def update_status(self, request, id=None): """Update workflow execution status.""" diff --git a/workers/file_processing/tasks.py b/workers/file_processing/tasks.py index 764953834b..1cf6b78a04 100644 --- a/workers/file_processing/tasks.py +++ b/workers/file_processing/tasks.py @@ -675,6 +675,67 @@ def _refactored_pre_create_file_executions( return context +def _file_item_name(file_item: Any) -> str | None: + """File name from a batch file item in any of its formats.""" + name: str | None = None + if isinstance(file_item, (list, tuple)) and len(file_item) == 2: + name = file_item[0] + elif isinstance(file_item, dict): + name = file_item.get("file_name") + return name + + +def _is_execution_stopped(api_client: InternalAPIClient, execution_id: str) -> bool: + """True when the execution was stopped (aborted) while this batch runs. + + Fails open: if the check itself fails, the next file is processed. + """ + try: + status = api_client.get_workflow_execution_status(execution_id) + except Exception as error: + logger.warning( + "[exec:%s] status check before the next file failed (%s); " + "processing it anyway", + execution_id, + error, + ) + return False + stopped: bool = status == ExecutionStatus.STOPPED.value + return stopped + + +def _mark_file_stopped(api_client: InternalAPIClient, file_execution_id: str) -> None: + """Mark one unprocessed file row STOPPED; a failed write is logged, not raised.""" + try: + api_client.update_file_execution_status( + file_execution_id=file_execution_id, + status=ExecutionStatus.STOPPED.value, + error_message="Not processed: execution was stopped", + ) + except Exception as error: + logger.warning( + "Failed to mark file execution %s STOPPED: %s", file_execution_id, error + ) + + +def _stop_remaining_files( + context: WorkflowContextData, remaining: list[Any], result: FileBatchResult +) -> None: + """Count the batch's unprocessed files as failed and mark their rows STOPPED.""" + api_client = context.organization_context.api_client + for file_item in remaining: + result.increment_failure() + name = _file_item_name(file_item) + pre_created = context.pre_created_file_executions.get(name) if name else None + if pre_created: + _mark_file_stopped(api_client, pre_created.id) + logger.warning( + "[exec:%s] Execution was stopped; skipped the batch's %s remaining file(s)", + context.execution_id, + len(remaining), + ) + + def _process_individual_files(context: WorkflowContextData) -> WorkflowContextData: """Process each file individually through the workflow. @@ -703,6 +764,11 @@ def _process_individual_files(context: WorkflowContextData) -> WorkflowContextDa # Process each file - handle list, tuple, and dictionary formats for file_number, file_item in enumerate(files, 1): + # Stop between files once the run is aborted. + if _is_execution_stopped(api_client, context.execution_id): + _stop_remaining_files(context, files[file_number - 1 :], result) + break + # Handle Django list format (from asdict serialization), tuple format, and dictionary format if isinstance(file_item, list): # Django backend format after asdict(): [file_name, file_hash_dict] diff --git a/workers/shared/api/internal_client.py b/workers/shared/api/internal_client.py index e832c609d0..6f6d186990 100644 --- a/workers/shared/api/internal_client.py +++ b/workers/shared/api/internal_client.py @@ -497,6 +497,14 @@ def get_workflow_execution( response_dict = response return WorkflowExecutionResponse.from_api_response(response_dict) + def get_workflow_execution_status( + self, execution_id: str | uuid.UUID, organization_id: str | None = None + ) -> str | None: + """Current execution status, without the execution context. Raises on failure.""" + return self.execution_client.get_workflow_execution_status( + execution_id, organization_id + ) + @with_cache( CacheType.WORKFLOW, lambda self, workflow_id, organization_id=None: str(workflow_id), diff --git a/workers/shared/clients/execution_client.py b/workers/shared/clients/execution_client.py index ef98f8daec..d738b53e77 100644 --- a/workers/shared/clients/execution_client.py +++ b/workers/shared/clients/execution_client.py @@ -143,6 +143,19 @@ def get_workflow_execution( status_code=getattr(e, "status_code", None), ) + def get_workflow_execution_status( + self, execution_id: str | uuid.UUID, organization_id: str | None = None + ) -> str | None: + """Current status of an execution, from the lightweight status endpoint. + + Raises on a failed request; the caller decides how to fail. + """ + response = self.get( + self._build_url("workflow_execution", f"{execution_id}/status/"), + organization_id=organization_id, + ) + return response.get("status") if isinstance(response, dict) else None + def get_workflow_definition( self, workflow_id: str | uuid.UUID, organization_id: str | None = None ) -> ExecutionResponse: diff --git a/workers/tests/test_abort_between_files.py b/workers/tests/test_abort_between_files.py new file mode 100644 index 0000000000..4e40ad439f --- /dev/null +++ b/workers/tests/test_abort_between_files.py @@ -0,0 +1,216 @@ +"""A running batch stops between files once its execution is STOPPED (UN-4318). + +Pins: the status is checked before each file; only STOPPED (abort) stops the +batch, with the remaining files counted as failed and their rows marked +STOPPED; any other status, or a failed check, processes the file (fail open). +Also pins the client URL and the facade delegation. +""" + +from __future__ import annotations + +import types +from typing import Any +from unittest import mock + +import pytest +from unstract.core.data_models import ExecutionStatus, FileBatchResult + +from file_processing.tasks import _file_item_name, _process_individual_files + +_TASKS = "file_processing.tasks" + + +def _context(api_client: Any, n_files: int = 3) -> Any: + files = { + f"file_{i}": {"file_name": f"f{i}.pdf", "file_path": f"/in/f{i}.pdf"} + for i in range(n_files) + } + pre_created = { + f"f{i}.pdf": mock.Mock( + id=f"fe-{i}", + file_hash=mock.Mock(is_manualreview_required=False), + ) + for i in range(n_files) + } + return types.SimpleNamespace( + files=files, + execution_id="exec-1", + workflow_id="wf-1", + pre_created_file_executions=pre_created, + organization_context=types.SimpleNamespace(api_client=api_client), + get_setting=lambda key, default=None: default, + metadata={ + "file_data": mock.Mock(), + "workflow_execution": {}, + "result": FileBatchResult(), + "successful_files_for_manual_review": [], + "total_files": n_files, + "workflow_logger": None, + }, + ) + + +def _run(api_client: Any, n_files: int = 3) -> tuple[FileBatchResult, Any]: + """Run the loop with file processing stubbed to succeed.""" + context = _context(api_client, n_files) + + def _succeeded( + _result: Any, _name: Any, _start: Any, result: Any, *args: Any, **kwargs: Any + ) -> None: + result.increment_success() + + with ( + mock.patch(f"{_TASKS}._process_file") as process_file, + mock.patch( + f"{_TASKS}._handle_file_processing_result", side_effect=_succeeded + ), + mock.patch(f"{_TASKS}.log_file_processing_start"), + mock.patch(f"{_TASKS}.log_file_info"), + ): + _process_individual_files(context) + return context.metadata["result"], process_file + + +def _stopped_ids(api_client: Any) -> list[str]: + return [ + c.kwargs["file_execution_id"] + for c in api_client.update_file_execution_status.call_args_list + if c.kwargs["status"] == ExecutionStatus.STOPPED.value + ] + + +def test_stops_after_the_status_flips_to_stopped() -> None: + api_client = mock.Mock() + api_client.get_workflow_execution_status.side_effect = [ + ExecutionStatus.EXECUTING.value, + ExecutionStatus.STOPPED.value, + ] + + result, process_file = _run(api_client) + + assert process_file.call_count == 1 + assert api_client.get_workflow_execution_status.call_count == 2 + assert _stopped_ids(api_client) == ["fe-1", "fe-2"] + assert (result.successful_files, result.failed_files) == (1, 2) + + +def test_stopped_before_the_first_file_processes_nothing() -> None: + api_client = mock.Mock() + api_client.get_workflow_execution_status.return_value = ( + ExecutionStatus.STOPPED.value + ) + + result, process_file = _run(api_client) + + process_file.assert_not_called() + assert _stopped_ids(api_client) == ["fe-0", "fe-1", "fe-2"] + assert (result.successful_files, result.failed_files) == (0, 3) + + +def test_a_failed_stopped_write_does_not_stop_the_rest() -> None: + api_client = mock.Mock() + api_client.get_workflow_execution_status.return_value = ( + ExecutionStatus.STOPPED.value + ) + api_client.update_file_execution_status.side_effect = RuntimeError("500") + + result, process_file = _run(api_client) + + process_file.assert_not_called() + assert _stopped_ids(api_client) == ["fe-0", "fe-1", "fe-2"] + assert (result.successful_files, result.failed_files) == (0, 3) + + +def test_error_mid_batch_keeps_processing_the_files() -> None: + """Only an abort stops a running batch; ERROR keeps today's behaviour.""" + api_client = mock.Mock() + api_client.get_workflow_execution_status.side_effect = [ + ExecutionStatus.EXECUTING.value, + ExecutionStatus.ERROR.value, + ExecutionStatus.ERROR.value, + ] + + result, process_file = _run(api_client) + + assert process_file.call_count == 3 + assert _stopped_ids(api_client) == [] + assert (result.successful_files, result.failed_files) == (3, 0) + + +@pytest.mark.parametrize( + "status", + [ + ExecutionStatus.PENDING.value, + ExecutionStatus.EXECUTING.value, + ExecutionStatus.ERROR.value, + ExecutionStatus.COMPLETED.value, + None, + ], +) +def test_any_other_status_processes_every_file(status: str | None) -> None: + api_client = mock.Mock() + api_client.get_workflow_execution_status.return_value = status + + result, process_file = _run(api_client) + + assert process_file.call_count == 3 + assert _stopped_ids(api_client) == [] + assert (result.successful_files, result.failed_files) == (3, 0) + + +def test_a_failed_check_processes_the_file_anyway() -> None: + api_client = mock.Mock() + api_client.get_workflow_execution_status.side_effect = RuntimeError("404") + + result, process_file = _run(api_client) + + assert process_file.call_count == 3 + assert _stopped_ids(api_client) == [] + assert result.successful_files == 3 + + +@pytest.mark.parametrize( + ("item", "name"), + [ + (["a.pdf", {}], "a.pdf"), + (("b.pdf", {}), "b.pdf"), + ({"file_name": "c.pdf"}, "c.pdf"), + (["too", "many", "parts"], None), + (object(), None), + ], +) +def test_file_item_name_reads_every_format(item: Any, name: str | None) -> None: + assert _file_item_name(item) == name + + +def test_client_reads_the_status_endpoint() -> None: + from shared.clients.execution_client import ExecutionAPIClient + + client = ExecutionAPIClient.__new__(ExecutionAPIClient) + body = {"id": "exec-1", "status": "STOPPED"} + with mock.patch.object(client, "get", return_value=body) as get: + assert client.get_workflow_execution_status("exec-1", "org-1") == "STOPPED" + get.assert_called_once_with( + "v1/workflow-execution/exec-1/status/", organization_id="org-1" + ) + + +def test_client_returns_none_for_a_non_dict_body() -> None: + from shared.clients.execution_client import ExecutionAPIClient + + client = ExecutionAPIClient.__new__(ExecutionAPIClient) + with mock.patch.object(client, "get", return_value=None): + assert client.get_workflow_execution_status("exec-1") is None + + +def test_facade_delegates_to_the_execution_client() -> None: + from shared.api.internal_client import InternalAPIClient + + client = InternalAPIClient.__new__(InternalAPIClient) + client.execution_client = mock.MagicMock() + client.execution_client.get_workflow_execution_status.return_value = "ERROR" + + assert client.get_workflow_execution_status("exec-1", "org-1") == "ERROR" + client.execution_client.get_workflow_execution_status.assert_called_once_with( + "exec-1", "org-1" + ) From 493fe0186a182bf0e30be9c74dcaa25f01dd2828 Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Thu, 5 Nov 2026 15:34:25 +0530 Subject: [PATCH 2/8] UN-4318 [FEAT] Finish aborted runs properly and fix review findings - Callback keeps STOPPED, so an aborted run is no longer reported as completed. - Abort sets the pipeline status, sends notifications and writes file counts, so a run with a dead worker is still reported. - Skipped batches of a stopped run count down and clear their active-file markers, so the callback cleans up and a re-run keeps every file. - A synchronous API call for an aborted run returns 422. - One id failing answers `failed` instead of a 500 for the whole request. - Add tests for org scoping, the row lock, per-id commits and the status route. - Correct the README and API schema, and document the known limits. Co-Authored-By: Claude Opus 5.5 --- backend/api_v2/api_deployment_views.py | 6 +- .../api_v2/tests/test_deployment_aborted.py | 61 ++++++ backend/platform_audit/README.md | 2 +- backend/workflow_manager/execution/README.md | 111 +++++++---- backend/workflow_manager/execution/abort.py | 49 ++++- backend/workflow_manager/execution/enum.py | 1 + .../execution/serializer/stuck_abort.py | 3 +- .../execution/tests/stuck_abort_base.py | 4 +- .../execution/tests/test_abort_executions.py | 176 ++++++++++++++++- .../tests/test_abort_transactions.py | 128 ++++++++++++ .../tests/test_execution_contract.py | 47 ++++- .../execution/tests/test_stuck_executions.py | 15 ++ backend/workflow_manager/execution/urls.py | 16 +- .../execution/views/stuck_abort.py | 23 +-- .../file_execution/serializers.py | 2 + workers/callback/tasks.py | 38 +++- workers/file_processing/tasks.py | 95 +++++++-- workers/queue_backend/pg_barrier.py | 20 +- .../infrastructure/logging/workflow_logger.py | 13 +- workers/tests/test_abort_between_files.py | 125 ++++++++++-- workers/tests/test_abort_finishes_the_run.py | 185 ++++++++++++++++++ 21 files changed, 989 insertions(+), 131 deletions(-) create mode 100644 backend/api_v2/tests/test_deployment_aborted.py create mode 100644 backend/workflow_manager/execution/tests/test_abort_transactions.py create mode 100644 workers/tests/test_abort_finishes_the_run.py diff --git a/backend/api_v2/api_deployment_views.py b/backend/api_v2/api_deployment_views.py index d1b9ac543f..7f18a79b9b 100644 --- a/backend/api_v2/api_deployment_views.py +++ b/backend/api_v2/api_deployment_views.py @@ -29,6 +29,7 @@ from utils.hubspot_notify import notify_hubspot_event from utils.pagination import CustomPagination from workflow_manager.workflow_v2.dto import ExecutionResponse +from workflow_manager.workflow_v2.enums import ExecutionStatus from workflow_manager.workflow_v2.models.execution import WorkflowExecution from api_v2.api_deployment_dto_registry import ApiDeploymentDTORegistry @@ -171,7 +172,10 @@ def post( # Determine response status based on execution result execution_status = response.get("execution_status", "") - has_error = response.get("error") or execution_status == "ERROR" + has_error = response.get("error") or execution_status in ( + ExecutionStatus.ERROR.value, + ExecutionStatus.STOPPED.value, + ) if has_error: # Check for tool not found in registry error - return 500 Internal Server Error diff --git a/backend/api_v2/tests/test_deployment_aborted.py b/backend/api_v2/tests/test_deployment_aborted.py new file mode 100644 index 0000000000..5cf0c0b122 --- /dev/null +++ b/backend/api_v2/tests/test_deployment_aborted.py @@ -0,0 +1,61 @@ +"""A synchronous API deployment call reports an aborted run as an error (UN-4318).""" + +from __future__ import annotations + +from unittest.mock import patch + +from account_v2.models import Organization +from django.core.files.uploadedfile import SimpleUploadedFile +from django.test import TestCase +from rest_framework.response import Response +from rest_framework.test import APIRequestFactory +from utils.user_context import UserContext +from workflow_manager.workflow_v2.models.workflow import Workflow + +from api_v2.api_deployment_views import DeploymentExecution +from api_v2.models import APIDeployment, APIKey + +ORG_ID = "org-aborted" + + +class SyncCallForAbortedRunTest(TestCase): + def setUp(self) -> None: + Organization.objects.create( + name=ORG_ID, display_name="Org", organization_id=ORG_ID + ) + UserContext.set_organization_identifier(ORG_ID) + workflow = Workflow.objects.create(workflow_name="wf", is_active=True) + self.api = APIDeployment.objects.create(api_name="api", workflow=workflow) + self.key = APIKey.objects.create(api=self.api) + + def _post_with_status(self, execution_status: str) -> Response: + request = APIRequestFactory().post( + f"/deployment/api/{ORG_ID}/api/", + {"files": SimpleUploadedFile("doc.txt", b"hello")}, + format="multipart", + HTTP_AUTHORIZATION=f"Bearer {self.key.api_key}", + ) + with ( + patch( + "api_v2.api_deployment_views.APIDeploymentRateLimiter.check_and_acquire", + return_value=(True, {}), + ), + patch( + "api_v2.api_deployment_views.DeploymentHelper.execute_workflow", + return_value={"execution_status": execution_status}, + ), + ): + return DeploymentExecution.as_view()( + request, org_name=ORG_ID, api_name="api" + ) + + def test_status_codes(self) -> None: + for execution_status, code in ( + ("STOPPED", 422), + ("ERROR", 422), + ("COMPLETED", 200), + ): + with self.subTest(execution_status=execution_status): + self.assertEqual( + self._post_with_status(execution_status).status_code, code + ) diff --git a/backend/platform_audit/README.md b/backend/platform_audit/README.md index 3030b01336..fbfd252244 100644 --- a/backend/platform_audit/README.md +++ b/backend/platform_audit/README.md @@ -197,7 +197,7 @@ Expected emitters outside this PR: | Task | Action | Notes | |---|---|---| | PB-T05 (UN-4360) | `key.expiry_set` | target per key type; cloud review keys may register their own `review_key.*` actions | -| PB-T08 (UN-4318) | `execution.aborted` | target `execution`; actor is the owner or admin who aborted | +| PB-T08 (UN-4318) | `execution.aborted` | target `execution` | ## Action vocabulary diff --git a/backend/workflow_manager/execution/README.md b/backend/workflow_manager/execution/README.md index 2515c59285..97c16cbdc1 100644 --- a/backend/workflow_manager/execution/README.md +++ b/backend/workflow_manager/execution/README.md @@ -1,9 +1,10 @@ # Executions: stuck runs and abort Finding executions that stopped making progress, and stopping running ones -(PB-T08, UN-4318). This README and the `@extend_schema` annotations in -`views/stuck_abort.py` and `views/execution.py` are the published contract. -Changes to anything below are contract changes. +(PB-T08, UN-4318). This README is the published contract, backed by the +`@extend_schema` annotations in `views/stuck_abort.py` and `views/execution.py`. +`stuck_since` is documented here only: the generated schema shows the plain list +row. Changes to anything below are contract changes. - [Endpoints](#endpoints) - [What counts as stuck](#what-counts-as-stuck) @@ -15,7 +16,7 @@ Changes to anything below are contract changes. ## Endpoints -Mounted under the organisation path: `/api/v1/unstract/{org_id}/execution/`. +Paths are relative to the organisation path: `/api/v1/unstract/{org_id}/`. | Method | Path | Returns | |---|---|---| @@ -25,7 +26,9 @@ Mounted under the organisation path: `/api/v1/unstract/{org_id}/execution/`. Everything here is additive. The existing list, detail and log endpoints keep their paths, parameters, response fields and permissions. Without `stuck`, or -with `stuck=false`, the list response is the same as before. +with `stuck=false`, the list response is the same as before. Aborted runs are the +one new case the existing endpoints see; see +[How the work stops](#how-the-work-stops). ## What counts as stuck @@ -87,7 +90,8 @@ is ignored, as if `stuck` were absent. Visibility is the same as the list's. ], "aborted": 1, "already_finished": 1, - "not_permitted": 1 + "not_permitted": 1, + "failed": 0 } ``` @@ -96,49 +100,69 @@ is ignored, as if `stuck` were absent. Visibility is the same as the list's. | `aborted` | The run was `PENDING` or `EXECUTING` and is now `STOPPED` | `STOPPED` | | `already_finished` | The run was already `COMPLETED`, `ERROR` or `STOPPED`; nothing changed | its status | | `not_permitted` | The caller may not abort it, or the id does not exist | `null` | +| `failed` | The abort hit an unexpected error for this id; retry it | `null` | Unknown ids and ids the caller cannot abort get the same answer, so the endpoint does not reveal which ids exist. **Who may abort.** The owner of what started the run: an OWNER of its API deployment or pipeline, or of its workflow when the run has no pipeline. Org -admins and service accounts (platform API keys) may abort any run in the -organisation. Shared users who are not owners get `not_permitted`. +admins and service accounts (platform API keys with write access) may abort any +run in the organisation. Shared users who are not owners get `not_permitted`. -**What abort does,** in one transaction per id: +**What abort does.** In one transaction per id: 1. Sets the execution to `STOPPED` with the error `Execution aborted`, which also frees an API deployment's rate-limit slot and sends the existing - "was stopped" inbox notice. -2. Sets its `PENDING` and `EXECUTING` file rows to `STOPPED`. Finished file - rows, file history, destination writes and usage are kept. -3. Records an `execution.aborted` event in the platform audit log, with - target `execution`, the workflow name, the summary `was ` - and `previous_status`, `pipeline_id` and `entity` in its metadata. - -The pipeline's last-run status is not changed. - -**How the work stops.** Queued batches of a stopped run are skipped when a -worker picks them up. A running batch checks the status before each file and -stops once it is `STOPPED`: its remaining files are marked `STOPPED` and -counted as failed. The file -being processed when the abort lands finishes, and its own result replaces -`STOPPED` on its row. Later status writes from workers cannot undo the abort on -runs using the PG queue. + "was stopped" inbox notice, unless the inbox item already shows a newer run. +2. Sets its `PENDING` and `EXECUTING` file rows to `STOPPED` with the same error, + and stamps their execution time. Finished file rows, file history, destination + writes and usage are kept. +3. Writes the run's file counts: finished files are successful, and every other + file is failed. + +Once that commits, it: + +1. Sets the pipeline's last-run status to failure and sends its notifications + with status `FAILURE`, or sends an API deployment's notifications with status + `STOPPED`. An inactive pipeline gets neither. +2. Records an `execution.aborted` event in the platform audit log, with target + `execution`, the workflow name, the summary `was ` and + `previous_status`, `pipeline_id` and `entity` in its metadata. This is best + effort: a lost event does not fail the abort and is not retried. + +### How the work stops + +- Queued batches of a stopped run are skipped when a worker picks them up. They + clear their files' active-file markers and still count down on the barrier. +- A running batch checks the status before each file and stops once it is + `STOPPED`: its remaining files are marked `STOPPED` and counted as failed. +- The file being processed when the abort lands finishes, and its own result + replaces `STOPPED` on its row. +- The last batch fires the aggregating callback. It keeps the status `STOPPED`, + saves the final file counts, commits usage for finished files, publishes the + final log and cleans up the execution directory. It does not send notifications + again. +- A caller waiting on a synchronous API deployment `POST` gets 422 for an aborted + run, the same as polling with `GET`. +- Later status writes from workers cannot undo the abort, unless the run's + `queue_message_id` is empty. ### Errors -Errors use the project's standard envelope (drf-standardized-errors). +400 uses the project's standard envelope (drf-standardized-errors). 401 and 403 +come from the auth middleware as `{"message": "..."}`. | Status | When | |---|---| | 400 | `ids` missing, empty, longer than 100, or holding a value that is not a UUID | -| 401 / 403 | Not authenticated | +| 401 | Not authenticated | +| 403 | Authenticated but refused: a user outside the organisation, or a platform API key without write access | -An unexpected failure part-way returns 500. Each id is aborted in its own -transaction, so the aborts made before it are kept (unless `ATOMIC_REQUESTS` is -on, which rolls the whole request back). Retrying is safe: ids already aborted -answer `already_finished`. +A failure on one id answers `failed` for that id, and the other ids still run. +Each id is aborted in its own transaction, so the others are kept (unless +`ATOMIC_REQUESTS` is on, which can roll the whole request back). Retrying is safe: +ids already aborted answer `already_finished`. ## Deploying @@ -147,14 +171,25 @@ internal `GET /internal/v1/workflow-execution//status/` endpoint. An older backend answers 404, and the worker then processes the file anyway (fail open), as it did before this change. +Expose abort to users only once the workers are updated too. An older worker's +callback reports an aborted run with a finished file as completed. + ## Limits -- The file being processed when the abort lands is not interrupted. A worker - that is alive but hung keeps its claim until the reaper cleans it up. -- Runs from before the PG queue (`queue_message_id` empty) are not protected - from a late worker status write. -- If the batch that was running was the last one, the aggregating callback - still runs. It can update the pipeline's last-run status and fire webhooks, - while the execution stays `STOPPED`. +- The file being processed when the abort lands is not interrupted. A hung worker + is covered in [`pg_queue/README.md`](../../../workers/queue_backend/pg_queue/README.md). +- **Re-running right after an abort can process a file twice.** The file still + being processed no longer counts as active, so a manual or scheduled re-run can + pick it up while the old worker finishes it: double LLM spend, and two writes to + a database destination. Wait for that file to finish before re-running. +- **A dead worker's run never gets its callback.** The pipeline status, + notifications and file counts are still set by the abort, but usage for its + finished files is not committed and its execution directory, which can hold + uploaded documents, is not cleaned up. It stays in storage until something else + deletes it. The reaper's `ERROR` path has the same gap. +- A run whose `queue_message_id` is empty (from before the PG queue, or whose + handle write failed) is not protected from a late worker status write. +- An abort that lands while the callback is already running stays `STOPPED`, but + the callback can still report the status it worked out. - A running batch stops between files only for `STOPPED`. A run that becomes `ERROR` or `COMPLETED` mid-batch keeps processing its files, as before. diff --git a/backend/workflow_manager/execution/abort.py b/backend/workflow_manager/execution/abort.py index c8dcf12d7d..365b9f0297 100644 --- a/backend/workflow_manager/execution/abort.py +++ b/backend/workflow_manager/execution/abort.py @@ -2,12 +2,14 @@ import logging from dataclasses import dataclass +from functools import partial from typing import Any from uuid import UUID from api_v2.models import APIDeployment from django.db import transaction -from django.db.models import Q, QuerySet +from django.db.models import Count, DurationField, ExpressionWrapper, F, Q, QuerySet +from django.db.models.functions import Extract, Now from permissions.roles import ResourceRole from pipeline_v2.models import Pipeline from platform_audit.actions import Action, TargetType @@ -18,6 +20,7 @@ from workflow_manager.execution.entity import entity_annotation from workflow_manager.execution.enum import AbortOutcome from workflow_manager.file_execution.models import WorkflowFileExecution +from workflow_manager.utils.pipeline_utils import PipelineUtils from workflow_manager.workflow_v2.enums import ExecutionStatus from workflow_manager.workflow_v2.models import Workflow, WorkflowExecution @@ -29,7 +32,7 @@ @dataclass(frozen=True) class AbortResult: - """Outcome for one id. ``status`` is ``None`` when not permitted.""" + """Outcome for one id. ``status`` is ``None`` when not permitted or failed.""" id: UUID outcome: AbortOutcome @@ -57,13 +60,17 @@ def already_finished(self) -> int: def not_permitted(self) -> int: return self._count(AbortOutcome.NOT_PERMITTED) + @property + def failed(self) -> int: + return self._count(AbortOutcome.FAILED) + def abort_executions(request: Any, ids: list[UUID]) -> AbortReport: """Abort each id the caller may abort; one transaction per id.""" unique_ids = list(dict.fromkeys(ids)) permitted = _abortable_entities(request.user, unique_ids) results = [ - _abort_one(request, execution_id, permitted[execution_id]) + _abort_or_fail(request, execution_id, permitted[execution_id]) if execution_id in permitted else AbortResult(execution_id, AbortOutcome.NOT_PERMITTED, None) for execution_id in unique_ids @@ -108,6 +115,15 @@ def owned(model: Any) -> QuerySet: ) +def _abort_or_fail(request: Any, execution_id: UUID, entity: str | None) -> AbortResult: + """Abort one id, so its failure cannot lose the other ids' outcomes.""" + try: + return _abort_one(request, execution_id, entity) + except Exception: + logger.exception("Failed to abort execution %s", execution_id) + return AbortResult(execution_id, AbortOutcome.FAILED, None) + + def _abort_one(request: Any, execution_id: UUID, entity: str | None) -> AbortResult: """Stop one execution under its row lock, then record the audit event.""" with transaction.atomic(): @@ -128,22 +144,47 @@ def _abort_one(request: Any, execution_id: UUID, entity: str | None) -> AbortRes status=ExecutionStatus.STOPPED, error=ABORT_ERROR_MESSAGE, row_locked=True ) _stop_open_files(execution) + _write_file_counts(execution) + # Pipeline status and webhook; no callback sends them if the worker is dead. + transaction.on_commit( + partial( + PipelineUtils.update_pipeline_status, execution.pipeline_id, execution + ), + robust=True, + ) _record_abort(request, execution, previous_status, entity) return AbortResult(execution_id, AbortOutcome.ABORTED, execution.status) def _stop_open_files(execution: WorkflowExecution) -> None: """Mark file rows not yet terminal STOPPED; finished rows keep their status.""" + elapsed = ExpressionWrapper(Now() - F("created_at"), output_field=DurationField()) stopped = ( WorkflowFileExecution.objects.filter(workflow_execution=execution) .exclude(status__in=ExecutionStatus.terminal_values()) - .update(status=ExecutionStatus.STOPPED.value, execution_error=ABORT_ERROR_MESSAGE) + .update( + status=ExecutionStatus.STOPPED.value, + execution_error=ABORT_ERROR_MESSAGE, + execution_time=Extract(elapsed, "epoch"), + ) ) logger.info( "Aborted execution %s; stopped %s open file execution(s)", execution.id, stopped ) +def _write_file_counts(execution: WorkflowExecution) -> None: + """Save the run's file counts; anything that did not finish counts as failed.""" + counts = WorkflowFileExecution.objects.filter(workflow_execution=execution).aggregate( + total=Count("pk"), + completed=Count("pk", filter=Q(status=ExecutionStatus.COMPLETED.value)), + ) + execution.total_files = max(execution.total_files, counts["total"]) + execution.successful_files = counts["completed"] + execution.failed_files = execution.total_files - counts["completed"] + execution.save(update_fields=["total_files", "successful_files", "failed_files"]) + + def _record_abort( request: Any, execution: WorkflowExecution, previous_status: str, entity: str | None ) -> None: diff --git a/backend/workflow_manager/execution/enum.py b/backend/workflow_manager/execution/enum.py index ba6f8fa1a7..48faa92b28 100644 --- a/backend/workflow_manager/execution/enum.py +++ b/backend/workflow_manager/execution/enum.py @@ -14,3 +14,4 @@ class AbortOutcome(StrEnum): ABORTED = "aborted" ALREADY_FINISHED = "already_finished" NOT_PERMITTED = "not_permitted" + FAILED = "failed" diff --git a/backend/workflow_manager/execution/serializer/stuck_abort.py b/backend/workflow_manager/execution/serializer/stuck_abort.py index c75cc25802..611c75729f 100644 --- a/backend/workflow_manager/execution/serializer/stuck_abort.py +++ b/backend/workflow_manager/execution/serializer/stuck_abort.py @@ -39,7 +39,7 @@ class ExecutionAbortResultSerializer(serializers.Serializer): id = serializers.UUIDField() outcome = serializers.ChoiceField(choices=[o.value for o in AbortOutcome]) status = serializers.CharField( - allow_null=True, help_text="The run's status; null when not_permitted." + allow_null=True, help_text="The run's status; null when not_permitted or failed." ) @@ -48,3 +48,4 @@ class ExecutionAbortResponseSerializer(serializers.Serializer): aborted = serializers.IntegerField() already_finished = serializers.IntegerField() not_permitted = serializers.IntegerField() + failed = serializers.IntegerField() diff --git a/backend/workflow_manager/execution/tests/stuck_abort_base.py b/backend/workflow_manager/execution/tests/stuck_abort_base.py index 72d2107f03..f36af73054 100644 --- a/backend/workflow_manager/execution/tests/stuck_abort_base.py +++ b/backend/workflow_manager/execution/tests/stuck_abort_base.py @@ -10,6 +10,7 @@ from unittest.mock import patch from api_v2.models import APIDeployment +from django.test import override_settings from django.utils import timezone from permissions.roles import ResourceRole from pipeline_v2.models import Pipeline @@ -24,11 +25,12 @@ "tenant_account_v2.organization_member_service." "OrganizationMemberService.is_user_organization_admin" ) -# Past the default 9000 s threshold, and well inside it. +# Past the 9000 s threshold set below, and well inside it. OLD = timedelta(hours=3) RECENT = timedelta(minutes=1) +@override_settings(STUCK_EXECUTION_THRESHOLD_SECONDS=9000) class StuckAbortTestBase(GroupSharingTestBase): def setUp(self) -> None: super().setUp() diff --git a/backend/workflow_manager/execution/tests/test_abort_executions.py b/backend/workflow_manager/execution/tests/test_abort_executions.py index 35b1cff59f..cb00fd3aef 100644 --- a/backend/workflow_manager/execution/tests/test_abort_executions.py +++ b/backend/workflow_manager/execution/tests/test_abort_executions.py @@ -1,9 +1,14 @@ """``POST execution/abort/`` (UN-4318): owner-or-admin, per-id outcomes, audit.""" import uuid +from datetime import timedelta from typing import Any from unittest.mock import MagicMock, patch +from account_v2.models import Organization +from django.db import connection +from django.test.utils import CaptureQueriesContext +from django.utils import timezone from permissions.roles import ResourceRole from pipeline_v2.models import Pipeline from platform_audit.actions import Action, TargetType @@ -11,11 +16,15 @@ from tenant_account_v2.tests import _add_viewers, _make_user from workflow_manager.execution.abort import ABORT_ERROR_MESSAGE, MAX_ABORT_IDS +from workflow_manager.execution.tests import test_update_status_locking as locking from workflow_manager.execution.tests.stuck_abort_base import RECENT, StuckAbortTestBase from workflow_manager.execution.views import ExecutionAbortView +from workflow_manager.file_execution.models import WorkflowFileExecution +from workflow_manager.file_execution.serializers import FileCentricExecutionSerializer from workflow_manager.internal_views import WorkflowExecutionInternalViewSet from workflow_manager.workflow_v2.enums import ExecutionStatus from workflow_manager.workflow_v2.models.execution import WorkflowExecution +from workflow_manager.workflow_v2.models.workflow import Workflow _abort = ExecutionAbortView.as_view() @@ -42,6 +51,8 @@ def test_abort_stops_the_run_and_its_open_files_only(self) -> None: done = self.file(execution, ExecutionStatus.COMPLETED, idle=RECENT) running = self.file(execution, ExecutionStatus.EXECUTING, idle=RECENT) queued = self.file(execution, ExecutionStatus.PENDING, idle=RECENT) + failed = self.file(execution, ExecutionStatus.ERROR, idle=RECENT) + WorkflowFileExecution.objects.filter(pk=failed.pk).update(execution_error="boom") response = self._abort(self.owner, execution.id) @@ -58,11 +69,146 @@ def test_abort_stops_the_run_and_its_open_files_only(self) -> None: (done, ExecutionStatus.COMPLETED), (running, ExecutionStatus.STOPPED), (queued, ExecutionStatus.STOPPED), + (failed, ExecutionStatus.ERROR), ): row.refresh_from_db() self.assertEqual(row.status, expected.value) - done.refresh_from_db() self.assertIsNone(done.execution_error) + self.assertEqual(failed.execution_error, "boom") + self.assertEqual(running.execution_error, ABORT_ERROR_MESSAGE) + self.assertEqual(queued.execution_error, ABORT_ERROR_MESSAGE) + + def test_stopped_file_row_shows_the_abort_and_a_fixed_time(self) -> None: + execution = self.execution(idle=RECENT) + row = self.file(execution, ExecutionStatus.EXECUTING, idle=RECENT) + started = timezone.now() - timedelta(hours=1) + WorkflowFileExecution.objects.filter(pk=row.pk).update(created_at=started) + + self._abort(self.owner, execution.id) + + row.refresh_from_db() + self.assertAlmostEqual(row.execution_time, 3600, delta=60) + self.assertEqual( + FileCentricExecutionSerializer().get_status_msg(row), ABORT_ERROR_MESSAGE + ) + + def test_abort_writes_the_file_counts_counting_the_unfinished_as_failed(self) -> None: + for known_total, total in ((5, 5), (0, 3)): + with self.subTest(known_total=known_total): + execution = self.execution(idle=RECENT) + WorkflowExecution.objects.filter(pk=execution.pk).update( + total_files=known_total + ) + self.file(execution, ExecutionStatus.COMPLETED, idle=RECENT) + self.file(execution, ExecutionStatus.ERROR, idle=RECENT) + self.file(execution, ExecutionStatus.EXECUTING, idle=RECENT) + + self._abort(self.owner, execution.id) + + execution.refresh_from_db() + self.assertEqual( + ( + execution.total_files, + execution.successful_files, + execution.failed_files, + ), + (total, 1, total - 1), + ) + + def test_abort_releases_the_slot_stamps_the_time_and_notifies_the_inbox(self) -> None: + execution = self.execution(self.deployment()) + started = timezone.now() - timedelta(hours=1) + WorkflowExecution.objects.filter(pk=execution.pk).update(created_at=started) + with ( + patch( + "api_v2.rate_limiter.APIDeploymentRateLimiter.release_slot" + ) as release, + patch("inbox.events.notify_execution_failed") as notify, + ): + self._abort(self.owner, execution.id) + release.assert_called_once_with(self.org.organization_id, str(execution.id)) + notify.assert_called_once_with(execution.id) + execution.refresh_from_db() + self.assertAlmostEqual(execution.execution_time, 3600, delta=60) + + def test_abort_marks_the_pipeline_failed_and_notifies(self) -> None: + pipeline = self.pipeline() + Pipeline.objects.filter(pk=pipeline.pk).update(active=True) + execution = self.execution(pipeline) + with patch( + "pipeline_v2.pipeline_processor.PipelineProcessor._send_notification" + ) as send: + outcome = self._outcomes(self._abort(self.owner, execution.id)) + pipeline.refresh_from_db() + self.assertEqual(outcome[str(execution.id)][0], "aborted") + self.assertEqual(pipeline.last_run_status, Pipeline.PipelineStatus.FAILURE) + send.assert_called_once() + self.assertEqual(send.call_args.kwargs["execution_id"], execution.id) + + def test_abort_of_a_paused_pipeline_run_still_aborts(self) -> None: + pipeline = self.pipeline() + execution = self.execution(pipeline) + with patch( + "pipeline_v2.pipeline_processor.PipelineProcessor._send_notification" + ) as send: + outcome = self._outcomes(self._abort(self.owner, execution.id)) + pipeline.refresh_from_db() + self.assertEqual(outcome[str(execution.id)][0], "aborted") + self.assertNotEqual(pipeline.last_run_status, Pipeline.PipelineStatus.FAILURE) + send.assert_not_called() + + def test_abort_of_a_deployment_run_notifies_the_deployment(self) -> None: + execution = self.execution(self.deployment()) + with patch( + "workflow_manager.utils.pipeline_utils.APIDeploymentUtils.send_notification" + ) as send: + self._abort(self.owner, execution.id) + send.assert_called_once() + self.assertEqual(send.call_args.kwargs["workflow_execution"].id, execution.id) + + def test_finished_run_announces_nothing(self) -> None: + execution = self.execution(self.pipeline(), status=ExecutionStatus.COMPLETED) + with patch( + "workflow_manager.execution.abort.PipelineUtils.update_pipeline_status" + ) as update: + self._abort(self.owner, execution.id) + update.assert_not_called() + + def test_abort_takes_one_row_lock(self) -> None: + execution = self.execution() + # The inbox notice takes its own lock on commit. + with ( + patch("inbox.events.notify_execution_failed"), + CaptureQueriesContext(connection) as queries, + ): + self._abort(self.owner, execution.id) + lock_statements: Any = locking._lock_statements + _, no_key, full = lock_statements(queries.captured_queries) + self.assertEqual(len(no_key), 1) + self.assertEqual(full, []) + + def test_aborting_twice_answers_already_finished(self) -> None: + execution = self.execution() + self._abort(self.owner, execution.id) + outcome = self._outcomes(self._abort(self.owner, execution.id)) + self.assertEqual(outcome[str(execution.id)], ("already_finished", "STOPPED")) + self.assertEqual(PlatformAuditEvent.objects.count(), 1) + + def test_one_failing_id_answers_failed_in_a_200(self) -> None: + bad, good = self.execution(), self.execution() + + def stop(execution: WorkflowExecution) -> None: + if execution.pk == bad.pk: + raise RuntimeError("boom") + + with patch("workflow_manager.execution.abort._stop_open_files", side_effect=stop): + response = self._abort(self.owner, bad.id, good.id) + + self.assertEqual( + self._outcomes(response), + {str(bad.id): ("failed", None), str(good.id): ("aborted", "STOPPED")}, + ) + self.assertEqual((response.data["aborted"], response.data["failed"]), (1, 1)) def test_abort_is_audited(self) -> None: deployment = self.deployment() @@ -140,6 +286,31 @@ def _run_owned_by(self, user: Any) -> WorkflowExecution: _add_viewers(deployment, self.owner) return self.execution(deployment) + def test_another_orgs_run_is_not_permitted_for_admin_or_service_account(self) -> None: + other_org = Organization.objects.create( + name="org-b", display_name="Org B", organization_id="org-b" + ) + other_workflow = Workflow.objects.create( + workflow_name="wf-b", organization=other_org, created_by=self.owner + ) + foreign = WorkflowExecution.objects.create( + workflow=other_workflow, + execution_mode=WorkflowExecution.Mode.QUEUE, + execution_method=WorkflowExecution.Method.DIRECT, + execution_type=WorkflowExecution.Type.COMPLETE, + status=ExecutionStatus.EXECUTING.value, + queue_message_id=1, + ) + service = _make_user("svc@example.com", is_service_account=True) + + for user in (self.admin, service): + with self.subTest(user=user.email): + outcome = self._outcomes(self._abort(user, foreign.id)) + self.assertEqual(outcome[str(foreign.id)], ("not_permitted", None)) + + self.assertEqual(self._status(foreign), ExecutionStatus.EXECUTING.value) + self.assertFalse(PlatformAuditEvent.objects.exists()) + def test_shared_viewer_is_not_permitted(self) -> None: deployment = self.deployment() _add_viewers(deployment, self.outsider) @@ -192,7 +363,8 @@ def test_duplicate_ids_are_answered_once(self) -> None: self.assertEqual(PlatformAuditEvent.objects.count(), 1) def test_invalid_bodies_are_400(self) -> None: - too_many = [str(uuid.uuid4()) for _ in range(MAX_ABORT_IDS + 1)] + self.assertEqual(MAX_ABORT_IDS, 100) + too_many = [str(uuid.uuid4()) for _ in range(101)] bodies: list[dict[str, Any]] = [ {}, {"ids": []}, diff --git a/backend/workflow_manager/execution/tests/test_abort_transactions.py b/backend/workflow_manager/execution/tests/test_abort_transactions.py new file mode 100644 index 0000000000..c9d619b220 --- /dev/null +++ b/backend/workflow_manager/execution/tests/test_abort_transactions.py @@ -0,0 +1,128 @@ +"""Abort on real transactions: the row lock and one commit per id (UN-4318).""" + +import threading +from typing import Any +from unittest.mock import patch + +from account_v2.models import Organization +from django.db import transaction +from django.test import TransactionTestCase +from permissions.roles import ResourceRole +from platform_audit.models import PlatformAuditEvent +from rest_framework.test import APIRequestFactory, force_authenticate +from tenant_account_v2.tests import _make_user +from utils.user_context import UserContext + +from workflow_manager.execution.tests import test_update_status_locking as locking +from workflow_manager.execution.views import ExecutionAbortView +from workflow_manager.workflow_v2.enums import ExecutionStatus +from workflow_manager.workflow_v2.models.execution import WorkflowExecution +from workflow_manager.workflow_v2.models.workflow import Workflow + +_ABORT = "workflow_manager.execution.abort" +_abort = ExecutionAbortView.as_view() + + +class AbortTransactionTests(TransactionTestCase): + def setUp(self) -> None: + self.org = Organization.objects.create( + name="org-a", display_name="Org A", organization_id="org-a" + ) + UserContext.set_organization_identifier(self.org.organization_id) + self.owner = _make_user("owner@example.com") + self.workflow = Workflow.objects.create( + workflow_name="wf-1", organization=self.org, created_by=self.owner + ) + self.workflow.memberships.create(user=self.owner, role=ResourceRole.OWNER) + + def _execution(self) -> WorkflowExecution: + execution: WorkflowExecution = WorkflowExecution.objects.create( + workflow=self.workflow, + execution_mode=WorkflowExecution.Mode.QUEUE, + execution_method=WorkflowExecution.Method.DIRECT, + execution_type=WorkflowExecution.Type.COMPLETE, + status=ExecutionStatus.EXECUTING.value, + queue_message_id=1, + ) + return execution + + def _post(self, *ids: Any) -> dict[str, tuple[str, str | None]]: + body = {"ids": [str(i) for i in ids]} + request = APIRequestFactory().post("/", body, format="json") + force_authenticate(request, user=self.owner) + response = _abort(request) + response.render() + self.assertEqual(response.status_code, 200, response.data) + return {r["id"]: (r["outcome"], r["status"]) for r in response.data["results"]} + + def _status(self, execution: WorkflowExecution) -> str: + execution.refresh_from_db() + status: str = execution.status + return status + + def test_abort_outside_any_transaction_succeeds(self) -> None: + execution = self._execution() + self.assertEqual( + self._post(execution.id), {str(execution.id): ("aborted", "STOPPED")} + ) + self.assertEqual(self._status(execution), ExecutionStatus.STOPPED.value) + + def test_a_failure_part_way_keeps_the_others_and_rolls_back_its_own(self) -> None: + first, bad, last = self._execution(), self._execution(), self._execution() + + def stop(execution: WorkflowExecution) -> None: + if execution.pk == bad.pk: + raise RuntimeError("boom") + + with ( + patch(f"{_ABORT}._stop_open_files", side_effect=stop), + patch(f"{_ABORT}.PipelineUtils.update_pipeline_status") as update, + ): + outcomes = self._post(first.id, bad.id, last.id) + + self.assertEqual( + outcomes, + { + str(first.id): ("aborted", "STOPPED"), + str(bad.id): ("failed", None), + str(last.id): ("aborted", "STOPPED"), + }, + ) + self.assertEqual(self._status(first), ExecutionStatus.STOPPED.value) + self.assertEqual(self._status(bad), ExecutionStatus.EXECUTING.value) + self.assertEqual(self._status(last), ExecutionStatus.STOPPED.value) + self.assertEqual(PlatformAuditEvent.objects.count(), 2) + self.assertEqual(update.call_count, 2) + + def test_a_run_finishing_during_the_abort_is_not_overwritten(self) -> None: + execution = self._execution() + locked, release = threading.Event(), threading.Event() + concurrent: Any = locking.ConcurrentStatusWriteTests + + def finish_the_run() -> None: + with transaction.atomic(): + WorkflowExecution.objects.select_for_update().get(pk=execution.pk) + locked.set() + release.wait(locking.WAIT_SECONDS) + WorkflowExecution.objects.filter(pk=execution.pk).update( + status=ExecutionStatus.COMPLETED.value + ) + + def abort() -> dict[str, tuple[str, str | None]]: + UserContext.set_organization_identifier(self.org.organization_id) + return self._post(execution.id) + + holder, _ = concurrent._run(finish_the_run) + self.assertTrue(locked.wait(locking.WAIT_SECONDS)) + thread, box = concurrent._run(abort) + concurrent._wait_until_a_backend_waits_on_a_lock(thread, box) + release.set() + holder.join(locking.WAIT_SECONDS) + thread.join(locking.WAIT_SECONDS) + + self.assertFalse(thread.is_alive()) + self.assertEqual( + box.get("result"), {str(execution.id): ("already_finished", "COMPLETED")} + ) + self.assertEqual(self._status(execution), ExecutionStatus.COMPLETED.value) + self.assertFalse(PlatformAuditEvent.objects.exists()) diff --git a/backend/workflow_manager/execution/tests/test_execution_contract.py b/backend/workflow_manager/execution/tests/test_execution_contract.py index d6baccdc76..e305c7a990 100644 --- a/backend/workflow_manager/execution/tests/test_execution_contract.py +++ b/backend/workflow_manager/execution/tests/test_execution_contract.py @@ -4,8 +4,9 @@ import uuid from typing import Any +from account_v2.models import Organization from django.test import SimpleTestCase -from django.urls import resolve, reverse +from django.urls import Resolver404, resolve, reverse from drf_spectacular.generators import SchemaGenerator from pipeline_v2.models import Pipeline from rest_framework.test import APIRequestFactory @@ -15,6 +16,7 @@ from workflow_manager.execution.views import ExecutionViewSet from workflow_manager.internal_views import WorkflowExecutionInternalViewSet from workflow_manager.workflow_v2.enums import ExecutionStatus +from workflow_manager.workflow_v2.models.execution import WorkflowExecution _URLCONF = "workflow_manager.execution.urls" _list = ExecutionViewSet.as_view({"get": "list"}) @@ -58,6 +60,8 @@ def test_list_and_detail_keep_their_keys(self) -> None: def test_stuck_false_is_the_same_response_as_no_parameter(self) -> None: self.execution() self.execution(status=ExecutionStatus.COMPLETED) + # A zero time renders as seconds since creation, which would tick between calls. + WorkflowExecution.objects.update(execution_time=1) plain = self.call(_list, self.owner) for value in ("false", "0", ""): with self.subTest(stuck=value): @@ -81,10 +85,17 @@ def test_entity_filter_still_selects_each_entity(self) -> None: def test_stuck_rows_add_only_stuck_since(self) -> None: self.execution() - listed = self.call(_list, self.owner, data={"stuck": "true"}) - self.assertEqual( - set(listed.data["results"][0]), _EXECUTION_KEYS | {"stuck_since"} - ) + for value in ("true", "1", "TRUE"): + with self.subTest(stuck=value): + listed = self.call(_list, self.owner, data={"stuck": value}) + self.assertEqual( + set(listed.data["results"][0]), _EXECUTION_KEYS | {"stuck_since"} + ) + + def test_unrecognised_stuck_value_is_ignored(self) -> None: + self.execution(status=ExecutionStatus.COMPLETED) + listed = self.call(_list, self.owner, data={"stuck": "garbage"}) + self.assertEqual(set(listed.data["results"][0]), _EXECUTION_KEYS) class InternalStatusActionTests(StuckAbortTestBase): @@ -103,7 +114,11 @@ def test_returns_the_status(self) -> None: def test_other_org_unknown_and_malformed_ids_are_404(self) -> None: execution = self.execution() + Organization.objects.create( + name="org-b", display_name="Org B", organization_id="org-b" + ) cases = ( + (execution.id, "org-b"), (execution.id, "some-other-org"), (execution.id, None), (uuid.uuid4(), self.org.organization_id), @@ -115,6 +130,13 @@ def test_other_org_unknown_and_malformed_ids_are_404(self) -> None: self.assertEqual(response.status_code, 404) +class WorkerStatusRouteTests(StuckAbortTestBase): + def test_worker_status_url_reaches_the_status_action(self) -> None: + match = resolve(f"/internal/v1/workflow-execution/{uuid.uuid4()}/status/") + self.assertEqual(match.func.actions, {"get": "execution_status"}) + self.assertIs(match.func.cls, WorkflowExecutionInternalViewSet) + + class RoutesUnchangedTests(SimpleTestCase): def test_existing_names_reverse_to_the_same_paths(self) -> None: pk = uuid.uuid4() @@ -131,6 +153,15 @@ def test_existing_names_reverse_to_the_same_paths(self) -> None: with self.subTest(name=name): self.assertEqual(reverse(name, urlconf=_URLCONF, kwargs=kwargs), path) + def test_new_paths_take_no_format_suffix(self) -> None: + for path in ("/abort.json", "/stuck-count.json"): + with self.subTest(path=path), self.assertRaises(Resolver404): + resolve(path, urlconf=_URLCONF) + self.assertEqual( + resolve(f"/{uuid.uuid4()}.json", urlconf=_URLCONF).url_name, + "execution-detail", + ) + def test_new_paths_do_not_shadow_detail(self) -> None: self.assertEqual( resolve("/stuck-count/", urlconf=_URLCONF).url_name, @@ -152,16 +183,16 @@ def test_schema_documents_the_new_operations(self) -> None: paths = schema["paths"] count = paths["/stuck-count/"]["get"] self.assertEqual(count["operationId"], "execution_stuck_count") - self.assertTrue({"200", "403"} <= set(count["responses"])) + self.assertTrue({"200", "401", "403"} <= set(count["responses"])) abort = paths["/abort/"]["post"] self.assertEqual(abort["operationId"], "execution_abort") - self.assertTrue({"200", "400", "403"} <= set(abort["responses"])) + self.assertTrue({"200", "400", "401", "403"} <= set(abort["responses"])) components = schema["components"]["schemas"] self.assertIn("ids", components["ExecutionAbortRequest"]["properties"]) self.assertEqual( set(components["ExecutionAbortResponse"]["properties"]), - {"results", "aborted", "already_finished", "not_permitted"}, + {"results", "aborted", "already_finished", "not_permitted", "failed"}, ) self.assertEqual( set(components["StuckExecutionCount"]["properties"]), diff --git a/backend/workflow_manager/execution/tests/test_stuck_executions.py b/backend/workflow_manager/execution/tests/test_stuck_executions.py index 3ddfadc0fd..b338f66367 100644 --- a/backend/workflow_manager/execution/tests/test_stuck_executions.py +++ b/backend/workflow_manager/execution/tests/test_stuck_executions.py @@ -62,6 +62,21 @@ def test_run_without_files_is_stuck_since_its_own_last_write(self) -> None: execution.modified_at, ) + def test_run_newer_than_its_files_is_stuck_since_its_own_last_write(self) -> None: + execution = self.execution(idle=OLD - timedelta(minutes=10)) + self.file(execution, ExecutionStatus.EXECUTING, idle=OLD) + rows = self._stuck_rows(self.owner) + self.assertEqual( + parse_datetime(rows[str(execution.id)]["stuck_since"]), + execution.modified_at, + ) + + @override_settings(STUCK_EXECUTION_THRESHOLD_SECONDS=3600) + def test_cutoff_is_the_threshold(self) -> None: + just_over = self.execution(idle=timedelta(minutes=61)) + self.execution(idle=timedelta(minutes=59)) + self.assertEqual(list(self._stuck_rows(self.owner)), [str(just_over.id)]) + def test_recent_file_activity_is_not_stuck(self) -> None: execution = self.execution() self.file(execution, ExecutionStatus.EXECUTING, idle=RECENT) diff --git a/backend/workflow_manager/execution/urls.py b/backend/workflow_manager/execution/urls.py index 1684fded05..435c86ecc5 100644 --- a/backend/workflow_manager/execution/urls.py +++ b/backend/workflow_manager/execution/urls.py @@ -29,11 +29,13 @@ execution_log_export, name="execution-log-export", ), - path( - "stuck-count/", - StuckExecutionCountView.as_view(), - name="execution-stuck-count", - ), - path("abort/", ExecutionAbortView.as_view(), name="execution-abort"), ] -) +) + [ + # Outside the suffix patterns: these views take no format argument. + path( + "stuck-count/", + StuckExecutionCountView.as_view(), + name="execution-stuck-count", + ), + path("abort/", ExecutionAbortView.as_view(), name="execution-abort"), +] diff --git a/backend/workflow_manager/execution/views/stuck_abort.py b/backend/workflow_manager/execution/views/stuck_abort.py index d8ac769ae1..76ac21cb50 100644 --- a/backend/workflow_manager/execution/views/stuck_abort.py +++ b/backend/workflow_manager/execution/views/stuck_abort.py @@ -1,10 +1,7 @@ """Stuck-execution count and abort. See ../README.md for the contract.""" -from drf_spectacular.utils import OpenApiExample, extend_schema -from drf_standardized_errors.openapi_serializers import ( - ErrorResponse403Serializer, - ValidationErrorResponseSerializer, -) +from drf_spectacular.utils import OpenApiExample, OpenApiResponse, extend_schema +from drf_standardized_errors.openapi_serializers import ValidationErrorResponseSerializer from rest_framework.permissions import IsAuthenticated from rest_framework.request import Request from rest_framework.response import Response @@ -20,6 +17,13 @@ from workflow_manager.workflow_v2.models import WorkflowExecution _TAGS = ["execution"] +# The auth middleware answers these before DRF runs, as {"message": ...}. +_AUTH_ERRORS = { + 401: OpenApiResponse(description='Not authenticated: {"message": ...}.'), + 403: OpenApiResponse( + description='Authenticated but refused access: {"message": ...}.' + ), +} class StuckExecutionCountView(APIView): @@ -29,7 +33,7 @@ class StuckExecutionCountView(APIView): operation_id="execution_stuck_count", tags=_TAGS, description="Stuck executions visible to the caller, per entity.", - responses={200: StuckExecutionCountSerializer, 403: ErrorResponse403Serializer}, + responses={200: StuckExecutionCountSerializer, **_AUTH_ERRORS}, examples=[ OpenApiExample( "Counts", @@ -53,15 +57,12 @@ class ExecutionAbortView(APIView): @extend_schema( operation_id="execution_abort", tags=_TAGS, - description=( - "Stop running executions. Only the owner or an org admin may abort; " - "every id gets its own outcome in a 200." - ), + description="Stop running executions; every id gets its own outcome in a 200.", request=ExecutionAbortRequestSerializer, responses={ 200: ExecutionAbortResponseSerializer, 400: ValidationErrorResponseSerializer, - 403: ErrorResponse403Serializer, + **_AUTH_ERRORS, }, ) def post(self, request: Request) -> Response: diff --git a/backend/workflow_manager/file_execution/serializers.py b/backend/workflow_manager/file_execution/serializers.py index a40f80b04b..59b62aa4f0 100644 --- a/backend/workflow_manager/file_execution/serializers.py +++ b/backend/workflow_manager/file_execution/serializers.py @@ -34,6 +34,8 @@ def get_status_msg(self, obj: FileExecution) -> dict[str, any] | None: return self.INIT_STATUS_MSG elif obj.status == ExecutionStatus.ERROR.value: return obj.execution_error + elif obj.status == ExecutionStatus.STOPPED.value and obj.execution_error: + return obj.execution_error # Use pre-annotated field from viewset queryset to avoid N+1 queries if hasattr(obj, "latest_log_data") and obj.latest_log_data: diff --git a/workers/callback/tasks.py b/workers/callback/tasks.py index cceb0aef04..b2024d0084 100644 --- a/workers/callback/tasks.py +++ b/workers/callback/tasks.py @@ -255,6 +255,13 @@ def _get_performance_stats() -> dict: return stats +def _was_stopped(execution_context: dict[str, Any] | None) -> bool: + """True when the run was already STOPPED (aborted) when this callback fetched it.""" + execution = (execution_context or {}).get("execution") or {} + stopped: bool = execution.get("status") == ExecutionStatus.STOPPED.value + return stopped + + def _determine_execution_status_unified( file_batch_results: list[dict[str, Any]], api_client: InternalAPIClient, @@ -324,6 +331,14 @@ def _determine_execution_status_unified( total_files = aggregated_results.get("total_files", 0) successful_files = aggregated_results.get("successful_files", 0) + # An abort ended this run, whatever its file counts say. + if _was_stopped(execution_context): + logger.warning( + f"Execution {execution_id} was stopped - {successful_files} of " + f"{total_files} files finished" + ) + return aggregated_results, ExecutionStatus.STOPPED.value, expected_files + # Step 5: Unified status determination with timeout failure detection # Detect timeout failures: expected files but processed none (likely SoftTimeLimitExceeded) @@ -605,6 +620,15 @@ def _handle_notifications_unified( Returns: Notification result dictionary """ + # The abort already sent these, so a run with a dead worker notifies too. + if status == ExecutionStatus.STOPPED.value: + logger.info(f"Execution {execution_id} was stopped - notifications already sent") + return { + "status": "skipped", + "reason": "sent_by_abort", + "message": "Notifications were sent when the execution was aborted", + } + try: if not pipeline_id: logger.warning("No pipeline_id provided - skipping notifications") @@ -1305,11 +1329,12 @@ def _track_subscription_usage_if_available( """Track subscription usage for successful file executions if plugin available. This is a non-blocking operation - errors are logged but do not fail the workflow. - Only tracks usage for completed workflows with successful file executions. + Only tracks usage for completed or stopped workflows with successful file + executions; files that finished before an abort are billed. Args: context: Callback execution context with organization and execution IDs - execution_status: Final execution status (should be COMPLETED for tracking) + execution_status: Final execution status (COMPLETED or STOPPED for tracking) successful_files: Successful file count from the aggregated batch results, used only to flag a fetch that came back without file-execution rows @@ -1317,8 +1342,11 @@ def _track_subscription_usage_if_available( dict with status, committed_count, message if tracking attempted None if plugin not available or workflow not completed """ - # Only track for completed workflows - if execution_status != ExecutionStatus.COMPLETED.value: + # Only track for workflows whose finished files stand + if execution_status not in ( + ExecutionStatus.COMPLETED.value, + ExecutionStatus.STOPPED.value, + ): logger.debug( f"Skipping subscription tracking for execution {context.execution_id} " f"with status {execution_status}" @@ -1954,6 +1982,7 @@ def _publish_final_workflow_ui_logs( total_files=total_files, successful_files=successful_files, failed_files=failed_files, + stopped=execution_status == ExecutionStatus.STOPPED.value, ) logger.info( @@ -2028,6 +2057,7 @@ def _publish_final_workflow_ui_logs_api( total_files=total_files, successful_files=successful_files, failed_files=failed_files, + stopped=execution_status == ExecutionStatus.STOPPED.value, ) logger.info( diff --git a/workers/file_processing/tasks.py b/workers/file_processing/tasks.py index 1cf6b78a04..41c348de91 100644 --- a/workers/file_processing/tasks.py +++ b/workers/file_processing/tasks.py @@ -19,6 +19,7 @@ # Import shared worker infrastructure from shared.api import InternalAPIClient +from shared.clients.base_client import InternalAPIClientError # Import from shared worker modules from shared.constants import Account @@ -240,7 +241,8 @@ class _TerminalExecutionSkip(Exception): The resurrection this prevents exists because the **reaper** can mark an in-flight execution terminal (and tear its barrier down) while a batch for it - is still in flight or waiting to be redelivered. + is still in flight or waiting to be redelivered. An abort can mark it STOPPED + the same way, but leaves the barrier armed. This is a validate-first *narrowing*, not a transactional guarantee: the status read and the later ``EXECUTING`` write are not atomic, so a reaper @@ -289,7 +291,9 @@ def _raise_if_execution_terminal( ) -def _terminal_skip_result(batch_data: FileBatchData) -> dict[str, Any]: +def _terminal_skip_result( + batch_data: FileBatchData, status: str | None = None +) -> dict[str, Any]: """Batch result for a batch skipped because its execution is already terminal. The files were never attempted; they are counted as failed (not left @@ -297,8 +301,9 @@ def _terminal_skip_result(batch_data: FileBatchData) -> dict[str, Any]: so reporting ``failed=total`` keeps a consumer that ever aggregates this result from rendering the files as perpetually in-progress. The ``SKIPPED_TERMINAL_EXECUTION_KEY`` marker tells ``run_batch_with_barrier`` to - bypass the barrier decrement (a terminal execution's barrier is by definition - already torn down). + bypass the barrier decrement (the reaper already tore that barrier down). A + STOPPED run's barrier is still armed, so it gets no marker: the batch counts + down and the last one fires the callback, which cleans up. """ n_files = len(batch_data.files) result = BatchExecutionResult( @@ -308,10 +313,30 @@ def _terminal_skip_result(batch_data: FileBatchData) -> dict[str, Any]: execution_time=0.0, organization_id=batch_data.file_data.organization_id, ).to_dict() - result[SKIPPED_TERMINAL_EXECUTION_KEY] = True + if status != ExecutionStatus.STOPPED.value: + result[SKIPPED_TERMINAL_EXECUTION_KEY] = True return result +def _clear_active_file_markers(batch_data: FileBatchData) -> None: + """Clear a skipped batch's active-file markers, which pre-create never reached. + + Without this, a re-run leaves these files out until the markers expire. + """ + from shared.workflow.execution.active_file_manager import cleanup_active_file_cache + + uuids = [ + uuid + for item in batch_data.files + if (uuid := _file_item_hash(item).get("provider_file_uuid")) + ] + cleanup_active_file_cache( + provider_file_uuids=uuids, + workflow_id=batch_data.file_data.workflow_id, + logger_instance=logger, + ) + + def _run_batch_stages( file_batch_data: dict[str, Any], celery_task_id: str ) -> dict[str, Any]: @@ -333,10 +358,11 @@ def _run_batch_stages( logger.warning( f"[exec:{skip.execution_id}] Skipping batch of " f"{len(batch_data.files)} file(s): execution already terminal " - f"({skip.status}) — stale/redelivered after reaper-recovery; not " - f"reprocessing." + f"({skip.status}); not processing." ) - return _terminal_skip_result(batch_data) + if skip.status == ExecutionStatus.STOPPED.value: + _clear_active_file_markers(batch_data) + return _terminal_skip_result(batch_data, skip.status) # Step 3: Handle manual review logic # context = _handle_manual_review_logic(context) @@ -685,14 +711,29 @@ def _file_item_name(file_item: Any) -> str | None: return name +def _file_item_hash(file_item: Any) -> dict[str, Any]: + """Hash data from a batch file item in any of its formats; empty if none.""" + data = file_item + if isinstance(file_item, (list, tuple)) and len(file_item) == 2: + data = file_item[1] + return data if isinstance(data, dict) else {} + + +def _file_item_identifier(file_item: Any) -> str | None: + """The ``uuid:path`` id the duplicate skip lists use; None without both parts.""" + data = _file_item_hash(file_item) + uuid, path = data.get("provider_file_uuid"), data.get("file_path") + return f"{uuid}:{path}" if uuid and path else None + + def _is_execution_stopped(api_client: InternalAPIClient, execution_id: str) -> bool: """True when the execution was stopped (aborted) while this batch runs. - Fails open: if the check itself fails, the next file is processed. + Fails open on an API failure: the next file is processed. """ try: status = api_client.get_workflow_execution_status(execution_id) - except Exception as error: + except InternalAPIClientError as error: logger.warning( "[exec:%s] status check before the next file failed (%s); " "processing it anyway", @@ -700,30 +741,44 @@ def _is_execution_stopped(api_client: InternalAPIClient, execution_id: str) -> b error, ) return False + if not status: + logger.warning( + "[exec:%s] status check before the next file returned no status; " + "processing it anyway", + execution_id, + ) stopped: bool = status == ExecutionStatus.STOPPED.value return stopped def _mark_file_stopped(api_client: InternalAPIClient, file_execution_id: str) -> None: - """Mark one unprocessed file row STOPPED; a failed write is logged, not raised.""" - try: - api_client.update_file_execution_status( - file_execution_id=file_execution_id, - status=ExecutionStatus.STOPPED.value, - error_message="Not processed: execution was stopped", - ) - except Exception as error: + """Mark one unprocessed file row STOPPED, logging a failed write.""" + # An APIResponse, whatever the facade's annotation says. + response: Any = api_client.update_file_execution_status( + file_execution_id=file_execution_id, + status=ExecutionStatus.STOPPED.value, + error_message="Not processed: execution was stopped", + ) + if not response.success: logger.warning( - "Failed to mark file execution %s STOPPED: %s", file_execution_id, error + "Failed to mark file execution %s STOPPED: %s", + file_execution_id, + response.error, ) def _stop_remaining_files( context: WorkflowContextData, remaining: list[Any], result: FileBatchResult ) -> None: - """Count the batch's unprocessed files as failed and mark their rows STOPPED.""" + """Count the batch's unprocessed files as failed and mark their rows STOPPED. + + Files already completed in this execution stay uncounted, as in the main loop. + """ api_client = context.organization_context.api_client + already_completed = (context.metadata or {}).get("skipped_already_completed", []) for file_item in remaining: + if _file_item_identifier(file_item) in already_completed: + continue result.increment_failure() name = _file_item_name(file_item) pre_created = context.pre_created_file_executions.get(name) if name else None diff --git a/workers/queue_backend/pg_barrier.py b/workers/queue_backend/pg_barrier.py index 08e0a1d464..6102f4c8bc 100644 --- a/workers/queue_backend/pg_barrier.py +++ b/workers/queue_backend/pg_barrier.py @@ -1305,11 +1305,12 @@ def _mark_execution_error_on_abort( ) -# A batch whose execution is already terminal returns its result with this -# marker set. run_batch_with_barrier bypasses the barrier decrement for it: the -# reaper has by definition already torn the barrier down, so decrementing would -# only log a spurious "decrement found no row" ERROR (false alert noise), and an -# already-terminal execution must not have its aggregating callback fired. +# A batch skipped for a run the reaper already finished returns its result with +# this marker set. run_batch_with_barrier bypasses the barrier decrement for it: +# the reaper already tore the barrier down, so decrementing would only log a +# spurious "decrement found no row" ERROR (false alert noise), and that run must +# not have its aggregating callback fired. An aborted (STOPPED) run keeps its +# barrier, so its batches do not set this. SKIPPED_TERMINAL_EXECUTION_KEY = "skipped_terminal_execution" @@ -1380,11 +1381,10 @@ def run_batch_with_barrier( try: result = work_fn() if result.get(SKIPPED_TERMINAL_EXECUTION_KEY): - # The batch's execution is already terminal, so the reaper has by - # definition already torn the barrier down. Bypass the decrement — - # decrementing a gone barrier only logs a spurious "decrement found - # no row" ERROR (false alert noise) — and do NOT abort (nothing to - # tear down; and an already-terminal execution must not have its + # The reaper already finished this run and tore its barrier down. + # Bypass the decrement — decrementing a gone barrier only logs a + # spurious "decrement found no row" ERROR (false alert noise) — and + # do NOT abort (nothing to tear down; and that run must not have its # aggregating callback re-fired). logger.info( f"[exec:{execution_id}] batch {batch_index} skipped — execution " diff --git a/workers/shared/infrastructure/logging/workflow_logger.py b/workers/shared/infrastructure/logging/workflow_logger.py index 1f7f572232..e2107a5521 100644 --- a/workers/shared/infrastructure/logging/workflow_logger.py +++ b/workers/shared/infrastructure/logging/workflow_logger.py @@ -469,6 +469,7 @@ def publish_final_workflow_logs( total_files: int, successful_files: int, failed_files: int, + stopped: bool = False, ) -> None: """Publish final workflow execution logs. @@ -479,14 +480,20 @@ def publish_final_workflow_logs( total_files: Total number of files processed successful_files: Number of successfully processed files failed_files: Number of files that failed processing + stopped: The run was aborted, so it did not execute successfully """ # Publish workflow end status to UI status bar self.publish_update_log(LogState.END_WORKFLOW, "1", LogComponent.STATUS_BAR) # Publish workflow completion status - self.publish_update_log( - LogState.SUCCESS, "Executed successfully", LogComponent.WORKFLOW - ) + if stopped: + self.publish_update_log( + LogState.ERROR, "Execution stopped", LogComponent.WORKFLOW + ) + else: + self.publish_update_log( + LogState.SUCCESS, "Executed successfully", LogComponent.WORKFLOW + ) # Publish detailed execution summary summary_message = ( diff --git a/workers/tests/test_abort_between_files.py b/workers/tests/test_abort_between_files.py index 4e40ad439f..978e18bb5b 100644 --- a/workers/tests/test_abort_between_files.py +++ b/workers/tests/test_abort_between_files.py @@ -2,7 +2,7 @@ Pins: the status is checked before each file; only STOPPED (abort) stops the batch, with the remaining files counted as failed and their rows marked -STOPPED; any other status, or a failed check, processes the file (fail open). +STOPPED; any other status, or a failed API call, processes the file (fail open). Also pins the client URL and the facade delegation. """ @@ -13,18 +13,35 @@ from unittest import mock import pytest +from celery.exceptions import SoftTimeLimitExceeded +from shared.api.internal_client import InternalAPIClient +from shared.clients.base_client import APIRequestError, AuthenticationError +from shared.data.models import APIResponse from unstract.core.data_models import ExecutionStatus, FileBatchResult from file_processing.tasks import _file_item_name, _process_individual_files _TASKS = "file_processing.tasks" +FORMATS = ["list", "tuple", "dict"] -def _context(api_client: Any, n_files: int = 3) -> Any: - files = { - f"file_{i}": {"file_name": f"f{i}.pdf", "file_path": f"/in/f{i}.pdf"} - for i in range(n_files) +def _client() -> Any: + """A client that rejects a misnamed method or argument.""" + return mock.create_autospec(InternalAPIClient, instance=True) + + +def _item(index: int, file_format: str) -> Any: + name = f"f{index}.pdf" + data = { + "file_name": name, + "file_path": f"/in/{name}", + "provider_file_uuid": f"uuid-{index}", } + return {"list": [name, data], "tuple": (name, data), "dict": data}[file_format] + + +def _context(api_client: Any, n_files: int = 3, file_format: str = "dict") -> Any: + files = {f"file_{i}": _item(i, file_format) for i in range(n_files)} pre_created = { f"f{i}.pdf": mock.Mock( id=f"fe-{i}", @@ -50,9 +67,14 @@ def _context(api_client: Any, n_files: int = 3) -> Any: ) -def _run(api_client: Any, n_files: int = 3) -> tuple[FileBatchResult, Any]: +def _run( + api_client: Any, n_files: int = 3, file_format: str = "dict" +) -> tuple[FileBatchResult, Any]: + return _run_context(_context(api_client, n_files, file_format)) + + +def _run_context(context: Any) -> tuple[FileBatchResult, Any]: """Run the loop with file processing stubbed to succeed.""" - context = _context(api_client, n_files) def _succeeded( _result: Any, _name: Any, _start: Any, result: Any, *args: Any, **kwargs: Any @@ -79,14 +101,15 @@ def _stopped_ids(api_client: Any) -> list[str]: ] -def test_stops_after_the_status_flips_to_stopped() -> None: - api_client = mock.Mock() +@pytest.mark.parametrize("file_format", FORMATS) +def test_stops_after_the_status_flips_to_stopped(file_format: str) -> None: + api_client = _client() api_client.get_workflow_execution_status.side_effect = [ ExecutionStatus.EXECUTING.value, ExecutionStatus.STOPPED.value, ] - result, process_file = _run(api_client) + result, process_file = _run(api_client, file_format=file_format) assert process_file.call_count == 1 assert api_client.get_workflow_execution_status.call_count == 2 @@ -95,7 +118,7 @@ def test_stops_after_the_status_flips_to_stopped() -> None: def test_stopped_before_the_first_file_processes_nothing() -> None: - api_client = mock.Mock() + api_client = _client() api_client.get_workflow_execution_status.return_value = ( ExecutionStatus.STOPPED.value ) @@ -107,23 +130,61 @@ def test_stopped_before_the_first_file_processes_nothing() -> None: assert (result.successful_files, result.failed_files) == (0, 3) -def test_a_failed_stopped_write_does_not_stop_the_rest() -> None: - api_client = mock.Mock() +def test_a_failed_stopped_write_is_logged_and_the_rest_still_written() -> None: + api_client = _client() api_client.get_workflow_execution_status.return_value = ( ExecutionStatus.STOPPED.value ) - api_client.update_file_execution_status.side_effect = RuntimeError("500") + api_client.update_file_execution_status.return_value = APIResponse( + success=False, error="500" + ) - result, process_file = _run(api_client) + with mock.patch(f"{_TASKS}.logger") as logger: + result, process_file = _run(api_client) process_file.assert_not_called() assert _stopped_ids(api_client) == ["fe-0", "fe-1", "fe-2"] assert (result.successful_files, result.failed_files) == (0, 3) + warned = [ + c.args[1] + for c in logger.warning.call_args_list + if c.args[0].startswith("Failed to mark") + ] + assert warned == ["fe-0", "fe-1", "fe-2"] + + +def test_a_file_without_a_row_is_counted_but_not_written() -> None: + api_client = _client() + api_client.get_workflow_execution_status.return_value = ( + ExecutionStatus.STOPPED.value + ) + context = _context(api_client, file_format="list") + del context.pre_created_file_executions["f1.pdf"] + + result, _ = _run_context(context) + + assert _stopped_ids(api_client) == ["fe-0", "fe-2"] + assert result.failed_files == 3 + + +@pytest.mark.parametrize("file_format", FORMATS) +def test_a_file_already_completed_in_this_run_is_not_counted(file_format: str) -> None: + api_client = _client() + api_client.get_workflow_execution_status.return_value = ( + ExecutionStatus.STOPPED.value + ) + context = _context(api_client, file_format=file_format) + context.metadata["skipped_already_completed"] = ["uuid-1:/in/f1.pdf"] + + result, _ = _run_context(context) + + assert _stopped_ids(api_client) == ["fe-0", "fe-2"] + assert result.failed_files == 2 def test_error_mid_batch_keeps_processing_the_files() -> None: """Only an abort stops a running batch; ERROR keeps today's behaviour.""" - api_client = mock.Mock() + api_client = _client() api_client.get_workflow_execution_status.side_effect = [ ExecutionStatus.EXECUTING.value, ExecutionStatus.ERROR.value, @@ -148,7 +209,7 @@ def test_error_mid_batch_keeps_processing_the_files() -> None: ], ) def test_any_other_status_processes_every_file(status: str | None) -> None: - api_client = mock.Mock() + api_client = _client() api_client.get_workflow_execution_status.return_value = status result, process_file = _run(api_client) @@ -158,9 +219,24 @@ def test_any_other_status_processes_every_file(status: str | None) -> None: assert (result.successful_files, result.failed_files) == (3, 0) -def test_a_failed_check_processes_the_file_anyway() -> None: - api_client = mock.Mock() - api_client.get_workflow_execution_status.side_effect = RuntimeError("404") +def test_a_missing_status_is_logged_and_the_file_processed() -> None: + api_client = _client() + api_client.get_workflow_execution_status.return_value = None + + with mock.patch(f"{_TASKS}.logger") as logger: + _, process_file = _run(api_client) + + assert process_file.call_count == 3 + warnings = [c.args[0] for c in logger.warning.call_args_list] + assert sum("returned no status" in w for w in warnings) == 3 + + +@pytest.mark.parametrize( + "error", [APIRequestError("404"), AuthenticationError("401")], ids=["404", "401"] +) +def test_a_failed_api_call_processes_the_file_anyway(error: Exception) -> None: + api_client = _client() + api_client.get_workflow_execution_status.side_effect = error result, process_file = _run(api_client) @@ -169,6 +245,15 @@ def test_a_failed_check_processes_the_file_anyway() -> None: assert result.successful_files == 3 +@pytest.mark.parametrize("error", [AttributeError("bug"), SoftTimeLimitExceeded()]) +def test_other_errors_in_the_check_are_not_hidden(error: Exception) -> None: + api_client = _client() + api_client.get_workflow_execution_status.side_effect = error + + with pytest.raises(type(error)): + _run(api_client) + + @pytest.mark.parametrize( ("item", "name"), [ diff --git a/workers/tests/test_abort_finishes_the_run.py b/workers/tests/test_abort_finishes_the_run.py new file mode 100644 index 0000000000..03fd147287 --- /dev/null +++ b/workers/tests/test_abort_finishes_the_run.py @@ -0,0 +1,185 @@ +"""An aborted run still finishes: skipped batches count down, and the callback +keeps STOPPED, bills finished files and does not notify again (UN-4318). +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any +from unittest import mock + +import callback.tasks as cb +import pytest +from shared.infrastructure.logging.workflow_logger import WorkerWorkflowLogger +from unstract.core.data_models import ExecutionStatus +from unstract.workflow_execution.enums import LogState + +from file_processing.tasks import ( + _run_batch_stages, + _terminal_skip_result, + _TerminalExecutionSkip, +) +from queue_backend import pg_barrier +from queue_backend.pg_barrier import SKIPPED_TERMINAL_EXECUTION_KEY + +STOPPED = ExecutionStatus.STOPPED.value +ERROR = ExecutionStatus.ERROR.value +COMPLETED = ExecutionStatus.COMPLETED.value +_TASKS = "file_processing.tasks" +_CLEANUP = "shared.workflow.execution.active_file_manager.cleanup_active_file_cache" + + +def _batch_data() -> Any: + def data(index: int) -> dict[str, str]: + return {"file_name": f"f{index}.pdf", "provider_file_uuid": f"uuid-{index}"} + + return SimpleNamespace( + files=[["f0.pdf", data(0)], ("f1.pdf", data(1)), data(2), {"file_name": "x"}], + file_data=SimpleNamespace(organization_id="org-1", workflow_id="wf-1"), + ) + + +def _skip_batch(status: str) -> tuple[dict[str, Any], Any]: + with ( + mock.patch(f"{_TASKS}._validate_and_parse_batch_data", return_value=_batch_data()), + mock.patch( + f"{_TASKS}._setup_execution_context", + side_effect=_TerminalExecutionSkip("exec-1", status), + ), + mock.patch(_CLEANUP) as cleanup, + ): + return _run_batch_stages({"any": "payload"}, "task-1"), cleanup + + +def test_a_stopped_runs_skipped_batch_clears_its_markers_and_counts_down() -> None: + result, cleanup = _skip_batch(STOPPED) + + assert SKIPPED_TERMINAL_EXECUTION_KEY not in result + assert result["failed_files"] == 4 + cleanup.assert_called_once_with( + provider_file_uuids=["uuid-0", "uuid-1", "uuid-2"], + workflow_id="wf-1", + logger_instance=mock.ANY, + ) + + +@pytest.mark.parametrize("status", [ERROR, COMPLETED]) +def test_a_run_the_reaper_finished_keeps_the_marker(status: str) -> None: + result, cleanup = _skip_batch(status) + + assert result[SKIPPED_TERMINAL_EXECUTION_KEY] is True + cleanup.assert_not_called() + + +@pytest.mark.parametrize(("status", "decrements"), [(STOPPED, 1), (ERROR, 0)]) +def test_only_a_stopped_runs_batch_decrements(status: str, decrements: int) -> None: + result = _terminal_skip_result(_batch_data(), status) + context: Any = {"execution_id": "e", "batch_index": 0, "callback_descriptor": {}} + with ( + mock.patch.object(pg_barrier, "claim_batch", return_value=True), + mock.patch.object(pg_barrier, "_barrier_pg_decrement") as decrement, + mock.patch.object(pg_barrier, "_abort_barrier_in_body") as abort, + ): + pg_barrier.run_batch_with_barrier(context, lambda: result) + + assert decrement.call_count == decrements + abort.assert_not_called() + + +def _status_for(run_status: str, successful: int, failed: int) -> str: + counts = { + "total_files": successful + failed, + "successful_files": successful, + "failed_files": failed, + "total_files_processed": successful + failed, + } + with ( + mock.patch.object(cb, "aggregate_file_batch_results", return_value=counts), + mock.patch.object( + cb.WallClockTimeCalculator, "calculate_execution_time", return_value=1.0 + ), + ): + _, status, _ = cb._determine_execution_status_unified( + file_batch_results=[], + api_client=mock.Mock(), + execution_id="e1", + organization_id="org-1", + execution_context={"execution": {"status": run_status}}, + ) + return status + + +@pytest.mark.parametrize(("successful", "failed"), [(1, 9), (0, 10), (10, 0)]) +def test_a_stopped_run_stays_stopped_whatever_the_counts( + successful: int, failed: int +) -> None: + assert _status_for(STOPPED, successful, failed) == STOPPED + + +@pytest.mark.parametrize( + ("successful", "failed", "expected"), [(1, 9, COMPLETED), (0, 10, ERROR)] +) +def test_a_running_run_is_still_judged_by_its_counts( + successful: int, failed: int, expected: str +) -> None: + executing = ExecutionStatus.EXECUTING.value + assert _status_for(executing, successful, failed) == expected + + +def test_a_missing_context_is_judged_by_its_counts() -> None: + assert cb._was_stopped(None) is False + assert cb._was_stopped({}) is False + + +@pytest.mark.parametrize(("status", "calls"), [(STOPPED, 0), (COMPLETED, 1)]) +def test_a_stopped_run_is_not_notified_twice(status: str, calls: int) -> None: + with mock.patch.object(cb, "handle_status_notifications") as notify: + result = cb._handle_notifications_unified( + api_client=mock.Mock(), + status=status, + organization_id="org-1", + execution_id="e1", + pipeline_id="p1", + ) + + assert notify.call_count == calls + assert result["status"] == ("skipped" if status == STOPPED else "completed") + + +@pytest.mark.parametrize(("status", "commits"), [(STOPPED, 1), (ERROR, 0)]) +def test_a_stopped_run_bills_only_its_finished_files(status: str, commits: int) -> None: + context = cb.CallbackContext() + context.execution_id = "e1" + context.organization_id = "org-1" + context.file_executions = [ + {"id": "done", "status": COMPLETED}, + {"id": "unfinished", "status": STOPPED}, + ] + plugin = mock.MagicMock() + with mock.patch("client_plugin_registry.get_client_plugin", return_value=plugin): + cb._track_subscription_usage_if_available(context, status) + + commit = plugin.commit_batch_subscription_usage + assert commit.call_count == commits + if commits: + assert commit.call_args.kwargs["file_execution_ids"] == ["done"] + + +@pytest.mark.parametrize( + ("stopped", "state", "message"), + [ + (True, LogState.ERROR, "Execution stopped"), + (False, LogState.SUCCESS, "Executed successfully"), + ], +) +def test_the_final_log_tells_a_stopped_run_apart( + stopped: bool, state: LogState, message: str +) -> None: + workflow_logger = WorkerWorkflowLogger.__new__(WorkerWorkflowLogger) + with ( + mock.patch.object(workflow_logger, "publish_update_log") as update, + mock.patch.object(workflow_logger, "publish_log"), + ): + workflow_logger.publish_final_workflow_logs(3, 1, 2, stopped=stopped) + + assert mock.call(state, message, mock.ANY) in update.call_args_list From 6dc7e1e1cec00e36de1615d4ca51d2f5ec475b60 Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Thu, 5 Nov 2026 15:37:59 +0530 Subject: [PATCH 3/8] UN-4318 [FEAT] Log an unknown status from the status check Co-Authored-By: Claude Opus 5.5 --- workers/file_processing/tasks.py | 5 +++-- workers/tests/test_abort_between_files.py | 9 +++++---- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/workers/file_processing/tasks.py b/workers/file_processing/tasks.py index 41c348de91..72634a1317 100644 --- a/workers/file_processing/tasks.py +++ b/workers/file_processing/tasks.py @@ -741,11 +741,12 @@ def _is_execution_stopped(api_client: InternalAPIClient, execution_id: str) -> b error, ) return False - if not status: + if status not in {s.value for s in ExecutionStatus}: logger.warning( - "[exec:%s] status check before the next file returned no status; " + "[exec:%s] status check before the next file returned %r; " "processing it anyway", execution_id, + status, ) stopped: bool = status == ExecutionStatus.STOPPED.value return stopped diff --git a/workers/tests/test_abort_between_files.py b/workers/tests/test_abort_between_files.py index 978e18bb5b..d96d7584ec 100644 --- a/workers/tests/test_abort_between_files.py +++ b/workers/tests/test_abort_between_files.py @@ -219,16 +219,17 @@ def test_any_other_status_processes_every_file(status: str | None) -> None: assert (result.successful_files, result.failed_files) == (3, 0) -def test_a_missing_status_is_logged_and_the_file_processed() -> None: +@pytest.mark.parametrize(("status", "warnings"), [(None, 3), ("RENAMED", 3), ("ERROR", 0)]) +def test_a_missing_or_unknown_status_is_logged(status: Any, warnings: int) -> None: api_client = _client() - api_client.get_workflow_execution_status.return_value = None + api_client.get_workflow_execution_status.return_value = status with mock.patch(f"{_TASKS}.logger") as logger: _, process_file = _run(api_client) assert process_file.call_count == 3 - warnings = [c.args[0] for c in logger.warning.call_args_list] - assert sum("returned no status" in w for w in warnings) == 3 + logged = [c.args[0] for c in logger.warning.call_args_list] + assert sum("status check before the next file returned" in m for m in logged) == warnings @pytest.mark.parametrize( From a75bb860eb23c8b9606a86aa0da8f8fb19d8694f Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Thu, 5 Nov 2026 18:00:31 +0530 Subject: [PATCH 4/8] UN-4318 [FEAT] Fix second review round: reporting, re-runs and crashes after abort - Report the abort with a named hook, so a failed notification can't turn a committed abort into `failed` or lose its audit event. - Only a pipeline's newest run writes its status; paused pipelines are reported. The callback leaves a stopped run's pipeline status and notifications to abort. - An abort landing during the callback is reported as STOPPED, not COMPLETED. - A batch that fails after the abort counts down, so the callback still runs. - A re-run ignores a stopped run's active-file markers, so it keeps those files before any worker picks up the queued batch. - GET filters a stopped run's results like POST, and the API spec names STOPPED. - Correct the barrier comments and pin the callback's final log argument. Co-Authored-By: Claude Opus 5.5 --- backend/api_v2/api_deployment_views.py | 4 + backend/api_v2/openapi_schema.py | 12 +- .../api_v2/tests/test_deployment_aborted.py | 56 ++++- backend/workflow_manager/execution/README.md | 26 ++- backend/workflow_manager/execution/abort.py | 44 +++- .../execution/tests/test_abort_executions.py | 61 +++--- .../tests/test_abort_transactions.py | 29 ++- specs/docstudio-oss.json | 6 +- workers/callback/tasks.py | 57 +++-- workers/file_processing/tasks.py | 62 +++--- workers/queue_backend/pg_barrier.py | 17 +- workers/shared/clients/execution_client.py | 5 +- workers/shared/processing/filter_pipeline.py | 29 ++- workers/tests/test_abort_finishes_the_run.py | 201 ++++++++++++++++-- 14 files changed, 480 insertions(+), 129 deletions(-) diff --git a/backend/api_v2/api_deployment_views.py b/backend/api_v2/api_deployment_views.py index 7f18a79b9b..45b09e57dc 100644 --- a/backend/api_v2/api_deployment_views.py +++ b/backend/api_v2/api_deployment_views.py @@ -257,6 +257,10 @@ def get( response_status = status.HTTP_422_UNPROCESSABLE_ENTITY if execution_status_value == CeleryTaskState.COMPLETED.value: response_status = status.HTTP_200_OK + if execution_status_value in ( + CeleryTaskState.COMPLETED.value, + ExecutionStatus.STOPPED.value, + ): DeploymentHelper.process_completed_execution( response=response, deployment_execution_dto=deployment_execution_dto, diff --git a/backend/api_v2/openapi_schema.py b/backend/api_v2/openapi_schema.py index a82e92739d..5178c490b3 100644 --- a/backend/api_v2/openapi_schema.py +++ b/backend/api_v2/openapi_schema.py @@ -246,7 +246,9 @@ def _get_error_response_examples(self) -> list[OpenApiExample]: "later call for that execution answers 406. Keep the payload of the call " "that returns it — it cannot be fetched again.\n\n" "A still-running execution answers 422 carrying its current `status`, so a " - "polling loop should treat 422 as the normal reply and stop on 200." + "polling loop should treat 422 as the normal reply and stop on 200, or when " + "`status` is `ERROR` or `STOPPED`. Those also answer 422 but the run has ended, " + "and that read is the only one that carries its results." ) @@ -263,7 +265,9 @@ def _get_error_response_examples(self) -> list[OpenApiExample]: responses={ 200: ExecuteResponse, 422: OpenApiResponse( - ExecuteResponse, description="The execution finished with an error." + ExecuteResponse, + description="The execution finished with an error, or was aborted " + "(`STOPPED`).", ), 500: OpenApiResponse( ExecuteResponse, @@ -289,8 +293,8 @@ def _get_error_response_examples(self) -> list[OpenApiExample]: ), 422: OpenApiResponse( StatusResponse, - description="The execution is still running, or it finished with an " - "error; read `status` to tell them apart.", + description="The execution is still running, finished with an error, " + "or was aborted (`STOPPED`); read `status` to tell them apart.", ), 500: OpenApiResponse( StatusResponse, diff --git a/backend/api_v2/tests/test_deployment_aborted.py b/backend/api_v2/tests/test_deployment_aborted.py index 5cf0c0b122..9a7caee0c2 100644 --- a/backend/api_v2/tests/test_deployment_aborted.py +++ b/backend/api_v2/tests/test_deployment_aborted.py @@ -1,4 +1,4 @@ -"""A synchronous API deployment call reports an aborted run as an error (UN-4318).""" +"""API deployment calls report an aborted run as an error (UN-4318).""" from __future__ import annotations @@ -10,6 +10,8 @@ from rest_framework.response import Response from rest_framework.test import APIRequestFactory from utils.user_context import UserContext +from workflow_manager.workflow_v2.dto import ExecutionResponse +from workflow_manager.workflow_v2.models.execution import WorkflowExecution from workflow_manager.workflow_v2.models.workflow import Workflow from api_v2.api_deployment_views import DeploymentExecution @@ -24,8 +26,8 @@ def setUp(self) -> None: name=ORG_ID, display_name="Org", organization_id=ORG_ID ) UserContext.set_organization_identifier(ORG_ID) - workflow = Workflow.objects.create(workflow_name="wf", is_active=True) - self.api = APIDeployment.objects.create(api_name="api", workflow=workflow) + self.workflow = Workflow.objects.create(workflow_name="wf", is_active=True) + self.api = APIDeployment.objects.create(api_name="api", workflow=self.workflow) self.key = APIKey.objects.create(api=self.api) def _post_with_status(self, execution_status: str) -> Response: @@ -59,3 +61,51 @@ def test_status_codes(self) -> None: self.assertEqual( self._post_with_status(execution_status).status_code, code ) + + def _get_stopped_result(self) -> Response: + item = { + "file": "doc.txt", + "metadata": {"workflow_id": "w", "highlight_data": [1]}, + "result": { + "output": {"a": 1}, + "metadata": {"inner": 1}, + "metrics": {"llm": 1}, + }, + } + execution = WorkflowExecution.objects.create( + workflow=self.workflow, + pipeline_id=self.api.id, + execution_mode=WorkflowExecution.Mode.QUEUE, + execution_method=WorkflowExecution.Method.DIRECT, + execution_type=WorkflowExecution.Type.COMPLETE, + status="STOPPED", + ) + stopped = ExecutionResponse( + workflow_id="w", + execution_id="e", + execution_status="STOPPED", + result=[item], + ) + request = APIRequestFactory().get( + f"/deployment/api/{ORG_ID}/api/", + {"execution_id": str(execution.id)}, + HTTP_AUTHORIZATION=f"Bearer {self.key.api_key}", + ) + helper = "api_v2.api_deployment_views.DeploymentHelper" + with ( + patch(f"{helper}.get_execution_status", return_value=stopped), + patch(f"{helper}._enrich_result_with_workflow_metadata"), + ): + return DeploymentExecution.as_view()( + request, org_name=ORG_ID, api_name="api" + ) + + def test_get_filters_a_stopped_runs_results_like_post(self) -> None: + response = self._get_stopped_result() + + self.assertEqual(response.status_code, 422) + item = response.data["message"][0] + self.assertNotIn("highlight_data", item["metadata"]) + self.assertNotIn("metrics", item["result"]) + self.assertNotIn("metadata", item["result"]) + self.assertEqual(item["result"]["output"], {"a": 1}) diff --git a/backend/workflow_manager/execution/README.md b/backend/workflow_manager/execution/README.md index 97c16cbdc1..00fc27daed 100644 --- a/backend/workflow_manager/execution/README.md +++ b/backend/workflow_manager/execution/README.md @@ -123,9 +123,11 @@ run in the organisation. Shared users who are not owners get `not_permitted`. Once that commits, it: -1. Sets the pipeline's last-run status to failure and sends its notifications - with status `FAILURE`, or sends an API deployment's notifications with status - `STOPPED`. An inactive pipeline gets neither. +1. Sends the notifications of the run's pipeline (status `FAILURE`) or API + deployment (status `STOPPED`), paused pipelines included. A failure here is + logged and does not change the answer. +1. Sets the pipeline's last-run status to failure, unless the pipeline has a newer + run. 2. Records an `execution.aborted` event in the platform audit log, with target `execution`, the workflow name, the summary `was ` and `previous_status`, `pipeline_id` and `entity` in its metadata. This is best @@ -133,18 +135,20 @@ Once that commits, it: ### How the work stops -- Queued batches of a stopped run are skipped when a worker picks them up. They - clear their files' active-file markers and still count down on the barrier. +- Queued batches of a stopped run are skipped when a worker picks them up, and + still count down on the barrier. So does a batch that fails after the abort. +- A re-run ignores the active-file markers of a stopped run, so it keeps those + files even if no worker has picked up the queued batch yet. - A running batch checks the status before each file and stops once it is `STOPPED`: its remaining files are marked `STOPPED` and counted as failed. - The file being processed when the abort lands finishes, and its own result replaces `STOPPED` on its row. - The last batch fires the aggregating callback. It keeps the status `STOPPED`, saves the final file counts, commits usage for finished files, publishes the - final log and cleans up the execution directory. It does not send notifications - again. + final log and cleans up the execution directory. It leaves the pipeline status + and notifications to the abort. - A caller waiting on a synchronous API deployment `POST` gets 422 for an aborted - run, the same as polling with `GET`. + run, the same as polling with `GET`. Both apply the same result filters. - Later status writes from workers cannot undo the abort, unless the run's `queue_message_id` is empty. @@ -189,7 +193,9 @@ callback reports an aborted run with a finished file as completed. deletes it. The reaper's `ERROR` path has the same gap. - A run whose `queue_message_id` is empty (from before the PG queue, or whose handle write failed) is not protected from a late worker status write. -- An abort that lands while the callback is already running stays `STOPPED`, but - the callback can still report the status it worked out. +- The file being processed when the abort lands is billed once it finishes, but an + API caller can't get its result: the first read after the abort returns the + results so far and discards them, and later reads answer 406. Its processing cost + was already spent, the same as for a caller who never collects a result. - A running batch stops between files only for `STOPPED`. A run that becomes `ERROR` or `COMPLETED` mid-batch keeps processing its files, as before. diff --git a/backend/workflow_manager/execution/abort.py b/backend/workflow_manager/execution/abort.py index 365b9f0297..1b74dd4141 100644 --- a/backend/workflow_manager/execution/abort.py +++ b/backend/workflow_manager/execution/abort.py @@ -2,16 +2,18 @@ import logging from dataclasses import dataclass -from functools import partial from typing import Any from uuid import UUID from api_v2.models import APIDeployment +from api_v2.utils import APIDeploymentUtils from django.db import transaction from django.db.models import Count, DurationField, ExpressionWrapper, F, Q, QuerySet from django.db.models.functions import Extract, Now +from django.utils import timezone from permissions.roles import ResourceRole from pipeline_v2.models import Pipeline +from pipeline_v2.notification import PipelineNotification from platform_audit.actions import Action, TargetType from platform_audit.recorder import record_platform_event from tenant_account_v2.organization_member_service import OrganizationMemberService @@ -20,7 +22,6 @@ from workflow_manager.execution.entity import entity_annotation from workflow_manager.execution.enum import AbortOutcome from workflow_manager.file_execution.models import WorkflowFileExecution -from workflow_manager.utils.pipeline_utils import PipelineUtils from workflow_manager.workflow_v2.enums import ExecutionStatus from workflow_manager.workflow_v2.models import Workflow, WorkflowExecution @@ -145,13 +146,8 @@ def _abort_one(request: Any, execution_id: UUID, entity: str | None) -> AbortRes ) _stop_open_files(execution) _write_file_counts(execution) - # Pipeline status and webhook; no callback sends them if the worker is dead. - transaction.on_commit( - partial( - PipelineUtils.update_pipeline_status, execution.pipeline_id, execution - ), - robust=True, - ) + # After commit, so a rolled-back abort reports nothing. + transaction.on_commit(lambda: _report_abort(execution), robust=True) _record_abort(request, execution, previous_status, entity) return AbortResult(execution_id, AbortOutcome.ABORTED, execution.status) @@ -185,6 +181,36 @@ def _write_file_counts(execution: WorkflowExecution) -> None: execution.save(update_fields=["total_files", "successful_files", "failed_files"]) +def _report_abort(execution: WorkflowExecution) -> None: + """Report the ended run to its pipeline or API deployment. + + The abort does this because a dead worker's run gets no callback. + """ + pipeline = Pipeline.objects.filter(pk=execution.pipeline_id).first() + if pipeline is not None: + _report_to_pipeline(pipeline, execution) + return + deployment = APIDeployment.objects.filter(pk=execution.pipeline_id).first() + if deployment is not None: + APIDeploymentUtils.send_notification(api=deployment, workflow_execution=execution) + + +def _report_to_pipeline(pipeline: Pipeline, execution: WorkflowExecution) -> None: + """Notify subscribers; the pipeline's own status follows only its newest run.""" + pipeline.last_run_status = Pipeline.PipelineStatus.FAILURE + newer_runs = WorkflowExecution.objects.filter( + pipeline_id=execution.pipeline_id, created_at__gt=execution.created_at + ) + if not newer_runs.exists(): + pipeline.last_run_time = timezone.now() + pipeline.save(update_fields=["last_run_status", "last_run_time"]) + PipelineNotification( + pipeline=pipeline, + execution_id=execution.id, + error_message=execution.error_message, + ).send() + + def _record_abort( request: Any, execution: WorkflowExecution, previous_status: str, entity: str | None ) -> None: diff --git a/backend/workflow_manager/execution/tests/test_abort_executions.py b/backend/workflow_manager/execution/tests/test_abort_executions.py index cb00fd3aef..10d7f95e2f 100644 --- a/backend/workflow_manager/execution/tests/test_abort_executions.py +++ b/backend/workflow_manager/execution/tests/test_abort_executions.py @@ -131,36 +131,47 @@ def test_abort_releases_the_slot_stamps_the_time_and_notifies_the_inbox(self) -> execution.refresh_from_db() self.assertAlmostEqual(execution.execution_time, 3600, delta=60) + def _abort_with_notification_patched(self, execution: WorkflowExecution) -> Any: + with patch("workflow_manager.execution.abort.PipelineNotification") as notify: + self._abort(self.owner, execution.id) + notify.return_value.send.assert_called_once() + return notify.call_args.kwargs + def test_abort_marks_the_pipeline_failed_and_notifies(self) -> None: - pipeline = self.pipeline() - Pipeline.objects.filter(pk=pipeline.pk).update(active=True) - execution = self.execution(pipeline) - with patch( - "pipeline_v2.pipeline_processor.PipelineProcessor._send_notification" - ) as send: - outcome = self._outcomes(self._abort(self.owner, execution.id)) - pipeline.refresh_from_db() - self.assertEqual(outcome[str(execution.id)][0], "aborted") - self.assertEqual(pipeline.last_run_status, Pipeline.PipelineStatus.FAILURE) - send.assert_called_once() - self.assertEqual(send.call_args.kwargs["execution_id"], execution.id) + for active in (True, False): + with self.subTest(active=active): + pipeline = self.pipeline() + Pipeline.objects.filter(pk=pipeline.pk).update(active=active) + execution = self.execution(pipeline) - def test_abort_of_a_paused_pipeline_run_still_aborts(self) -> None: + sent = self._abort_with_notification_patched(execution) + + pipeline.refresh_from_db() + self.assertEqual( + pipeline.last_run_status, Pipeline.PipelineStatus.FAILURE + ) + self.assertIsNotNone(pipeline.last_run_time) + self.assertEqual(sent["execution_id"], execution.id) + + def test_abort_of_an_old_run_notifies_but_keeps_the_pipeline_status(self) -> None: pipeline = self.pipeline() - execution = self.execution(pipeline) - with patch( - "pipeline_v2.pipeline_processor.PipelineProcessor._send_notification" - ) as send: - outcome = self._outcomes(self._abort(self.owner, execution.id)) + old_run = self.execution(pipeline) + self.execution(pipeline, status=ExecutionStatus.COMPLETED) + Pipeline.objects.filter(pk=pipeline.pk).update( + last_run_status=Pipeline.PipelineStatus.SUCCESS + ) + + sent = self._abort_with_notification_patched(old_run) + + self.assertEqual(sent["pipeline"].last_run_status, Pipeline.PipelineStatus.FAILURE) pipeline.refresh_from_db() - self.assertEqual(outcome[str(execution.id)][0], "aborted") - self.assertNotEqual(pipeline.last_run_status, Pipeline.PipelineStatus.FAILURE) - send.assert_not_called() + self.assertEqual(pipeline.last_run_status, Pipeline.PipelineStatus.SUCCESS) + self.assertIsNone(pipeline.last_run_time) def test_abort_of_a_deployment_run_notifies_the_deployment(self) -> None: execution = self.execution(self.deployment()) with patch( - "workflow_manager.utils.pipeline_utils.APIDeploymentUtils.send_notification" + "workflow_manager.execution.abort.APIDeploymentUtils.send_notification" ) as send: self._abort(self.owner, execution.id) send.assert_called_once() @@ -168,11 +179,9 @@ def test_abort_of_a_deployment_run_notifies_the_deployment(self) -> None: def test_finished_run_announces_nothing(self) -> None: execution = self.execution(self.pipeline(), status=ExecutionStatus.COMPLETED) - with patch( - "workflow_manager.execution.abort.PipelineUtils.update_pipeline_status" - ) as update: + with patch("workflow_manager.execution.abort._report_abort") as report: self._abort(self.owner, execution.id) - update.assert_not_called() + report.assert_not_called() def test_abort_takes_one_row_lock(self) -> None: execution = self.execution() diff --git a/backend/workflow_manager/execution/tests/test_abort_transactions.py b/backend/workflow_manager/execution/tests/test_abort_transactions.py index c9d619b220..2d4bcf530a 100644 --- a/backend/workflow_manager/execution/tests/test_abort_transactions.py +++ b/backend/workflow_manager/execution/tests/test_abort_transactions.py @@ -5,6 +5,7 @@ from unittest.mock import patch from account_v2.models import Organization +from api_v2.models import APIDeployment from django.db import transaction from django.test import TransactionTestCase from permissions.roles import ResourceRole @@ -35,9 +36,10 @@ def setUp(self) -> None: ) self.workflow.memberships.create(user=self.owner, role=ResourceRole.OWNER) - def _execution(self) -> WorkflowExecution: + def _execution(self, pipeline_id: Any = None) -> WorkflowExecution: execution: WorkflowExecution = WorkflowExecution.objects.create( workflow=self.workflow, + pipeline_id=pipeline_id, execution_mode=WorkflowExecution.Mode.QUEUE, execution_method=WorkflowExecution.Method.DIRECT, execution_type=WorkflowExecution.Type.COMPLETE, @@ -76,7 +78,7 @@ def stop(execution: WorkflowExecution) -> None: with ( patch(f"{_ABORT}._stop_open_files", side_effect=stop), - patch(f"{_ABORT}.PipelineUtils.update_pipeline_status") as update, + patch(f"{_ABORT}._report_abort") as report, ): outcomes = self._post(first.id, bad.id, last.id) @@ -92,7 +94,28 @@ def stop(execution: WorkflowExecution) -> None: self.assertEqual(self._status(bad), ExecutionStatus.EXECUTING.value) self.assertEqual(self._status(last), ExecutionStatus.STOPPED.value) self.assertEqual(PlatformAuditEvent.objects.count(), 2) - self.assertEqual(update.call_count, 2) + self.assertEqual(report.call_count, 2) + + def test_a_failing_report_keeps_the_abort_and_its_audit_event(self) -> None: + deployment = APIDeployment.objects.create( + api_name="api-1", + workflow=self.workflow, + organization=self.org, + created_by=self.owner, + ) + deployment.memberships.create(user=self.owner, role=ResourceRole.OWNER) + execution = self._execution(deployment.id) + + with patch( + f"{_ABORT}.APIDeploymentUtils.send_notification", + side_effect=RuntimeError("boom"), + ) as send: + outcomes = self._post(execution.id) + + send.assert_called_once() + self.assertEqual(outcomes, {str(execution.id): ("aborted", "STOPPED")}) + self.assertEqual(self._status(execution), ExecutionStatus.STOPPED.value) + self.assertEqual(PlatformAuditEvent.objects.count(), 1) def test_a_run_finishing_during_the_abort_is_not_overwritten(self) -> None: execution = self._execution() diff --git a/specs/docstudio-oss.json b/specs/docstudio-oss.json index 08b71f8e15..534ede8228 100644 --- a/specs/docstudio-oss.json +++ b/specs/docstudio-oss.json @@ -639,7 +639,7 @@ }, "/deployment/api/{org_name}/{api_name}/": { "get": { - "description": "Reads a previously started execution, taking the same deployment key that ran it. The read is one-shot: the first call that observes a completed execution acknowledges it and the stored result is discarded, so every later call for that execution answers 406. Keep the payload of the call that returns it \u2014 it cannot be fetched again.\n\nA still-running execution answers 422 carrying its current `status`, so a polling loop should treat 422 as the normal reply and stop on 200.", + "description": "Reads a previously started execution, taking the same deployment key that ran it. The read is one-shot: the first call that observes a completed execution acknowledges it and the stored result is discarded, so every later call for that execution answers 406. Keep the payload of the call that returns it \u2014 it cannot be fetched again.\n\nA still-running execution answers 422 carrying its current `status`, so a polling loop should treat 422 as the normal reply and stop on 200, or when `status` is `ERROR` or `STOPPED`. Those also answer 422 but the run has ended, and that read is the only one that carries its results.", "operationId": "status", "parameters": [ { @@ -833,7 +833,7 @@ } } }, - "description": "The execution is still running, or it finished with an error; read `status` to tell them apart." + "description": "The execution is still running, finished with an error, or was aborted (`STOPPED`); read `status` to tell them apart." }, "500": { "content": { @@ -1043,7 +1043,7 @@ } } }, - "description": "The execution finished with an error." + "description": "The execution finished with an error, or was aborted (`STOPPED`)." }, "429": { "content": { diff --git a/workers/callback/tasks.py b/workers/callback/tasks.py index b2024d0084..667d61468f 100644 --- a/workers/callback/tasks.py +++ b/workers/callback/tasks.py @@ -376,6 +376,15 @@ def _determine_execution_status_unified( return aggregated_results, final_status, expected_files +def _stopped_meanwhile(response: Any) -> bool: + """True when the backend refused the write because an abort got there first.""" + data = getattr(response, "data", None) or {} + return ( + data.get("reason") == "already_final" + and data.get("new_status") == ExecutionStatus.STOPPED.value + ) + + def _update_execution_status_unified( api_client: InternalAPIClient, execution_id: str, @@ -410,16 +419,27 @@ def _update_execution_status_unified( successful_files = aggregated_results.get("successful_files", 0) failed_files = aggregated_results.get("failed_files", 0) + def post_status(status: str, error: str | None) -> Any: + return api_client.update_workflow_execution_status( + execution_id=execution_id, + status=status, + total_files=total_files, + successful_files=successful_files, + failed_files=failed_files, + organization_id=organization_id, + error_message=error, + ) + # Make the unified API call - api_client.update_workflow_execution_status( - execution_id=execution_id, - status=final_status, - total_files=total_files, - successful_files=successful_files, - failed_files=failed_files, - organization_id=organization_id, - error_message=error_message, - ) + response = post_status(final_status, error_message) + if _stopped_meanwhile(response): + logger.warning( + f"Execution {execution_id} was stopped during this callback; " + f"reporting STOPPED instead of {final_status}" + ) + final_status = ExecutionStatus.STOPPED.value + # The refused write saved no file counts. + post_status(final_status, None) # Return consistent result structure update_result = { @@ -493,6 +513,15 @@ def _handle_pipeline_updates_unified( "pipeline_id": context.pipeline_id, } + # The abort wrote it, and only if this is the pipeline's newest run. + if final_status == ExecutionStatus.STOPPED.value: + return { + "status": "skipped", + "reason": "handled_by_abort", + "message": "The abort handles a stopped run's pipeline status", + "pipeline_id": context.pipeline_id, + } + # Non-API workflows (ETL/TASK/APP) need pipeline status updates if not context.pipeline_id: logger.info("No pipeline_id provided - skipping pipeline status update") @@ -620,13 +649,13 @@ def _handle_notifications_unified( Returns: Notification result dictionary """ - # The abort already sent these, so a run with a dead worker notifies too. + # The abort handles these, so a run with a dead worker notifies too. if status == ExecutionStatus.STOPPED.value: - logger.info(f"Execution {execution_id} was stopped - notifications already sent") + logger.info(f"Execution {execution_id} was stopped - the abort handles notifying") return { "status": "skipped", - "reason": "sent_by_abort", - "message": "Notifications were sent when the execution was aborted", + "reason": "handled_by_abort", + "message": "The abort handles a stopped run's notifications", } try: @@ -1537,6 +1566,7 @@ def _process_batch_callback_core( organization_id=context.organization_id, error_message=None, ) + execution_status = execution_update_result["final_status"] # Handle pipeline updates using unified function (non-API deployment) pipeline_result = _handle_pipeline_updates_unified( context, execution_status, is_api_deployment=False @@ -1803,6 +1833,7 @@ def process_batch_callback_api( aggregated_results=aggregated_results, organization_id=organization_id, ) + execution_status = execution_update_result["final_status"] # Create minimal context for unified pipeline handling context = CallbackContext() diff --git a/workers/file_processing/tasks.py b/workers/file_processing/tasks.py index 72634a1317..52b53f2d32 100644 --- a/workers/file_processing/tasks.py +++ b/workers/file_processing/tasks.py @@ -242,7 +242,7 @@ class _TerminalExecutionSkip(Exception): The resurrection this prevents exists because the **reaper** can mark an in-flight execution terminal (and tear its barrier down) while a batch for it is still in flight or waiting to be redelivered. An abort can mark it STOPPED - the same way, but leaves the barrier armed. + the same way. This is a validate-first *narrowing*, not a transactional guarantee: the status read and the later ``EXECUTING`` write are not atomic, so a reaper @@ -301,40 +301,52 @@ def _terminal_skip_result( so reporting ``failed=total`` keeps a consumer that ever aggregates this result from rendering the files as perpetually in-progress. The ``SKIPPED_TERMINAL_EXECUTION_KEY`` marker tells ``run_batch_with_barrier`` to - bypass the barrier decrement (the reaper already tore that barrier down). A - STOPPED run's barrier is still armed, so it gets no marker: the batch counts + bypass the barrier decrement. A STOPPED run gets no marker: the batch counts down and the last one fires the callback, which cleans up. """ - n_files = len(batch_data.files) - result = BatchExecutionResult( + result = _all_failed_result( + len(batch_data.files), batch_data.file_data.organization_id + ) + if status != ExecutionStatus.STOPPED.value: + result[SKIPPED_TERMINAL_EXECUTION_KEY] = True + return result + + +def _all_failed_result(n_files: int, organization_id: str | None) -> dict[str, Any]: + """Batch result that counts every file as failed.""" + result: dict[str, Any] = BatchExecutionResult( total_files=n_files, successful_files=0, failed_files=n_files, execution_time=0.0, - organization_id=batch_data.file_data.organization_id, + organization_id=organization_id, ).to_dict() - if status != ExecutionStatus.STOPPED.value: - result[SKIPPED_TERMINAL_EXECUTION_KEY] = True return result -def _clear_active_file_markers(batch_data: FileBatchData) -> None: - """Clear a skipped batch's active-file markers, which pre-create never reached. +def _run_stages_under_barrier( + file_batch_data: dict[str, Any], celery_task_id: str +) -> dict[str, Any]: + """Run the batch stages; a stopped run's failed batch still counts down. - Without this, a re-run leaves these files out until the markers expire. + Otherwise the barrier is torn down and the run's callback never fires. """ - from shared.workflow.execution.active_file_manager import cleanup_active_file_cache - - uuids = [ - uuid - for item in batch_data.files - if (uuid := _file_item_hash(item).get("provider_file_uuid")) - ] - cleanup_active_file_cache( - provider_file_uuids=uuids, - workflow_id=batch_data.file_data.workflow_id, - logger_instance=logger, - ) + try: + return _run_batch_stages(file_batch_data, celery_task_id) + except Exception: + file_data = file_batch_data.get("file_data") or {} + execution_id = file_data.get("execution_id") + organization_id = file_data.get("organization_id") + api_client = create_api_client(organization_id) + if not _is_execution_stopped(api_client, execution_id): + raise + logger.exception( + "[exec:%s] Batch failed after the run was stopped; counting it down", + execution_id, + ) + return _all_failed_result( + len(file_batch_data.get("files") or []), organization_id + ) def _run_batch_stages( @@ -360,8 +372,6 @@ def _run_batch_stages( f"{len(batch_data.files)} file(s): execution already terminal " f"({skip.status}); not processing." ) - if skip.status == ExecutionStatus.STOPPED.value: - _clear_active_file_markers(batch_data) return _terminal_skip_result(batch_data, skip.status) # Step 3: Handle manual review logic @@ -424,7 +434,7 @@ def _process_file_batch_core( # the barrier in-body / self-chain the callback. return run_batch_with_barrier( barrier_context, - lambda: _run_batch_stages(file_batch_data, celery_task_id), + lambda: _run_stages_under_barrier(file_batch_data, celery_task_id), ) diff --git a/workers/queue_backend/pg_barrier.py b/workers/queue_backend/pg_barrier.py index 6102f4c8bc..a85e6ac541 100644 --- a/workers/queue_backend/pg_barrier.py +++ b/workers/queue_backend/pg_barrier.py @@ -1305,12 +1305,10 @@ def _mark_execution_error_on_abort( ) -# A batch skipped for a run the reaper already finished returns its result with -# this marker set. run_batch_with_barrier bypasses the barrier decrement for it: -# the reaper already tore the barrier down, so decrementing would only log a -# spurious "decrement found no row" ERROR (false alert noise), and that run must -# not have its aggregating callback fired. An aborted (STOPPED) run keeps its -# barrier, so its batches do not set this. +# A batch skipped because its execution is terminal returns its result with this +# marker set, and run_batch_with_barrier bypasses the barrier decrement for it, so +# that run's aggregating callback does not fire. Set for every terminal status +# except STOPPED: an aborted run's batches count down so the callback can finish it. SKIPPED_TERMINAL_EXECUTION_KEY = "skipped_terminal_execution" @@ -1381,11 +1379,8 @@ def run_batch_with_barrier( try: result = work_fn() if result.get(SKIPPED_TERMINAL_EXECUTION_KEY): - # The reaper already finished this run and tore its barrier down. - # Bypass the decrement — decrementing a gone barrier only logs a - # spurious "decrement found no row" ERROR (false alert noise) — and - # do NOT abort (nothing to tear down; and that run must not have its - # aggregating callback re-fired). + # Bypass the decrement and do NOT abort: this run's aggregating + # callback must not fire. logger.info( f"[exec:{execution_id}] batch {batch_index} skipped — execution " f"already terminal; barrier decrement bypassed." diff --git a/workers/shared/clients/execution_client.py b/workers/shared/clients/execution_client.py index 0914d4543c..abc225bc4d 100644 --- a/workers/shared/clients/execution_client.py +++ b/workers/shared/clients/execution_client.py @@ -344,7 +344,10 @@ def update_workflow_execution_status( ) # Convert dict response to consistent APIResponse - return convert_dict_response(response, APIResponse) + api_response = convert_dict_response(response, APIResponse) + # The body is flat, and a refused write says why only there. + api_response.data = response + return api_response def recover_stuck_pg_executions( self, diff --git a/workers/shared/processing/filter_pipeline.py b/workers/shared/processing/filter_pipeline.py index 7a47bc7d42..a1a0b76750 100644 --- a/workers/shared/processing/filter_pipeline.py +++ b/workers/shared/processing/filter_pipeline.py @@ -11,6 +11,7 @@ from ..api.internal_client import InternalAPIClient from ..cache.cache_backends import RedisCacheBackend +from ..clients.base_client import InternalAPIClientError from ..infrastructure.logging import WorkerLogger from ..workflow.execution.active_file_manager import ActiveFileManager @@ -610,6 +611,25 @@ def apply( return filtered + @staticmethod + def _marker_run_is_live( + run_id: str, api_client: Any, live_runs: dict[str, bool] + ) -> bool: + """False for a stopped run, whose queued files no longer hold their markers. + + A failed status read keeps the marker in force. + """ + if run_id not in live_runs: + try: + status = api_client.get_workflow_execution_status(run_id) + except InternalAPIClientError as error: + logger.warning( + "[ActiveFileFilter] Status read for run %s failed: %s", run_id, error + ) + status = None + live_runs[run_id] = status != ExecutionStatus.STOPPED.value + return live_runs[run_id] + def _check_active_files_batch( self, file_identifiers: dict[str, dict[str, Any]], @@ -640,6 +660,7 @@ def _check_active_files_batch( # 1. Check Redis cache for active files using precise cache keys try: cache = RedisCacheBackend() + live_runs: dict[str, bool] = {} for identifier in identifiers_to_check: data = file_identifiers[identifier] uuid = data["uuid"] @@ -653,7 +674,13 @@ def _check_active_files_batch( if cached_data and isinstance(cached_data, dict): cached_exec_id = cached_data.get("data", {}).get("execution_id") - if cached_exec_id and cached_exec_id != execution_id: + if ( + cached_exec_id + and cached_exec_id != execution_id + and self._marker_run_is_live( + cached_exec_id, api_client, live_runs + ) + ): active_identifiers.add(identifier) logger.debug( f"[ActiveFileFilter] File {identifier} active in cache (exec: {cached_exec_id})" diff --git a/workers/tests/test_abort_finishes_the_run.py b/workers/tests/test_abort_finishes_the_run.py index 03fd147287..105533881b 100644 --- a/workers/tests/test_abort_finishes_the_run.py +++ b/workers/tests/test_abort_finishes_the_run.py @@ -1,5 +1,6 @@ -"""An aborted run still finishes: skipped batches count down, and the callback -keeps STOPPED, bills finished files and does not notify again (UN-4318). +"""An aborted run still finishes: skipped and failed batches count down, and the +callback keeps STOPPED, bills finished files and leaves the pipeline status and +notifications to the abort. A re-run ignores the stopped run's markers (UN-4318). """ from __future__ import annotations @@ -10,12 +11,16 @@ import callback.tasks as cb import pytest +from shared.clients.base_client import APIRequestError +from shared.clients.execution_client import ExecutionAPIClient from shared.infrastructure.logging.workflow_logger import WorkerWorkflowLogger +from shared.processing.filter_pipeline import ActiveFileFilter from unstract.core.data_models import ExecutionStatus from unstract.workflow_execution.enums import LogState from file_processing.tasks import ( _run_batch_stages, + _run_stages_under_barrier, _terminal_skip_result, _TerminalExecutionSkip, ) @@ -26,7 +31,7 @@ ERROR = ExecutionStatus.ERROR.value COMPLETED = ExecutionStatus.COMPLETED.value _TASKS = "file_processing.tasks" -_CLEANUP = "shared.workflow.execution.active_file_manager.cleanup_active_file_cache" +_REFUSAL = {"status": "skipped", "reason": "already_final", "new_status": STOPPED} def _batch_data() -> Any: @@ -39,36 +44,27 @@ def data(index: int) -> dict[str, str]: ) -def _skip_batch(status: str) -> tuple[dict[str, Any], Any]: +def _skip_batch(status: str) -> dict[str, Any]: with ( mock.patch(f"{_TASKS}._validate_and_parse_batch_data", return_value=_batch_data()), mock.patch( f"{_TASKS}._setup_execution_context", side_effect=_TerminalExecutionSkip("exec-1", status), ), - mock.patch(_CLEANUP) as cleanup, ): - return _run_batch_stages({"any": "payload"}, "task-1"), cleanup + return _run_batch_stages({"any": "payload"}, "task-1") -def test_a_stopped_runs_skipped_batch_clears_its_markers_and_counts_down() -> None: - result, cleanup = _skip_batch(STOPPED) +def test_a_stopped_runs_skipped_batch_counts_down() -> None: + result = _skip_batch(STOPPED) assert SKIPPED_TERMINAL_EXECUTION_KEY not in result assert result["failed_files"] == 4 - cleanup.assert_called_once_with( - provider_file_uuids=["uuid-0", "uuid-1", "uuid-2"], - workflow_id="wf-1", - logger_instance=mock.ANY, - ) @pytest.mark.parametrize("status", [ERROR, COMPLETED]) -def test_a_run_the_reaper_finished_keeps_the_marker(status: str) -> None: - result, cleanup = _skip_batch(status) - - assert result[SKIPPED_TERMINAL_EXECUTION_KEY] is True - cleanup.assert_not_called() +def test_other_terminal_statuses_keep_the_marker(status: str) -> None: + assert _skip_batch(status)[SKIPPED_TERMINAL_EXECUTION_KEY] is True @pytest.mark.parametrize(("status", "decrements"), [(STOPPED, 1), (ERROR, 0)]) @@ -86,6 +82,49 @@ def test_only_a_stopped_runs_batch_decrements(status: str, decrements: int) -> N abort.assert_not_called() +_PAYLOAD = { + "files": [["f0.pdf", {}], ["f1.pdf", {}]], + "file_data": {"execution_id": "e", "organization_id": "org-1"}, +} + + +def _failing_batch(status: Any) -> Any: + client = mock.Mock() + client.get_workflow_execution_status.side_effect = [status] + return ( + mock.patch(f"{_TASKS}._run_batch_stages", side_effect=RuntimeError("boom")), + mock.patch(f"{_TASKS}.create_api_client", return_value=client), + ) + + +def test_a_batch_failing_after_the_abort_counts_down() -> None: + stages, client = _failing_batch(STOPPED) + context: Any = {"execution_id": "e", "batch_index": 0, "callback_descriptor": {}} + with ( + stages, + client, + mock.patch.object(pg_barrier, "claim_batch", return_value=True), + mock.patch.object(pg_barrier, "_barrier_pg_decrement") as decrement, + mock.patch.object(pg_barrier, "_abort_barrier_in_body") as abort, + ): + pg_barrier.run_batch_with_barrier( + context, lambda: _run_stages_under_barrier(_PAYLOAD, "task-1") + ) + + assert decrement.call_args.args[0]["failed_files"] == 2 + assert SKIPPED_TERMINAL_EXECUTION_KEY not in decrement.call_args.args[0] + abort.assert_not_called() + + +@pytest.mark.parametrize( + "status", [ExecutionStatus.EXECUTING.value, APIRequestError("500")], ids=["up", "down"] +) +def test_a_batch_failing_in_a_live_run_still_raises(status: Any) -> None: + stages, client = _failing_batch(status) + with stages, client, pytest.raises(RuntimeError, match="boom"): + _run_stages_under_barrier(_PAYLOAD, "task-1") + + def _status_for(run_status: str, successful: int, failed: int) -> str: counts = { "total_files": successful + failed, @@ -131,6 +170,60 @@ def test_a_missing_context_is_judged_by_its_counts() -> None: assert cb._was_stopped({}) is False +def test_an_abort_after_the_callbacks_fetch_is_reported_as_stopped() -> None: + api_client = mock.Mock() + api_client.update_workflow_execution_status.side_effect = [ + SimpleNamespace(data=_REFUSAL), + SimpleNamespace(data={"status": "updated"}), + ] + + result = cb._update_execution_status_unified( + api_client, "e1", COMPLETED, {"total_files": 2, "successful_files": 2}, "org-1" + ) + + calls = api_client.update_workflow_execution_status.call_args_list + assert [c.kwargs["status"] for c in calls] == [COMPLETED, STOPPED] + assert calls[1].kwargs["successful_files"] == 2 + assert calls[1].kwargs["error_message"] is None + assert result["final_status"] == STOPPED + + +def test_an_accepted_write_keeps_its_status() -> None: + api_client = mock.Mock() + api_client.update_workflow_execution_status.return_value = SimpleNamespace( + data={"status": "updated"} + ) + + result = cb._update_execution_status_unified(api_client, "e1", COMPLETED, {}, "org-1") + + assert api_client.update_workflow_execution_status.call_count == 1 + assert result["final_status"] == COMPLETED + + +def test_the_client_keeps_the_body_of_a_refused_write() -> None: + client = ExecutionAPIClient.__new__(ExecutionAPIClient) + with mock.patch.object(client, "post", return_value=dict(_REFUSAL)): + response = client.update_workflow_execution_status("e1", COMPLETED) + + assert response.success is True + assert cb._stopped_meanwhile(response) is True + + +@pytest.mark.parametrize(("status", "calls"), [(STOPPED, 0), (COMPLETED, 1)]) +def test_only_the_abort_writes_a_stopped_runs_pipeline_status( + status: str, calls: int +) -> None: + context = cb.CallbackContext() + context.pipeline_id = "3f0d6c1e-5a0b-4a39-9a7e-0d6f3c2b1a90" + context.execution_id = "e1" + context.organization_id = "org-1" + context.api_client = mock.Mock() + with mock.patch.object(cb, "_update_pipeline_directly") as update: + cb._handle_pipeline_updates_unified(context, status, is_api_deployment=False) + + assert update.call_count == calls + + @pytest.mark.parametrize(("status", "calls"), [(STOPPED, 0), (COMPLETED, 1)]) def test_a_stopped_run_is_not_notified_twice(status: str, calls: int) -> None: with mock.patch.object(cb, "handle_status_notifications") as notify: @@ -143,7 +236,10 @@ def test_a_stopped_run_is_not_notified_twice(status: str, calls: int) -> None: ) assert notify.call_count == calls - assert result["status"] == ("skipped" if status == STOPPED else "completed") + if status == STOPPED: + assert result["reason"] == "handled_by_abort" + else: + assert result["status"] == "completed" @pytest.mark.parametrize(("status", "commits"), [(STOPPED, 1), (ERROR, 0)]) @@ -165,6 +261,73 @@ def test_a_stopped_run_bills_only_its_finished_files(status: str, commits: int) assert commit.call_args.kwargs["file_execution_ids"] == ["done"] +@pytest.mark.parametrize( + ("publish", "creator"), + [ + (cb._publish_final_workflow_ui_logs, "create_for_general_workflow"), + (cb._publish_final_workflow_ui_logs_api, "create_for_api_workflow"), + ], + ids=["general", "api"], +) +@pytest.mark.parametrize(("status", "stopped"), [(STOPPED, True), (COMPLETED, False)]) +def test_both_callbacks_pass_the_stop_to_the_final_log( + publish: Any, creator: str, status: str, stopped: bool +) -> None: + context = cb.CallbackContext() + context.execution_id = "e1" + context.organization_id = "org-1" + context.pipeline_id = "p1" + context.api_client = mock.Mock() + context.api_client.get_workflow_execution.return_value = SimpleNamespace( + success=True, data={"aggregated_usage_cost": 0} + ) + workflow_logger = mock.Mock() + with mock.patch.object( + cb.WorkerWorkflowLogger, creator, return_value=workflow_logger + ): + publish(context=context, aggregated_results={}, execution_status=status) + + final = workflow_logger.publish_final_workflow_logs + assert final.call_args.kwargs["stopped"] is stopped + + +def _active_files(status: Any) -> tuple[set[str], Any]: + api_client = mock.Mock() + api_client.get_workflow_execution_status.side_effect = [status] + api_client.check_files_active_processing.return_value = SimpleNamespace( + success=True, data={} + ) + cache = mock.Mock() + cache.get.return_value = {"data": {"execution_id": "old-run"}} + files = { + "u1:/a": {"uuid": "u1", "path": "/a"}, + "u2:/b": {"uuid": "u2", "path": "/b"}, + } + with mock.patch( + "shared.processing.filter_pipeline.RedisCacheBackend", return_value=cache + ): + found = ActiveFileFilter()._check_active_files_batch( + files, "wf-1", "new-run", api_client + ) + return found, api_client + + +def test_a_rerun_ignores_a_stopped_runs_markers() -> None: + found, api_client = _active_files(STOPPED) + + assert found == set() + api_client.get_workflow_execution_status.assert_called_once_with("old-run") + + +@pytest.mark.parametrize( + "status", [ExecutionStatus.EXECUTING.value, APIRequestError("404")], ids=["live", "down"] +) +def test_other_runs_markers_stay_in_force(status: Any) -> None: + found, _ = _active_files(status) + + assert found == {"u1:/a", "u2:/b"} + + @pytest.mark.parametrize( ("stopped", "state", "message"), [ From d742282b86e844f8800495dd44b3355ad2eabe85 Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Tue, 10 Nov 2026 20:08:39 +0530 Subject: [PATCH 5/8] UN-4318 [FEAT] Fix third review round: stopped-run counts, failed batches and MCP status - The backend counts a stopped run's file rows on a STOPPED write, so a failed callback or a crashed batch can't lower its saved counts. - A failed batch marks its run ERROR and counts down when the backend refuses because the run is STOPPED, so an abort landing just before the failure no longer loses the callback. - The abort's report logs failures with the run id, and a failed pipeline save still sends the notifications. - The MCP getExecutionStatus tool filters a stopped run's result and names STOPPED as a final status. - A re-run logs each stopped-run marker it ignores. - Tests drive both callbacks and the batch core; README, docstrings and log formatting corrected. Co-Authored-By: Claude Opus 5.5 --- backend/mcp_server/registry.py | 6 +- .../mcp_server/tests/test_tool_execution.py | 14 ++ backend/mcp_server/tools/execution.py | 7 +- backend/mcp_server/tools/info.py | 2 +- backend/workflow_manager/execution/README.md | 32 +++-- backend/workflow_manager/execution/abort.py | 52 +++++-- .../execution/tests/test_abort_executions.py | 42 ++++++ .../tests/test_abort_transactions.py | 12 +- backend/workflow_manager/internal_views.py | 8 +- workers/callback/tasks.py | 32 ++--- workers/file_processing/tasks.py | 56 ++++++-- workers/shared/clients/execution_client.py | 13 +- workers/shared/processing/filter_pipeline.py | 12 +- workers/tests/test_abort_finishes_the_run.py | 134 ++++++++++++++++-- 14 files changed, 331 insertions(+), 91 deletions(-) diff --git a/backend/mcp_server/registry.py b/backend/mcp_server/registry.py index e3e16bf711..7a93beddbb 100644 --- a/backend/mcp_server/registry.py +++ b/backend/mcp_server/registry.py @@ -266,8 +266,8 @@ def build_deployment_registry() -> MCPToolRegistry: "Fetch the status and, once available, the result of an " "extraction previously started by `extractDocument`. Pass the " "`execution_id` that call returned.\n\n" - "Returns an `execution_status` of PENDING, EXECUTING, COMPLETED " - "or ERROR. Poll while it is PENDING or EXECUTING, leaving a few " + "Returns an `execution_status` of PENDING, EXECUTING, COMPLETED, " + "ERROR or STOPPED. Poll while it is PENDING or EXECUTING, leaving a few " "seconds between calls." ), input_schema=get_execution_status_schema(), @@ -715,7 +715,7 @@ def build_platform_registry() -> MCPToolRegistry: "extractDocument. Free to call.\n\n" "If extractDocument already returned execution_status " "COMPLETED, the result is in that response and no polling is " - "needed. Otherwise poll this until COMPLETED or ERROR, " + "needed. Otherwise poll this until COMPLETED, ERROR or STOPPED, " "pausing a few seconds between calls." ), input_schema=platform_execution_status_schema(), diff --git a/backend/mcp_server/tests/test_tool_execution.py b/backend/mcp_server/tests/test_tool_execution.py index 94c2872e82..917e409698 100644 --- a/backend/mcp_server/tests/test_tool_execution.py +++ b/backend/mcp_server/tests/test_tool_execution.py @@ -314,6 +314,20 @@ def test_completed_execution_is_enriched_and_returned(self) -> None: assert result["result"][0]["result"]["output"]["total"] == 42 assert result["execution_id"] == EXECUTION_ID + def test_stopped_execution_is_enriched_too(self) -> None: + response = ExecutionResponse( + workflow_id="wf", + execution_id=EXECUTION_ID, + execution_status="STOPPED", + result=[{"file": "invoice.pdf", "result": {"output": {"total": 42}}}], + ) + + result, process = self._run(response, include_metrics=True) + + process.assert_called_once() + assert process.call_args.kwargs["include_metrics"] is True + assert result["execution_status"] == "STOPPED" + def test_pending_execution_is_not_enriched(self) -> None: response = ExecutionResponse( workflow_id="wf", diff --git a/backend/mcp_server/tools/execution.py b/backend/mcp_server/tools/execution.py index 5ac468091d..d2bca3eee8 100644 --- a/backend/mcp_server/tools/execution.py +++ b/backend/mcp_server/tools/execution.py @@ -14,6 +14,7 @@ from tags.serializers import TagParamsSerializer from utils.enums import CeleryTaskState from workflow_manager.workflow_v2.dto import ExecutionResponse +from workflow_manager.workflow_v2.enums import ExecutionStatus from workflow_manager.workflow_v2.exceptions import ExecutionDoesNotExistError from workflow_manager.workflow_v2.models.execution import WorkflowExecution @@ -375,7 +376,11 @@ def get_execution_status( ), } - if response.execution_status == CeleryTaskState.COMPLETED.value: + # A stopped run returns its results so far, filtered the same way. + if response.execution_status in ( + CeleryTaskState.COMPLETED.value, + ExecutionStatus.STOPPED.value, + ): deployment_execution_dto = DeploymentExecutionDTO( api=context.api, api_key=context.api_key ) diff --git a/backend/mcp_server/tools/info.py b/backend/mcp_server/tools/info.py index e5bbbc2807..68a69bdefc 100644 --- a/backend/mcp_server/tools/info.py +++ b/backend/mcp_server/tools/info.py @@ -50,7 +50,7 @@ def read_me_first(context: MCPContext) -> dict[str, Any]: "3. If the response has execution_status COMPLETED, the result is " "already in the response — you are done.", "4. Otherwise poll getExecutionStatus with the returned " - "execution_id until the status is COMPLETED or ERROR, pausing a " + "execution_id until the status is COMPLETED, ERROR or STOPPED, pausing a " "few seconds between polls.", ], "notes": [ diff --git a/backend/workflow_manager/execution/README.md b/backend/workflow_manager/execution/README.md index 00fc27daed..82cbd68e76 100644 --- a/backend/workflow_manager/execution/README.md +++ b/backend/workflow_manager/execution/README.md @@ -3,8 +3,7 @@ Finding executions that stopped making progress, and stopping running ones (PB-T08, UN-4318). This README is the published contract, backed by the `@extend_schema` annotations in `views/stuck_abort.py` and `views/execution.py`. -`stuck_since` is documented here only: the generated schema shows the plain list -row. Changes to anything below are contract changes. +Changes to anything below are contract changes. - [Endpoints](#endpoints) - [What counts as stuck](#what-counts-as-stuck) @@ -118,17 +117,18 @@ run in the organisation. Shared users who are not owners get `not_permitted`. 2. Sets its `PENDING` and `EXECUTING` file rows to `STOPPED` with the same error, and stamps their execution time. Finished file rows, file history, destination writes and usage are kept. -3. Writes the run's file counts: finished files are successful, and every other - file is failed. +3. Writes the run's file counts from its file rows: completed files are + successful, and every other file is failed. Once that commits, it: -1. Sends the notifications of the run's pipeline (status `FAILURE`) or API - deployment (status `STOPPED`), paused pipelines included. A failure here is - logged and does not change the answer. 1. Sets the pipeline's last-run status to failure, unless the pipeline has a newer run. -2. Records an `execution.aborted` event in the platform audit log, with target +2. Sends the notifications of the run's pipeline (status `FAILURE`) or API + deployment (status `STOPPED`), paused pipelines included. They go out even if + step 1 fails. A failure in either step is logged with the run's id and does + not change the answer. +3. Records an `execution.aborted` event in the platform audit log, with target `execution`, the workflow name, the summary `was ` and `previous_status`, `pipeline_id` and `entity` in its metadata. This is best effort: a lost event does not fail the abort and is not retried. @@ -136,7 +136,8 @@ Once that commits, it: ### How the work stops - Queued batches of a stopped run are skipped when a worker picks them up, and - still count down on the barrier. So does a batch that fails after the abort. + still count down on the barrier. So does a batch that fails after the abort, + unless the backend can't be reached when it fails. - A re-run ignores the active-file markers of a stopped run, so it keeps those files even if no worker has picked up the queued batch yet. - A running batch checks the status before each file and stops once it is @@ -144,7 +145,8 @@ Once that commits, it: - The file being processed when the abort lands finishes, and its own result replaces `STOPPED` on its row. - The last batch fires the aggregating callback. It keeps the status `STOPPED`, - saves the final file counts, commits usage for finished files, publishes the + saves the final file counts (the backend counts a stopped run's file rows, + whatever the worker sends), commits usage for completed files, publishes the final log and cleans up the execution directory. It leaves the pipeline status and notifications to the abort. - A caller waiting on a synchronous API deployment `POST` gets 422 for an aborted @@ -173,10 +175,11 @@ ids already aborted answer `already_finished`. Deploy the backend before the workers. Workers read the status from the internal `GET /internal/v1/workflow-execution//status/` endpoint. An older backend answers 404, and the worker then processes the file anyway (fail -open), as it did before this change. +open), as it did before this change. Until the backend is updated, each file +also logs an `ERROR` and a `WARNING` for the failed check. Expose abort to users only once the workers are updated too. An older worker's -callback reports an aborted run with a finished file as completed. +callback reports an aborted run with a completed file as completed. ## Limits @@ -186,9 +189,10 @@ callback reports an aborted run with a finished file as completed. being processed no longer counts as active, so a manual or scheduled re-run can pick it up while the old worker finishes it: double LLM spend, and two writes to a database destination. Wait for that file to finish before re-running. -- **A dead worker's run never gets its callback.** The pipeline status, +- **A dead worker's run never gets its callback**, and neither does a run whose + batch fails while the backend can't be reached. The pipeline status, notifications and file counts are still set by the abort, but usage for its - finished files is not committed and its execution directory, which can hold + completed files is not committed and its execution directory, which can hold uploaded documents, is not cleaned up. It stays in storage until something else deletes it. The reaper's `ERROR` path has the same gap. - A run whose `queue_message_id` is empty (from before the PG queue, or whose diff --git a/backend/workflow_manager/execution/abort.py b/backend/workflow_manager/execution/abort.py index 1b74dd4141..bdf0105a3e 100644 --- a/backend/workflow_manager/execution/abort.py +++ b/backend/workflow_manager/execution/abort.py @@ -145,7 +145,7 @@ def _abort_one(request: Any, execution_id: UUID, entity: str | None) -> AbortRes status=ExecutionStatus.STOPPED, error=ABORT_ERROR_MESSAGE, row_locked=True ) _stop_open_files(execution) - _write_file_counts(execution) + write_file_counts(execution) # After commit, so a rolled-back abort reports nothing. transaction.on_commit(lambda: _report_abort(execution), robust=True) _record_abort(request, execution, previous_status, entity) @@ -169,8 +169,8 @@ def _stop_open_files(execution: WorkflowExecution) -> None: ) -def _write_file_counts(execution: WorkflowExecution) -> None: - """Save the run's file counts; anything that did not finish counts as failed.""" +def write_file_counts(execution: WorkflowExecution) -> None: + """Save the run's file counts from its rows; anything not completed is failed.""" counts = WorkflowFileExecution.objects.filter(workflow_execution=execution).aggregate( total=Count("pk"), completed=Count("pk", filter=Q(status=ExecutionStatus.COMPLETED.value)), @@ -186,17 +186,44 @@ def _report_abort(execution: WorkflowExecution) -> None: The abort does this because a dead worker's run gets no callback. """ - pipeline = Pipeline.objects.filter(pk=execution.pipeline_id).first() - if pipeline is not None: - _report_to_pipeline(pipeline, execution) - return - deployment = APIDeployment.objects.filter(pk=execution.pipeline_id).first() - if deployment is not None: - APIDeploymentUtils.send_notification(api=deployment, workflow_execution=execution) + try: + pipeline = Pipeline.objects.filter(pk=execution.pipeline_id).first() + if pipeline is not None: + _report_to_pipeline(pipeline, execution) + return + deployment = APIDeployment.objects.filter(pk=execution.pipeline_id).first() + if deployment is not None: + APIDeploymentUtils.send_notification( + api=deployment, workflow_execution=execution + ) + except Exception: + logger.exception( + "Failed to report aborted execution %s to %s", + execution.id, + execution.pipeline_id, + ) def _report_to_pipeline(pipeline: Pipeline, execution: WorkflowExecution) -> None: """Notify subscribers; the pipeline's own status follows only its newest run.""" + try: + _save_pipeline_status(pipeline, execution) + except Exception: + # Subscribers are still notified. + logger.exception( + "Failed to save pipeline %s status for aborted execution %s", + pipeline.id, + execution.id, + ) + PipelineNotification( + pipeline=pipeline, + execution_id=execution.id, + error_message=execution.error_message, + ).send() + + +def _save_pipeline_status(pipeline: Pipeline, execution: WorkflowExecution) -> None: + """Set the last-run status to failure, unless the pipeline has a newer run.""" pipeline.last_run_status = Pipeline.PipelineStatus.FAILURE newer_runs = WorkflowExecution.objects.filter( pipeline_id=execution.pipeline_id, created_at__gt=execution.created_at @@ -204,11 +231,6 @@ def _report_to_pipeline(pipeline: Pipeline, execution: WorkflowExecution) -> Non if not newer_runs.exists(): pipeline.last_run_time = timezone.now() pipeline.save(update_fields=["last_run_status", "last_run_time"]) - PipelineNotification( - pipeline=pipeline, - execution_id=execution.id, - error_message=execution.error_message, - ).send() def _record_abort( diff --git a/backend/workflow_manager/execution/tests/test_abort_executions.py b/backend/workflow_manager/execution/tests/test_abort_executions.py index 10d7f95e2f..9d123ccf8e 100644 --- a/backend/workflow_manager/execution/tests/test_abort_executions.py +++ b/backend/workflow_manager/execution/tests/test_abort_executions.py @@ -153,6 +153,17 @@ def test_abort_marks_the_pipeline_failed_and_notifies(self) -> None: self.assertIsNotNone(pipeline.last_run_time) self.assertEqual(sent["execution_id"], execution.id) + def test_a_failed_pipeline_save_still_notifies(self) -> None: + execution = self.execution(self.pipeline()) + with ( + patch.object(Pipeline, "save", side_effect=RuntimeError("boom")), + self.assertLogs("workflow_manager.execution.abort", level="ERROR") as logs, + ): + sent = self._abort_with_notification_patched(execution) + + self.assertEqual(sent["execution_id"], execution.id) + self.assertIn(str(execution.id), logs.output[0]) + def test_abort_of_an_old_run_notifies_but_keeps_the_pipeline_status(self) -> None: pipeline = self.pipeline() old_run = self.execution(pipeline) @@ -403,3 +414,34 @@ def test_later_completion_cannot_undo_the_abort(self) -> None: self.assertEqual(out.data.get("reason"), "already_final") self.assertEqual(self._status(execution), ExecutionStatus.STOPPED.value) + + def test_a_stopped_write_saves_the_counts_of_the_file_rows(self) -> None: + execution = self.execution(self.pipeline()) + for status in (ExecutionStatus.COMPLETED, ExecutionStatus.COMPLETED): + self.file(execution, status, idle=RECENT) + self.file(execution, ExecutionStatus.EXECUTING, idle=RECENT) + self._abort(self.owner, execution.id) + + # The worker reports zeros, as a failed callback's retry does. + view = WorkflowExecutionInternalViewSet() + request = MagicMock() + request.data = { + "status": ExecutionStatus.STOPPED.value, + "total_files": 0, + "successful_files": 0, + "failed_files": 0, + } + fresh = WorkflowExecution.objects.get(pk=execution.pk) + with patch.object(view, "get_object", return_value=fresh): + out = view.update_status(request, id=str(execution.id)) + + self.assertEqual(out.data.get("status"), "updated") + execution.refresh_from_db() + self.assertEqual( + ( + execution.total_files, + execution.successful_files, + execution.failed_files, + ), + (3, 2, 1), + ) diff --git a/backend/workflow_manager/execution/tests/test_abort_transactions.py b/backend/workflow_manager/execution/tests/test_abort_transactions.py index 2d4bcf530a..750141564b 100644 --- a/backend/workflow_manager/execution/tests/test_abort_transactions.py +++ b/backend/workflow_manager/execution/tests/test_abort_transactions.py @@ -106,13 +106,17 @@ def test_a_failing_report_keeps_the_abort_and_its_audit_event(self) -> None: deployment.memberships.create(user=self.owner, role=ResourceRole.OWNER) execution = self._execution(deployment.id) - with patch( - f"{_ABORT}.APIDeploymentUtils.send_notification", - side_effect=RuntimeError("boom"), - ) as send: + with ( + patch( + f"{_ABORT}.APIDeploymentUtils.send_notification", + side_effect=RuntimeError("boom"), + ) as send, + self.assertLogs(_ABORT, level="ERROR") as logs, + ): outcomes = self._post(execution.id) send.assert_called_once() + self.assertIn(str(execution.id), logs.output[0]) self.assertEqual(outcomes, {str(execution.id): ("aborted", "STOPPED")}) self.assertEqual(self._status(execution), ExecutionStatus.STOPPED.value) self.assertEqual(PlatformAuditEvent.objects.count(), 1) diff --git a/backend/workflow_manager/internal_views.py b/backend/workflow_manager/internal_views.py index 72f2ca1415..a56e3116d4 100644 --- a/backend/workflow_manager/internal_views.py +++ b/backend/workflow_manager/internal_views.py @@ -645,7 +645,13 @@ def update_status(self, request, id=None): # either way: Django's foreign keys are DEFERRABLE INITIALLY DEFERRED, # so an open insert holds no lock on this row until it commits, and # FOR UPDATE never waited for it.) - self._update_file_aggregates(execution, validated_data) + if status_enum == ExecutionStatus.STOPPED: + from workflow_manager.execution.abort import write_file_counts + + # A stopped run's counts come from its file rows, not the worker. + write_file_counts(execution) + else: + self._update_file_aggregates(execution, validated_data) # Reaper recovery: cascade the terminal status to any non-terminal # file executions so the execution and its files agree (atomic with diff --git a/workers/callback/tasks.py b/workers/callback/tasks.py index 667d61468f..b86d6d64cc 100644 --- a/workers/callback/tasks.py +++ b/workers/callback/tasks.py @@ -12,6 +12,7 @@ # Import shared worker infrastructure from shared.api import InternalAPIClient +from shared.clients.execution_client import refused_as_stopped # Import from shared worker modules from shared.enums import PipelineType @@ -56,7 +57,7 @@ class CallbackContext: """Container for callback processing context data.""" - def __init__(self): + def __init__(self) -> None: self.task_id: str = "" self.execution_id: str = "" self.pipeline_id: str | None = None @@ -334,8 +335,10 @@ def _determine_execution_status_unified( # An abort ended this run, whatever its file counts say. if _was_stopped(execution_context): logger.warning( - f"Execution {execution_id} was stopped - {successful_files} of " - f"{total_files} files finished" + "Execution %s was stopped - %s of %s files completed", + execution_id, + successful_files, + total_files, ) return aggregated_results, ExecutionStatus.STOPPED.value, expected_files @@ -376,15 +379,6 @@ def _determine_execution_status_unified( return aggregated_results, final_status, expected_files -def _stopped_meanwhile(response: Any) -> bool: - """True when the backend refused the write because an abort got there first.""" - data = getattr(response, "data", None) or {} - return ( - data.get("reason") == "already_final" - and data.get("new_status") == ExecutionStatus.STOPPED.value - ) - - def _update_execution_status_unified( api_client: InternalAPIClient, execution_id: str, @@ -432,10 +426,12 @@ def post_status(status: str, error: str | None) -> Any: # Make the unified API call response = post_status(final_status, error_message) - if _stopped_meanwhile(response): + if refused_as_stopped(response): logger.warning( - f"Execution {execution_id} was stopped during this callback; " - f"reporting STOPPED instead of {final_status}" + "Execution %s was stopped during this callback; " + "reporting STOPPED instead of %s", + execution_id, + final_status, ) final_status = ExecutionStatus.STOPPED.value # The refused write saved no file counts. @@ -651,7 +647,9 @@ def _handle_notifications_unified( """ # The abort handles these, so a run with a dead worker notifies too. if status == ExecutionStatus.STOPPED.value: - logger.info(f"Execution {execution_id} was stopped - the abort handles notifying") + logger.info( + "Execution %s was stopped - the abort handles notifying", execution_id + ) return { "status": "skipped", "reason": "handled_by_abort", @@ -1359,7 +1357,7 @@ def _track_subscription_usage_if_available( This is a non-blocking operation - errors are logged but do not fail the workflow. Only tracks usage for completed or stopped workflows with successful file - executions; files that finished before an abort are billed. + executions; files that completed before an abort are billed. Args: context: Callback execution context with organization and execution IDs diff --git a/workers/file_processing/tasks.py b/workers/file_processing/tasks.py index 52b53f2d32..3d3437c223 100644 --- a/workers/file_processing/tasks.py +++ b/workers/file_processing/tasks.py @@ -20,6 +20,7 @@ # Import shared worker infrastructure from shared.api import InternalAPIClient from shared.clients.base_client import InternalAPIClientError +from shared.clients.execution_client import refused_as_stopped # Import from shared worker modules from shared.constants import Account @@ -241,8 +242,7 @@ class _TerminalExecutionSkip(Exception): The resurrection this prevents exists because the **reaper** can mark an in-flight execution terminal (and tear its barrier down) while a batch for it - is still in flight or waiting to be redelivered. An abort can mark it STOPPED - the same way. + is still in flight or waiting to be redelivered. This is a validate-first *narrowing*, not a transactional guarantee: the status read and the later ``EXECUTING`` write are not atomic, so a reaper @@ -335,10 +335,10 @@ def _run_stages_under_barrier( return _run_batch_stages(file_batch_data, celery_task_id) except Exception: file_data = file_batch_data.get("file_data") or {} - execution_id = file_data.get("execution_id") - organization_id = file_data.get("organization_id") + execution_id = file_data.get("execution_id") or "" + organization_id = file_data.get("organization_id") or "" api_client = create_api_client(organization_id) - if not _is_execution_stopped(api_client, execution_id): + if not _stopped_before_error_mark(api_client, execution_id, organization_id): raise logger.exception( "[exec:%s] Batch failed after the run was stopped; counting it down", @@ -349,6 +349,31 @@ def _run_stages_under_barrier( ) +def _stopped_before_error_mark( + api_client: InternalAPIClient, execution_id: str, organization_id: str | None +) -> bool: + """Mark the run ERROR; True when the backend refuses because it is STOPPED. + + The backend orders this write and an abort under the run's row lock, so an + abort landing just before the batch failed is not missed. + """ + try: + response = api_client.update_workflow_execution_status( + execution_id=execution_id, + status=ExecutionStatus.ERROR.value, + error_message="Batch failed", + organization_id=organization_id, + ) + except InternalAPIClientError as error: + logger.warning( + "[exec:%s] Marking the failed batch's run ERROR failed: %s", + execution_id, + error, + ) + return False + return refused_as_stopped(response) + + def _run_batch_stages( file_batch_data: dict[str, Any], celery_task_id: str ) -> dict[str, Any]: @@ -361,16 +386,18 @@ def _run_batch_stages( # Step 1: Validate and parse input data batch_data = _validate_and_parse_batch_data(file_batch_data) - # Step 2: Setup execution context. The terminal guard inside it may skip a - # stale/redelivered batch for an already-terminal execution — see + # Step 2: Setup execution context. The terminal guard inside it skips a + # batch whose execution is already terminal — see # :class:`_TerminalExecutionSkip`. try: context = _setup_execution_context(batch_data, celery_task_id) except _TerminalExecutionSkip as skip: logger.warning( - f"[exec:{skip.execution_id}] Skipping batch of " - f"{len(batch_data.files)} file(s): execution already terminal " - f"({skip.status}); not processing." + "[exec:%s] Skipping batch of %s file(s): execution already terminal " + "(%s); not processing.", + skip.execution_id, + len(batch_data.files), + skip.status, ) return _terminal_skip_result(batch_data, skip.status) @@ -518,8 +545,7 @@ def _setup_execution_context( Raises: ValueError: If required context fields are missing - _TerminalExecutionSkip: If the execution is already terminal (a - stale/redelivered batch after reaper-recovery). + _TerminalExecutionSkip: If the execution is already terminal. """ # Extract context using dataclass file_data = batch_data.file_data @@ -559,9 +585,9 @@ def _setup_execution_context( execution_context = execution_response.data workflow_execution = execution_context.get("execution", {}) - # Validate-first terminal guard — skip a stale/redelivered batch - # for an already-terminal execution, before we touch status or process - # anything. See :class:`_TerminalExecutionSkip`. + # Validate-first terminal guard — skip a batch whose execution is already + # terminal, before we touch status or process anything. See + # :class:`_TerminalExecutionSkip`. _raise_if_execution_terminal(workflow_execution, execution_id) # Set LOG_EVENTS_ID in StateStore for WebSocket messaging (critical for UI logs) diff --git a/workers/shared/clients/execution_client.py b/workers/shared/clients/execution_client.py index abc225bc4d..9c13474c97 100644 --- a/workers/shared/clients/execution_client.py +++ b/workers/shared/clients/execution_client.py @@ -13,7 +13,7 @@ import logging import uuid -from typing import Any +from typing import Any, cast from uuid import UUID from unstract.core.data_models import ExecutionStatus @@ -88,6 +88,15 @@ def wrapper(*args, **kwargs): logger = logging.getLogger(__name__) +def refused_as_stopped(response: Any) -> bool: + """True when a status write was refused because an abort got there first.""" + data = getattr(response, "data", None) or {} + return bool( + data.get("reason") == "already_final" + and data.get("new_status") == ExecutionStatus.STOPPED.value + ) + + class ExecutionAPIClient(BaseAPIClient): """Specialized API client for workflow execution operations. @@ -344,7 +353,7 @@ def update_workflow_execution_status( ) # Convert dict response to consistent APIResponse - api_response = convert_dict_response(response, APIResponse) + api_response = cast(APIResponse, convert_dict_response(response, APIResponse)) # The body is flat, and a refused write says why only there. api_response.data = response return api_response diff --git a/workers/shared/processing/filter_pipeline.py b/workers/shared/processing/filter_pipeline.py index a1a0b76750..79bd901218 100644 --- a/workers/shared/processing/filter_pipeline.py +++ b/workers/shared/processing/filter_pipeline.py @@ -530,7 +530,7 @@ def get_name(self) -> str: class ActiveFileFilter(FileFilter): """Filter files that are currently being processed by other executions.""" - def __init__(self): + def __init__(self) -> None: self._cache_checked: set[str] = ( set() ) # Track composite identifiers we've already checked @@ -613,7 +613,7 @@ def apply( @staticmethod def _marker_run_is_live( - run_id: str, api_client: Any, live_runs: dict[str, bool] + run_id: str, identifier: str, api_client: Any, live_runs: dict[str, bool] ) -> bool: """False for a stopped run, whose queued files no longer hold their markers. @@ -628,6 +628,12 @@ def _marker_run_is_live( ) status = None live_runs[run_id] = status != ExecutionStatus.STOPPED.value + if not live_runs[run_id]: + logger.info( + "[ActiveFileFilter] Ignoring the marker of stopped run %s for file %s", + run_id, + identifier, + ) return live_runs[run_id] def _check_active_files_batch( @@ -678,7 +684,7 @@ def _check_active_files_batch( cached_exec_id and cached_exec_id != execution_id and self._marker_run_is_live( - cached_exec_id, api_client, live_runs + cached_exec_id, identifier, api_client, live_runs ) ): active_identifiers.add(identifier) diff --git a/workers/tests/test_abort_finishes_the_run.py b/workers/tests/test_abort_finishes_the_run.py index 105533881b..99ac0daad7 100644 --- a/workers/tests/test_abort_finishes_the_run.py +++ b/workers/tests/test_abort_finishes_the_run.py @@ -1,10 +1,11 @@ """An aborted run still finishes: skipped and failed batches count down, and the -callback keeps STOPPED, bills finished files and leaves the pipeline status and +callback keeps STOPPED, bills completed files and leaves the pipeline status and notifications to the abort. A re-run ignores the stopped run's markers (UN-4318). """ from __future__ import annotations +from contextlib import ExitStack from types import SimpleNamespace from typing import Any from unittest import mock @@ -12,13 +13,14 @@ import callback.tasks as cb import pytest from shared.clients.base_client import APIRequestError -from shared.clients.execution_client import ExecutionAPIClient +from shared.clients.execution_client import ExecutionAPIClient, refused_as_stopped from shared.infrastructure.logging.workflow_logger import WorkerWorkflowLogger from shared.processing.filter_pipeline import ActiveFileFilter from unstract.core.data_models import ExecutionStatus from unstract.workflow_execution.enums import LogState from file_processing.tasks import ( + _process_file_batch_core, _run_batch_stages, _run_stages_under_barrier, _terminal_skip_result, @@ -32,6 +34,7 @@ COMPLETED = ExecutionStatus.COMPLETED.value _TASKS = "file_processing.tasks" _REFUSAL = {"status": "skipped", "reason": "already_final", "new_status": STOPPED} +_UPDATED = {"status": "updated"} def _batch_data() -> Any: @@ -88,40 +91,43 @@ def test_only_a_stopped_runs_batch_decrements(status: str, decrements: int) -> N } -def _failing_batch(status: Any) -> Any: +def _failing_batch(write: Any) -> tuple[Any, Any, Any]: + """A batch whose stages raise; ``write`` answers its run's ERROR mark.""" client = mock.Mock() - client.get_workflow_execution_status.side_effect = [status] + client.update_workflow_execution_status.side_effect = [write] return ( mock.patch(f"{_TASKS}._run_batch_stages", side_effect=RuntimeError("boom")), mock.patch(f"{_TASKS}.create_api_client", return_value=client), + client, ) def test_a_batch_failing_after_the_abort_counts_down() -> None: - stages, client = _failing_batch(STOPPED) + """Also covers an abort landing after the batch's last status check.""" + stages, factory, client = _failing_batch(SimpleNamespace(data=_REFUSAL)) context: Any = {"execution_id": "e", "batch_index": 0, "callback_descriptor": {}} with ( stages, - client, + factory, mock.patch.object(pg_barrier, "claim_batch", return_value=True), mock.patch.object(pg_barrier, "_barrier_pg_decrement") as decrement, mock.patch.object(pg_barrier, "_abort_barrier_in_body") as abort, ): - pg_barrier.run_batch_with_barrier( - context, lambda: _run_stages_under_barrier(_PAYLOAD, "task-1") - ) + _process_file_batch_core(mock.Mock(), _PAYLOAD, context) + mark = client.update_workflow_execution_status.call_args.kwargs + assert (mark["execution_id"], mark["status"]) == ("e", ERROR) assert decrement.call_args.args[0]["failed_files"] == 2 assert SKIPPED_TERMINAL_EXECUTION_KEY not in decrement.call_args.args[0] abort.assert_not_called() @pytest.mark.parametrize( - "status", [ExecutionStatus.EXECUTING.value, APIRequestError("500")], ids=["up", "down"] + "write", [SimpleNamespace(data=_UPDATED), APIRequestError("500")], ids=["up", "down"] ) -def test_a_batch_failing_in_a_live_run_still_raises(status: Any) -> None: - stages, client = _failing_batch(status) - with stages, client, pytest.raises(RuntimeError, match="boom"): +def test_a_batch_failing_in_a_live_run_still_raises(write: Any) -> None: + stages, factory, _ = _failing_batch(write) + with stages, factory, pytest.raises(RuntimeError, match="boom"): _run_stages_under_barrier(_PAYLOAD, "task-1") @@ -206,7 +212,102 @@ def test_the_client_keeps_the_body_of_a_refused_write() -> None: response = client.update_workflow_execution_status("e1", COMPLETED) assert response.success is True - assert cb._stopped_meanwhile(response) is True + assert refused_as_stopped(response) is True + + +_DOWNSTREAM = ( + "_handle_pipeline_updates_unified", + "_track_subscription_usage_if_available", + "_publish_final_workflow_ui_logs", + "_publish_final_workflow_ui_logs_api", + "_cleanup_execution_resources", + "_handle_notifications_unified", +) + + +def _refusing_client() -> Any: + """Refuses the callback's COMPLETED write, then accepts STOPPED.""" + api_client = mock.Mock() + api_client.update_workflow_execution_status.side_effect = [ + SimpleNamespace(data=_REFUSAL), + SimpleNamespace(data=_UPDATED), + ] + api_client.get_workflow_execution.return_value = SimpleNamespace( + success=True, + data={"execution": {"status": "EXECUTING", "workflow_id": "wf"}, "workflow": {}}, + ) + return api_client + + +def _run_callback(run: Any) -> dict[str, Any]: + """Run a callback whose COMPLETED write is refused; return the stubbed steps.""" + determined = ({"total_files": 1, "successful_files": 1}, COMPLETED, 0) + with ExitStack() as stack: + stack.enter_context( + mock.patch.object( + cb, "_determine_execution_status_unified", return_value=determined + ) + ) + stack.enter_context(mock.patch.object(cb, "_initialize_performance_managers")) + steps = { + name: stack.enter_context(mock.patch.object(cb, name)) for name in _DOWNSTREAM + } + run() + return steps + + +def _statuses(steps: dict[str, Any]) -> set[Any]: + """Every status the callback handed to its later steps.""" + seen = set() + for step in steps.values(): + for call in step.call_args_list: + seen.update(call.args[1:2]) + seen.update( + value + for key, value in call.kwargs.items() + if key in ("final_status", "execution_status", "status") + ) + return seen + + +def test_the_general_callback_carries_on_as_stopped() -> None: + context = cb.CallbackContext() + context.execution_id = "e1" + context.organization_id = "org-1" + context.api_client = _refusing_client() + context.execution_status = "EXECUTING" + + def run() -> None: + with mock.patch.object(cb, "_extract_callback_parameters", return_value=context): + cb._process_batch_callback_core(mock.Mock(), []) + + steps = _run_callback(run) + + assert steps["_handle_pipeline_updates_unified"].called + assert steps["_handle_notifications_unified"].called + assert _statuses(steps) == {STOPPED} + + +def test_the_api_callback_carries_on_as_stopped() -> None: + api_client = _refusing_client() + + def run() -> None: + with ( + mock.patch.object(cb, "create_api_client", return_value=api_client), + mock.patch.object( + cb, "_fetch_pipeline_data_simplified", return_value=("api", "API") + ), + ): + # Celery binds the task itself as ``self``. + cb.process_batch_callback_api( + [], execution_id="e1", organization_id="org-1", pipeline_id="p1" + ) + + steps = _run_callback(run) + + assert steps["_handle_pipeline_updates_unified"].called + assert steps["_handle_notifications_unified"].called + assert _statuses(steps) == {STOPPED} @pytest.mark.parametrize(("status", "calls"), [(STOPPED, 0), (COMPLETED, 1)]) @@ -313,10 +414,13 @@ def _active_files(status: Any) -> tuple[set[str], Any]: def test_a_rerun_ignores_a_stopped_runs_markers() -> None: - found, api_client = _active_files(STOPPED) + with mock.patch("shared.processing.filter_pipeline.logger") as logger: + found, api_client = _active_files(STOPPED) assert found == set() api_client.get_workflow_execution_status.assert_called_once_with("old-run") + ignored = [c.args[2] for c in logger.info.call_args_list if "stopped run" in c.args[0]] + assert ignored == ["u1:/a", "u2:/b"] @pytest.mark.parametrize( From cfeb9f4389128f27866140f0db9045b863edfb6d Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Tue, 10 Nov 2026 21:53:29 +0530 Subject: [PATCH 6/8] UN-4318 [FEAT] Protect STOPPED on every run and match the failed batch's ERROR write - update_status locks every run and refuses another status over STOPPED, including scheduled and API-triggered pipeline runs with no queue handle - A failed batch's ERROR write uses the barrier's text and file cascade - A stopped run's final log uses the counts the backend saved - MCP instructions treat COMPLETED, ERROR and STOPPED as final - Tests for runs without a handle, the recount and a COMPLETED refusal Co-Authored-By: Claude Opus 5.5 --- backend/mcp_server/registry.py | 18 ++-- backend/mcp_server/tools/info.py | 5 +- backend/mcp_server/tools/platform.py | 4 +- backend/workflow_manager/execution/README.md | 9 +- .../execution/tests/stuck_abort_base.py | 8 +- .../execution/tests/test_abort_executions.py | 60 +++++++---- .../workflow_manager/internal_serializers.py | 2 + backend/workflow_manager/internal_views.py | 102 +++++++++--------- workers/callback/tasks.py | 36 +++++-- workers/file_processing/tasks.py | 21 ++-- workers/tests/test_abort_finishes_the_run.py | 28 ++++- 11 files changed, 179 insertions(+), 114 deletions(-) diff --git a/backend/mcp_server/registry.py b/backend/mcp_server/registry.py index 7a93beddbb..2aaad626be 100644 --- a/backend/mcp_server/registry.py +++ b/backend/mcp_server/registry.py @@ -55,9 +55,9 @@ class MCPTool: consumes_result: True when *reading* is itself destructive — the result store is one-shot, so a successful read acknowledges the execution and the payload cannot be fetched again. ``getExecutionStatus`` on - a COMPLETED execution is the case: it neither writes configuration - nor spends money, so ``writes``/``billable`` are both false, but - calling it twice loses data. Without this flag the derivation below + a COMPLETED or STOPPED execution is the case: it neither writes + configuration nor spends money, so ``writes``/``billable`` are both + false, but calling it twice loses data. Without this flag the derivation below would advertise ``readOnlyHint: true`` and tell an agent it is safe to poll freely — the opposite of the truth. preflight: Optional ``preflight(context, **arguments)`` run *before* the @@ -272,8 +272,8 @@ def build_deployment_registry() -> MCPToolRegistry: ), input_schema=get_execution_status_schema(), handler=get_execution_status, - # Reading a COMPLETED execution acknowledges it and the result - # store is one-shot, so the payload cannot be fetched again. + # Reading a COMPLETED or STOPPED execution acknowledges it and the + # result store is one-shot, so the payload cannot be fetched again. # Not `writes` (no configuration changes) and not `billable` # (no spend), but emphatically not read-only either. consumes_result=True, @@ -714,14 +714,14 @@ def build_platform_registry() -> MCPToolRegistry: "Poll for the result of an extraction started by " "extractDocument. Free to call.\n\n" "If extractDocument already returned execution_status " - "COMPLETED, the result is in that response and no polling is " - "needed. Otherwise poll this until COMPLETED, ERROR or STOPPED, " + "COMPLETED, ERROR or STOPPED, that response is final and no " + "polling is needed. Otherwise poll this until one of those, " "pausing a few seconds between calls." ), input_schema=platform_execution_status_schema(), handler=get_platform_execution_status, - # Reading a COMPLETED execution acknowledges it and the result - # store is one-shot, so the payload cannot be fetched again. + # Reading a COMPLETED or STOPPED execution acknowledges it and the + # result store is one-shot, so the payload cannot be fetched again. # Not `writes` (no configuration changes) and not `billable` # (no spend), but emphatically not read-only either. consumes_result=True, diff --git a/backend/mcp_server/tools/info.py b/backend/mcp_server/tools/info.py index 68a69bdefc..aa32278e62 100644 --- a/backend/mcp_server/tools/info.py +++ b/backend/mcp_server/tools/info.py @@ -47,8 +47,9 @@ def read_me_first(context: MCPContext) -> dict[str, Any]: "1. Call getApiInfo to confirm the deployment is active and see " "what it is meant to process.", "2. Call extractDocument with the URL(s) of the document(s).", - "3. If the response has execution_status COMPLETED, the result is " - "already in the response — you are done.", + "3. If the response has execution_status COMPLETED, ERROR or " + "STOPPED, it is final and any result is already in it — you are " + "done.", "4. Otherwise poll getExecutionStatus with the returned " "execution_id until the status is COMPLETED, ERROR or STOPPED, pausing a " "few seconds between polls.", diff --git a/backend/mcp_server/tools/platform.py b/backend/mcp_server/tools/platform.py index c5826f8bb1..5c61d36316 100644 --- a/backend/mcp_server/tools/platform.py +++ b/backend/mcp_server/tools/platform.py @@ -178,8 +178,8 @@ def platform_read_me_first(context: PlatformMCPContext) -> dict[str, Any]: "to_run_an_extraction": ( "Call listApiDeployments for an api_name, then extractDocument " "with that name and your document URLs. If the response is not " - "already COMPLETED, poll getExecutionStatus with the returned " - "execution_id.\n\n" + "already COMPLETED, ERROR or STOPPED, poll getExecutionStatus with " + "the returned execution_id.\n\n" "The same extraction is also available on a deployment-scoped " "server at /deployment/api///mcp, which takes that " "one deployment's own API key and can reach nothing else. Both " diff --git a/backend/workflow_manager/execution/README.md b/backend/workflow_manager/execution/README.md index 82cbd68e76..9322bbe69f 100644 --- a/backend/workflow_manager/execution/README.md +++ b/backend/workflow_manager/execution/README.md @@ -137,7 +137,7 @@ Once that commits, it: - Queued batches of a stopped run are skipped when a worker picks them up, and still count down on the barrier. So does a batch that fails after the abort, - unless the backend can't be reached when it fails. + unless its status write fails. - A re-run ignores the active-file markers of a stopped run, so it keeps those files even if no worker has picked up the queued batch yet. - A running batch checks the status before each file and stops once it is @@ -151,8 +151,7 @@ Once that commits, it: and notifications to the abort. - A caller waiting on a synchronous API deployment `POST` gets 422 for an aborted run, the same as polling with `GET`. Both apply the same result filters. -- Later status writes from workers cannot undo the abort, unless the run's - `queue_message_id` is empty. +- Later status writes from workers cannot undo the abort. ### Errors @@ -190,13 +189,11 @@ callback reports an aborted run with a completed file as completed. pick it up while the old worker finishes it: double LLM spend, and two writes to a database destination. Wait for that file to finish before re-running. - **A dead worker's run never gets its callback**, and neither does a run whose - batch fails while the backend can't be reached. The pipeline status, + batch fails and then can't write its status. The pipeline status, notifications and file counts are still set by the abort, but usage for its completed files is not committed and its execution directory, which can hold uploaded documents, is not cleaned up. It stays in storage until something else deletes it. The reaper's `ERROR` path has the same gap. -- A run whose `queue_message_id` is empty (from before the PG queue, or whose - handle write failed) is not protected from a late worker status write. - The file being processed when the abort lands is billed once it finishes, but an API caller can't get its result: the first read after the abort returns the results so far and discards them, and later reads answer 406. Its processing cost diff --git a/backend/workflow_manager/execution/tests/stuck_abort_base.py b/backend/workflow_manager/execution/tests/stuck_abort_base.py index f36af73054..a3fe305381 100644 --- a/backend/workflow_manager/execution/tests/stuck_abort_base.py +++ b/backend/workflow_manager/execution/tests/stuck_abort_base.py @@ -68,8 +68,12 @@ def execution( resource: Any = None, status: ExecutionStatus = ExecutionStatus.EXECUTING, idle: timedelta = OLD, + queue_message_id: int | None = 1, ) -> WorkflowExecution: - """A PG-transport run whose last own write was ``idle`` ago.""" + """A run whose last own write was ``idle`` ago. + + Scheduled and API-triggered pipeline runs have no ``queue_message_id``. + """ execution: WorkflowExecution = WorkflowExecution.objects.create( workflow=self.workflow, pipeline_id=resource.id if resource else None, @@ -77,7 +81,7 @@ def execution( execution_method=WorkflowExecution.Method.DIRECT, execution_type=WorkflowExecution.Type.COMPLETE, status=status.value, - queue_message_id=1, + queue_message_id=queue_message_id, ) self.backdate(WorkflowExecution, execution.pk, idle) execution.refresh_from_db() diff --git a/backend/workflow_manager/execution/tests/test_abort_executions.py b/backend/workflow_manager/execution/tests/test_abort_executions.py index 9d123ccf8e..2408c34760 100644 --- a/backend/workflow_manager/execution/tests/test_abort_executions.py +++ b/backend/workflow_manager/execution/tests/test_abort_executions.py @@ -401,39 +401,55 @@ def test_requires_authentication(self) -> None: response = _abort(request) self.assertIn(response.status_code, (401, 403)) - def test_later_completion_cannot_undo_the_abort(self) -> None: - execution = self.execution(self.pipeline()) - self._abort(self.owner, execution.id) - + def test_later_writes_cannot_undo_the_abort(self) -> None: + # Scheduled and API-triggered pipeline runs have no queue_message_id. + for queue_message_id in (1, None): + for status in (ExecutionStatus.COMPLETED, ExecutionStatus.ERROR): + with self.subTest(queue_message_id=queue_message_id, status=status): + execution = self.execution( + self.pipeline(), queue_message_id=queue_message_id + ) + self._abort(self.owner, execution.id) + + out = self._write_status(execution, {"status": status.value}) + + self.assertEqual(out.data.get("reason"), "already_final") + self.assertEqual(out.data.get("new_status"), "STOPPED") + self.assertEqual( + self._status(execution), ExecutionStatus.STOPPED.value + ) + + @staticmethod + def _write_status(execution: WorkflowExecution, data: dict[str, Any]) -> Any: + """A worker's status write through the internal API.""" view = WorkflowExecutionInternalViewSet() request = MagicMock() - request.data = {"status": ExecutionStatus.COMPLETED.value} + request.data = data fresh = WorkflowExecution.objects.get(pk=execution.pk) with patch.object(view, "get_object", return_value=fresh): - out = view.update_status(request, id=str(execution.id)) - - self.assertEqual(out.data.get("reason"), "already_final") - self.assertEqual(self._status(execution), ExecutionStatus.STOPPED.value) + return view.update_status(request, id=str(execution.id)) def test_a_stopped_write_saves_the_counts_of_the_file_rows(self) -> None: execution = self.execution(self.pipeline()) for status in (ExecutionStatus.COMPLETED, ExecutionStatus.COMPLETED): self.file(execution, status, idle=RECENT) - self.file(execution, ExecutionStatus.EXECUTING, idle=RECENT) + in_flight = self.file(execution, ExecutionStatus.EXECUTING, idle=RECENT) self._abort(self.owner, execution.id) + # The file in progress at the abort finishes afterwards. + WorkflowFileExecution.objects.filter(pk=in_flight.pk).update( + status=ExecutionStatus.COMPLETED.value + ) # The worker reports zeros, as a failed callback's retry does. - view = WorkflowExecutionInternalViewSet() - request = MagicMock() - request.data = { - "status": ExecutionStatus.STOPPED.value, - "total_files": 0, - "successful_files": 0, - "failed_files": 0, - } - fresh = WorkflowExecution.objects.get(pk=execution.pk) - with patch.object(view, "get_object", return_value=fresh): - out = view.update_status(request, id=str(execution.id)) + out = self._write_status( + execution, + { + "status": ExecutionStatus.STOPPED.value, + "total_files": 0, + "successful_files": 0, + "failed_files": 0, + }, + ) self.assertEqual(out.data.get("status"), "updated") execution.refresh_from_db() @@ -443,5 +459,5 @@ def test_a_stopped_write_saves_the_counts_of_the_file_rows(self) -> None: execution.successful_files, execution.failed_files, ), - (3, 2, 1), + (3, 3, 0), ) diff --git a/backend/workflow_manager/internal_serializers.py b/backend/workflow_manager/internal_serializers.py index 9127763873..fbe9e54b53 100644 --- a/backend/workflow_manager/internal_serializers.py +++ b/backend/workflow_manager/internal_serializers.py @@ -96,6 +96,8 @@ class Meta: "status", "result_acknowledged", "total_files", + "successful_files", + "failed_files", "error_message", "attempts", "execution_time", diff --git a/backend/workflow_manager/internal_views.py b/backend/workflow_manager/internal_views.py index a56e3116d4..fbf56a641e 100644 --- a/backend/workflow_manager/internal_views.py +++ b/backend/workflow_manager/internal_views.py @@ -570,70 +570,54 @@ def update_status(self, request, id=None): # with failed_files=None, which is_failure_run() reads as a success and # silently bypasses notify_on_failures subscribers. with transaction.atomic(): - row_locked = False - # Concurrency guard for PG-transport executions, identified by the - # durable ``queue_message_id`` marker — NOT a re-evaluated Flipt flag. - # Keying on the row's own transport keeps the invariant correct even - # if the rollout flag is toggled off or Flipt is unreachable while a - # PG execution is still finalizing (per the "transport carried on the - # row, never re-read from Flipt" rule). - # - # Protect COMPLETED and STOPPED — the two statuses that assert a - # deliberate final outcome. COMPLETED is set exclusively by a - # successful callback (its redelivery is skipped, see - # _callback_already_ran); STOPPED records an explicit user/operator - # stop. Under the PG path's high concurrency ~15 call sites write - # status via unlocked read-modify-write, so a stale writer can commit - # AFTER one of these and clobber it (lost update → execution stranded - # in EXECUTING, a confusing success→failed flip, or a straggler - # callback erasing a user's stop). Lock the row and refuse to overwrite - # either with a different status. + # Lock the row and refuse to overwrite a deliberate final outcome: + # COMPLETED (set only by a successful callback) and STOPPED (an + # explicit stop). Under high concurrency ~15 call sites write status + # via unlocked read-modify-write, so a stale writer can commit AFTER + # one of these and clobber it (lost update → execution stranded in + # EXECUTING, a confusing success→failed flip, or a straggler callback + # erasing a user's stop). See _protected_statuses for which rows + # protect which status. # # ERROR is deliberately NOT protected: other paths (upstream error, # the reaper) can set it BEFORE the first real callback, which then # legitimately corrects the row to COMPLETED when the files actually # succeeded — blocking that would freeze a genuinely successful run at - # a premature ERROR. Celery executions (queue_message_id NULL) keep the - # legacy path. + # a premature ERROR. # # ONE lock per call: the row read here is handed to update_execution() - # with row_locked=True, so it does not lock the same row a second time. The - # Celery path takes no lock here and lets update_execution() take its - # own — still one. FOR NO KEY UPDATE (no_key=True) keeps writers - # mutually exclusive without blocking the FOR KEY SHARE that every - # child-row INSERT takes on this row for its foreign-key check. - if execution.queue_message_id is not None: - execution = WorkflowExecution.objects.select_for_update( - no_key=True - ).get(pk=execution.pk) - row_locked = True - protected = { - ExecutionStatus.COMPLETED.value, - ExecutionStatus.STOPPED.value, - } - if ( - execution.status in protected - and status_enum.value != execution.status - ): - logger.warning( - f"update_status: refusing to overwrite protected status " - f"'{execution.status}' with '{status_enum.value}' " - f"for execution {id}" - ) - return Response( - { - "status": "skipped", - "reason": "already_final", - "execution_id": str(execution.id), - "new_status": execution.status, - } - ) + # with row_locked=True, so it does not lock the same row a second + # time. FOR NO KEY UPDATE (no_key=True) keeps writers mutually + # exclusive without blocking the FOR KEY SHARE that every child-row + # INSERT takes on this row for its foreign-key check. + execution = WorkflowExecution.objects.select_for_update(no_key=True).get( + pk=execution.pk + ) + if ( + execution.status in self._protected_statuses(execution) + and status_enum.value != execution.status + ): + logger.warning( + "update_status: refusing to overwrite protected status " + "'%s' with '%s' for execution %s", + execution.status, + status_enum.value, + id, + ) + return Response( + { + "status": "skipped", + "reason": "already_final", + "execution_id": str(execution.id), + "new_status": execution.status, + } + ) execution.update_execution( status=status_enum, error=error_message, increment_attempt=increment_attempt, - row_locked=row_locked, + row_locked=True, ) # File aggregates are not handled by update_execution; persist @@ -900,6 +884,20 @@ def _recover_one_stuck_pg_execution(self, exec_id, cutoff, terminal) -> str: logger.exception("recover_stuck_pg_executions: failed to recover %s", exec_id) return "failed" + @staticmethod + def _protected_statuses(execution: WorkflowExecution) -> set[str]: + """Statuses a different status write must not overwrite on this row. + + Any run can be stopped, so STOPPED is protected on every row. COMPLETED + stays protected only on rows carrying the PG dispatch marker + ``queue_message_id``, keyed on the row and never on a re-read Flipt flag. + """ + from workflow_manager.workflow_v2.enums import ExecutionStatus + + if execution.queue_message_id is None: + return {ExecutionStatus.STOPPED.value} + return {ExecutionStatus.COMPLETED.value, ExecutionStatus.STOPPED.value} + @staticmethod def _cascade_terminal_files(execution, status_enum, error_message) -> None: """Mark this execution's non-terminal (PENDING/EXECUTING) file executions to diff --git a/workers/callback/tasks.py b/workers/callback/tasks.py index b86d6d64cc..79029faaf4 100644 --- a/workers/callback/tasks.py +++ b/workers/callback/tasks.py @@ -1955,6 +1955,26 @@ def process_batch_callback_api( logger.warning("api_client.close() failed during cleanup: %s", e) +def _final_log_counts( + aggregated_results: dict[str, Any], + execution_data: dict[str, Any] | None, + execution_status: str, +) -> tuple[int, int, int]: + """Total, successful and failed files for the final log line. + + A stopped run logs the counts the backend saved from its file rows. + """ + counts = aggregated_results + if execution_status == ExecutionStatus.STOPPED.value: + saved = (execution_data or {}).get("execution") or {} + counts = {**aggregated_results, **saved} + return ( + counts.get("total_files") or 0, + counts.get("successful_files") or 0, + counts.get("failed_files") or 0, + ) + + def _publish_final_workflow_ui_logs( context: "CallbackContext", aggregated_results: dict[str, Any], @@ -1971,11 +1991,6 @@ def _publish_final_workflow_ui_logs( execution_status: Final execution status """ try: - # Extract file statistics from aggregated results - total_files = aggregated_results.get("total_files", 0) - successful_files = aggregated_results.get("successful_files", 0) - failed_files = aggregated_results.get("failed_files", 0) - # Get execution data to extract cost information (with cost calculation) execution_response = context.api_client.get_workflow_execution( context.execution_id, include_cost=True, file_execution=False @@ -1988,6 +2003,9 @@ def _publish_final_workflow_ui_logs( # Cost data is at the top level of response when include_cost=True aggregated_usage_cost = execution_response.data.get("aggregated_usage_cost") + total_files, successful_files, failed_files = _final_log_counts( + aggregated_results, execution_response.data, execution_status + ) # Create workflow logger for UI feedback # Use general workflow logger since this is called from general workflow callback @@ -2046,11 +2064,6 @@ def _publish_final_workflow_ui_logs_api( execution_status: Final execution status """ try: - # Extract file statistics from aggregated results - total_files = aggregated_results.get("total_files", 0) - successful_files = aggregated_results.get("successful_files", 0) - failed_files = aggregated_results.get("failed_files", 0) - # Get execution data to extract cost information (with cost calculation) execution_response = context.api_client.get_workflow_execution( context.execution_id, include_cost=True, file_execution=False @@ -2063,6 +2076,9 @@ def _publish_final_workflow_ui_logs_api( # Cost data is at the top level of response when include_cost=True aggregated_usage_cost = execution_response.data.get("aggregated_usage_cost") + total_files, successful_files, failed_files = _final_log_counts( + aggregated_results, execution_response.data, execution_status + ) # Create workflow logger for UI feedback # Use API workflow logger since this is called from API workflow callback diff --git a/workers/file_processing/tasks.py b/workers/file_processing/tasks.py index 3d3437c223..c808805d42 100644 --- a/workers/file_processing/tasks.py +++ b/workers/file_processing/tasks.py @@ -325,7 +325,7 @@ def _all_failed_result(n_files: int, organization_id: str | None) -> dict[str, A def _run_stages_under_barrier( - file_batch_data: dict[str, Any], celery_task_id: str + file_batch_data: dict[str, Any], celery_task_id: str, batch_index: int ) -> dict[str, Any]: """Run the batch stages; a stopped run's failed batch still counts down. @@ -338,7 +338,9 @@ def _run_stages_under_barrier( execution_id = file_data.get("execution_id") or "" organization_id = file_data.get("organization_id") or "" api_client = create_api_client(organization_id) - if not _stopped_before_error_mark(api_client, execution_id, organization_id): + if not _stopped_before_error_mark( + api_client, execution_id, organization_id, batch_index + ): raise logger.exception( "[exec:%s] Batch failed after the run was stopped; counting it down", @@ -350,19 +352,24 @@ def _run_stages_under_barrier( def _stopped_before_error_mark( - api_client: InternalAPIClient, execution_id: str, organization_id: str | None + api_client: InternalAPIClient, + execution_id: str, + organization_id: str | None, + batch_index: int, ) -> bool: """Mark the run ERROR; True when the backend refuses because it is STOPPED. The backend orders this write and an abort under the run's row lock, so an - abort landing just before the batch failed is not missed. + abort landing just before the batch failed is not missed. The write matches + the barrier's own ERROR mark, so the two agree. """ try: response = api_client.update_workflow_execution_status( execution_id=execution_id, status=ExecutionStatus.ERROR.value, - error_message="Batch failed", + error_message=f"[pg-barrier-abort] batch {batch_index} failed.", organization_id=organization_id, + cascade_terminal_files=True, ) except InternalAPIClientError as error: logger.warning( @@ -461,7 +468,9 @@ def _process_file_batch_core( # the barrier in-body / self-chain the callback. return run_batch_with_barrier( barrier_context, - lambda: _run_stages_under_barrier(file_batch_data, celery_task_id), + lambda: _run_stages_under_barrier( + file_batch_data, celery_task_id, int(barrier_context["batch_index"]) + ), ) diff --git a/workers/tests/test_abort_finishes_the_run.py b/workers/tests/test_abort_finishes_the_run.py index 99ac0daad7..409c04e189 100644 --- a/workers/tests/test_abort_finishes_the_run.py +++ b/workers/tests/test_abort_finishes_the_run.py @@ -34,6 +34,7 @@ COMPLETED = ExecutionStatus.COMPLETED.value _TASKS = "file_processing.tasks" _REFUSAL = {"status": "skipped", "reason": "already_final", "new_status": STOPPED} +_COMPLETED_REFUSAL = {**_REFUSAL, "new_status": COMPLETED} _UPDATED = {"status": "updated"} @@ -105,7 +106,7 @@ def _failing_batch(write: Any) -> tuple[Any, Any, Any]: def test_a_batch_failing_after_the_abort_counts_down() -> None: """Also covers an abort landing after the batch's last status check.""" stages, factory, client = _failing_batch(SimpleNamespace(data=_REFUSAL)) - context: Any = {"execution_id": "e", "batch_index": 0, "callback_descriptor": {}} + context: Any = {"execution_id": "e", "batch_index": 3, "callback_descriptor": {}} with ( stages, factory, @@ -117,18 +118,27 @@ def test_a_batch_failing_after_the_abort_counts_down() -> None: mark = client.update_workflow_execution_status.call_args.kwargs assert (mark["execution_id"], mark["status"]) == ("e", ERROR) + # The same write as the barrier's own mark (see test_pg_barrier). + assert mark["error_message"] == "[pg-barrier-abort] batch 3 failed." + assert mark["cascade_terminal_files"] is True assert decrement.call_args.args[0]["failed_files"] == 2 assert SKIPPED_TERMINAL_EXECUTION_KEY not in decrement.call_args.args[0] abort.assert_not_called() @pytest.mark.parametrize( - "write", [SimpleNamespace(data=_UPDATED), APIRequestError("500")], ids=["up", "down"] + "write", + [ + SimpleNamespace(data=_UPDATED), + SimpleNamespace(data=_COMPLETED_REFUSAL), + APIRequestError("500"), + ], + ids=["up", "completed", "down"], ) def test_a_batch_failing_in_a_live_run_still_raises(write: Any) -> None: stages, factory, _ = _failing_batch(write) with stages, factory, pytest.raises(RuntimeError, match="boom"): - _run_stages_under_barrier(_PAYLOAD, "task-1") + _run_stages_under_barrier(_PAYLOAD, "task-1", 0) def _status_for(run_status: str, successful: int, failed: int) -> str: @@ -450,3 +460,15 @@ def test_the_final_log_tells_a_stopped_run_apart( workflow_logger.publish_final_workflow_logs(3, 1, 2, stopped=stopped) assert mock.call(state, message, mock.ANY) in update.call_args_list + + +@pytest.mark.parametrize( + ("status", "expected"), [(STOPPED, (3, 3, 0)), (COMPLETED, (3, 2, 1))] +) +def test_a_stopped_runs_final_log_uses_the_saved_counts( + status: str, expected: tuple[int, int, int] +) -> None: + worker = {"total_files": 3, "successful_files": 2, "failed_files": 1} + saved = {"execution": {"total_files": 3, "successful_files": 3, "failed_files": 0}} + + assert cb._final_log_counts(worker, saved, status) == expected From e0006b5f07db5f28d42cef2d782cb447e8c98f99 Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Wed, 11 Nov 2026 11:52:15 +0530 Subject: [PATCH 7/8] UN-4318 [FEAT] Keep STOPPED against backend writes and leave a stopped run's pipeline to the abort - The model sends a stopped row through the guarded write whatever its queue handle, so staging or dispatch failures cannot replace an abort; a refused COMPLETED keeps the abort's counts. - A run that finds no files skips the pipeline SUCCESS write once it reads as stopped. - The refusal reason is a shared constant in unstract.core. - Tests cover the callbacks' logged counts and the internal payload's counts. - README Limits names runs aborted before a worker picks them up. Co-Authored-By: Claude Opus 5.5 --- backend/mcp_server/registry.py | 6 +-- backend/workflow_manager/execution/README.md | 5 +- .../execution/tests/test_abort_executions.py | 24 ++++++++++ .../tests/test_execution_contract.py | 18 +++++++ backend/workflow_manager/internal_views.py | 3 +- .../workflow_manager/workflow_v2/execution.py | 36 +++++++------- .../workflow_v2/models/execution.py | 23 +++++---- .../core/src/unstract/core/data_models.py | 3 ++ workers/general/tasks.py | 48 +++++++++++-------- workers/shared/clients/execution_client.py | 4 +- workers/tests/test_abort_finishes_the_run.py | 39 ++++++++++++--- 11 files changed, 149 insertions(+), 60 deletions(-) diff --git a/backend/mcp_server/registry.py b/backend/mcp_server/registry.py index 2aaad626be..81338e1ae0 100644 --- a/backend/mcp_server/registry.py +++ b/backend/mcp_server/registry.py @@ -55,7 +55,7 @@ class MCPTool: consumes_result: True when *reading* is itself destructive — the result store is one-shot, so a successful read acknowledges the execution and the payload cannot be fetched again. ``getExecutionStatus`` on - a COMPLETED or STOPPED execution is the case: it neither writes + a finished execution is the case: it neither writes configuration nor spends money, so ``writes``/``billable`` are both false, but calling it twice loses data. Without this flag the derivation below would advertise ``readOnlyHint: true`` and tell an agent it is safe @@ -272,7 +272,7 @@ def build_deployment_registry() -> MCPToolRegistry: ), input_schema=get_execution_status_schema(), handler=get_execution_status, - # Reading a COMPLETED or STOPPED execution acknowledges it and the + # Reading a finished execution acknowledges it and the # result store is one-shot, so the payload cannot be fetched again. # Not `writes` (no configuration changes) and not `billable` # (no spend), but emphatically not read-only either. @@ -720,7 +720,7 @@ def build_platform_registry() -> MCPToolRegistry: ), input_schema=platform_execution_status_schema(), handler=get_platform_execution_status, - # Reading a COMPLETED or STOPPED execution acknowledges it and the + # Reading a finished execution acknowledges it and the # result store is one-shot, so the payload cannot be fetched again. # Not `writes` (no configuration changes) and not `billable` # (no spend), but emphatically not read-only either. diff --git a/backend/workflow_manager/execution/README.md b/backend/workflow_manager/execution/README.md index 9322bbe69f..c236b5f8a9 100644 --- a/backend/workflow_manager/execution/README.md +++ b/backend/workflow_manager/execution/README.md @@ -151,7 +151,7 @@ Once that commits, it: and notifications to the abort. - A caller waiting on a synchronous API deployment `POST` gets 422 for an aborted run, the same as polling with `GET`. Both apply the same result filters. -- Later status writes from workers cannot undo the abort. +- Later status writes, from workers or the backend, cannot undo the abort. ### Errors @@ -189,7 +189,8 @@ callback reports an aborted run with a completed file as completed. pick it up while the old worker finishes it: double LLM spend, and two writes to a database destination. Wait for that file to finish before re-running. - **A dead worker's run never gets its callback**, and neither does a run whose - batch fails and then can't write its status. The pipeline status, + batch fails and then can't write its status, or a workflow or ETL run aborted + before a worker picks it up (the worker then skips it). The pipeline status, notifications and file counts are still set by the abort, but usage for its completed files is not committed and its execution directory, which can hold uploaded documents, is not cleaned up. It stays in storage until something else diff --git a/backend/workflow_manager/execution/tests/test_abort_executions.py b/backend/workflow_manager/execution/tests/test_abort_executions.py index 2408c34760..023a0e8d64 100644 --- a/backend/workflow_manager/execution/tests/test_abort_executions.py +++ b/backend/workflow_manager/execution/tests/test_abort_executions.py @@ -1,6 +1,7 @@ """``POST execution/abort/`` (UN-4318): owner-or-admin, per-id outcomes, audit.""" import uuid +from collections.abc import Callable from datetime import timedelta from typing import Any from unittest.mock import MagicMock, patch @@ -23,6 +24,7 @@ from workflow_manager.file_execution.serializers import FileCentricExecutionSerializer from workflow_manager.internal_views import WorkflowExecutionInternalViewSet from workflow_manager.workflow_v2.enums import ExecutionStatus +from workflow_manager.workflow_v2.execution import WorkflowExecutionServiceHelper from workflow_manager.workflow_v2.models.execution import WorkflowExecution from workflow_manager.workflow_v2.models.workflow import Workflow @@ -419,6 +421,28 @@ def test_later_writes_cannot_undo_the_abort(self) -> None: self._status(execution), ExecutionStatus.STOPPED.value ) + def test_backend_writes_before_dispatch_cannot_undo_the_abort(self) -> None: + # Staging or dispatch fails, or stages nothing, after a PENDING run's abort. + helper = WorkflowExecutionServiceHelper + writes: dict[str, Callable[[str], Any]] = { + "error": lambda pk: helper.update_execution_err(pk, "staging failed"), + "completed": lambda pk: helper.update_execution_completed(pk, 2, 2), + } + for name, write in writes.items(): + with self.subTest(write=name): + execution = self.execution( + self.deployment(), + status=ExecutionStatus.PENDING, + queue_message_id=None, + ) + self._abort(self.owner, execution.id) + + write(str(execution.pk)) + + execution.refresh_from_db() + self.assertEqual(execution.status, ExecutionStatus.STOPPED.value) + self.assertEqual((execution.total_files, execution.failed_files), (0, 0)) + @staticmethod def _write_status(execution: WorkflowExecution, data: dict[str, Any]) -> Any: """A worker's status write through the internal API.""" diff --git a/backend/workflow_manager/execution/tests/test_execution_contract.py b/backend/workflow_manager/execution/tests/test_execution_contract.py index e305c7a990..04bae87d54 100644 --- a/backend/workflow_manager/execution/tests/test_execution_contract.py +++ b/backend/workflow_manager/execution/tests/test_execution_contract.py @@ -22,6 +22,7 @@ _list = ExecutionViewSet.as_view({"get": "list"}) _detail = ExecutionViewSet.as_view({"get": "retrieve"}) _status = WorkflowExecutionInternalViewSet.as_view({"get": "execution_status"}) +_retrieve = WorkflowExecutionInternalViewSet.as_view({"get": "retrieve"}) # The default row, pinned: a new key here is a contract change. _EXECUTION_KEYS = { @@ -130,6 +131,23 @@ def test_other_org_unknown_and_malformed_ids_are_404(self) -> None: self.assertEqual(response.status_code, 404) +class InternalExecutionPayloadTests(StuckAbortTestBase): + def test_carries_the_saved_file_counts(self) -> None: + # A stopped run's final log uses these counts. + execution = self.execution(status=ExecutionStatus.STOPPED) + WorkflowExecution.objects.filter(pk=execution.pk).update( + total_files=3, successful_files=3, failed_files=0 + ) + request = APIRequestFactory().get("/", {"file_execution": "false"}) + request.organization_id = self.org.organization_id + + response = _retrieve(request, id=str(execution.id)) + + self.assertEqual(response.status_code, 200, response.data) + saved = response.data["execution"] + self.assertEqual((saved["successful_files"], saved["failed_files"]), (3, 0)) + + class WorkerStatusRouteTests(StuckAbortTestBase): def test_worker_status_url_reaches_the_status_action(self) -> None: match = resolve(f"/internal/v1/workflow-execution/{uuid.uuid4()}/status/") diff --git a/backend/workflow_manager/internal_views.py b/backend/workflow_manager/internal_views.py index fbf56a641e..1d26d41c45 100644 --- a/backend/workflow_manager/internal_views.py +++ b/backend/workflow_manager/internal_views.py @@ -20,6 +20,7 @@ # Import new dataclasses for WorkflowDefinitionAPIView from unstract.core.data_models import ( + ALREADY_FINAL_REASON, ConnectionType, ConnectorInstanceData, WorkflowDefinitionResponseData, @@ -607,7 +608,7 @@ def update_status(self, request, id=None): return Response( { "status": "skipped", - "reason": "already_final", + "reason": ALREADY_FINAL_REASON, "execution_id": str(execution.id), "new_status": execution.status, } diff --git a/backend/workflow_manager/workflow_v2/execution.py b/backend/workflow_manager/workflow_v2/execution.py index dd177e50d3..22731c5f46 100644 --- a/backend/workflow_manager/workflow_v2/execution.py +++ b/backend/workflow_manager/workflow_v2/execution.py @@ -175,8 +175,8 @@ def update_execution( increment_attempt: bool = False, ) -> None: # Delegate to the model method, which owns the single, robust terminal-one-way - # guard (atomic select_for_update, PG-scoped, field-scoped writes). Duplicating - # the guard here risked drift + its own TOCTOU/dropped-field bugs. + # guard (atomic select_for_update, PG and STOPPED rows, field-scoped writes). + # Duplicating the guard here risked drift + its own TOCTOU/dropped-field bugs. execution = WorkflowExecution.objects.get(pk=self.execution_id) execution.update_execution( status=status, error=error, increment_attempt=increment_attempt @@ -405,23 +405,23 @@ def update_execution_completed( execution = WorkflowExecution.objects.get(pk=execution_id) # One transaction, so on-commit hooks of the status write see the counts. with transaction.atomic(): - # Routed through the model method so the PG terminal-one-way guard - # applies; the legacy transport has no such guard. + # Routed through the model method so its guard applies. execution.update_execution(status=ExecutionStatus.COMPLETED) - execution.total_files = total_files - execution.successful_files = 0 - execution.failed_files = failed_files - # Field-scoped for the columns, but save() still republishes this - # object's status to the execution cache. Safe only because the row - # is undispatched here, so no worker can have moved it on. - execution.save( - update_fields=[ - "total_files", - "successful_files", - "failed_files", - "modified_at", - ] - ) + # A refused write (e.g. the run was aborted) keeps the stored counts. + if execution.status == ExecutionStatus.COMPLETED.value: + execution.total_files = total_files + execution.successful_files = 0 + execution.failed_files = failed_files + # Field-scoped, but save() still republishes this object's status + # to the cache; safe under the lock this block's write holds. + execution.save( + update_fields=[ + "total_files", + "successful_files", + "failed_files", + "modified_at", + ] + ) # The guard may have refused the status change without raising; re-read # so the returned row reflects what is actually stored. execution.refresh_from_db() diff --git a/backend/workflow_manager/workflow_v2/models/execution.py b/backend/workflow_manager/workflow_v2/models/execution.py index 536fda11f2..9ffdff8f61 100644 --- a/backend/workflow_manager/workflow_v2/models/execution.py +++ b/backend/workflow_manager/workflow_v2/models/execution.py @@ -460,7 +460,11 @@ def update_execution( ) if locked is None: return - if locked["queue_message_id"] is None: + # A stopped row takes the guarded write too, so nothing replaces the stop. + if ( + locked["queue_message_id"] is None + and locked["status"] != ExecutionStatus.STOPPED.value + ): should_release_rate_limit = self._apply_legacy_update( status, error, increment_attempt ) @@ -492,8 +496,9 @@ def _apply_legacy_update( error: str | None, increment_attempt: bool, ) -> bool: - """Celery-path (queue_message_id NULL) — unchanged full save(); runs inside the - caller's locked transaction. Returns whether to release the rate-limit slot. + """Celery-path (queue_message_id NULL, not STOPPED) — unchanged full save(); + runs inside the caller's locked transaction. Returns whether to release the + rate-limit slot. """ should_release = False if status is not None: @@ -520,12 +525,12 @@ def _apply_pg_guarded_update( error: str | None, increment_attempt: bool, ) -> bool: - """PG-path terminal-one-way guarded write, inside the caller's row lock. Refuses - a revert of a protected-terminal (COMPLETED/STOPPED) status — ERROR stays - correctable — while still applying error / increment_attempt. Field-scoped - (``update_fields``) so a stale object can never clobber counters or the marker, - and the post-save cache publishes the status actually persisted. Returns whether - to release the rate-limit slot. + """Terminal-one-way guarded write for PG and STOPPED rows, inside the caller's + row lock. Refuses a revert of a protected-terminal (COMPLETED/STOPPED) status + — ERROR stays correctable — while still applying error / increment_attempt. + Field-scoped (``update_fields``) so a stale object can never clobber counters or + the marker, and the post-save cache publishes the status actually persisted. + Returns whether to release the rate-limit slot. """ committed = locked["status"] should_release = False diff --git a/unstract/core/src/unstract/core/data_models.py b/unstract/core/src/unstract/core/data_models.py index 2f4f58a54d..d3e7167632 100644 --- a/unstract/core/src/unstract/core/data_models.py +++ b/unstract/core/src/unstract/core/data_models.py @@ -582,6 +582,9 @@ def _protected_terminal_values(cls) -> frozenset[str]: ExecutionStatus.protected_terminal_values = classmethod(_protected_terminal_values) +# Reason the backend gives when it refuses a status write over a final status. +ALREADY_FINAL_REASON = "already_final" + # Add the is_completed method as a class method def _is_completed(cls, status: str) -> bool: diff --git a/workers/general/tasks.py b/workers/general/tasks.py index d1ffe422f0..4a4eaf6b78 100644 --- a/workers/general/tasks.py +++ b/workers/general/tasks.py @@ -17,6 +17,7 @@ # Import shared worker infrastructure using new structure from shared.api import InternalAPIClient +from shared.clients.execution_client import refused_as_stopped # Import worker-specific data models from shared.data.models import ( @@ -781,25 +782,7 @@ def _execute_general_workflow( successful_files=0, failed_files=0, total_time=0.0 ) - # Complete immediately with no files - try: - api_client.update_workflow_execution_status( - execution_id=execution_id, - status=ExecutionStatus.COMPLETED.value, - total_files=0, - ) - except Exception as status_error: - logger.warning(f"Failed to update execution status: {status_error}") - - # Update pipeline status if needed - if pipeline_id: - try: - api_client.update_pipeline_status( - pipeline_id=pipeline_id, - status=PipelineStatus.SUCCESS.value, - ) - except Exception as pipeline_error: - logger.warning(f"Failed to update pipeline status: {pipeline_error}") + _complete_without_files(api_client, execution_id, pipeline_id) execution_time = time.time() - start_time response = WorkerTaskResponse.success_response( @@ -881,6 +864,33 @@ def _execute_general_workflow( raise +def _complete_without_files( + api_client: InternalAPIClient, execution_id: str, pipeline_id: str | None +) -> None: + """Mark a run with no files COMPLETED and its pipeline SUCCESS. + + A run aborted meanwhile stays STOPPED, and the abort sets its pipeline status. + """ + stopped = False + try: + reply = api_client.update_workflow_execution_status( + execution_id=execution_id, + status=ExecutionStatus.COMPLETED.value, + total_files=0, + ) + stopped = refused_as_stopped(reply) + except Exception as status_error: + logger.warning("Failed to update execution status: %s", status_error) + if not pipeline_id or stopped: + return + try: + api_client.update_pipeline_status( + pipeline_id=pipeline_id, status=PipelineStatus.SUCCESS.value + ) + except Exception as pipeline_error: + logger.warning("Failed to update pipeline status: %s", pipeline_error) + + def _orchestrate_file_processing_general( api_client: InternalAPIClient, workflow_id: str, diff --git a/workers/shared/clients/execution_client.py b/workers/shared/clients/execution_client.py index 9c13474c97..30bbdb8e80 100644 --- a/workers/shared/clients/execution_client.py +++ b/workers/shared/clients/execution_client.py @@ -16,7 +16,7 @@ from typing import Any, cast from uuid import UUID -from unstract.core.data_models import ExecutionStatus +from unstract.core.data_models import ALREADY_FINAL_REASON, ExecutionStatus from ..data.models import ( BatchOperationRequest, @@ -92,7 +92,7 @@ def refused_as_stopped(response: Any) -> bool: """True when a status write was refused because an abort got there first.""" data = getattr(response, "data", None) or {} return bool( - data.get("reason") == "already_final" + data.get("reason") == ALREADY_FINAL_REASON and data.get("new_status") == ExecutionStatus.STOPPED.value ) diff --git a/workers/tests/test_abort_finishes_the_run.py b/workers/tests/test_abort_finishes_the_run.py index 409c04e189..a797eb4645 100644 --- a/workers/tests/test_abort_finishes_the_run.py +++ b/workers/tests/test_abort_finishes_the_run.py @@ -380,26 +380,53 @@ def test_a_stopped_run_bills_only_its_finished_files(status: str, commits: int) ], ids=["general", "api"], ) -@pytest.mark.parametrize(("status", "stopped"), [(STOPPED, True), (COMPLETED, False)]) +@pytest.mark.parametrize( + ("status", "stopped", "counts"), + [(STOPPED, True, (3, 3, 0)), (COMPLETED, False, (3, 2, 1))], +) def test_both_callbacks_pass_the_stop_to_the_final_log( - publish: Any, creator: str, status: str, stopped: bool + publish: Any, creator: str, status: str, stopped: bool, counts: tuple[int, int, int] ) -> None: context = cb.CallbackContext() context.execution_id = "e1" context.organization_id = "org-1" context.pipeline_id = "p1" context.api_client = mock.Mock() + saved = {"total_files": 3, "successful_files": 3, "failed_files": 0} context.api_client.get_workflow_execution.return_value = SimpleNamespace( - success=True, data={"aggregated_usage_cost": 0} + success=True, data={"aggregated_usage_cost": 0, "execution": saved} ) + worker = {"total_files": 3, "successful_files": 2, "failed_files": 1} workflow_logger = mock.Mock() with mock.patch.object( cb.WorkerWorkflowLogger, creator, return_value=workflow_logger ): - publish(context=context, aggregated_results={}, execution_status=status) + publish(context=context, aggregated_results=worker, execution_status=status) + + final = workflow_logger.publish_final_workflow_logs.call_args.kwargs + assert final["stopped"] is stopped + logged = (final["total_files"], final["successful_files"], final["failed_files"]) + assert logged == counts + + +@pytest.mark.parametrize( + ("reply", "pipeline_writes"), + [(_REFUSAL, 0), (_UPDATED, 1), (_COMPLETED_REFUSAL, 1)], + ids=["stopped", "live", "completed"], +) +def test_a_run_with_no_files_leaves_a_stopped_pipeline_to_the_abort( + reply: dict[str, Any], pipeline_writes: int +) -> None: + from general import tasks as general_tasks + + api_client = mock.Mock() + api_client.update_workflow_execution_status.return_value = SimpleNamespace( + data=reply + ) + + general_tasks._complete_without_files(api_client, "exec-1", "pipe-1") - final = workflow_logger.publish_final_workflow_logs - assert final.call_args.kwargs["stopped"] is stopped + assert api_client.update_pipeline_status.call_count == pipeline_writes def _active_files(status: Any) -> tuple[set[str], Any]: From 71d701b2c3e2b1e54ca601864820b4a62ef29e9c Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Wed, 11 Nov 2026 12:54:00 +0530 Subject: [PATCH 8/8] UN-4318 [FEAT] Report an aborted run once and keep its abort message - Callers whose status write the abort refused now act on the stored status: no second webhook or pipeline write, no ERROR answer for a STOPPED run, and no "completed" log for a stopped run with no files. - A stopped run keeps the error text "Execution aborted". - Comments and the entry-skip log name runs aborted before pickup. - Tests cover each caller, including the no-files call site. Co-Authored-By: Claude Opus 5.5 --- backend/api_v2/deployment_helper.py | 12 +- .../api_v2/tests/test_deployment_helper.py | 39 +++++++ backend/workflow_manager/execution/README.md | 4 +- .../execution/tests/test_abort_executions.py | 20 ++++ .../workflow_v2/models/execution.py | 8 +- .../workflow_v2/workflow_helper.py | 7 +- .../core/src/unstract/core/data_models.py | 2 +- workers/general/tasks.py | 53 ++++----- workers/tests/test_abort_finishes_the_run.py | 105 ++++++++++++++++++ 9 files changed, 216 insertions(+), 34 deletions(-) diff --git a/backend/api_v2/deployment_helper.py b/backend/api_v2/deployment_helper.py index a04fd79ff8..76585cc91a 100644 --- a/backend/api_v2/deployment_helper.py +++ b/backend/api_v2/deployment_helper.py @@ -296,10 +296,13 @@ def execute_workflow( ) except Exception as error: # Isolate the DB error-marking so cleanup still runs if it raises. + stored_status = ExecutionStatus.ERROR.value try: - WorkflowExecutionServiceHelper.update_execution_err( + marked = WorkflowExecutionServiceHelper.update_execution_err( str(execution_id), str(error) ) + # An aborted run stays STOPPED; answer with what is stored. + stored_status = marked.status if marked else stored_status except Exception: logger.exception(f"Failed to mark execution {execution_id} as ERROR") @@ -317,7 +320,7 @@ def execute_workflow( ExecutionResponse( workflow_id=workflow_id, execution_id=execution_id, - execution_status=ExecutionStatus.ERROR.value, + execution_status=stored_status, error=str(error), ) ).data @@ -360,6 +363,11 @@ def execute_workflow( f"Failed to acknowledge results for execution {execution_id}" ) + # An aborted run was already reported by the abort. + if ( + execution is not None + and execution.status == ExecutionStatus.COMPLETED.value + ): try: # Terminalising here bypasses the usual notifier, so subscribers # would otherwise hear nothing. diff --git a/backend/api_v2/tests/test_deployment_helper.py b/backend/api_v2/tests/test_deployment_helper.py index 2e8c3ec7d3..f3f37f8634 100644 --- a/backend/api_v2/tests/test_deployment_helper.py +++ b/backend/api_v2/tests/test_deployment_helper.py @@ -92,6 +92,24 @@ def test_staging_failure_cleanup_survives_db_marking_error(collaborators) -> Non collaborators["DestinationConnector"].delete_api_storage_dir.assert_called_once() +def test_staging_failure_after_an_abort_answers_stopped( + collaborators: dict[str, MagicMock], +) -> None: + """The answer matches the stored status, which the abort keeps STOPPED.""" + collaborators["WorkflowExecutionServiceHelper"].update_execution_err.return_value = ( + MagicMock(status="STOPPED") + ) + + response = dh.DeploymentHelper.execute_workflow( + organization_name="org", + api=_api(), + file_objs=[], + timeout=-1, + ) + + assert response["execution_status"] == "STOPPED" + + @pytest.fixture def staging_rejects_everything(): """Patch execute_workflow's collaborators; staging returns no dispatchable files.""" @@ -195,6 +213,27 @@ def test_all_files_rejected_acknowledges_and_notifies( ) +def test_an_aborted_run_is_not_reported_again( + staging_rejects_everything: dict[str, MagicMock], +) -> None: + """The abort already notified subscribers of a run it stopped.""" + mocks = staging_rejects_everything + stopped_row = mocks[ + "WorkflowExecutionServiceHelper" + ].update_execution_completed.return_value + stopped_row.status = "STOPPED" + + response = dh.DeploymentHelper.execute_workflow( + organization_name="org", + api=_api(), + file_objs=[MagicMock()], + timeout=-1, + ) + + mocks["PipelineUtils"].update_pipeline_status.assert_not_called() + assert response["execution_status"] == "STOPPED" + + def test_notification_survives_a_failed_acknowledgement( staging_rejects_everything, ) -> None: diff --git a/backend/workflow_manager/execution/README.md b/backend/workflow_manager/execution/README.md index c236b5f8a9..51b6ee50db 100644 --- a/backend/workflow_manager/execution/README.md +++ b/backend/workflow_manager/execution/README.md @@ -151,7 +151,9 @@ Once that commits, it: and notifications to the abort. - A caller waiting on a synchronous API deployment `POST` gets 422 for an aborted run, the same as polling with `GET`. Both apply the same result filters. -- Later status writes, from workers or the backend, cannot undo the abort. +- Later status writes, from workers or the backend, cannot undo the abort or + replace its error. Code whose write is refused leaves the pipeline status and + notifications to the abort. ### Errors diff --git a/backend/workflow_manager/execution/tests/test_abort_executions.py b/backend/workflow_manager/execution/tests/test_abort_executions.py index 023a0e8d64..38622294a0 100644 --- a/backend/workflow_manager/execution/tests/test_abort_executions.py +++ b/backend/workflow_manager/execution/tests/test_abort_executions.py @@ -27,6 +27,7 @@ from workflow_manager.workflow_v2.execution import WorkflowExecutionServiceHelper from workflow_manager.workflow_v2.models.execution import WorkflowExecution from workflow_manager.workflow_v2.models.workflow import Workflow +from workflow_manager.workflow_v2.workflow_helper import WorkflowHelper _abort = ExecutionAbortView.as_view() @@ -442,6 +443,25 @@ def test_backend_writes_before_dispatch_cannot_undo_the_abort(self) -> None: execution.refresh_from_db() self.assertEqual(execution.status, ExecutionStatus.STOPPED.value) self.assertEqual((execution.total_files, execution.failed_files), (0, 0)) + self.assertEqual(execution.error_message, ABORT_ERROR_MESSAGE) + + def test_a_failed_dispatch_after_the_abort_answers_stopped(self) -> None: + execution = self.execution( + self.deployment(), status=ExecutionStatus.PENDING, queue_message_id=None + ) + self._abort(self.owner, execution.id) + + with patch.object( + WorkflowHelper, "_dispatch_orchestrator_task", side_effect=RuntimeError("down") + ): + response = WorkflowHelper.execute_workflow_async( + workflow_id=str(self.workflow.id), + execution_id=str(execution.id), + hash_values_of_files={}, + ) + + self.assertEqual(response.execution_status, ExecutionStatus.STOPPED.value) + self.assertEqual(self._status(execution), ExecutionStatus.STOPPED.value) @staticmethod def _write_status(execution: WorkflowExecution, data: dict[str, Any]) -> Any: diff --git a/backend/workflow_manager/workflow_v2/models/execution.py b/backend/workflow_manager/workflow_v2/models/execution.py index 9ffdff8f61..da1d2db49f 100644 --- a/backend/workflow_manager/workflow_v2/models/execution.py +++ b/backend/workflow_manager/workflow_v2/models/execution.py @@ -435,7 +435,7 @@ def update_execution( # self.queue_message_id — unless the caller passes row_locked=True, i.e. # this instance was itself read under the lock (see the docstring): a PG row snapshotted before dispatch recorded its # handle, mis-routed to an unlocked legacy full save(), would write NULL - # back into queue_message_id and permanently disable every guard. Under + # back into queue_message_id and permanently drop the COMPLETED guard. Under # this lock a concurrent _record_dispatch_handle serializes after us. # # no_key=True takes FOR NO KEY UPDATE rather than FOR UPDATE. Status @@ -527,7 +527,8 @@ def _apply_pg_guarded_update( ) -> bool: """Terminal-one-way guarded write for PG and STOPPED rows, inside the caller's row lock. Refuses a revert of a protected-terminal (COMPLETED/STOPPED) status - — ERROR stays correctable — while still applying error / increment_attempt. + — ERROR stays correctable — while still applying error / increment_attempt + (a STOPPED row keeps its error text). Field-scoped (``update_fields``) so a stale object can never clobber counters or the marker, and the post-save cache publishes the status actually persisted. Returns whether to release the rate-limit slot. @@ -540,7 +541,8 @@ def _apply_pg_guarded_update( ExecutionStatus(status), committed, error, increment_attempt ) update_fields.extend(status_fields) - if error: + # A stopped run keeps the abort's message. + if error and committed != ExecutionStatus.STOPPED.value: self.error_message = error[:EXECUTION_ERROR_LENGTH] update_fields.append("error_message") if increment_attempt: diff --git a/backend/workflow_manager/workflow_v2/workflow_helper.py b/backend/workflow_manager/workflow_v2/workflow_helper.py index e6ceb8d6c8..2b323fd458 100644 --- a/backend/workflow_manager/workflow_v2/workflow_helper.py +++ b/backend/workflow_manager/workflow_v2/workflow_helper.py @@ -783,7 +783,9 @@ def execute_workflow_async( return ExecutionResponse( workflow_id, execution_id, ExecutionStatus.EXECUTING.value ) - WorkflowExecutionServiceHelper.update_execution_err(execution_id, str(error)) + marked = WorkflowExecutionServiceHelper.update_execution_err( + execution_id, str(error) + ) logger.exception( f"Error while enqueuing async job for WF '{workflow_id}', " f"execution '{execution_id}'" @@ -791,7 +793,8 @@ def execute_workflow_async( return ExecutionResponse( workflow_id, execution_id, - ExecutionStatus.ERROR.value, + # An aborted run stays STOPPED. + marked.status if marked else ExecutionStatus.ERROR.value, error=str(error), ) diff --git a/unstract/core/src/unstract/core/data_models.py b/unstract/core/src/unstract/core/data_models.py index d3e7167632..70424621e2 100644 --- a/unstract/core/src/unstract/core/data_models.py +++ b/unstract/core/src/unstract/core/data_models.py @@ -582,7 +582,7 @@ def _protected_terminal_values(cls) -> frozenset[str]: ExecutionStatus.protected_terminal_values = classmethod(_protected_terminal_values) -# Reason the backend gives when it refuses a status write over a final status. +# Reason the backend gives when it refuses a status write over a protected status. ALREADY_FINAL_REASON = "already_final" diff --git a/workers/general/tasks.py b/workers/general/tasks.py index 4a4eaf6b78..69637ee4fb 100644 --- a/workers/general/tasks.py +++ b/workers/general/tasks.py @@ -305,12 +305,14 @@ def async_execute_bin_general( # to test) a redelivery could reach here, win a fresh claim, and re-arm + # re-dispatch a finished execution — duplicate (costly) destination # writes when use_file_history=False. Skip if already terminal, keeping - # the freshly-won claim so it re-establishes the tombstone. + # the freshly-won claim so it re-establishes the tombstone. A run aborted + # before pickup is skipped here too. if current_status in ExecutionStatus.terminal_values(): logger.warning( - f"Execution {execution_id} already terminal ({current_status}) " - f"at orchestration entry — skipping re-orchestration (a " - f"redelivery of a finished execution)." + "Execution %s already terminal (%s) at orchestration entry — " + "skipping (a redelivery, or a run aborted before pickup).", + execution_id, + current_status, ) return { "status": "skipped_terminal_execution", @@ -425,14 +427,15 @@ def async_execute_bin_general( try: with InternalAPIClient(config) as api_client: api_client.set_organization_context(schema_name) - api_client.update_workflow_execution_status( + reply = api_client.update_workflow_execution_status( execution_id=execution_id, status=ExecutionStatus.ERROR.value, error_message=str(e), ) - # CRITICAL FIX: Also update pipeline status to FAILED for consistency - if pipeline_id: + # CRITICAL FIX: Also update pipeline status to FAILED for consistency. + # An aborted run's pipeline status and notifications are the abort's. + if pipeline_id and not refused_as_stopped(reply): try: api_client.update_pipeline_status( pipeline_id=pipeline_id, @@ -777,12 +780,11 @@ def _execute_general_workflow( "💤 No files to process - all files have been filtered out (already processed, active, or no matches)", ) - # Send completion log to UI - workflow_logger.publish_execution_complete( - successful_files=0, failed_files=0, total_time=0.0 - ) - - _complete_without_files(api_client, execution_id, pipeline_id) + if _complete_without_files(api_client, execution_id, pipeline_id): + # Send completion log to UI + workflow_logger.publish_execution_complete( + successful_files=0, failed_files=0, total_time=0.0 + ) execution_time = time.time() - start_time response = WorkerTaskResponse.success_response( @@ -866,29 +868,30 @@ def _execute_general_workflow( def _complete_without_files( api_client: InternalAPIClient, execution_id: str, pipeline_id: str | None -) -> None: +) -> bool: """Mark a run with no files COMPLETED and its pipeline SUCCESS. - A run aborted meanwhile stays STOPPED, and the abort sets its pipeline status. + Returns False for a run aborted meanwhile: it stays STOPPED, and the abort + sets its pipeline status. """ - stopped = False try: reply = api_client.update_workflow_execution_status( execution_id=execution_id, status=ExecutionStatus.COMPLETED.value, total_files=0, ) - stopped = refused_as_stopped(reply) + if refused_as_stopped(reply): + return False except Exception as status_error: logger.warning("Failed to update execution status: %s", status_error) - if not pipeline_id or stopped: - return - try: - api_client.update_pipeline_status( - pipeline_id=pipeline_id, status=PipelineStatus.SUCCESS.value - ) - except Exception as pipeline_error: - logger.warning("Failed to update pipeline status: %s", pipeline_error) + if pipeline_id: + try: + api_client.update_pipeline_status( + pipeline_id=pipeline_id, status=PipelineStatus.SUCCESS.value + ) + except Exception as pipeline_error: + logger.warning("Failed to update pipeline status: %s", pipeline_error) + return True def _orchestrate_file_processing_general( diff --git a/workers/tests/test_abort_finishes_the_run.py b/workers/tests/test_abort_finishes_the_run.py index a797eb4645..4fc18747dd 100644 --- a/workers/tests/test_abort_finishes_the_run.py +++ b/workers/tests/test_abort_finishes_the_run.py @@ -429,6 +429,111 @@ def test_a_run_with_no_files_leaves_a_stopped_pipeline_to_the_abort( assert api_client.update_pipeline_status.call_count == pipeline_writes +def _run_with_no_files(reply: dict[str, Any]) -> tuple[Any, Any]: + from general import tasks as general_tasks + + api_client = mock.Mock(organization_id="org-1") + api_client.update_workflow_execution_status.return_value = SimpleNamespace( + data=reply + ) + source = mock.Mock() + source.list_files_from_source.return_value = ({}, 0) + workflow_logger = mock.Mock() + with ExitStack() as stack: + for name in ("create_organization_context", "WorkflowContextData"): + stack.enter_context(mock.patch.object(general_tasks, name)) + stack.enter_context( + mock.patch.object(general_tasks, "WorkerSourceConnector", return_value=source) + ) + stack.enter_context( + mock.patch.object( + general_tasks.WorkerWorkflowLogger, + "create_for_general_workflow", + return_value=workflow_logger, + ) + ) + general_tasks._execute_general_workflow( + api_client=api_client, + execution_context={"execution": {"id": "exec-1", "workflow_id": "wf-1"}}, + file_batch_results=[], + pipeline_id="pipe-1", + execution_mode=None, + use_file_history=True, + scheduled=False, + schema_name="org-1", + ) + return api_client, workflow_logger + + +@pytest.mark.parametrize( + ("reply", "reported"), [(_REFUSAL, 0), (_UPDATED, 1)], ids=["stopped", "live"] +) +def test_a_stopped_run_with_no_files_reports_nothing( + reply: dict[str, Any], reported: int +) -> None: + api_client, workflow_logger = _run_with_no_files(reply) + + assert api_client.update_pipeline_status.call_count == reported + assert workflow_logger.publish_execution_complete.call_count == reported + + +def _fail_orchestration(reply: dict[str, Any]) -> tuple[Any, Any]: + from general import tasks as general_tasks + + task = general_tasks.async_execute_bin_general + run = task + while hasattr(run, "__wrapped__"): + run = run.__wrapped__ + client = mock.Mock() + client.get_workflow_execution.return_value = SimpleNamespace( + success=False, error="boom" + ) + error_client = mock.MagicMock() + error_client.__enter__.return_value = error_client + error_client.update_workflow_execution_status.return_value = SimpleNamespace( + data=reply + ) + notify = mock.Mock() + last_attempt = SimpleNamespace( + request=SimpleNamespace(id="t-1", retries=task.max_retries), + max_retries=task.max_retries, + ) + with ExitStack() as stack: + stack.enter_context( + mock.patch.dict( + run.__globals__, + { + "_should_skip_duplicate_orchestration": lambda *_: False, + "release_orchestration_claim": lambda _eid: None, + "InternalAPIClient": lambda _config: error_client, + "notify_execution_failure": notify, + }, + ) + ) + stack.enter_context( + mock.patch.object( + general_tasks.WorkerExecutionContext, + "setup_execution_context", + return_value=(mock.Mock(), client), + ) + ) + with pytest.raises(Exception, match="boom"): + run(last_attempt, "org-1", "wf-1", "exec-1", {}, pipeline_id="pipe-1") + return error_client, notify + + +@pytest.mark.parametrize( + ("reply", "reported"), [(_REFUSAL, 0), (_UPDATED, 1)], ids=["stopped", "live"] +) +def test_a_failed_orchestration_leaves_a_stopped_run_to_the_abort( + reply: dict[str, Any], reported: int +) -> None: + error_client, notify = _fail_orchestration(reply) + + assert error_client.update_pipeline_status.call_count == reported + assert notify.call_count == reported + + def _active_files(status: Any) -> tuple[set[str], Any]: api_client = mock.Mock() api_client.get_workflow_execution_status.side_effect = [status]