50 lines
1.6 KiB
Python
50 lines
1.6 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Dict
|
|
import threading
|
|
|
|
# Simple in-process cancel flags storage (per pipeline_id)
|
|
# Thread-safe for FastAPI workers in same process
|
|
_cancel_flags: Dict[str, bool] = {}
|
|
# Mode of cancellation per pipeline: "graceful" (default) or "abort"
|
|
_cancel_modes: Dict[str, str] = {}
|
|
_lock = threading.Lock()
|
|
|
|
|
|
def request_cancel(pipeline_id: str, mode: str = "graceful") -> None:
|
|
"""Set cancel flag for given pipeline id with an optional mode.
|
|
|
|
mode:
|
|
- "graceful": do not interrupt in-flight operations, stop before next step
|
|
- "abort": attempt to cancel in-flight operations immediately
|
|
"""
|
|
pid = str(pipeline_id or "pipeline_editor")
|
|
m = str(mode or "graceful").lower().strip()
|
|
if m not in {"graceful", "abort"}:
|
|
m = "graceful"
|
|
with _lock:
|
|
_cancel_flags[pid] = True
|
|
_cancel_modes[pid] = m
|
|
|
|
|
|
def clear_cancel(pipeline_id: str) -> None:
|
|
"""Clear cancel flag for given pipeline id."""
|
|
pid = str(pipeline_id or "pipeline_editor")
|
|
with _lock:
|
|
_cancel_flags.pop(pid, None)
|
|
_cancel_modes.pop(pid, None)
|
|
|
|
|
|
def is_cancelled(pipeline_id: str) -> bool:
|
|
"""Check cancel flag for given pipeline id."""
|
|
pid = str(pipeline_id or "pipeline_editor")
|
|
with _lock:
|
|
return bool(_cancel_flags.get(pid, False))
|
|
|
|
|
|
def get_cancel_mode(pipeline_id: str) -> str:
|
|
"""Return current cancel mode for given pipeline id: 'graceful' or 'abort' (default graceful)."""
|
|
pid = str(pipeline_id or "pipeline_editor")
|
|
with _lock:
|
|
m = _cancel_modes.get(pid)
|
|
return m if m in {"graceful", "abort"} else "graceful" |