diff --git a/openrag/api/routers/admin/indexing.py b/openrag/api/routers/admin/indexing.py index 9006ba177..078075918 100644 --- a/openrag/api/routers/admin/indexing.py +++ b/openrag/api/routers/admin/indexing.py @@ -30,6 +30,7 @@ validate_metadata, ) from api.routers.admin.task_logs import collect_task_logs +from core.models.catalog import TERMINAL_TASK_STATES from core.utils.exceptions import OpenRAGError from core.utils.filename import sanitize_filename from core.utils.log_tail import app_log_file @@ -557,5 +558,11 @@ async def cancel_task( ): cancelled = await service.cancel_task(task_id) if not cancelled: + task_state = await service.get_task_state(task_id) + if task_state in TERMINAL_TASK_STATES: + raise HTTPException( + status.HTTP_409_CONFLICT, + f"Task {task_id} is already {task_state.lower()} and cannot be cancelled", + ) raise HTTPException(404, f"No ObjectRef stored for task {task_id}") return {"message": f"Cancellation signal sent for task {task_id}"} diff --git a/openrag/core/indexing/dispatcher.py b/openrag/core/indexing/dispatcher.py index 03ed0f001..ab6b619cf 100644 --- a/openrag/core/indexing/dispatcher.py +++ b/openrag/core/indexing/dispatcher.py @@ -82,8 +82,7 @@ async def get_task_error(self, task_id: str) -> str | None: async def cancel_task(self, task_id: str) -> bool: """Cancel a running/queued task. - Returns ``False`` when no object ref is stored for ``task_id`` - (the caller maps that to a 404), ``True`` once the cancel signal - has been sent. + Returns ``False`` when the task cannot be cancelled, ``True`` once + the cancel signal has been sent. """ ... diff --git a/openrag/core/models/__init__.py b/openrag/core/models/__init__.py index 1ddb25e67..dd6bd1a64 100644 --- a/openrag/core/models/__init__.py +++ b/openrag/core/models/__init__.py @@ -1,6 +1,6 @@ """Domain models — pure Pydantic, no infrastructure imports.""" -from .catalog import DocumentRecord, DocumentStatus, IndexationJob, JobStatus +from .catalog import TERMINAL_TASK_STATES, DocumentRecord, DocumentStatus, IndexationJob, JobStatus from .chunk import Chunk, ChunkType from .contextualization import ContextualizedQuery from .conversation import Conversation, Message @@ -13,6 +13,7 @@ from .workspace import Workspace __all__ = [ + "TERMINAL_TASK_STATES", "ApiKey", "Chunk", "ChunkType", diff --git a/openrag/core/models/catalog.py b/openrag/core/models/catalog.py index 721047f26..c650c5049 100644 --- a/openrag/core/models/catalog.py +++ b/openrag/core/models/catalog.py @@ -20,6 +20,13 @@ class DocumentStatus(str, Enum): CANCELLED = "CANCELLED" +# States a task/document cannot leave once reached. Single source of truth for +# the cancellation guard shared by the indexing route, dispatcher, and +# TaskStateManager — keep those in sync via this constant rather than +# re-declaring the set. +TERMINAL_TASK_STATES = frozenset({DocumentStatus.COMPLETED, DocumentStatus.FAILED, DocumentStatus.CANCELLED}) + + class JobStatus(str, Enum): QUEUED = "QUEUED" RUNNING = "RUNNING" diff --git a/openrag/services/workers/dispatcher.py b/openrag/services/workers/dispatcher.py index 8a38c99dc..7622c53de 100644 --- a/openrag/services/workers/dispatcher.py +++ b/openrag/services/workers/dispatcher.py @@ -237,11 +237,19 @@ async def cancel_task(self, task_id: str) -> bool: if obj_ref is None: return False - ray.cancel(obj_ref["ref"], recursive=True) - await self._call( - self._tsm.set_state.remote(task_id, "CANCELLED"), - task_description=f"set_state({task_id})", + # Atomically claim the cancellation first: if ray.cancel() ran before + # this and the RPC then failed, a killed worker could never report + # back and the task would be stuck active forever (a zombie). + # TaskStateManager keeps CANCELLED sticky, so a worker that starts in + # this small window cannot report active/success after the cancel claim. + cancelled = await self._call( + self._tsm.set_cancelled_if_active.remote(task_id), + task_description=f"set_cancelled_if_active({task_id})", ) + if not cancelled: + return False + + ray.cancel(obj_ref["ref"], recursive=True) return True diff --git a/openrag/services/workers/task_state.py b/openrag/services/workers/task_state.py index 7e5bc6329..a5fb6258d 100644 --- a/openrag/services/workers/task_state.py +++ b/openrag/services/workers/task_state.py @@ -5,6 +5,7 @@ from typing import Any import ray +from core.models.catalog import TERMINAL_TASK_STATES, DocumentStatus try: from core.config import load_config as _load_config @@ -46,6 +47,8 @@ async def _ensure_task(self, task_id: str) -> TaskInfo: async def set_state(self, task_id: str, state: str) -> None: async with self.lock: info = await self._ensure_task(task_id) + if info.state == DocumentStatus.CANCELLED and state != DocumentStatus.CANCELLED: + return info.state = state @ray.method(concurrency_group="set") @@ -65,6 +68,15 @@ async def set_failed_if_not_cancelled(self, task_id: str, tb_str: str) -> bool: info.error = tb_str return True + @ray.method(concurrency_group="set") + async def set_cancelled_if_active(self, task_id: str) -> bool: + async with self.lock: + info = self.tasks.get(task_id) + if info is None or info.state in TERMINAL_TASK_STATES: + return False + info.state = "CANCELLED" + return True + @ray.method(concurrency_group="set") async def set_details( self, diff --git a/tests/unit/services/workers/test_dispatcher.py b/tests/unit/services/workers/test_dispatcher.py index 39e3db6f5..ac5ca5cd8 100644 --- a/tests/unit/services/workers/test_dispatcher.py +++ b/tests/unit/services/workers/test_dispatcher.py @@ -61,6 +61,7 @@ def _workspace_repo() -> MagicMock: def _task_state_manager() -> MagicMock: tsm = MagicMock() tsm.set_state = _remote_mock() + tsm.set_cancelled_if_active = _remote_mock(True) tsm.set_details = _remote_mock() tsm.set_object_ref = _remote_mock() tsm.get_state = _remote_mock("SERIALIZING") @@ -195,7 +196,6 @@ async def test_cancel_task_uses_stored_pool_object_ref() -> None: ref = object() tsm = _task_state_manager() tsm.get_object_ref.remote = AsyncMock(return_value={"ref": ref}) - tsm.get_state.remote = AsyncMock(return_value="SERIALIZING") dispatcher = WorkerDispatcher( pool=_pool_with_ref(object()), task_state_manager=tsm, @@ -209,18 +209,18 @@ async def test_cancel_task_uses_stored_pool_object_ref() -> None: result = await dispatcher.cancel_task("task-1") assert result is True + tsm.set_cancelled_if_active.remote.assert_called_once_with("task-1") cancel.assert_called_once_with(ref, recursive=True) - tsm.set_state.remote.assert_called_once_with("task-1", "CANCELLED") @pytest.mark.asyncio -async def test_cancel_task_marks_cancelled_even_if_worker_finished_first() -> None: +async def test_cancel_task_does_not_cancel_terminal_task() -> None: from services.workers.dispatcher import WorkerDispatcher ref = object() tsm = _task_state_manager() tsm.get_object_ref.remote = AsyncMock(return_value={"ref": ref}) - tsm.get_state.remote = AsyncMock(return_value="COMPLETED") + tsm.set_cancelled_if_active.remote = AsyncMock(return_value=False) dispatcher = WorkerDispatcher( pool=_pool_with_ref(object()), task_state_manager=tsm, @@ -230,11 +230,12 @@ async def test_cancel_task_marks_cancelled_even_if_worker_finished_first() -> No collection="default", ) - with patch("ray.cancel"): + with patch("ray.cancel") as cancel: result = await dispatcher.cancel_task("task-1") - assert result is True - tsm.set_state.remote.assert_called_once_with("task-1", "CANCELLED") + assert result is False + tsm.set_cancelled_if_active.remote.assert_called_once_with("task-1") + cancel.assert_not_called() @pytest.mark.asyncio diff --git a/tests/unit/services/workers/test_task_state.py b/tests/unit/services/workers/test_task_state.py new file mode 100644 index 000000000..881bba8b1 --- /dev/null +++ b/tests/unit/services/workers/test_task_state.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +from typing import Any + +import pytest +from services.workers.task_state import TaskStateManager + + +def _task_state_manager() -> Any: + return TaskStateManager.__ray_metadata__.modified_class() + + +@pytest.mark.asyncio +async def test_cancelled_state_is_not_overwritten_by_worker_transitions() -> None: + manager = _task_state_manager() + + await manager.set_state("task-1", "QUEUED") + assert await manager.set_cancelled_if_active("task-1") is True + + await manager.set_state("task-1", "SERIALIZING") + await manager.set_state("task-1", "COMPLETED") + + assert await manager.get_state("task-1") == "CANCELLED" diff --git a/ui/src/components/shared/data-table.test.tsx b/ui/src/components/shared/data-table.test.tsx index 936bd5a6b..9ec126a30 100644 --- a/ui/src/components/shared/data-table.test.tsx +++ b/ui/src/components/shared/data-table.test.tsx @@ -119,4 +119,18 @@ describe("DataTable selection", () => { fireEvent.click(screen.getAllByRole("checkbox")[0]); expect(screen.queryByText(/selected/)).toBeNull(); }); + + it("disables select-all when the visible rows cannot be selected", () => { + render( + false} + getRowId={(r) => String(r.id)} + />, + ); + + expect(screen.getByRole("checkbox", { name: /select visible rows/i }).hasAttribute("disabled")).toBe(true); + }); }); diff --git a/ui/src/components/shared/data-table.tsx b/ui/src/components/shared/data-table.tsx index 463673257..71fb6668f 100644 --- a/ui/src/components/shared/data-table.tsx +++ b/ui/src/components/shared/data-table.tsx @@ -31,6 +31,8 @@ interface BaseDataTableProps { initialSorting?: SortingState; /** Render a leading checkbox column. */ enableSelection?: boolean; + /** Optional row-level selection guard for pages with state-dependent bulk actions. */ + canSelectRow?: (row: TData) => boolean; /** Stable row id (e.g. file_id) so selection survives a data refetch. */ getRowId?: (row: TData) => string; } @@ -54,6 +56,7 @@ export function DataTable({ pageSize = 10, initialSorting = [], enableSelection = false, + canSelectRow, getRowId, rowSelection: controlledRowSelection, onRowSelectionChange, @@ -71,20 +74,25 @@ export function DataTable({ const selectColumn: ColumnDef = { id: "__select", enableSorting: false, - header: ({ table }) => ( - table.toggleAllPageRowsSelected(!!value)} - aria-label="Select visible rows" - /> - ), + header: ({ table }) => { + const hasSelectablePageRows = table.getRowModel().rows.some((row) => row.getCanSelect()); + return ( + table.toggleAllPageRowsSelected(!!value)} + disabled={!hasSelectablePageRows} + aria-label="Select visible rows" + /> + ); + }, cell: ({ row }) => ( row.toggleSelected(!!value)} + disabled={!row.getCanSelect()} aria-label="Select row" /> ), @@ -103,7 +111,9 @@ export function DataTable({ onColumnFiltersChange: setColumnFilters, onPaginationChange: setPagination, onRowSelectionChange: setRowSelection, - enableRowSelection: enableSelection, + enableRowSelection: enableSelection + ? (row) => (canSelectRow ? canSelectRow(row.original) : true) + : false, getRowId, // Keep the user on their current page when the rows change underneath them // (e.g. deleting a row refetches the list). TanStack resets pageIndex to 0 diff --git a/ui/src/pages/admin/jobs/list.test.tsx b/ui/src/pages/admin/jobs/list.test.tsx index ee5142310..2c0d5e6fa 100644 --- a/ui/src/pages/admin/jobs/list.test.tsx +++ b/ui/src/pages/admin/jobs/list.test.tsx @@ -2,10 +2,18 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { MemoryRouter } from "react-router-dom"; +import { toast } from "sonner"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import { getQueueInfo, listTasks } from "@/lib/api/jobs"; +import { cancelTask, getQueueInfo, listTasks, type TaskState } from "@/lib/api/jobs"; import JobListPage from "./list"; +vi.mock("sonner", () => ({ + toast: { + success: vi.fn(), + error: vi.fn(), + }, +})); + const permissions = vi.hoisted(() => ({ isAdmin: true })); vi.mock("@/lib/permissions", () => ({ @@ -14,15 +22,22 @@ vi.mock("@/lib/permissions", () => ({ }), })); -vi.mock("@/lib/api/jobs", () => ({ - getQueueInfo: vi.fn(), - listTasks: vi.fn(), -})); +vi.mock("@/lib/api/jobs", async () => { + const actual = await vi.importActual("@/lib/api/jobs"); + return { + ...actual, + cancelTask: vi.fn(), + getQueueInfo: vi.fn(), + listTasks: vi.fn(), + }; +}); +const cancelTaskMock = vi.mocked(cancelTask); const getQueueInfoMock = vi.mocked(getQueueInfo); const listTasksMock = vi.mocked(listTasks); +const toastErrorMock = vi.mocked(toast.error); -const task = (task_id: string, state: "COMPLETED" | "FAILED", filename: string, partition = "docs") => ({ +const task = (task_id: string, state: TaskState, filename: string, partition = "docs") => ({ task_id, state, details: { @@ -54,6 +69,7 @@ describe("JobListPage filters", () => { beforeEach(() => { permissions.isAdmin = true; vi.clearAllMocks(); + cancelTaskMock.mockResolvedValue({ message: "Cancellation signal sent" }); getQueueInfoMock.mockResolvedValue({ workers: { total_slots: 4, pool_size: 2, max_per_actor: 2 }, tasks: { @@ -90,6 +106,70 @@ describe("JobListPage filters", () => { expect(await screen.findByText("failed.pdf")).not.toBeNull(); }); + it("bulk-cancels only selected active jobs", async () => { + listTasksMock.mockResolvedValue({ + tasks: [ + task("queued-task", "QUEUED", "queued.pdf"), + task("completed-task", "COMPLETED", "completed.pdf"), + ], + }); + + renderJobs(); + + expect(await screen.findByText("queued.pdf")).not.toBeNull(); + expect(screen.getByText("completed.pdf")).not.toBeNull(); + + const rowCheckboxes = screen.getAllByRole("checkbox", { name: /select row/i }); + expect(rowCheckboxes[0].hasAttribute("disabled")).toBe(false); + expect(rowCheckboxes[1].hasAttribute("disabled")).toBe(true); + + await userEvent.click(screen.getByRole("checkbox", { name: /select visible rows/i })); + expect(screen.getByText("1 selected")).not.toBeNull(); + + await userEvent.click(screen.getByRole("button", { name: /cancel selected/i })); + await userEvent.click(screen.getByRole("button", { name: /confirm/i })); + + await waitFor(() => expect(cancelTaskMock).toHaveBeenCalledWith("queued-task")); + expect(cancelTaskMock).not.toHaveBeenCalledWith("completed-task"); + }); + + it("rechecks selected jobs before bulk cancelling", async () => { + let tasks = [task("queued-task", "QUEUED", "queued.pdf")]; + listTasksMock.mockImplementation(async () => ({ tasks })); + + renderJobs(); + + expect(await screen.findByText("queued.pdf")).not.toBeNull(); + + await userEvent.click(screen.getByRole("checkbox", { name: /select row/i })); + expect(screen.getByText("1 selected")).not.toBeNull(); + + tasks = [task("queued-task", "COMPLETED", "queued.pdf")]; + + await userEvent.click(screen.getByRole("button", { name: /cancel selected/i })); + await userEvent.click(screen.getByRole("button", { name: /confirm/i })); + + await waitFor(() => expect(toastErrorMock).toHaveBeenCalledWith("1 task(s) were no longer active")); + expect(cancelTaskMock).not.toHaveBeenCalled(); + }); + + it("clears selected jobs when the search changes", async () => { + listTasksMock.mockResolvedValue({ + tasks: [task("queued-task", "QUEUED", "queued.pdf"), task("other-task", "QUEUED", "other.pdf")], + }); + + renderJobs(); + + expect(await screen.findByText("queued.pdf")).not.toBeNull(); + + await userEvent.click(screen.getAllByRole("checkbox", { name: /select row/i })[0]); + expect(screen.getByText("1 selected")).not.toBeNull(); + + await userEvent.type(screen.getByPlaceholderText("Search jobs..."), "other"); + + await waitFor(() => expect(screen.queryByText("1 selected")).toBeNull()); + }); + it("keeps long job values constrained while exposing full names", async () => { const longTaskId = "task-" + "1234567890".repeat(5); const longFilename = "benchmark-run-with-a-very-long-document-name-that-should-not-stretch-the-table.pdf"; diff --git a/ui/src/pages/admin/jobs/list.tsx b/ui/src/pages/admin/jobs/list.tsx index 5b8fcac54..f47475ccb 100644 --- a/ui/src/pages/admin/jobs/list.tsx +++ b/ui/src/pages/admin/jobs/list.tsx @@ -1,18 +1,20 @@ import { useEffect, useMemo, useState } from "react"; import { Link } from "react-router-dom"; -import { useQuery } from "@tanstack/react-query"; -import type { ColumnDef } from "@tanstack/react-table"; -import { RefreshCw, Search } from "lucide-react"; +import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; +import type { ColumnDef, RowSelectionState } from "@tanstack/react-table"; +import { Ban, RefreshCw, Search } from "lucide-react"; +import { toast } from "sonner"; import { usePermissions } from "@/lib/permissions"; import { PageHeader } from "@/components/shared/page-header"; import { DataTable } from "@/components/shared/data-table"; import { StatusBadge } from "@/components/shared/status-badge"; +import { ConfirmDialog } from "@/components/shared/confirm-dialog"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Tabs, TabsList, TabsTrigger, TabsContent } from "@/components/ui/tabs"; -import { getQueueInfo, listTasks, type QueueInfo, type TaskListItem } from "@/lib/api/jobs"; +import { cancelTask, getQueueInfo, isActiveState, listTasks, type QueueInfo, type TaskListItem } from "@/lib/api/jobs"; // OpenRag exposes per-file indexing tasks (TaskStateManager), not batch "jobs". const STATUS_TABS = ["ALL", "ACTIVE", "COMPLETED", "FAILED", "CANCELLED"] as const; @@ -111,10 +113,13 @@ const columns: ColumnDef[] = [ export default function JobListPage() { const { isAdmin } = usePermissions(); + const queryClient = useQueryClient(); const [statusTab, setStatusTab] = useState("ALL"); const [search, setSearch] = useState(""); const [debouncedSearch, setDebouncedSearch] = useState(""); + const [rowSelection, setRowSelection] = useState({}); const [manualRefreshing, setManualRefreshing] = useState(false); + const taskStatusFilter = statusTab === "ALL" ? undefined : statusTab === "ACTIVE" ? "active" : statusTab; useEffect(() => { const timeout = window.setTimeout(() => setDebouncedSearch(search), JOB_SEARCH_DEBOUNCE_MS); @@ -123,7 +128,7 @@ export default function JobListPage() { const tasksQuery = useQuery({ queryKey: ["tasks", statusTab], - queryFn: () => listTasks(statusTab === "ALL" ? undefined : statusTab === "ACTIVE" ? "active" : statusTab), + queryFn: () => listTasks(taskStatusFilter), refetchInterval: JOBS_REFETCH_INTERVAL_MS, // poll — OpenRag has no task SSE }); const queueInfoQuery = useQuery({ @@ -146,11 +151,37 @@ export default function JobListPage() { ); }); }, [tasks, debouncedSearch]); + const selectedActiveTasks = useMemo( + () => filteredTasks.filter((task) => rowSelection[task.task_id] && isActiveState(task.state)), + [filteredTasks, rowSelection], + ); + + const bulkCancelMutation = useMutation({ + mutationFn: async (taskIds: string[]) => { + const latestTasks = await listTasks(taskStatusFilter); + const selectedIds = new Set(taskIds); + const activeTaskIds = latestTasks.tasks + .filter((task) => selectedIds.has(task.task_id) && isActiveState(task.state)) + .map((task) => task.task_id); + const results = await Promise.allSettled(activeTaskIds.map((taskId) => cancelTask(taskId))); + const ok = results.filter((result) => result.status === "fulfilled").length; + return { ok, failed: results.length - ok, skipped: taskIds.length - activeTaskIds.length }; + }, + onSuccess: ({ ok, failed, skipped }) => { + if (ok) toast.success(`${ok} task(s) cancelled`); + if (failed) toast.error(`${failed} task(s) failed to cancel`); + if (skipped) toast.error(`${skipped} task(s) were no longer active`); + setRowSelection({}); + queryClient.invalidateQueries({ queryKey: ["tasks"] }); + }, + onError: (error: Error) => toast.error(`Bulk cancel failed: ${error.message}`), + }); const handleStatusTabChange = (value: string) => { setStatusTab(value); setSearch(""); setDebouncedSearch(""); + setRowSelection({}); }; return ( @@ -158,13 +189,16 @@ export default function JobListPage() { -
+
setSearch(e.target.value)} + onChange={(e) => { + setSearch(e.target.value); + setRowSelection({}); + }} className="pl-9" />
@@ -172,6 +206,23 @@ export default function JobListPage() {

{filteredTasks.length} job{filteredTasks.length === 1 ? "" : "s"}

+ {selectedActiveTasks.length > 0 && ( + <> + bulkCancelMutation.mutate(selectedActiveTasks.map((task) => task.task_id))} + > + + +

+ {selectedActiveTasks.length} selected +

+ + )} {isAdmin && ( ) : ( // Remounting resets pagination for a new tab/search context; sort state resets with it. - + isActiveState(task.state)} + getRowId={(task) => task.task_id} + rowSelection={rowSelection} + onRowSelectionChange={setRowSelection} + /> )} ))}