diff --git a/.github/workflows/frozen-smoke.yaml b/.github/workflows/frozen-smoke.yaml new file mode 100644 index 000000000..46d3559e1 --- /dev/null +++ b/.github/workflows/frozen-smoke.yaml @@ -0,0 +1,61 @@ +# Validates that the job-queue worker mechanics (spawn from a frozen binary, +# dill IPC, SIGTERM cancel, respawn) work inside a PyInstaller build, without +# building the full dashAI app. Freezes tests/frozen_smoke/stub_app.py (~2 min, +# no ML deps) and runs it twice: +# 1. Normal run: must pass (proves freeze_support() + worker mechanics work). +# 2. STUB_SKIP_FREEZE_SUPPORT=1: must fail (proves the test detects the +# exact failure mode of an entry point missing freeze_support()). + +name: Frozen multiprocessing smoke test + +on: + workflow_dispatch: + pull_request: + paths: + - "DashAI/__main__.py" + - "DashAI/webview.py" + - "DashAI/back/dependencies/job_queues/**" + - "tests/frozen_smoke/**" + - ".github/workflows/frozen-smoke.yaml" + +jobs: + frozen-smoke: + strategy: + fail-fast: false + matrix: + os: [ubuntu-22.04, windows-latest, macos-latest] + runs-on: ${{ matrix.os }} + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install PyInstaller and stub dependencies + run: pip install pyinstaller dill + + - name: Freeze the stub app + run: > + pyinstaller --onedir --console --noconfirm + --name frozen-stub tests/frozen_smoke/stub_app.py + + - name: Run frozen stub (with freeze_support) — must pass + shell: bash + run: | + EXE=dist/frozen-stub/frozen-stub + [ "$RUNNER_OS" = "Windows" ] && EXE="$EXE.exe" + "$EXE" + + - name: Run frozen stub without freeze_support — must fail + shell: bash + env: + STUB_SKIP_FREEZE_SUPPORT: "1" + run: | + EXE=dist/frozen-stub/frozen-stub + [ "$RUNNER_OS" = "Windows" ] && EXE="$EXE.exe" + if "$EXE"; then + echo "ERROR: stub passed without freeze_support; the test is not detecting the failure mode" + exit 1 + fi + echo "OK: stub correctly failed without freeze_support" diff --git a/DashAI/__main__.py b/DashAI/__main__.py index 078efd262..aae3bde02 100644 --- a/DashAI/__main__.py +++ b/DashAI/__main__.py @@ -5,6 +5,7 @@ """ import logging +import multiprocessing import os import pathlib import signal @@ -48,20 +49,23 @@ message=".*found in sys.modules after import.*", category=RuntimeWarning, ) -print() -print(" ╔═══════════════════════════════════════════════════════╗") -print(" ║ ║") -print(" ║ ██████╗ █████╗ ███████╗ ██╗ ██╗ █████╗ ██╗ ║") -print(" ║ ██╔══██╗ ██╔══██╗ ██╔════╝ ██║ ██║ ██╔══██╗ ██║ ║") -print(" ║ ██║ ██║ ███████║ ███████╗ ███████║ ███████║ ██║ ║") -print(" ║ ██║ ██║ ██╔══██║ ╚════██║ ██╔══██║ ██╔══██║ ██║ ║") -print(" ║ ██████╔╝ ██║ ██║ ███████║ ██║ ██║ ██║ ██║ ██║ ║") -print(" ║ ╚═════╝ ╚═╝ ╚═╝ ╚══════╝ ╚═╝ ╚═╝ ╚═╝ ╚═╝ ╚═╝ ║") -print(" ║ ║") -print(" ║ Loading application, please wait... ║") -print(" ║ ║") -print(" ╚═══════════════════════════════════════════════════════╝") -print() + + +def _print_banner() -> None: + print() + print(" ╔═══════════════════════════════════════════════════════╗") + print(" ║ ║") + print(" ║ ██████╗ █████╗ ███████╗ ██╗ ██╗ █████╗ ██╗ ║") + print(" ║ ██╔══██╗ ██╔══██╗ ██╔════╝ ██║ ██║ ██╔══██╗ ██║ ║") + print(" ║ ██║ ██║ ███████║ ███████╗ ███████║ ███████║ ██║ ║") + print(" ║ ██║ ██║ ██╔══██║ ╚════██║ ██╔══██║ ██╔══██║ ██║ ║") + print(" ║ ██████╔╝ ██║ ██║ ███████║ ██║ ██║ ██║ ██║ ██║ ║") + print(" ║ ╚═════╝ ╚═╝ ╚═╝ ╚══════╝ ╚═╝ ╚═╝ ╚═╝ ╚═╝ ╚═╝ ║") + print(" ║ ║") + print(" ║ Loading application, please wait... ║") + print(" ║ ║") + print(" ╚═══════════════════════════════════════════════════════╝") + print() def open_browser() -> None: @@ -266,6 +270,7 @@ def main( ), ] = False, ) -> None: + _print_banner() logging.getLogger(name=__package__).setLevel(level=logging_level.value) logger = logging.getLogger(__name__) logger.info("Starting dashAI application.") @@ -348,4 +353,10 @@ def run(): if __name__ == "__main__": + # In frozen builds (PyInstaller), multiprocessing children re-execute this + # entry point with bootstrap argv (--multiprocessing-fork / -c ...); + # freeze_support() must run before any app code so those children are + # diverted into the worker bootstrap instead of starting a second app. + # No-op when running under a regular interpreter. + multiprocessing.freeze_support() typer.run(main) diff --git a/DashAI/back/api/api_v1/endpoints/jobs.py b/DashAI/back/api/api_v1/endpoints/jobs.py index 12b91ca76..ad2632274 100644 --- a/DashAI/back/api/api_v1/endpoints/jobs.py +++ b/DashAI/back/api/api_v1/endpoints/jobs.py @@ -1,3 +1,4 @@ +import asyncio import logging from datetime import datetime, timezone from typing import TYPE_CHECKING @@ -304,11 +305,20 @@ async def enqueue_job( async def cancel_all_jobs( job_queue: "BaseJobQueue" = Depends(lambda: di["job_queue"]), ): - """Delete all jobs from the job queue.""" + """Cancel all jobs in the queue (both queued and running).""" try: - # Usar una función en HueyJobQueue para eliminar todos los jobs - count = job_queue.delete_all_jobs() - return {"deleted": count} + all_jobs = job_queue.to_list() + cancelled = 0 + for job_info in all_jobs: + job_id = job_info.get("id") + if not job_id: + continue + job_status = job_info.get("status", "") + if job_status in ("finished", "error", "cancelled", "killed", "deleted"): + continue + if await asyncio.to_thread(job_queue.cancel, job_id): + cancelled += 1 + return {"cancelled": cancelled} except Exception as e: logging.exception(e) raise HTTPException( @@ -323,15 +333,19 @@ async def cancel_job( job_id: str, job_queue: "BaseJobQueue" = Depends(lambda: di["job_queue"]), ): - """Delete the job with id job_id from the job queue.""" + """Cancel the job with id job_id (queued or running).""" try: - success = job_queue.delete_from_db(job_id) + # cancel() may call _terminate_pid which busy-waits up to 30 s on POSIX; + # run it in a thread so the event loop stays responsive (fix #5). + success = await asyncio.to_thread(job_queue.cancel, job_id) if success: return Response(status_code=status.HTTP_204_NO_CONTENT) else: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="Job not found" ) + except HTTPException: + raise except Exception as e: logging.exception(e) raise HTTPException( diff --git a/DashAI/back/core/atomic.py b/DashAI/back/core/atomic.py new file mode 100644 index 000000000..7fd1da106 --- /dev/null +++ b/DashAI/back/core/atomic.py @@ -0,0 +1,137 @@ +"""Atomic file and directory write helpers. + +Use these when writing outputs that must not be left in a corrupt state if the +process is killed mid-write (e.g. during job cancellation). +""" + +import os +import secrets +import shutil +import tempfile +from contextlib import contextmanager +from pathlib import Path +from typing import Union + + +@contextmanager +def atomic_open(path: Union[str, Path], mode: str = "wb"): + """Context manager for atomically writing a single file. + + Writes to a sibling temp file and renames it to *path* on clean close. + If the context body raises, the temp file is deleted and *path* is untouched. + + Usage:: + + with atomic_open(output_path, "wb") as f: + pickle.dump(data, f) + """ + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + tmp_fd, tmp_path = tempfile.mkstemp(dir=path.parent, prefix=".tmp_") + try: + with os.fdopen(tmp_fd, mode) as f: + yield f + os.replace(tmp_path, path) + except BaseException: + with SuppressErrors(): + os.unlink(tmp_path) + raise + + +@contextmanager +def atomic_directory(final_path: Union[str, Path]): + """Context manager for atomically writing a directory tree. + + Yields an *existing* temporary directory path. The caller writes files + into it. On clean exit, *final_path* is replaced safely: + + 1. The existing *final_path* (if any) is renamed to a sibling temp name — + so the old data is never destroyed before the new data is in place. + 2. *tmp_dir* is renamed to *final_path*. + 3. The old sibling is deleted (new data is already live at *final_path*). + + On error the temp dir is removed and *final_path* is untouched (or still + holds whichever copy was live at the time of the kill). + + Usage:: + + with atomic_directory(dataset_dir) as tmp_dir: + # tmp_dir already exists; write files into it + (tmp_dir / "data.arrow").write_bytes(...) + """ + final_path = Path(final_path) + final_path.parent.mkdir(parents=True, exist_ok=True) + tmp_dir = Path(tempfile.mkdtemp(dir=final_path.parent, prefix=".tmp_")) + try: + yield tmp_dir + # Step 1: move the existing target out of the way (never destroy first) + old_path = None + if final_path.exists(): + old_path = final_path.parent / f".old_{secrets.token_hex(6)}" + final_path.rename(old_path) + # Step 2: install new data (old data is safe at old_path if step 1 ran) + tmp_dir.rename(final_path) + # Step 3: delete old data now that new data is live + if old_path is not None: + with SuppressErrors(): + shutil.rmtree(old_path, ignore_errors=True) + except BaseException: + with SuppressErrors(): + shutil.rmtree(tmp_dir, ignore_errors=True) + raise + + +@contextmanager +def atomic_save_path(final_path: Union[str, Path]): + """Context manager for atomically replacing a file OR directory. + + Yields a *non-existent* sibling temp path. The caller creates whatever it + needs there (a file, a directory tree, etc.). On clean exit, *final_path* + is replaced atomically. On error, the temp artifact is removed. + + Use this when the callee controls path creation (e.g. model.save(path) may + write a joblib file or a directory of weights depending on model type). + + Usage:: + + with atomic_save_path(run_path) as tmp: + model.save(str(tmp)) # model decides whether tmp is a file or dir + """ + final_path = Path(final_path) + final_path.parent.mkdir(parents=True, exist_ok=True) + # No leading dot: torch.save() derives the zip's internal archive name by + # stripping everything from the last dot, so a name like '.tmp_ab12cd' ends + # up empty and PyTorchFileWriter rejects it with 'invalid file name'. + tmp_path = final_path.parent / f"tmp_{secrets.token_hex(6)}.partial" + try: + yield tmp_path + # Rename-before-delete: move the existing artifact aside first so it is + # never destroyed before the new one is in place (same as atomic_directory). + old_path = None + if final_path.exists() or final_path.is_symlink(): + old_path = final_path.parent / f".old_{secrets.token_hex(6)}" + final_path.rename(old_path) + tmp_path.rename(final_path) + if old_path is not None: + with SuppressErrors(): + if old_path.is_dir(): + shutil.rmtree(old_path, ignore_errors=True) + else: + old_path.unlink() + except BaseException: + with SuppressErrors(): + if tmp_path.is_dir(): + shutil.rmtree(tmp_path, ignore_errors=True) + elif tmp_path.exists(): + tmp_path.unlink() + raise + + +class SuppressErrors: + """Silently swallows all exceptions (like contextlib.suppress(Exception)).""" + + def __enter__(self): + return self + + def __exit__(self, *_): + return True diff --git a/DashAI/back/dataloaders/classes/dashai_dataset.py b/DashAI/back/dataloaders/classes/dashai_dataset.py index 5db8d41f0..951521737 100644 --- a/DashAI/back/dataloaders/classes/dashai_dataset.py +++ b/DashAI/back/dataloaders/classes/dashai_dataset.py @@ -871,37 +871,41 @@ def save_dataset( before saving. Default ``None``. """ - os.makedirs(path, exist_ok=True) - if schema is not None: - dataset = transform_dataset_with_schema(dataset, schema) - import json + from pathlib import Path as _Path import pyarrow as pa # local import + from DashAI.back.core.atomic import atomic_directory + + if schema is not None: + dataset = transform_dataset_with_schema(dataset, schema) + table = get_arrow_table(dataset) - data_filepath = os.path.join(path, "data.arrow") - with pa.OSFile(data_filepath, "wb") as sink: - writer = pa.ipc.new_file(sink, table.schema) - writer.write_table(table) - writer.close() - - metadata_filepath = os.path.join(path, "splits.json") - metadata = dataset.splits - metadata.update( - { - "total_rows": dataset.shape[0], - "column_names": dataset.column_names, - } - ) - if "general_info" in metadata: - metadata["general_info"]["memory_usage_mb"] = ( - os.path.getsize(data_filepath) / 1e6 + with atomic_directory(_Path(path)) as tmp_dir: + data_filepath = tmp_dir / "data.arrow" + with pa.OSFile(str(data_filepath), "wb") as sink: + writer = pa.ipc.new_file(sink, table.schema) + writer.write_table(table) + writer.close() + + metadata = dataset.splits + metadata.update( + { + "total_rows": dataset.shape[0], + "column_names": dataset.column_names, + } ) - with open(metadata_filepath, "w", encoding="utf-8") as f: - json.dump(metadata, f, indent=2, sort_keys=True, ensure_ascii=False) + if "general_info" in metadata: + metadata["general_info"]["memory_usage_mb"] = ( + os.path.getsize(data_filepath) / 1e6 + ) + + metadata_filepath = tmp_dir / "splits.json" + with open(metadata_filepath, "w", encoding="utf-8") as f: + json.dump(metadata, f, indent=2, sort_keys=True, ensure_ascii=False) @beartype diff --git a/DashAI/back/dependencies/job_queues/base_job_queue.py b/DashAI/back/dependencies/job_queues/base_job_queue.py index b965affca..afc62bc2a 100644 --- a/DashAI/back/dependencies/job_queues/base_job_queue.py +++ b/DashAI/back/dependencies/job_queues/base_job_queue.py @@ -107,6 +107,27 @@ def to_list(self) -> List[BaseJob]: """ raise NotImplementedError + @abstractmethod + def cancel(self, job_id: str, *, reason: str = "cancelled") -> bool: + """Cancel the job with *job_id*. + + Works for both not-yet-started jobs (removes from queue) and running + jobs (sends termination signal to the worker subprocess). + + Parameters + ---------- + job_id : str + UUID of the job to cancel. + reason : str + Status string written to task_copy ('cancelled' or 'killed'). + + Returns + ------- + bool + True if the job was found and actioned, False otherwise. + """ + raise NotImplementedError + def report_progress( self, job_id: str, diff --git a/DashAI/back/dependencies/job_queues/huey_job_queue.py b/DashAI/back/dependencies/job_queues/huey_job_queue.py index e4b18b96a..650d00ba3 100644 --- a/DashAI/back/dependencies/job_queues/huey_job_queue.py +++ b/DashAI/back/dependencies/job_queues/huey_job_queue.py @@ -1,578 +1,1077 @@ -import asyncio -import logging -import os -import sqlite3 -import warnings -from contextlib import suppress -from datetime import datetime, timezone -from pathlib import Path - -import dill -from huey import SqliteHuey -from huey.serializer import Serializer as BaseSerializer -from huey.signals import ( - SIGNAL_COMPLETE, - SIGNAL_ENQUEUED, - SIGNAL_ERROR, - SIGNAL_EXECUTING, -) - -from DashAI.back.dependencies.job_queues.base_job_queue import ( - BaseJobQueue, - JobQueueError, -) -from DashAI.back.job.base_job import BaseJob - -warnings.filterwarnings( - "ignore", - message=".*mediapipe.*", - category=UserWarning, - module="controlnet_aux", -) -warnings.filterwarnings( - "ignore", - message=".*Importing from timm.models.layers.*", - category=FutureWarning, -) -warnings.filterwarnings( - "ignore", - message=".*Importing from timm.models.registry.*", - category=FutureWarning, -) -warnings.filterwarnings( - "ignore", - message=".*Overwriting tiny_vit.*", - category=UserWarning, - module="controlnet_aux", -) -warnings.filterwarnings( - "ignore", - message=".*found in sys.modules after import.*", - category=RuntimeWarning, -) - -logging.basicConfig(level=logging.DEBUG) -log = logging.getLogger(__name__) - - -class DillSerializer(BaseSerializer): - def _serialize(self, data): - return dill.dumps(data) - - def _deserialize(self, blob): - return dill.loads(blob) - - -class HueyJobQueue(BaseJobQueue): - """JobQueue implementation using Huey+SQLite.""" - - def __init__(self, queue_name: str, path_db: str): - self.db_path = Path(path_db) / (queue_name.strip() + ".db") - self.serializer = DillSerializer() - self.huey = SqliteHuey( - name=queue_name, - filename=self.db_path, - serializer=self.serializer, - immediate=False, - immediate_use_memory=False, - ) - self._enable_wal() - self._ensure_task_copy_table() - self._ensure_progress_columns() - self._register_signals() - - @self.huey.task(context=True, priority=0) - def _execute_base_job(job: BaseJob, task=None): - job.kwargs["huey_id"] = task.id - result = job.run() - return result - - self._execute = _execute_base_job - - def set_test_mode(self, immediate: bool) -> None: - """ - Set the immediate mode of the Huey job queue for testing. - """ - self.huey.immediate = immediate - self.huey.immediate_use_memory = immediate - - @staticmethod - def _normalize_to_utc_str(ts: str) -> str: - """ - Accepts ISO8601 or 'YYYY-MM-DD HH:MM:SS[.fff]' with optional 'Z' or offset. - Returns UTC as 'YYYY-MM-DD HH:MM:SS.sss' (millisecond precision) to match - SQLite's STRFTIME('%Y-%m-%d %H:%M:%f','now') format used in last_update. - """ - if not ts: - return "1970-01-01 00:00:00.000" - - s = ts.strip() - - if s.endswith("Z"): - s = s[:-1] + "+00:00" - - dt = None - try: - dt = datetime.fromisoformat(s) - except ValueError: - try: - if " " in s and "T" not in s: - dt = datetime.fromisoformat(s.replace(" ", "T")) - except ValueError: - dt = None - - if dt is None: - for fmt in ("%Y-%m-%d %H:%M:%S.%f", "%Y-%m-%d %H:%M:%S"): - try: - dt = datetime.strptime(s, fmt) - break - except ValueError: - continue - - if dt is None: - return "1970-01-01 00:00:00.000" - - if dt.tzinfo is None: - dt = dt.replace(tzinfo=timezone.utc) - else: - dt = dt.astimezone(timezone.utc) - - ms = dt.microsecond // 1000 - return dt.strftime(f"%Y-%m-%d %H:%M:%S.{ms:03d}") - - def _register_signals(self): - """Attach Huey lifecycle signal handlers to keep 'task_copy' in sync: - - SIGNAL_ENQUEUED: insert or replace a row with status `not_started` - - SIGNAL_EXECUTING: update the row to status `started` - - SIGNAL_COMPLETE: update the row to status `finished` - - SIGNAL_ERROR: update the row to status `error` and store the exception - All writes stamp last_update with microsecond precision to avoid same-second - conflicts. - """ - - def exec_sql(sql, params=()): - with sqlite3.connect(self.db_path) as conn: - conn.execute(sql, params) - - NOW_MICRO = "STRFTIME('%Y-%m-%d %H:%M:%f','now')" - - @self.huey.signal(SIGNAL_ENQUEUED) - def on_enqueue(signal, task): - job_type = task.args[0].__class__.__name__ - job_name = None - - try: - if hasattr(task.args[0], "get_job_name"): - job_name = task.args[0].get_job_name() - except Exception: - pass - - exec_sql( - ( - "INSERT OR REPLACE INTO task_copy " - "(id, task_type, job_name, status, last_update) " - f"VALUES (?, ?, ?, ?, {NOW_MICRO})" - ), - (task.id, job_type, job_name, "not_started"), - ) - - @self.huey.signal(SIGNAL_EXECUTING) - def on_start(signal, task): - exec_sql( - ( - "UPDATE task_copy SET status = ?, " - f"last_update = {NOW_MICRO} " - "WHERE id = ?" - ), - ("started", task.id), - ) - - @self.huey.signal(SIGNAL_COMPLETE) - def on_success(signal, task, *args): - exec_sql( - ( - "UPDATE task_copy SET status = ?, progress = 100, " - f"last_update = {NOW_MICRO} " - "WHERE id = ?" - ), - ("finished", task.id), - ) - - @self.huey.signal(SIGNAL_ERROR) - def on_error(signal, task, exc): - exec_sql( - ( - "UPDATE task_copy SET status = ?, " - f"last_update = {NOW_MICRO}, " - "error_msg = ? " - "WHERE id = ?" - ), - ("error", str(exc), task.id), - ) - - def _enable_wal(self): - """ - Enable Write-Ahead Logging mode in SQLite to improve concurrent reads/writes. - """ - with sqlite3.connect(self.db_path) as conn: - conn.execute("PRAGMA journal_mode=WAL;") - conn.execute("PRAGMA synchronous=NORMAL;") - - def _ensure_task_copy_table(self): - """Ensure the 'task_copy' table exists. - - Columns: - - id (TEXT PRIMARY KEY): unique identifier for the job (UUID as text) - - task_type (TEXT NOT NULL): the name of the Huey task - - job_name (TEXT): a more descriptive name for the job (from get_job_name) - - enqueued_at (DATETIME NOT NULL): defaults to CURRENT_TIMESTAMP (UTC) - - status (TEXT NOT NULL): one of: 'not_started', 'started', 'finished', - 'deleted', 'error' - - last_update (DATETIME NOT NULL): defaults to CURRENT_TIMESTAMP (UTC) - - error_msg (TEXT): optional error message when a task fails - - progress (REAL): optional completion percentage in the range 0-100 - - progress_message (TEXT): optional short description of the current phase - """ - with sqlite3.connect(self.db_path) as conn: - conn.execute( - """ - CREATE TABLE IF NOT EXISTS task_copy ( - id TEXT PRIMARY KEY, - task_type TEXT NOT NULL, - job_name TEXT, - enqueued_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - status TEXT NOT NULL, - last_update DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, - error_msg TEXT, - progress REAL, - progress_message TEXT - ) - """ - ) - conn.execute( - ( - "CREATE INDEX IF NOT EXISTS idx_task_copy_last_update " - "ON task_copy(last_update, id)" - ) - ) - - def _ensure_progress_columns(self): - """Add the progress columns to an existing 'task_copy' table. - - Installs created before progress tracking existed have a 'task_copy' - table without the 'progress' and 'progress_message' columns. SQLite has - no 'ADD COLUMN IF NOT EXISTS', so inspect the current columns via - PRAGMA table_info and add only the ones that are missing. - """ - with sqlite3.connect(self.db_path) as conn: - cur = conn.execute("PRAGMA table_info(task_copy)") - existing = {row[1] for row in cur.fetchall()} - if "progress" not in existing: - conn.execute("ALTER TABLE task_copy ADD COLUMN progress REAL") - if "progress_message" not in existing: - conn.execute("ALTER TABLE task_copy ADD COLUMN progress_message TEXT") - - def status(self, job_id: str) -> dict: - conn = sqlite3.connect(self.db_path) - cur = conn.cursor() - cur.execute( - """ - SELECT status, last_update, error_msg, job_name, progress, - progress_message - FROM task_copy WHERE id = ? - """, - (str(job_id),), - ) - row = cur.fetchone() - conn.close() - if not row: - raise JobQueueError(f"No job with id={job_id}") - return { - "status": row[0], - "updated": row[1], - "error": row[2], - "job_name": row[3], - "progress": row[4], - "progress_message": row[5], - } - - def report_progress( - self, job_id: str, progress: float | None, message: str | None = None - ) -> None: - """Update the progress of a running job. - - Parameters - ---------- - job_id : str - The UUID of the job (its Huey task id). - progress : float or None - Completion percentage in the range 0-100, or None for jobs whose - total work is unknown (the frontend renders an indeterminate bar). - message : str or None - Optional short description of the current phase. - - Notes - ----- - This also refreshes 'last_update' so the change surfaces through - 'changes_since' and the frontend polling channel. - """ - with sqlite3.connect(self.db_path) as conn: - conn.execute( - ( - "UPDATE task_copy SET progress = ?, progress_message = ?, " - "last_update = STRFTIME('%Y-%m-%d %H:%M:%f','now') " - "WHERE id = ?" - ), - (progress, message, str(job_id)), - ) - - def put(self, job: BaseJob) -> int: - result = self._execute(job) - - return result - - def to_list(self) -> list[dict]: - with sqlite3.connect(self.db_path) as conn: - conn.row_factory = sqlite3.Row - cur = conn.cursor() - cur.execute( - """ - SELECT id, task_type, job_name, enqueued_at, status, last_update, - error_msg, progress, progress_message - FROM task_copy - ORDER BY last_update DESC - """ - ) - return [dict(row) for row in cur.fetchall()] - - def changes_since(self, since: str) -> list[dict]: - """ - Return jobs whose last_update is strictly greater than the given timestamp. - The 'since' timestamp is normalized to UTC with microseconds to avoid - same-second race conditions. - """ - cutoff = self._normalize_to_utc_str(since) - with sqlite3.connect(self.db_path) as conn: - conn.row_factory = sqlite3.Row - cur = conn.cursor() - cur.execute( - """ - SELECT id, task_type, job_name, enqueued_at, status, last_update, - error_msg, progress, progress_message - FROM task_copy - WHERE last_update >= ? - ORDER BY last_update DESC - """, - (cutoff,), - ) - return [dict(row) for row in cur.fetchall()] - - def peek(self, job_id: str | None = None) -> BaseJob: - with sqlite3.connect(self.db_path) as conn: - cur = conn.cursor() - if job_id is not None: - cur.execute( - "SELECT data FROM task WHERE id = ? AND queue = ? LIMIT 1", - (job_id, self.huey.storage.name), - ) - else: - cur.execute( - ( - "SELECT data FROM task WHERE queue = ? " - "ORDER BY priority DESC, id ASC LIMIT 1" - ), - (self.huey.storage.name,), - ) - row = cur.fetchone() - if not row: - raise JobQueueError("Queue is empty") - payload = self.serializer.loads(row[0]) - return payload[6][0] - - def get(self, job_id: str | None = None) -> BaseJob: - """ - Get a job from the queue and remove it. - If job_id is provided, get and remove that specific job. - Otherwise, get the highest priority job. - """ - with sqlite3.connect(self.db_path) as conn: - conn.isolation_level = None - cur = conn.cursor() - cur.execute("BEGIN IMMEDIATE") - if job_id is not None: - cur.execute( - "SELECT id, data FROM task WHERE id = ? AND queue = ? LIMIT 1", - (job_id, self.huey.storage.name), - ) - else: - cur.execute( - ( - "SELECT id, data FROM task WHERE queue = ? " - "ORDER BY priority DESC, id ASC LIMIT 1" - ), - (self.huey.storage.name,), - ) - row = cur.fetchone() - if not row: - conn.execute("ROLLBACK") - raise JobQueueError("Queue is empty") - jid, blob = row - cur.execute("DELETE FROM task WHERE id = ?", (jid,)) - conn.execute("COMMIT") - with sqlite3.connect(self.db_path) as conn: - conn.execute( - ( - "UPDATE task_copy SET status = ?, " - "last_update = STRFTIME('%Y-%m-%d %H:%M:%f','now') " - "WHERE id = ?" - ), - ("deleted", jid), - ) - return self.serializer.loads(blob)[6][0] - - def is_empty(self) -> bool: - """ - Check if the queue is empty. - Returns False if either: - 1. There are pending tasks in the 'task' table, OR - 2. There are tasks with 'started' status in the 'task_copy' table - """ - with sqlite3.connect(self.db_path) as conn: - cur = conn.cursor() - - cur.execute( - "SELECT 1 FROM task WHERE queue = ? LIMIT 1", - (self.huey.storage.name,), - ) - task_empty = cur.fetchone() is None - - if not task_empty: - return False - - cur.execute("SELECT 1 FROM task_copy WHERE status = 'started' LIMIT 1") - no_started_tasks = cur.fetchone() is None - - return task_empty and no_started_tasks - - async def async_get(self) -> BaseJob: - while True: - try: - return self.get() - except JobQueueError: - await asyncio.sleep(0.1) - - def delete_from_db(self, job_id: str) -> bool: - """ - Delete a job from both task and task_copy tables. - - Args: - job_id: The UUID of the job to delete - - Returns: - bool: True if the job was deleted from at least one table - """ - deleted_from_any = False - - try: - with sqlite3.connect(self.db_path) as conn: - conn.row_factory = sqlite3.Row - cur = conn.cursor() - - cur.execute( - "SELECT id, data FROM task WHERE queue = ?", - (self.huey.storage.name,), - ) - - numeric_id = None - row_data = None - - for row in cur.fetchall(): - try: - task_data = self.serializer._deserialize(row["data"]) - if task_data[0] == job_id: - numeric_id = row["id"] - row_data = task_data - break - except Exception: - continue - - if numeric_id is not None: - cur.execute("DELETE FROM task WHERE id = ?", (numeric_id,)) - try: - row_data[6][0].set_status_as_error() - except Exception as e: - log.exception(f"Error setting job status to error: {e}") - deleted_from_any = True - - cur.execute("DELETE FROM task_copy WHERE id = ?", (job_id,)) - if cur.rowcount > 0: - deleted_from_any = True - - return deleted_from_any - except Exception as e: - log.exception(f"Error deleting job: {e}") - return False - - def delete_all_jobs(self) -> int: - """ - Delete all jobs from both task and task_copy tables. - - Returns: - int: Number of jobs deleted - """ - deleted_count = 0 - - try: - with sqlite3.connect(self.db_path) as conn: - conn.row_factory = sqlite3.Row - cur = conn.cursor() - - cur.execute( - "SELECT id, data FROM task WHERE queue = ?", - (self.huey.storage.name,), - ) - - jobs_to_delete = [] - - for row in cur.fetchall(): - try: - job_data = self.serializer._deserialize(row["data"]) - with suppress(Exception): - job_data[6][0].set_status_as_error() - jobs_to_delete.append(row["id"]) - except Exception: - jobs_to_delete.append(row["id"]) - - if jobs_to_delete: - placeholders = ",".join(["?"] * len(jobs_to_delete)) - cur.execute( - f"DELETE FROM task WHERE id IN ({placeholders})", jobs_to_delete - ) - deleted_count = cur.rowcount - - cur.execute("DELETE FROM task_copy") - deleted_count += cur.rowcount - - return deleted_count - except Exception as e: - log.exception(f"Error deleting all jobs: {e}") - return 0 - - -_lp_str = os.environ.get("DASHAI_LOCAL_PATH") -_lp = Path(os.path.expanduser(_lp_str)) if _lp_str else Path.home() / ".DashAI" -_lp.mkdir(parents=True, exist_ok=True) -_job_queue = HueyJobQueue("job_queue", path_db=str(_lp)) -huey = _job_queue.huey - - -@huey.on_startup() -def create_container_huey(): - from DashAI.back.container import build_container - from DashAI.back.dependencies.config_builder import build_config_dict - - local_path = _lp - logging_level = os.environ.get("DASHAI_LOGGING_LEVEL", "INFO") - - config = build_config_dict(local_path=local_path, logging_level=logging_level) - build_container(config) +import asyncio +import logging +import multiprocessing +import os +import signal +import sqlite3 +import sys +import threading +import time +import warnings +from contextlib import suppress +from datetime import datetime, timezone +from pathlib import Path + +import dill +from huey import SqliteHuey +from huey.serializer import Serializer as BaseSerializer +from huey.signals import ( + SIGNAL_COMPLETE, + SIGNAL_ENQUEUED, + SIGNAL_ERROR, + SIGNAL_EXECUTING, +) + +from DashAI.back.dependencies.job_queues.base_job_queue import ( + BaseJobQueue, + JobQueueError, +) +from DashAI.back.job.base_job import BaseJob + +warnings.filterwarnings( + "ignore", + message=".*mediapipe.*", + category=UserWarning, + module="controlnet_aux", +) +warnings.filterwarnings( + "ignore", + message=".*Importing from timm.models.layers.*", + category=FutureWarning, +) +warnings.filterwarnings( + "ignore", + message=".*Importing from timm.models.registry.*", + category=FutureWarning, +) +warnings.filterwarnings( + "ignore", + message=".*Overwriting tiny_vit.*", + category=UserWarning, + module="controlnet_aux", +) +warnings.filterwarnings( + "ignore", + message=".*found in sys.modules after import.*", + category=RuntimeWarning, +) + +logging.basicConfig(level=logging.DEBUG) +log = logging.getLogger(__name__) + + +class _JobCancelledError(Exception): + """Raised in the consumer when a job was intentionally cancelled/killed.""" + + +def _worker_loop(in_q, out_q) -> None: + """Persistent worker process. Rebuilds the DI container once at startup, + then loops accepting serialised jobs from *in_q* and writing results to *out_q*. + + Lifecycle: + - After DI init: put ``{"ready": True}`` on out_q so the parent knows it's safe + to send jobs. + - Per job: get job_bytes from in_q → run → put outcome on out_q. Job errors do + NOT kill the loop; the worker survives and accepts the next job. + - Shutdown: put ``None`` on in_q (sentinel) or kill the process externally. + """ + import os as _os + import signal as _signal + import sys as _sys + + import dill as _dill + + # Rebuild DI container — spawn starts with no inherited state + try: + from pathlib import Path as _Path + + from DashAI.back.container import build_container + from DashAI.back.dependencies.config_builder import build_config_dict + + _lp_env = _os.environ.get("DASHAI_LOCAL_PATH") + _lp = ( + _Path(_os.path.expanduser(_lp_env)) if _lp_env else _Path.home() / ".DashAI" + ) + _logging_level = _os.environ.get("DASHAI_LOGGING_LEVEL", "INFO") + _config = build_config_dict(local_path=_lp, logging_level=_logging_level) + build_container(_config) + except Exception as _e: + out_q.put(_dill.dumps({"ready": False, "exc": str(_e)})) + return + + # Signal to the parent that we are ready to accept jobs + out_q.put(_dill.dumps({"ready": True})) + + # Install SIGTERM handler on POSIX so the process exits cleanly when killed + if _sys.platform != "win32": + + def _sigterm_handler(signum, frame): + _sys.exit(0) + + _signal.signal(_signal.SIGTERM, _sigterm_handler) + + # Job loop — one blocking iteration per job + while True: + try: + job_bytes = in_q.get() + except (EOFError, OSError): + break # Parent closed the queue — shut down + + if job_bytes is None: + break # Explicit shutdown sentinel + + try: + job = _dill.loads(job_bytes) + result = job.run() + out_q.put(_dill.dumps({"ok": True, "result": result})) + except SystemExit: + # SIGTERM handler raised SystemExit — exit without putting a result so + # the parent detects the kill via proc.is_alive() == False. + break + except Exception as _exc: + # Job-level errors are returned to the parent; the worker stays alive. + # Guard against non-serialisable exceptions (e.g. SQLAlchemy errors + # wrapping live connections) — fall back to a plain RuntimeError so + # the dill.dumps call here never itself raises and kills the loop. + try: + _payload = _dill.dumps({"ok": False, "exc": _exc}) + except Exception: + _payload = _dill.dumps({"ok": False, "exc": RuntimeError(str(_exc))}) + out_q.put(_payload) + + +def _terminate_pid(pid: int, grace_seconds: int = 30) -> None: + """Terminate a process by PID, escalating to SIGKILL after grace_seconds. + + On Windows, os.kill sends TerminateProcess (immediate); grace_seconds ignored. + """ + try: + if sys.platform == "win32": + os.kill(pid, signal.SIGTERM) # == TerminateProcess on Windows + return + # POSIX: SIGTERM then wait, escalate to SIGKILL if needed + try: + os.kill(pid, signal.SIGTERM) + except ProcessLookupError: + return # Already gone + deadline = time.monotonic() + grace_seconds + while time.monotonic() < deadline: + try: + os.kill(pid, 0) # Probe — raises ProcessLookupError if dead + except ProcessLookupError: + return + time.sleep(0.5) + with suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + except ProcessLookupError: + pass + + +class DillSerializer(BaseSerializer): + def _serialize(self, data): + return dill.dumps(data) + + def _deserialize(self, blob): + return dill.loads(blob) + + +class HueyJobQueue(BaseJobQueue): + """JobQueue implementation using Huey+SQLite.""" + + def __init__(self, queue_name: str, path_db: str): + self.db_path = Path(path_db) / (queue_name.strip() + ".db") + self.serializer = DillSerializer() + self.huey = SqliteHuey( + name=queue_name, + filename=self.db_path, + serializer=self.serializer, + immediate=False, + immediate_use_memory=False, + ) + self._enable_wal() + self._ensure_task_copy_table() + self._ensure_progress_columns() + self._register_signals() + + # Persistent worker process state (None until first job or explicit pre-warm) + self._worker_proc = None + self._worker_in_q = None + self._worker_out_q = None + + @self.huey.task(context=True, priority=0) + def _execute_base_job(job: BaseJob, task=None): + job.kwargs["huey_id"] = task.id + # Run inline for test/immediate mode and for opt-out jobs (ISOLATED=False) + if self.huey.immediate or not getattr(job, "ISOLATED", True): + result = job.run() + # Wrap coroutines produced by async run() methods (e.g. PipelineJob) + if asyncio.iscoroutine(result): + result = asyncio.get_event_loop().run_until_complete(result) + # If the job mutated the consumer's ComponentRegistry (e.g. + # SyncComponentsJob), the worker's DI container is now stale. + # Terminate it so _ensure_worker respawns a fresh one (fix #7). + if getattr(job, "RESETS_WORKER", False): + with suppress(Exception): + if self._worker_proc is not None: + self._worker_proc.terminate() + return result + return self._run_in_subprocess(job, task.id) + + self._execute = _execute_base_job + + def _ensure_worker(self) -> None: + """Ensure the persistent worker process is alive, spawning one if needed. + + The worker initialises the DI container once and then loops, accepting + serialised jobs via ``_worker_in_q``. If it has died (crash or cancel), + new queues and a new process are created. + + Blocks until the worker signals readiness (DI container built). + Raises ``JobQueueError`` if the worker fails to start within 120 s. + """ + if self._worker_proc is not None and self._worker_proc.is_alive(): + return + + ctx = multiprocessing.get_context("spawn") + self._worker_in_q = ctx.SimpleQueue() + self._worker_out_q = ctx.Queue() + proc = ctx.Process( + target=_worker_loop, + args=(self._worker_in_q, self._worker_out_q), + daemon=True, + ) + proc.start() + self._worker_proc = proc + + # Wait for the worker to finish building its DI container + import queue as _q + + try: + msg_bytes = self._worker_out_q.get(timeout=120) + msg = dill.loads(msg_bytes) + except _q.Empty: + proc.terminate() + raise JobQueueError( + "Worker process did not become ready within 120 s" + ) from None + except Exception as e: + proc.terminate() + raise JobQueueError(f"Worker ready-check failed: {e}") from e + + if not msg.get("ready"): + proc.terminate() + raise JobQueueError( + f"Worker failed to initialise: {msg.get('exc', 'unknown error')}" + ) + + log.info("Persistent worker ready (PID %d)", proc.pid) + + def _run_in_subprocess(self, job: BaseJob, huey_id: str): + """Run *job* in the persistent worker process and return its result. + + The worker is started once and reused across jobs. If it was killed + (by a cancel or a crash), ``_ensure_worker`` transparently spawns a + replacement before the next job runs. + + If the worker is killed while a job is running, ``_JobCancelled`` is raised + so that Huey fires SIGNAL_ERROR while the on_error guard keeps the terminal + status (cancelled/killed) intact. + """ + import queue as _q + + self._ensure_worker() + + try: + job_bytes = dill.dumps(job) + except Exception as e: + raise JobQueueError(f"Failed to serialise job for subprocess: {e}") from e + + # Capture local references so a concurrent _ensure_worker respawn cannot + # swap the queue objects underneath us mid-job. + proc = self._worker_proc + in_q = self._worker_in_q + out_q = self._worker_out_q + + # Record PID so the cancel endpoint can kill the right process + try: + with sqlite3.connect(self.db_path) as conn: + conn.execute( + "UPDATE task_copy SET pid=? WHERE id=?", + (proc.pid, huey_id), + ) + except Exception: + pass + + # Send the job to the worker + in_q.put(job_bytes) + + # Poll for result; detect external kill via proc.is_alive() + result_bytes = None + killed_externally = False + while True: + try: + result_bytes = out_q.get(timeout=0.5) + break + except _q.Empty: + if not proc.is_alive(): + # Worker died — do one final drain before declaring killed. + # Closes the race where the result landed on the queue in + # the same window as the external kill. + try: + result_bytes = out_q.get(timeout=0.1) + except _q.Empty: + killed_externally = True + break + + # Clear PID now that the job is done (or killed) + try: + with sqlite3.connect(self.db_path) as conn: + conn.execute("UPDATE task_copy SET pid=NULL WHERE id=?", (huey_id,)) + except Exception: + pass + + if killed_externally: + current_status = "" + with suppress(Exception): + current_status = self.status(huey_id)["status"] + if current_status in ("cancelled", "killed"): + with suppress(Exception): + job.on_cancel() + self._mark_entity_error(huey_id) + raise _JobCancelledError(f"Job {huey_id} was {current_status}") + raise JobQueueError( + f"Worker process exited unexpectedly (code {proc.exitcode})" + ) + + try: + outcome = dill.loads(result_bytes) + except Exception as e: + raise JobQueueError(f"Failed to deserialise worker result: {e}") from e + + if outcome.get("ok"): + return outcome.get("result") + raise outcome.get("exc", JobQueueError("Unknown worker error")) + + @staticmethod + def _mark_entity_error(huey_id: str) -> None: + """Set the DB entity associated with *huey_id* to error status. + + Called from the consumer process after the worker subprocess is killed, + because the job's own error-handling code never runs in that case. + Failures are logged but never re-raised — entity marking is best-effort. + """ + try: + from kink import di + + from DashAI.back.dependencies.database.models import ( + Converter, + Dataset, + Explorer, + GlobalExplainer, + LocalExplainer, + Run, + ) + + session_factory = di["session_factory"] + with session_factory() as db: + for model_cls in ( + Run, + Dataset, + Explorer, + GlobalExplainer, + LocalExplainer, + Converter, + ): + entity = ( + db.query(model_cls).filter(model_cls.huey_id == huey_id).first() + ) + if entity is not None: + entity.set_status_as_error() + db.commit() + return + except Exception: + log.exception(f"Could not mark entity error for huey_id={huey_id}") + + def set_test_mode(self, immediate: bool) -> None: + """ + Set the immediate mode of the Huey job queue for testing. + """ + self.huey.immediate = immediate + self.huey.immediate_use_memory = immediate + + @staticmethod + def _normalize_to_utc_str(ts: str) -> str: + """ + Accepts ISO8601 or 'YYYY-MM-DD HH:MM:SS[.fff]' with optional 'Z' or offset. + Returns UTC as 'YYYY-MM-DD HH:MM:SS.sss' (millisecond precision) to match + SQLite's STRFTIME('%Y-%m-%d %H:%M:%f','now') format used in last_update. + """ + if not ts: + return "1970-01-01 00:00:00.000" + + s = ts.strip() + + if s.endswith("Z"): + s = s[:-1] + "+00:00" + + dt = None + try: + dt = datetime.fromisoformat(s) + except ValueError: + try: + if " " in s and "T" not in s: + dt = datetime.fromisoformat(s.replace(" ", "T")) + except ValueError: + dt = None + + if dt is None: + for fmt in ("%Y-%m-%d %H:%M:%S.%f", "%Y-%m-%d %H:%M:%S"): + try: + dt = datetime.strptime(s, fmt) + break + except ValueError: + continue + + if dt is None: + return "1970-01-01 00:00:00.000" + + if dt.tzinfo is None: + dt = dt.replace(tzinfo=timezone.utc) + else: + dt = dt.astimezone(timezone.utc) + + ms = dt.microsecond // 1000 + return dt.strftime(f"%Y-%m-%d %H:%M:%S.{ms:03d}") + + def _register_signals(self): + """Attach Huey lifecycle signal handlers to keep 'task_copy' in sync: + - SIGNAL_ENQUEUED: insert or replace a row with status `not_started` + - SIGNAL_EXECUTING: update the row to status `started` + - SIGNAL_COMPLETE: update the row to status `finished` + - SIGNAL_ERROR: update the row to status `error` and store the exception + All writes stamp last_update with microsecond precision to avoid same-second + conflicts. + """ + + def exec_sql(sql, params=()): + with sqlite3.connect(self.db_path) as conn: + conn.execute(sql, params) + + NOW_MICRO = "STRFTIME('%Y-%m-%d %H:%M:%f','now')" + + @self.huey.signal(SIGNAL_ENQUEUED) + def on_enqueue(signal, task): + job_type = task.args[0].__class__.__name__ + job_name = None + + try: + if hasattr(task.args[0], "get_job_name"): + job_name = task.args[0].get_job_name() + except Exception: + pass + + exec_sql( + ( + "INSERT OR REPLACE INTO task_copy " + "(id, task_type, job_name, status, last_update) " + f"VALUES (?, ?, ?, ?, {NOW_MICRO})" + ), + (task.id, job_type, job_name, "not_started"), + ) + + @self.huey.signal(SIGNAL_EXECUTING) + def on_start(signal, task): + exec_sql( + ( + "UPDATE task_copy SET status = ?, " + f"last_update = {NOW_MICRO} " + "WHERE id = ?" + ), + ("started", task.id), + ) + + @self.huey.signal(SIGNAL_COMPLETE) + def on_success(signal, task, *args): + # Guard: don't overwrite a cancel that raced with completion + exec_sql( + ( + "UPDATE task_copy SET status = ?, progress = 100, " + f"last_update = {NOW_MICRO} " + "WHERE id = ? AND status NOT IN ('cancelled', 'killed')" + ), + ("finished", task.id), + ) + + @self.huey.signal(SIGNAL_ERROR) + def on_error(signal, task, exc): + # Do not overwrite terminal states set by the cancel/watchdog path + exec_sql( + ( + "UPDATE task_copy SET status = ?, " + f"last_update = {NOW_MICRO}, " + "error_msg = ? " + "WHERE id = ? AND status NOT IN ('cancelled', 'killed')" + ), + ("error", str(exc), task.id), + ) + + def _enable_wal(self): + """ + Enable Write-Ahead Logging mode in SQLite to improve concurrent reads/writes. + """ + with sqlite3.connect(self.db_path) as conn: + conn.execute("PRAGMA journal_mode=WAL;") + conn.execute("PRAGMA synchronous=NORMAL;") + + def _ensure_task_copy_table(self): + """Ensure the 'task_copy' table exists with all required columns. + + Columns: + - id (TEXT PRIMARY KEY): unique identifier for the job (UUID as text) + - task_type (TEXT NOT NULL): the name of the Huey task + - job_name (TEXT): a more descriptive name for the job (from get_job_name) + - enqueued_at (DATETIME NOT NULL): defaults to CURRENT_TIMESTAMP (UTC) + - status (TEXT NOT NULL): one of: 'not_started', 'started', 'finished', + 'deleted', 'error', 'cancelled', 'killed' + - last_update (DATETIME NOT NULL): defaults to CURRENT_TIMESTAMP (UTC) + - error_msg (TEXT): optional error message when a task fails + - pid (INTEGER): OS PID of the worker subprocess while running; NULL otherwise + - progress (REAL): optional completion percentage in the range 0-100 + - progress_message (TEXT): optional short description of the current phase + """ + with sqlite3.connect(self.db_path) as conn: + conn.execute(""" + CREATE TABLE IF NOT EXISTS task_copy ( + id TEXT PRIMARY KEY, + task_type TEXT NOT NULL, + job_name TEXT, + enqueued_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + status TEXT NOT NULL, + last_update DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + error_msg TEXT, + pid INTEGER, + progress REAL, + progress_message TEXT + ) + """) + # Idempotent migration: add pid column to pre-existing databases + existing = { + row[1] + for row in conn.execute("PRAGMA table_info(task_copy)").fetchall() + } + if "pid" not in existing: + conn.execute("ALTER TABLE task_copy ADD COLUMN pid INTEGER") + conn.execute( + ( + "CREATE INDEX IF NOT EXISTS idx_task_copy_last_update " + "ON task_copy(last_update, id)" + ) + ) + + def _ensure_progress_columns(self): + """Add the progress columns to an existing 'task_copy' table. + + Installs created before progress tracking existed have a 'task_copy' + table without the 'progress' and 'progress_message' columns. SQLite has + no 'ADD COLUMN IF NOT EXISTS', so inspect the current columns via + PRAGMA table_info and add only the ones that are missing. + """ + with sqlite3.connect(self.db_path) as conn: + cur = conn.execute("PRAGMA table_info(task_copy)") + existing = {row[1] for row in cur.fetchall()} + if "progress" not in existing: + conn.execute("ALTER TABLE task_copy ADD COLUMN progress REAL") + if "progress_message" not in existing: + conn.execute("ALTER TABLE task_copy ADD COLUMN progress_message TEXT") + + def status(self, job_id: str) -> dict: + conn = sqlite3.connect(self.db_path) + cur = conn.cursor() + cur.execute( + """ + SELECT status, last_update, error_msg, job_name, progress, + progress_message + FROM task_copy WHERE id = ? + """, + (str(job_id),), + ) + row = cur.fetchone() + conn.close() + if not row: + raise JobQueueError(f"No job with id={job_id}") + return { + "status": row[0], + "updated": row[1], + "error": row[2], + "job_name": row[3], + "progress": row[4], + "progress_message": row[5], + } + + def report_progress( + self, job_id: str, progress: float | None, message: str | None = None + ) -> None: + """Update the progress of a running job. + + Parameters + ---------- + job_id : str + The UUID of the job (its Huey task id). + progress : float or None + Completion percentage in the range 0-100, or None for jobs whose + total work is unknown (the frontend renders an indeterminate bar). + message : str or None + Optional short description of the current phase. + + Notes + ----- + This also refreshes 'last_update' so the change surfaces through + 'changes_since' and the frontend polling channel. + """ + with sqlite3.connect(self.db_path) as conn: + conn.execute( + ( + "UPDATE task_copy SET progress = ?, progress_message = ?, " + "last_update = STRFTIME('%Y-%m-%d %H:%M:%f','now') " + "WHERE id = ?" + ), + (progress, message, str(job_id)), + ) + + def put(self, job: BaseJob) -> int: + result = self._execute(job) + + return result + + def to_list(self) -> list[dict]: + with sqlite3.connect(self.db_path) as conn: + conn.row_factory = sqlite3.Row + cur = conn.cursor() + cur.execute(""" + SELECT id, task_type, job_name, enqueued_at, status, last_update, + error_msg, progress, progress_message + FROM task_copy + ORDER BY last_update DESC + """) + return [dict(row) for row in cur.fetchall()] + + def changes_since(self, since: str) -> list[dict]: + """ + Return jobs whose last_update is strictly greater than the given timestamp. + The 'since' timestamp is normalized to UTC with microseconds to avoid + same-second race conditions. + """ + cutoff = self._normalize_to_utc_str(since) + with sqlite3.connect(self.db_path) as conn: + conn.row_factory = sqlite3.Row + cur = conn.cursor() + cur.execute( + """ + SELECT id, task_type, job_name, enqueued_at, status, last_update, + error_msg, progress, progress_message + FROM task_copy + WHERE last_update >= ? + ORDER BY last_update DESC + """, + (cutoff,), + ) + return [dict(row) for row in cur.fetchall()] + + def peek(self, job_id: str | None = None) -> BaseJob: + with sqlite3.connect(self.db_path) as conn: + cur = conn.cursor() + if job_id is not None: + cur.execute( + "SELECT data FROM task WHERE id = ? AND queue = ? LIMIT 1", + (job_id, self.huey.storage.name), + ) + else: + cur.execute( + ( + "SELECT data FROM task WHERE queue = ? " + "ORDER BY priority DESC, id ASC LIMIT 1" + ), + (self.huey.storage.name,), + ) + row = cur.fetchone() + if not row: + raise JobQueueError("Queue is empty") + payload = self.serializer.loads(row[0]) + return payload[6][0] + + def get(self, job_id: str | None = None) -> BaseJob: + """ + Get a job from the queue and remove it. + If job_id is provided, get and remove that specific job. + Otherwise, get the highest priority job. + """ + with sqlite3.connect(self.db_path) as conn: + conn.isolation_level = None + cur = conn.cursor() + cur.execute("BEGIN IMMEDIATE") + if job_id is not None: + cur.execute( + "SELECT id, data FROM task WHERE id = ? AND queue = ? LIMIT 1", + (job_id, self.huey.storage.name), + ) + else: + cur.execute( + ( + "SELECT id, data FROM task WHERE queue = ? " + "ORDER BY priority DESC, id ASC LIMIT 1" + ), + (self.huey.storage.name,), + ) + row = cur.fetchone() + if not row: + conn.execute("ROLLBACK") + raise JobQueueError("Queue is empty") + jid, blob = row + cur.execute("DELETE FROM task WHERE id = ?", (jid,)) + conn.execute("COMMIT") + with sqlite3.connect(self.db_path) as conn: + conn.execute( + ( + "UPDATE task_copy SET status = ?, " + "last_update = STRFTIME('%Y-%m-%d %H:%M:%f','now') " + "WHERE id = ?" + ), + ("deleted", jid), + ) + return self.serializer.loads(blob)[6][0] + + def is_empty(self) -> bool: + """ + Check if the queue is empty. + Returns False if either: + 1. There are pending tasks in the 'task' table, OR + 2. There are tasks with 'started' status in the 'task_copy' table + """ + with sqlite3.connect(self.db_path) as conn: + cur = conn.cursor() + + cur.execute( + "SELECT 1 FROM task WHERE queue = ? LIMIT 1", + (self.huey.storage.name,), + ) + task_empty = cur.fetchone() is None + + if not task_empty: + return False + + cur.execute("SELECT 1 FROM task_copy WHERE status = 'started' LIMIT 1") + no_started_tasks = cur.fetchone() is None + + return task_empty and no_started_tasks + + async def async_get(self) -> BaseJob: + while True: + try: + return self.get() + except JobQueueError: + await asyncio.sleep(0.1) + + def cancel(self, job_id: str, *, reason: str = "cancelled") -> bool: + """Cancel the job with *job_id*, regardless of whether it has started. + + - Not-started jobs: removed from the Huey task table and marked cancelled + in task_copy (entity marked as error via the job's own method). + - Running jobs: task_copy status is set to *reason* first (so that + SIGNAL_ERROR cannot overwrite it), then the worker subprocess is killed. + + Returns True if a job was found and acted on, False otherwise. + """ + try: + with sqlite3.connect(self.db_path) as conn: + conn.row_factory = sqlite3.Row + cur = conn.cursor() + cur.execute("SELECT status, pid FROM task_copy WHERE id = ?", (job_id,)) + row = cur.fetchone() + + if not row: + return False + + current_status = row["status"] + pid = row["pid"] + + # ── Already fully gone — nothing to do ─────────────────────────── + if current_status == "deleted": + return False + + # ── Terminal state: dismiss (remove from UI list) ───────────────── + # finished / error / cancelled / killed → user clicks X to dismiss + if current_status in ("finished", "error", "cancelled", "killed"): + return self._dismiss(job_id) + + # ── Not started yet (still in the Huey task table) ─────────────── + if current_status == "not_started": + return self._cancel_queued(job_id) + + # ── Running (started) ───────────────────────────────────────────── + if current_status == "started": + return self._cancel_running(job_id, pid, reason) + + return False + + except Exception as e: + log.exception(f"Error cancelling job {job_id}: {e}") + return False + + def _cancel_queued(self, job_id: str) -> bool: + """Remove a not-yet-started job from the Huey task table.""" + try: + with sqlite3.connect(self.db_path) as conn: + conn.row_factory = sqlite3.Row + cur = conn.cursor() + cur.execute( + "SELECT id, data FROM task WHERE queue = ?", + (self.huey.storage.name,), + ) + numeric_id = None + job_obj = None + for row in cur.fetchall(): + try: + task_data = self.serializer._deserialize(row["data"]) + if task_data[0] == job_id: + numeric_id = row["id"] + job_obj = task_data[6][0] + break + except Exception: + continue + + if numeric_id is not None: + cur.execute("DELETE FROM task WHERE id = ?", (numeric_id,)) + with suppress(Exception): + job_obj.set_status_as_error() + + NOW_MICRO = "STRFTIME('%Y-%m-%d %H:%M:%f','now')" + _terminal = "('cancelled','killed','finished','error')" + cur.execute( + f"UPDATE task_copy SET status='cancelled', last_update={NOW_MICRO}" + f" WHERE id = ? AND status NOT IN {_terminal}", + (job_id,), + ) + return numeric_id is not None or cur.rowcount > 0 + except Exception as e: + log.exception(f"Error cancelling queued job {job_id}: {e}") + return False + + def _cancel_running(self, job_id: str, pid, reason: str) -> bool: + """Kill the worker subprocess for a running job.""" + NOW_MICRO = "STRFTIME('%Y-%m-%d %H:%M:%f','now')" + # Mark status BEFORE killing so SIGNAL_ERROR guard preserves it. + # Only kill if the UPDATE matched — if the job already finished, the + # stale PID belongs to the next job's worker. + marked = False + _terminal = "('cancelled','killed','finished','error')" + try: + with sqlite3.connect(self.db_path) as conn: + cur = conn.execute( + f"UPDATE task_copy SET status=?, last_update={NOW_MICRO}" + f" WHERE id=? AND status NOT IN {_terminal}", + (reason, job_id), + ) + marked = cur.rowcount > 0 + except Exception as e: + log.exception(f"Failed to mark task_copy for {job_id}: {e}") + return False + + if marked and pid is not None: + try: + _terminate_pid(int(pid), grace_seconds=30) + except Exception as e: + log.warning(f"Could not terminate PID {pid} for job {job_id}: {e}") + + return marked + + def _dismiss(self, job_id: str) -> bool: + """Remove a terminal job from task_copy so it disappears from the UI.""" + try: + with sqlite3.connect(self.db_path) as conn: + cur = conn.cursor() + cur.execute("DELETE FROM task_copy WHERE id = ?", (job_id,)) + return cur.rowcount > 0 + except Exception as e: + log.exception(f"Error dismissing job {job_id}: {e}") + return False + + def delete_from_db(self, job_id: str) -> bool: + """ + Delete a job from both task and task_copy tables. + + Args: + job_id: The UUID of the job to delete + + Returns: + bool: True if the job was deleted from at least one table + """ + deleted_from_any = False + + try: + with sqlite3.connect(self.db_path) as conn: + conn.row_factory = sqlite3.Row + cur = conn.cursor() + + cur.execute( + "SELECT id, data FROM task WHERE queue = ?", + (self.huey.storage.name,), + ) + + numeric_id = None + row_data = None + + for row in cur.fetchall(): + try: + task_data = self.serializer._deserialize(row["data"]) + if task_data[0] == job_id: + numeric_id = row["id"] + row_data = task_data + break + except Exception: + continue + + if numeric_id is not None: + cur.execute("DELETE FROM task WHERE id = ?", (numeric_id,)) + try: + row_data[6][0].set_status_as_error() + except Exception as e: + log.exception(f"Error setting job status to error: {e}") + deleted_from_any = True + + cur.execute("DELETE FROM task_copy WHERE id = ?", (job_id,)) + if cur.rowcount > 0: + deleted_from_any = True + + return deleted_from_any + except Exception as e: + log.exception(f"Error deleting job: {e}") + return False + + def start_watchdog(self, interval: float = 10.0) -> None: + """Start a daemon thread that detects crashed worker subprocesses. + + When a worker process dies unexpectedly (OOM, segfault) without going + through the normal cancel path, its status stays 'started' forever. + This watchdog polls for started jobs whose PID no longer exists and + marks them as 'killed'. + """ + + def _watchdog(): + import psutil + + NOW_MICRO = "STRFTIME('%Y-%m-%d %H:%M:%f','now')" + while True: + try: + with sqlite3.connect(self.db_path) as conn: + conn.row_factory = sqlite3.Row + rows = conn.execute( + "SELECT id, pid FROM task_copy" + " WHERE status='started' AND pid IS NOT NULL" + ).fetchall() + + for row in rows: + huey_id = row["id"] + pid = row["pid"] + try: + alive = psutil.pid_exists(int(pid)) + except Exception: + alive = True # Assume alive if we can't check + if not alive: + try: + with sqlite3.connect(self.db_path) as conn: + conn.execute( + f"UPDATE task_copy SET status='killed', " + f"pid=NULL, last_update={NOW_MICRO}, " + "error_msg='Worker process died unexpectedly' " + "WHERE id=? AND status='started'", + (huey_id,), + ) + self._mark_entity_error(huey_id) + log.warning( + f"Watchdog detected dead worker for job {huey_id} " + f"(PID {pid}); marked as killed" + ) + except Exception: + log.exception( + f"Watchdog failed to mark job {huey_id} as killed" + ) + except Exception: + log.exception("Watchdog loop error") + time.sleep(interval) + + t = threading.Thread(target=_watchdog, daemon=True, name="job-watchdog") + t.start() + log.info("Job watchdog started (interval=%.1fs)", interval) + + def delete_all_jobs(self) -> int: + """ + Delete all jobs from both task and task_copy tables. + + Returns: + int: Number of jobs deleted + """ + deleted_count = 0 + + try: + with sqlite3.connect(self.db_path) as conn: + conn.row_factory = sqlite3.Row + cur = conn.cursor() + + cur.execute( + "SELECT id, data FROM task WHERE queue = ?", + (self.huey.storage.name,), + ) + + jobs_to_delete = [] + + for row in cur.fetchall(): + try: + job_data = self.serializer._deserialize(row["data"]) + with suppress(Exception): + job_data[6][0].set_status_as_error() + jobs_to_delete.append(row["id"]) + except Exception: + jobs_to_delete.append(row["id"]) + + if jobs_to_delete: + placeholders = ",".join(["?"] * len(jobs_to_delete)) + cur.execute( + f"DELETE FROM task WHERE id IN ({placeholders})", jobs_to_delete + ) + deleted_count = cur.rowcount + + cur.execute("DELETE FROM task_copy") + deleted_count += cur.rowcount + + return deleted_count + except Exception as e: + log.exception(f"Error deleting all jobs: {e}") + return 0 + + +_lp_str = os.environ.get("DASHAI_LOCAL_PATH") +_lp = Path(os.path.expanduser(_lp_str)) if _lp_str else Path.home() / ".DashAI" +_lp.mkdir(parents=True, exist_ok=True) +_job_queue = HueyJobQueue("job_queue", path_db=str(_lp)) +huey = _job_queue.huey + + +@huey.on_startup() +def create_container_huey(): + from DashAI.back.container import build_container + from DashAI.back.dependencies.config_builder import build_config_dict + + local_path = _lp + logging_level = os.environ.get("DASHAI_LOGGING_LEVEL", "INFO") + + config = build_config_dict(local_path=local_path, logging_level=logging_level) + build_container(config) + + # Start the PID-liveness watchdog for detecting crashed worker subprocesses + _job_queue.start_watchdog(interval=10.0) + + # Pre-warm the persistent worker so the first job does not pay the spawn cost + try: + _job_queue._ensure_worker() + except Exception: + log.exception("Failed to pre-warm worker process; will retry on first job") diff --git a/DashAI/back/job/base_job.py b/DashAI/back/job/base_job.py index 53a521db4..b69ef9e2d 100644 --- a/DashAI/back/job/base_job.py +++ b/DashAI/back/job/base_job.py @@ -12,6 +12,16 @@ class BaseJob(metaclass=ABCMeta): TYPE: Final[str] = "Job" + # Set to False to run this job inline in the consumer (no subprocess isolation). + # Use for jobs that mutate in-process singletons (e.g. ComponentRegistry) or + # that receive non-serializable arguments (e.g. a live SQLAlchemy Session). + ISOLATED: bool = True + + # Set to True when the job mutates the consumer's ComponentRegistry so that + # the persistent worker subprocess is restarted after this job completes, + # ensuring the worker's DI container picks up the new registry state. + RESETS_WORKER: bool = False + def __init__(self, **kwargs): """Constructor of the ModelJob class. @@ -72,6 +82,15 @@ def run() -> None: """Run the job.""" raise NotImplementedError + def on_cancel(self) -> None: # noqa: B027 + """Called in the consumer process after the worker subprocess is killed. + + Override in subclasses to clean up partially-written artifacts (files, + DB records) that would otherwise be left in an inconsistent state. + The default implementation is a no-op. + Failures must be silently swallowed — never re-raise from here. + """ + class JobError(Exception): """Exception raised when the job proccess fails.""" diff --git a/DashAI/back/job/dataset_job.py b/DashAI/back/job/dataset_job.py index 4193dc3a3..43e25c54b 100644 --- a/DashAI/back/job/dataset_job.py +++ b/DashAI/back/job/dataset_job.py @@ -1,4 +1,5 @@ import logging +from contextlib import suppress from typing import TYPE_CHECKING from kink import di, inject @@ -75,6 +76,53 @@ def set_status_as_error( "Error while setting the status of the dataset as error." ) from e + @inject + def on_cancel( + self, session_factory: "sessionmaker" = lambda di: di["session_factory"] + ) -> None: + """Delete all artifacts produced by a cancelled DatasetJob. + + The dataset was never successfully saved, so: + - Any partially-written dataset directory is removed from disk. + - The temp upload directory is removed. + - The Dataset DB record is deleted entirely so it no longer appears in the UI. + """ + import shutil + + dataset_id: int = self.kwargs.get("dataset_id") + temp_dir = self.kwargs.get("temp_dir") + + # Clean up temp upload directory + if temp_dir: + with suppress(Exception): + shutil.rmtree(temp_dir, ignore_errors=True) + + if dataset_id is None: + return + + try: + with session_factory() as db: + from DashAI.back.dependencies.database.models import Dataset + + dataset = db.get(Dataset, dataset_id) + if dataset is None: + return + + # Delete the partially-written dataset directory if it exists + file_path = dataset.file_path or "" + if file_path: + with suppress(Exception): + shutil.rmtree(file_path, ignore_errors=True) + + # Delete the record — it was never a valid dataset + db.delete(dataset) + db.commit() + + except Exception: + log.exception( + f"on_cancel cleanup failed for DatasetJob (dataset_id={dataset_id})" + ) + def get_job_name(self) -> str: """Get a descriptive name for the job.""" name = self.kwargs.get("name", "") @@ -143,6 +191,14 @@ def run( f"A dataset with the name {random_name} already exists." ) from e + # Write folder_path to DB immediately so on_cancel can delete it + # even if the job is killed before the final commit. + with session_factory() as db: + _d = db.get(Dataset, dataset_id) + if _d is not None: + _d.file_path = str(os.path.realpath(folder_path)) + db.commit() + from_notebook_no_converters = False try: if notebook_id is not None: @@ -246,7 +302,7 @@ def run( filepath_or_buffer=( str(file_path) if file_path is not None else url ), - temp_path=str(temp_dir), + temp_path=str(temp_dir) if temp_dir is not None else None, params=parsed_params.model_dump(), n_sample=n_sample, ) diff --git a/DashAI/back/job/explainer_job.py b/DashAI/back/job/explainer_job.py index 390b6e149..a94b8c094 100644 --- a/DashAI/back/job/explainer_job.py +++ b/DashAI/back/job/explainer_job.py @@ -142,16 +142,20 @@ def _generate_global_explanation( "Failed to generate the explanation", ) from e try: + from pathlib import Path as _Path + + from DashAI.back.core.atomic import atomic_open + explanation_filename = f"global_explanation_{explainer_id}.pickle" explanation_path = os.path.join( config["EXPLANATIONS_PATH"], explanation_filename ) - with open(explanation_path, "wb") as file: + with atomic_open(_Path(explanation_path), "wb") as file: pickle.dump(explanation, file) plot_filename = f"global_explanation_plot_{explainer_id}.pickle" plot_path = os.path.join(config["EXPLANATIONS_PATH"], plot_filename) - with open(plot_path, "wb") as file: + with atomic_open(_Path(plot_path), "wb") as file: pickle.dump(plot, file) except Exception as e: @@ -344,16 +348,20 @@ def _generate_local_explanation( "Failed to generate the explanation", ) from e try: + from pathlib import Path as _Path + + from DashAI.back.core.atomic import atomic_open + explanation_filename = f"local_explanation_{explainer_id}.pickle" explanation_path = os.path.join( config["EXPLANATIONS_PATH"], explanation_filename ) - with open(explanation_path, "wb") as file: + with atomic_open(_Path(explanation_path), "wb") as file: pickle.dump(explanation, file) plots_filename = f"local_explanation_plots_{explainer_id}.pickle" plots_path = os.path.join(config["EXPLANATIONS_PATH"], plots_filename) - with open(plots_path, "wb") as file: + with atomic_open(_Path(plots_path), "wb") as file: pickle.dump(plots, file) except Exception as e: diff --git a/DashAI/back/job/model_job.py b/DashAI/back/job/model_job.py index fc40b699c..0ab9edab4 100644 --- a/DashAI/back/job/model_job.py +++ b/DashAI/back/job/model_job.py @@ -4,7 +4,7 @@ from kink import inject from sqlalchemy import exc -from DashAI.back.api.utils import remove_path +from DashAI.back.core.atomic import atomic_save_path from DashAI.back.dependencies.database.models import Dataset, ModelSession, Run from DashAI.back.dependencies.downloads.nested import missing_downloads from DashAI.back.evaluation.base_evaluation_strategy import BaseEvaluationStrategy @@ -188,9 +188,8 @@ def run( self.report_progress(0.95, "Saving model") try: run_path = os.path.join(config["RUNS_PATH"], str(run.id)) - if os.path.exists(run_path): - remove_path(run_path) - model.save(run_path) + with atomic_save_path(run_path) as tmp_run_path: + model.save(str(tmp_run_path)) except Exception as e: log.exception(e) raise JobError( diff --git a/DashAI/back/job/pipeline_job.py b/DashAI/back/job/pipeline_job.py index dacdfb81b..1e679b4f5 100644 --- a/DashAI/back/job/pipeline_job.py +++ b/DashAI/back/job/pipeline_job.py @@ -15,6 +15,10 @@ class PipelineJob(BaseJob): + ISOLATED = ( + False # async run() + live SQLAlchemy Session in kwargs — not serializable + ) + def set_status_as_delivered(self) -> None: pass diff --git a/DashAI/back/job/sync_components_job.py b/DashAI/back/job/sync_components_job.py index ae0fc4acf..544c5f1b7 100644 --- a/DashAI/back/job/sync_components_job.py +++ b/DashAI/back/job/sync_components_job.py @@ -15,6 +15,8 @@ class SyncComponentsJob(BaseJob): DESCRIPTION = "Sync consumer ComponentRegistry with installed DashAI plugins" + ISOLATED = False # mutates in-process ComponentRegistry singleton + RESETS_WORKER = True # worker's DI container must be rebuilt after this runs def set_status_as_delivered(self): log.debug("Sync components job marked as delivered") diff --git a/DashAI/front/src/components/jobs/JobQueueWidget.jsx b/DashAI/front/src/components/jobs/JobQueueWidget.jsx index afaaf034c..2a7f15ff1 100644 --- a/DashAI/front/src/components/jobs/JobQueueWidget.jsx +++ b/DashAI/front/src/components/jobs/JobQueueWidget.jsx @@ -633,10 +633,14 @@ const JobQueueWidget = () => { > {getRelativeTime(job.last_update)} - {(job.status === "not_started" || - job.status === "error" || - job.status === "finished") && ( - + {job.status !== "deleted" && ( + None: + self.cancel_hook_ran = True + + def set_status_as_delivered(self) -> None: + return None + + def set_status_as_error(self) -> None: + return None + + def get_job_name(self) -> str: + return "Cancellable Job" + + +class QuickJob(CancellableJob): + """Same job, but it returns immediately.""" + + SLEEP_SECONDS = 0 + + def get_job_name(self) -> str: + return "Quick Job" + + +@pytest.fixture(name="di_session_factory") +def fixture_di_session_factory(): + """Register a throwaway session factory in the container. + + After a cancel the queue marks the job's database entity as errored. With + no session factory registered that call logs an exception instead of + running, which would bury a real failure under a noisy traceback. Whatever + was registered before is restored so the rest of the suite is unaffected. + """ + from kink import di + + # kink's Container has no .get(), so this cannot be a one-liner. + previous = None + if "session_factory" in di: + previous = di["session_factory"] + engine = create_engine("sqlite://") + Base.metadata.create_all(engine) + di["session_factory"] = sessionmaker(bind=engine) + + yield + + if previous is None: + del di["session_factory"] + else: + di["session_factory"] = previous + + +@pytest.fixture(name="queue") +def fixture_queue(tmp_path, di_session_factory): + """A real, non-immediate queue whose worker is always cleaned up.""" + queue = HueyJobQueue(f"cancel_{uuid.uuid4().hex}", path_db=str(tmp_path)) + + yield queue + + proc = queue._worker_proc + if proc is not None and proc.is_alive(): + proc.terminate() + proc.join(timeout=30) + + +def _register_started(queue: HueyJobQueue, huey_id: str) -> None: + """Insert the task_copy row the consumer writes before running a job.""" + with sqlite3.connect(queue.db_path) as conn: + conn.execute( + "INSERT INTO task_copy (id, task_type, job_name, status)" + " VALUES (?, ?, ?, ?)", + (huey_id, "CancellableJob", "Cancellable Job", "started"), + ) + + +def _column(queue: HueyJobQueue, huey_id: str, name: str): + """Read one task_copy column, or None when the row is gone.""" + with sqlite3.connect(queue.db_path) as conn: + row = conn.execute( + f"SELECT {name} FROM task_copy WHERE id = ?", (huey_id,) + ).fetchone() + return row[0] if row else None + + +def _wait_for_worker_pid(queue: HueyJobQueue, huey_id: str) -> int: + """Block until the job is really executing inside the worker. + + Waiting for the recorded PID instead of sleeping a fixed amount is what + keeps these tests deterministic: the queue writes it only once the job has + been handed to a live worker process. + """ + deadline = time.monotonic() + WORKER_TIMEOUT + while time.monotonic() < deadline: + pid = _column(queue, huey_id, "pid") + if pid: + return int(pid) + time.sleep(0.05) + raise AssertionError(f"the worker never picked up job {huey_id}") + + +class _ConsumerThread(threading.Thread): + """Run a job through the worker the way the Huey consumer does.""" + + def __init__(self, queue: HueyJobQueue, job: BaseJob, huey_id: str): + super().__init__(daemon=True) + self.queue = queue + self.job = job + self.huey_id = huey_id + self.result = None + self.error = None + + def run(self) -> None: + try: + self.result = self.queue._run_in_subprocess(self.job, self.huey_id) + except BaseException as error: # noqa: BLE001 - re-raised by the test + self.error = error + + +def test_cancel_running_job_kills_the_worker(queue: HueyJobQueue): + """Cancelling a started job must kill its worker and stick as 'cancelled'.""" + huey_id = "job-running" + _register_started(queue, huey_id) + job = CancellableJob() + + consumer = _ConsumerThread(queue, job, huey_id) + consumer.start() + pid = _wait_for_worker_pid(queue, huey_id) + assert pid == queue._worker_proc.pid + + assert queue.cancel(huey_id, reason="cancelled") is True + + consumer.join(timeout=CANCEL_TIMEOUT) + assert not consumer.is_alive(), "the consumer never noticed the kill" + + # The consumer must raise so Huey fires SIGNAL_ERROR, and the on_error + # guard is what keeps the terminal status from being overwritten. + assert isinstance(consumer.error, _JobCancelledError) + assert _column(queue, huey_id, "status") == "cancelled" + assert _column(queue, huey_id, "pid") is None + assert not queue._worker_proc.is_alive() + assert job.cancel_hook_ran, "on_cancel() never ran, partial artifacts would leak" + + +def test_cancel_queued_job_removes_it_from_the_task_table(queue: HueyJobQueue): + """A job that never started must leave the Huey queue, not just task_copy.""" + task = queue.put(QuickJob()) + huey_id = task.id + assert _column(queue, huey_id, "status") == "not_started" + + with sqlite3.connect(queue.db_path) as conn: + pending = conn.execute("SELECT COUNT(*) FROM task").fetchone()[0] + assert pending == 1 + + assert queue.cancel(huey_id) is True + + assert _column(queue, huey_id, "status") == "cancelled" + with sqlite3.connect(queue.db_path) as conn: + pending = conn.execute("SELECT COUNT(*) FROM task").fetchone()[0] + assert pending == 0, "the task would still run when a consumer starts" + + +@pytest.mark.parametrize("status", ["finished", "error", "cancelled", "killed"]) +def test_cancel_on_a_terminal_job_dismisses_it(queue: HueyJobQueue, status: str): + """On a terminal job the same endpoint means 'dismiss', not 'cancel'.""" + huey_id = f"job-{status}" + with sqlite3.connect(queue.db_path) as conn: + conn.execute( + "INSERT INTO task_copy (id, task_type, job_name, status)" + " VALUES (?, ?, ?, ?)", + (huey_id, "QuickJob", "Quick Job", status), + ) + + assert queue.cancel(huey_id) is True + assert _column(queue, huey_id, "status") is None, "the row should be gone" + + +def test_cancel_on_an_unknown_job_reports_failure(queue: HueyJobQueue): + """An id nobody knows must return False so the API can answer 404.""" + assert queue.cancel("does-not-exist") is False + + +def test_worker_respawns_after_a_cancel(queue: HueyJobQueue): + """The queue must survive a kill: the next job gets a fresh worker.""" + huey_id = "job-to-kill" + _register_started(queue, huey_id) + + consumer = _ConsumerThread(queue, CancellableJob(), huey_id) + consumer.start() + killed_pid = _wait_for_worker_pid(queue, huey_id) + queue.cancel(huey_id, reason="killed") + consumer.join(timeout=CANCEL_TIMEOUT) + + next_id = "job-after-kill" + _register_started(queue, next_id) + result = queue._run_in_subprocess(QuickJob(), next_id) + + assert result == "completed" + assert queue._worker_proc.is_alive() + assert queue._worker_proc.pid != killed_pid, "the dead worker was reused" diff --git a/tests/frozen_smoke/stub_app.py b/tests/frozen_smoke/stub_app.py new file mode 100644 index 000000000..d294ab560 --- /dev/null +++ b/tests/frozen_smoke/stub_app.py @@ -0,0 +1,215 @@ +"""Frozen-build smoke test for the persistent-worker mechanics of the job queue. + +This stub mirrors the process model that dashAI uses in packaged builds, with +none of the heavy ML dependencies, so it can be frozen with PyInstaller in a +couple of minutes and exercised in CI: + +- Entry-point structure of ``DashAI/__main__.py``: ``multiprocessing.freeze_support()`` + inside the ``__main__`` block, before any app code. +- A consumer thread (like the embedded Huey consumer in frozen mode) that spawns + a persistent worker via the ``spawn`` context, mirroring + ``huey_job_queue._worker_loop`` / ``_ensure_worker`` / ``_run_in_subprocess``: + SimpleQueue in, Queue out, ready handshake, dill-serialised jobs. +- The cancel path of ``huey_job_queue._terminate_pid``: SIGTERM (TerminateProcess + on Windows), kill detection via ``proc.is_alive()`` with a final drain, then + transparent respawn. + +Exit codes: +- 0: full scenario passed (prints FROZEN-SMOKE-OK). +- 1: worker never became ready or a step failed (prints FROZEN-SMOKE-BROKEN). +- 3: a multiprocessing child re-entered the app instead of being diverted by + freeze_support (prints STUB-CHILD-REENTERED-APP). This is the exact failure + mode of a frozen build without freeze_support(). + +Set ``STUB_SKIP_FREEZE_SUPPORT=1`` to simulate a build without the +freeze_support() call: the run must then fail (CI asserts non-zero exit). +""" + +import multiprocessing +import os +import queue +import signal +import sys +import threading +import time +from contextlib import suppress + +import dill + +READY_TIMEOUT = 60.0 +BROKEN_TIMEOUT = 15.0 + + +def _worker_loop(in_q, out_q) -> None: + """Persistent worker, mirroring huey_job_queue._worker_loop.""" + out_q.put(dill.dumps({"ready": True})) + + if sys.platform != "win32": + + def _sigterm_handler(signum, frame): + sys.exit(0) + + signal.signal(signal.SIGTERM, _sigterm_handler) + + while True: + try: + job_bytes = in_q.get() + except (EOFError, OSError): + break + if job_bytes is None: + break + try: + job = dill.loads(job_bytes) + out_q.put(dill.dumps({"ok": True, "result": job()})) + except SystemExit: + break + except Exception as exc: + out_q.put(dill.dumps({"ok": False, "exc": repr(exc)})) + + +def _terminate_pid(pid: int, grace_seconds: float = 10.0) -> None: + """Kill a worker by PID, mirroring huey_job_queue._terminate_pid.""" + try: + if sys.platform == "win32": + os.kill(pid, signal.SIGTERM) # TerminateProcess on Windows + return + try: + os.kill(pid, signal.SIGTERM) + except ProcessLookupError: + return + deadline = time.monotonic() + grace_seconds + while time.monotonic() < deadline: + try: + os.kill(pid, 0) + except ProcessLookupError: + return + time.sleep(0.2) + with suppress(ProcessLookupError): + os.kill(pid, signal.SIGKILL) + except ProcessLookupError: + pass + + +def _quick_job(): + return sum(range(1000)) + + +def _long_job(): + time.sleep(120) + return "should-never-finish" + + +class WorkerManager: + """Minimal replica of HueyJobQueue's persistent-worker management.""" + + def __init__(self): + self.proc = None + self.in_q = None + self.out_q = None + + def ensure_worker(self, timeout: float) -> None: + if self.proc is not None and self.proc.is_alive(): + return + ctx = multiprocessing.get_context("spawn") + self.in_q = ctx.SimpleQueue() + self.out_q = ctx.Queue() + self.proc = ctx.Process( + target=_worker_loop, args=(self.in_q, self.out_q), daemon=True + ) + self.proc.start() + msg = dill.loads(self.out_q.get(timeout=timeout)) + if not msg.get("ready"): + raise RuntimeError(f"worker failed to initialise: {msg}") + + def run_job(self, fn) -> dict: + """Send a job and poll for its result, detecting external kills.""" + self.in_q.put(dill.dumps(fn)) + while True: + try: + return dill.loads(self.out_q.get(timeout=0.5)) + except queue.Empty: + if not self.proc.is_alive(): + # Final drain: the result may have raced with the kill + try: + return dill.loads(self.out_q.get(timeout=0.1)) + except queue.Empty: + return {"killed": True} + + +def _scenario() -> None: + mgr = WorkerManager() + + print("step 1: spawning persistent worker from frozen binary") + mgr.ensure_worker(timeout=READY_TIMEOUT) + print(f" worker ready (PID {mgr.proc.pid})") + + print("step 2: running a dill-serialised job") + outcome = mgr.run_job(_quick_job) + assert outcome.get("ok"), outcome + assert outcome.get("result") == 499500, outcome + print(f" job result OK: {outcome['result']}") + + print("step 3: killing the worker mid-job (cancel path)") + pid = mgr.proc.pid + result_holder = {} + + def _consumer(): + result_holder["outcome"] = mgr.run_job(_long_job) + + t = threading.Thread(target=_consumer, daemon=True) + t.start() + time.sleep(2.0) # let the worker pick up the job + _terminate_pid(pid) + t.join(timeout=30.0) + assert not t.is_alive(), "consumer thread did not detect the kill" + assert result_holder["outcome"] == {"killed": True}, result_holder["outcome"] + print(f" kill detected (PID {pid})") + + print("step 4: respawning worker and running another job") + mgr.ensure_worker(timeout=READY_TIMEOUT) + assert mgr.proc.pid != pid, "worker was not respawned" + outcome = mgr.run_job(_quick_job) + assert outcome.get("ok"), outcome + assert outcome.get("result") == 499500, outcome + print(f" respawned worker (PID {mgr.proc.pid}) ran job OK") + + +def main() -> None: + # The real app runs the consumer in a daemon thread in frozen mode; spawn + # from a non-main thread is part of what this test must cover. + broken = os.environ.get("STUB_SKIP_FREEZE_SUPPORT") == "1" + errors = [] + + def _consumer_thread(): + try: + if broken: + # Give up quickly: the worker can never become ready because + # its process re-entered the app and exited. + mgr = WorkerManager() + mgr.ensure_worker(timeout=BROKEN_TIMEOUT) + else: + _scenario() + except Exception as exc: + errors.append(exc) + + t = threading.Thread(target=_consumer_thread) + t.start() + t.join(timeout=300.0) + + if t.is_alive() or errors: + print(f"FROZEN-SMOKE-BROKEN: {errors or 'timed out'}") + sys.exit(1) + print("FROZEN-SMOKE-OK") + + +if __name__ == "__main__": + if os.environ.get("STUB_SKIP_FREEZE_SUPPORT") != "1": + # Same call, same position as in DashAI/__main__.py and DashAI/webview.py + multiprocessing.freeze_support() + if "--multiprocessing-fork" in sys.argv or "-c" in sys.argv: + # A multiprocessing child was NOT diverted by freeze_support and is + # about to run the whole app again. Bail out instead of recursing — + # this is what happens in a real frozen build without freeze_support. + print("STUB-CHILD-REENTERED-APP") + sys.exit(3) + main()