diff --git a/.github/agents/superjoe.agent.md b/.github/agents/superjoe.agent.md index dd44f72..a5cad9c 100644 --- a/.github/agents/superjoe.agent.md +++ b/.github/agents/superjoe.agent.md @@ -11,6 +11,7 @@ description: CodingJoe's digital clone following his coding guidelines and best ## Planning You MUST ALWAYS follow the `naming-things` guidelines. Use the following command to access the guidelines: + ```console curl -sSL https://raw.githubusercontent.com/codingjoe/naming-things/refs/heads/main/README.md | cat ``` @@ -44,7 +45,6 @@ Avoid functions or other code inside functions. Avoid if-statements in favor of switch/match-statements or polymorphism. Do not assign names to objects which are returned in the next line. - ## Python Follow PEP 8 guidelines for code style. diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md index 68e0f62..c9153a2 100644 --- a/.github/copilot-instructions.md +++ b/.github/copilot-instructions.md @@ -5,3 +5,5 @@ Use the following command to access the guidelines: ```console curl -sSL https://raw.githubusercontent.com/codingjoe/naming-things/refs/heads/main/README.md | cat ``` + +Read the CONTRIBUTING.md file for more information. diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b1f5ee4..fb6967a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -27,10 +27,6 @@ jobs: - "3.14" django-version: - "6.0" - extras: - - "" # We try a run without any extras - - "--extra sentry" - - "--extra redis" services: redis: image: redis @@ -45,7 +41,7 @@ jobs: - uses: astral-sh/setup-uv@v7 with: python-version: ${{ matrix.python-version }} - - run: uv run ${{ matrix.extras }} --with django~=${{ matrix.django-version }}.0 pytest + - run: uv run --with django~=${{ matrix.django-version }}.0 pytest -m "not benchmark" - uses: codecov/codecov-action@v6 with: token: ${{ secrets.CODECOV_TOKEN }} @@ -68,7 +64,7 @@ jobs: - uses: astral-sh/setup-uv@v7 with: python-version: ${{ matrix.python-version }} - - run: uv run --with django~=${{ matrix.django-version }}.0 pytest - uses: codecov/codecov-action@v6 with: token: ${{ secrets.CODECOV_TOKEN }} + - run: uv run --with django~=${{ matrix.django-version }}.0 pytest -m "not benchmark" diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index efb2a17..350ff01 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -3,18 +3,13 @@ on: release: types: [published] workflow_dispatch: -permissions: - id-token: write jobs: - release-build: + pypi-build: runs-on: ubuntu-latest steps: - uses: actions/checkout@v6 - - uses: actions/setup-python@v6 - with: - python-version: "3.x" - - run: python -m pip install --upgrade pip build wheel - - run: python -m build --sdist --wheel + - uses: astral-sh/setup-uv@v7 + - run: uvx --from build pyproject-build --sdist --wheel - uses: actions/upload-artifact@v7 with: name: release-dists @@ -22,7 +17,7 @@ jobs: pypi-publish: runs-on: ubuntu-latest needs: - - release-build + - pypi-build permissions: id-token: write steps: diff --git a/.gitignore b/.gitignore index 1a5596a..318e31b 100644 --- a/.gitignore +++ b/.gitignore @@ -166,7 +166,10 @@ cython_debug/ # Packaging -grinder/_version.py +threadmill/_version.py # uv uv.lock + +# pytest-benchmark +.benchmarks/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index e0096fc..8cd5471 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -33,6 +33,7 @@ repos: - mdformat-footnote - mdformat-gfm - mdformat-gfm-alerts + exclude: '.github/agents/' - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.15.11 hooks: diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..d942904 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,50 @@ +# Contributing + +When writing code, you MUST ALWAYS follow the [naming-things](https://github.com/codingjoe/naming-things/blob/main/README.md) guidelines. + + + +## Design Principles + +- Consistency – We never lose data, even if someone unplugs the power or network. +- Durability – We recover from any failures, even poorly written tasks. +- Utilization – We keep the CPU saturated with tasks, not with idle time or waiting for locks. + +## Testing + +We have unit tests, integration tests, and benchmarks. Avoid mocking if possible. + +To run the tests, use the following command: + +```bash +uv run pytest +``` + +To run only integration tests: + +```bash +uv run pytest -m integration +``` + +To run only integration benchmarks: + +```bash +uv run pytest -m "integration and benchmark" +``` + +Benchmarking snapshots are created automatically. +To compare your feature branch against the main branch, +run the test suite on main, followed by: + +``` +uv run pytest --benchmark-compare +``` + +Before your first commit, ensure that the pre-commit hooks are installed by running: + +```bash +uvx prek install +``` diff --git a/README.md b/README.md index a65ef26..a71f339 100644 --- a/README.md +++ b/README.md @@ -1,51 +1,137 @@ -# Django Grinder +# Threadmill

- - - Django Grinder: A queue agnostic worker for Django's task framework. + + + Django Grinder: A queue agnostic worker for Django's task framework.
- Documentation | - Issues | - Changelog | + Documentation | + Issues | + Changelog | Funding 💚

**A queue agnostic worker for Django's task framework.** -- self-healing workers -- graceful shutdown -- CPU, IO, or memory optimized workers +## Design Principles -[![PyPi Version](https://img.shields.io/pypi/v/django-grinder.svg)](https://pypi.python.org/pypi/django-grinder/) -[![Test Coverage](https://codecov.io/gh/codingjoe/django-grinder/branch/main/graph/badge.svg)](https://codecov.io/gh/codingjoe/django-grinder) -[![GitHub License](https://img.shields.io/github/license/codingjoe/django-grinder)](https://raw.githubusercontent.com/codingjoe/django-grinder/master/LICENSE) +- **Durability** – We recover from any failures, even poorly written tasks. +- **Consistency** – We never lose data, even if someone unplugs the power or network. +- **Utilization** – We keep the CPU saturated with tasks, not with idle time or waiting for locks. + +> [!WARNING] +> Threadmill requires a development version of Django and is in a preview stage. + +[![PyPi Version](https://img.shields.io/pypi/v/threadmill.svg)](https://pypi.python.org/pypi/threadmill/) +[![Test Coverage](https://codecov.io/gh/codingjoe/threadmill/branch/main/graph/badge.svg)](https://codecov.io/gh/codingjoe/threadmill) +[![GitHub License](https://img.shields.io/github/license/codingjoe/threadmill)](https://raw.githubusercontent.com/codingjoe/threadmill/master/LICENSE) ## Setup You need to have [Django's Task framework][django-tasks] setup properly. ```console -uv add django-grinder +uv add threadmill ``` -Add `grinder` to your `INSTALLED_APPS` in `settings.py`: +Add `threadmill` to your `INSTALLED_APPS` in `settings.py`: ```python # settings.py INSTALLED_APPS = [ - "grinder", + "threadmill", # ... ] ``` -Finally, you launch the scheduler in a separate process: +Finally, you launch the worker pool: + +```console +uv run manage.py threadmill +``` + +## Usage + +The workers are inspired by Gunicorn, and the CLI is very similar. + +### Utilization + +Depending on your workload, you can tweak the number of processes and threads. +Processes allow for parallel compute (no GIL) while threads are great for low-memory concurrent IO. + +```console +uv run manage.py threadmill --processes 4 --threads 2 +``` + +### Health + +If your tasks leak memory, you can recycle (restart) the workers after a certain number of tasks have been processed: + +```console +uv run manage.py threadmill --max-tasks 1000 --max-tasks-jitter 100 +``` + +This will restart the workers after 1000 tasks have been processed, with a random jitter of up to 100 tasks to avoid all workers restarting at the same time. + +Should a worker crash or be killed, the pool will automatically restart it. + +### Shutdown + +A graceful shutdown is possible with the `SIGTERM` or a keyboard interrupt. +All workers will finish the tasks they acquired and publish them. + +You can use `--exit-empty` to exit immediately after all tasks have been processed, +which might be useful for draining a one-off queue. + +### Task Backlog + +You can prefetch tasks from a queue to avoid IO latency bottlenecks. +However, this will increase the memory usage of the worker pool. ```console -uv run manage.py grinder +uv run manage.py threadmill --prefetch 100 ``` +### Task Timeouts + +> [!WARNING] +> Work in progress, this feature is not yet stable. + +Task timeouts are important to ensure the long-term health of your pool. +However, they need to be aligned with your queueing system's timeout settings. +The message queue needs to requeue a task that hasn't been acknowledged within the timeout. + +## Integration -[django-tasks]: https://docs.djangoproject.com/en/6.0/topics/tasks/ +> [!NOTE] +> This section is for people who want to integrate Threadmill into their queueing system. + +Threadmill is designed to be durable and requires a queueing system to support late acknowledgement. + +To use Threadmill, your backend will need to inherit from `threadmill.backends.AcknowledgeableTaskBackend` and implement the following methods: + +```python +class AcknowledgeableTaskBackend(BaseTaskBackend, ABC): + """Provide an interface for tasks queues to be processed by the executor.""" + + def acquire( + self, *queue_names: str, timeout: datetime.timedelta | None = None + ) -> TaskResult: + """ + Return and lock the next task to be processed without removing it from the queue. + + Args: + queue_names: The names of the queues to acquire tasks from. + timeout: The maximum time to wait for a task. If None, wait indefinitely. + + Raises: + TimeoutError: If no task is available within the specified timeout. + """ + raise NotImplementedError + + def acknowledge(self, task_result: TaskResult) -> None: + """Remove the task from the queue and publish the result.""" + raise NotImplementedError +``` diff --git a/images/logo-dark.svg b/images/logo-dark.svg index 7eee30e..e1f06bc 100644 --- a/images/logo-dark.svg +++ b/images/logo-dark.svg @@ -1,11 +1,17 @@ - + + + + + Django - - Grinder + + Threadmill A queue agnostic worker for Django's task framework. diff --git a/images/logo-light.svg b/images/logo-light.svg index c7afafb..af7e09c 100644 --- a/images/logo-light.svg +++ b/images/logo-light.svg @@ -1,13 +1,20 @@ - + - - + + + + + + Django - - Grinder + + Threadmill - - A queue agnostic worker for Django's task framework. + + A queue agnostic worker for Django's task framework. diff --git a/pyproject.toml b/pyproject.toml index a3f5796..8f58c57 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ requires = ["flit_core>=3.2", "flit_scm", "wheel"] build-backend = "flit_scm:buildapi" [project] -name = "django-grinder" +name = "threadmill" authors = [ { name = "Johannes Maron", email = "johannes@maron.family" }, ] @@ -33,39 +33,43 @@ classifiers = [ "Framework :: Django :: 6.0", ] requires-python = ">=3.12" -dependencies = ["django>=6.0"] +dependencies = ["django @ git+https://github.com/django/django.git@main#egg=django"] [project.urls] # https://packaging.python.org/en/latest/specifications/well-known-project-urls/#well-known-labels -Homepage = "https://github.com/codingjoe/django-grinder" -Changelog = "https://github.com/codingjoe/django-grinder/releases" -Source = "https://github.com/codingjoe/django-grinder" -Releasenotes = "https://github.com/codingjoe/django-grinder/releases/latest" -Documentation = "https://django-grinder.rtfd.io/" -Issues = "https://github.com/codingjoe/django-grinder/issues" +Homepage = "https://github.com/codingjoe/threadmill" +Changelog = "https://github.com/codingjoe/threadmill/releases" +Source = "https://github.com/codingjoe/threadmill" +Releasenotes = "https://github.com/codingjoe/threadmill/releases/latest" +Documentation = "https://github.com/codingjoe/threadmill/" +Issues = "https://github.com/codingjoe/threadmill/issues" Funding = "https://github.com/sponsors/codingjoe" [tool.flit.module] -name = "grinder" +name = "threadmill" [tool.setuptools_scm] -write_to = "grinder/_version.py" +write_to = "threadmill/_version.py" [tool.pytest.ini_options] minversion = "6.0" -addopts = "--cov --cov-report=xml --cov-report=term --tb=short -rxs" +addopts = "--cov --cov-report=xml --cov-report=term --tb=short -rxs --benchmark-autosave --benchmark-group-by=fullname --benchmark-min-rounds=10" testpaths = ["tests"] DJANGO_SETTINGS_MODULE = "tests.testapp.settings" +markers = [ + "benchmark: mark benchmark tests.", + "integration: mark integration tests.", +] [tool.coverage.run] -source = ["grinder"] +source = ["threadmill"] [tool.coverage.report] show_missing = true skip_covered = true [tool.ruff] -src = ["grinder", "tests"] +src = ["threadmill", "tests"] [tool.ruff.lint] select = [ @@ -97,6 +101,8 @@ dev = [ ] test = [ "pytest", + "pytest-benchmark", + "pytest-asyncio", "pytest-cov", "pytest-django", ] diff --git a/tests/test_backends.py b/tests/test_backends.py new file mode 100644 index 0000000..2c9ef4f --- /dev/null +++ b/tests/test_backends.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +import datetime + +import pytest +from threadmill.backends import AcknowledgeableTaskBackend + + +class BackendDouble(AcknowledgeableTaskBackend): + def enqueue(self, task): + return task + + +class TestAcknowledgeableTaskBackend: + def test_acquire__raise_not_implemented_error(self) -> None: + """Raise NotImplementedError for backend acquire API.""" + with pytest.raises(NotImplementedError): + BackendDouble(alias="default", params={}).acquire( + timeout=datetime.timedelta(seconds=1) + ) + + def test_acknowledge__raise_not_implemented_error(self) -> None: + """Raise NotImplementedError for backend acknowledge API.""" + with pytest.raises(NotImplementedError): + BackendDouble(alias="default", params={}).acknowledge(task_result=None) diff --git a/tests/test_command.py b/tests/test_command.py new file mode 100644 index 0000000..0947aa9 --- /dev/null +++ b/tests/test_command.py @@ -0,0 +1,135 @@ +from __future__ import annotations + +import argparse +import signal + +import pytest +from django.core.management import call_command +from django.tasks import default_task_backend +from threadmill.management.commands import threadmill + + +class TestKillSoftly: + def test_kill_softly__raise_keyboard_interrupt_with_signal_name(self): + """Raise KeyboardInterrupt with signal metadata in message.""" + with pytest.raises(KeyboardInterrupt, match="SIGINT"): + threadmill.kill_softly(signal.SIGINT, None) + + +class TestCommand: + def test_add_arguments__register_all_worker_options(self): + """Register command arguments for worker runtime configuration.""" + parser = argparse.ArgumentParser() + + threadmill.Command().add_arguments(parser) + parsed_arguments = parser.parse_args([]) + + assert parsed_arguments.backend == "default" + assert parsed_arguments.queues == ["default"] + assert parsed_arguments.threads == 1 + assert parsed_arguments.max_tasks == 0 + assert parsed_arguments.max_tasks_jitter == 0 + assert parsed_arguments.task_timeout == 3600.0 + + @pytest.mark.benchmark + def test_call_command__benchmark_compute( + self, + benchmark, + ): + """Benchmark command execution for one CPU intense task solved 100 times.""" + default_task_backend.reset() + benchmark.pedantic( + lambda: call_command( + "threadmill", + verbosity=0, + queues=["compute"], + exit_empty=True, + ), + rounds=1, + iterations=1, + warmup_rounds=0, + ) + + assert default_task_backend.solved_task_count == 100 + + @pytest.mark.benchmark + def test_call_command__benchmark_io( + self, + benchmark, + ): + """Benchmark command execution for one CPU intense task solved 100 times.""" + default_task_backend.reset() + benchmark.pedantic( + lambda: call_command( + "threadmill", + verbosity=0, + queues=["io"], + threads=6, + exit_empty=True, + ), + rounds=1, + iterations=1, + warmup_rounds=0, + ) + + assert default_task_backend.solved_task_count == 100 + + @pytest.mark.benchmark + def test_call_command__benchmark_compute_and_io( + self, + benchmark, + ): + """Benchmark command execution for one CPU intense task solved 100 times.""" + default_task_backend.reset() + benchmark.pedantic( + lambda: call_command( + "threadmill", + verbosity=0, + queues=["compute", "io"], + exit_empty=True, + threads=2, + ), + rounds=1, + iterations=1, + warmup_rounds=0, + ) + + assert default_task_backend.solved_task_count == 200 + + @pytest.mark.benchmark + def test_call_command__benchmark_memory_leak_recovery( + self, + benchmark, + ): + """Benchmark command execution for one CPU intense task solved 100 times.""" + default_task_backend.reset(1000) + benchmark.pedantic( + lambda: call_command( + "threadmill", + verbosity=0, + queues=["memory"], + exit_empty=True, + max_tasks=10, + ), + rounds=1, + iterations=1, + warmup_rounds=0, + ) + + @pytest.mark.benchmark + def test_call_command__benchmark_random_crash(self, benchmark): + """Benchmark command execution for one CPU intense task solved 100 times.""" + default_task_backend.reset() + benchmark.pedantic( + lambda: call_command( + "threadmill", + verbosity=0, + exit_empty=True, + ), + rounds=1, + iterations=1, + warmup_rounds=0, + ) + assert default_task_backend.issued_task_count == 100, ( + "All tasks should be issued." + ) diff --git a/tests/testapp/backends.py b/tests/testapp/backends.py new file mode 100644 index 0000000..f237f90 --- /dev/null +++ b/tests/testapp/backends.py @@ -0,0 +1,111 @@ +from queue import Empty + +from django.tasks import TaskResult, TaskResultStatus +from django.utils import timezone +from django.utils.module_loading import import_string +from threadmill.backends import AcknowledgeableTaskBackend + + +class GeneratingTaskBackend(AcknowledgeableTaskBackend): + solved_task_count = 0 + issued_task_count = 0 + target_task_count = 100 + supports_async_task = True + + def __init__(self, alias, params): + super().__init__(alias=alias, params=params) + self._queues = None + + def reset(self, task_count=100): + GeneratingTaskBackend.solved_task_count = 0 + GeneratingTaskBackend.issued_task_count = 0 + self._queues = { + "default": [ + TaskResult( + task=import_string("tests.testapp.tasks.random_crash"), + enqueued_at=timezone.now(), + status=TaskResultStatus.READY, + id=f"default-{i + 1}", + args=[], + kwargs={}, + worker_ids=[], + started_at=None, + finished_at=None, + errors=[], + backend=self.alias, + last_attempted_at=None, + ) + for i in range(task_count) + ], + "compute": [ + TaskResult( + task=import_string("tests.testapp.tasks.compute_workload"), + enqueued_at=timezone.now(), + status=TaskResultStatus.READY, + id=f"compute-{i + 1}", + args=[], + kwargs={}, + worker_ids=[], + started_at=None, + finished_at=None, + errors=[], + backend=self.alias, + last_attempted_at=None, + ) + for i in range(task_count) + ], + "io": [ + TaskResult( + task=import_string("tests.testapp.tasks.io_workload"), + enqueued_at=timezone.now(), + status=TaskResultStatus.READY, + id=f"io-{i + 1}", + args=[], + kwargs={}, + worker_ids=[], + started_at=None, + finished_at=None, + errors=[], + backend=self.alias, + last_attempted_at=None, + ) + for i in range(task_count) + ], + "memory": [ + TaskResult( + task=import_string("tests.testapp.tasks.memory_workload"), + enqueued_at=timezone.now(), + status=TaskResultStatus.READY, + id=f"memory-{i + 1}", + args=[], + kwargs={}, + worker_ids=[], + started_at=None, + finished_at=None, + errors=[], + backend=self.alias, + last_attempted_at=None, + ) + for i in range(task_count) + ], + } + + def enqueue(self, task): + return task + + def acquire(self, *queue_names, timeout=None): + if self._queues is None: + self.reset() + GeneratingTaskBackend.issued_task_count += 1 + queues = [self._queues[queue_name] for queue_name in queue_names] + try: + # pop from the longest queue first to simulate a more realistic scenario + for queue in sorted(queues, key=len, reverse=True): + return queue.pop(0) + except IndexError as e: + GeneratingTaskBackend.issued_task_count -= 1 + raise Empty("No more tasks to solve.") from e + raise Empty("No more tasks to solve.") + + def acknowledge(self, task_result: TaskResult) -> None: + GeneratingTaskBackend.solved_task_count += 1 diff --git a/tests/testapp/settings.py b/tests/testapp/settings.py index b626f9f..4884ec4 100644 --- a/tests/testapp/settings.py +++ b/tests/testapp/settings.py @@ -10,9 +10,10 @@ https://docs.djangoproject.com/en/4.2/ref/settings/ """ -import os from pathlib import Path +from django.tasks import DEFAULT_TASK_QUEUE_NAME + # Build paths inside the project like this: BASE_DIR / 'subdir'. BASE_DIR = Path(__file__).resolve().parent.parent @@ -38,7 +39,7 @@ "django.contrib.sessions", "django.contrib.messages", "django.contrib.staticfiles", - "grinder", + "threadmill", "tests.testapp", ] @@ -83,7 +84,12 @@ } } -TASKS = {"default": {"BACKEND": "django.tasks.backends.immediate.ImmediateBackend"}} +TASKS = { + "default": { + "BACKEND": "tests.testapp.backends.GeneratingTaskBackend", + "QUEUES": [DEFAULT_TASK_QUEUE_NAME, "compute", "io", "memory"], + }, +} # Password validation # https://docs.djangoproject.com/en/4.2/ref/settings/#auth-password-validators diff --git a/tests/testapp/tasks.py b/tests/testapp/tasks.py index 3b49002..87d912c 100644 --- a/tests/testapp/tasks.py +++ b/tests/testapp/tasks.py @@ -1,12 +1,55 @@ +import asyncio import logging +import random +import uuid -from grinder import cron from django.tasks import task logger = logging.getLogger(__name__) -@cron("*/5 * * * *") -@task -def my_task(): - logger.info("Hello World!") +@task(queue_name="compute") +def compute_workload(): + """Calculate the first 1000 prime numbers.""" + + def is_prime(number: int) -> bool: + if number < 2: + return False + if number in (2, 3): + return True + if number % 2 == 0: + return False + for divisor in range(3, int(number**0.5) + 1, 2): + if number % divisor == 0: + return False + return True + + prime_count = 0 + number = 2 + while prime_count < 100_000: + if is_prime(number): + prime_count += 1 + number += 1 + return prime_count + + +@task(queue_name="io") +async def io_workload(): + """Sleep for a random amount of time.""" + await asyncio.sleep(random.uniform(0.1, 0.5)) # noqa: S311 + + +leak = {} + + +@task(queue_name="memory") +def memory_workload(): + """Allocate and leak 100MB of memory.""" + leak[uuid.uuid4()] = "x" * 1024 * 1024 * 100 + + +@task() +def random_crash(): + """Raise a random exception.""" + if random.random() < 0.75: # noqa: S311 + exit(1) diff --git a/grinder/__init__.py b/threadmill/__init__.py similarity index 79% rename from grinder/__init__.py rename to threadmill/__init__.py index 5b666e3..5f32987 100644 --- a/grinder/__init__.py +++ b/threadmill/__init__.py @@ -4,3 +4,6 @@ __version__ = _version.version VERSION = _version.version_tuple + + +__all__ = ["VERSION", "__version__"] diff --git a/threadmill/backends.py b/threadmill/backends.py new file mode 100644 index 0000000..cb0ab7b --- /dev/null +++ b/threadmill/backends.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import datetime +from abc import ABC + +from django.tasks import TaskResult +from django.tasks.backends.base import BaseTaskBackend + + +class AcknowledgeableTaskBackend(BaseTaskBackend, ABC): + """Provide an interface for tasks queues to be processed by the executor.""" + + supports_async_task = True + supports_get_result = True + + def acquire( + self, *queue_names: str, timeout: datetime.timedelta | None = None + ) -> TaskResult: + """ + Return and lock the next task to be processed without removing it from the queue. + + Args: + queue_names: The names of the queues to acquire tasks from. + timeout: The maximum time to wait for a task. If None, wait indefinitely. + + Raises: + TimeoutError: If no task is available within the specified timeout. + queue.Empty: If no task is available and timeout is None. + """ + raise NotImplementedError + + def acknowledge(self, task_result: TaskResult) -> None: + """Remove the task from the queue and publish the result.""" + raise NotImplementedError diff --git a/threadmill/executor.py b/threadmill/executor.py new file mode 100644 index 0000000..a0a3510 --- /dev/null +++ b/threadmill/executor.py @@ -0,0 +1,307 @@ +"""Task worker executor implementation.""" + +from __future__ import annotations + +import asyncio +import dataclasses +import datetime +import logging +import multiprocessing +import random +import socket +import threading +import time +import typing +from concurrent.futures import ThreadPoolExecutor +from contextlib import suppress +from inspect import iscoroutinefunction +from multiprocessing.queues import JoinableQueue +from queue import Empty +from traceback import format_exception + +from django.tasks import TaskResult +from django.tasks.base import TaskContext, TaskError, TaskResultStatus +from django.tasks.signals import task_finished, task_started +from django.utils import timezone +from django.utils.json import normalize_json + +if typing.TYPE_CHECKING: + from .backends import AcknowledgeableTaskBackend + + +logger = multiprocessing.get_logger() +formatter = logging.Formatter( + "%(levelname)s: %(asctime)s - pid=%(process)s - %(message)s" +) +handler = logging.StreamHandler() +handler.setFormatter(formatter) +logger.addHandler(handler) +logger.setLevel(logging.INFO) + + +@dataclasses.dataclass(kw_only=True, slots=True) +class TaskExecutor: + """Consume tasks from shared joinable queues with process and thread pools.""" + + backend: AcknowledgeableTaskBackend + workers: int | None = None + threads: int = 1 + max_tasks: int = 0 + max_tasks_jitter: int = 0 + task_timeout: datetime.timedelta = datetime.timedelta(hours=1) + is_acquiring: bool = dataclasses.field(default=True, init=False) + is_publishing: bool = dataclasses.field(default=True, init=False) + worker_processes: list[WorkerProcess] = dataclasses.field( + default_factory=list, init=False + ) + process_count: int = dataclasses.field(init=False) + thread_count: int = dataclasses.field(init=False) + queues: tuple[str] + shared_task_queue: multiprocessing.JoinableQueue[TaskResult] = dataclasses.field( + init=False + ) + processed_task_queue: multiprocessing.JoinableQueue[TaskResult] = dataclasses.field( + init=False + ) + exit_empty: bool = False + + def __post_init__(self) -> None: + """Initialize derived orchestration fields and queues.""" + self.process_count = self.workers or max(multiprocessing.cpu_count() - 1, 1) + self.thread_count = max(self.threads, 1) + self.shared_task_queue = multiprocessing.JoinableQueue( + maxsize=self.process_count * self.thread_count, + ) + self.processed_task_queue = multiprocessing.JoinableQueue() + + def get_maximum_tasks_per_child(self) -> int | None: + """Return worker recycling limit based on config and thread count.""" + if self.max_tasks: + return ( + self.max_tasks + random.randint(0, self.max_tasks_jitter) # noqa: S311 + ) // self.thread_count + + def create_worker_process(self) -> WorkerProcess: + """Create and start a new worker process.""" + worker = WorkerProcess( + self.shared_task_queue, + self.processed_task_queue, + self.thread_count, + self.task_timeout, + self.get_maximum_tasks_per_child(), + ) + worker.start() + return worker + + def run(self) -> None: + """Start consuming tasks until shutdown is requested.""" + self.worker_processes = [ + self.create_worker_process() for _ in range(self.process_count) + ] + threads = [ + threading.Thread(target=self.acknowledge_tasks, daemon=True), + threading.Thread(target=self.maintain_worker_pool, daemon=True), + threading.Thread(target=self.acquire_tasks, daemon=True), + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + def acquire_tasks(self) -> None: + """Buffer tasks in shared task queue.""" + while self.is_acquiring: + try: + work = self.backend.acquire(*self.queues) + except Empty: + if self.exit_empty: + logger.info("No more tasks to solve. Shutting down.") + self.shutdown() + return + time.sleep(0.01) + else: + self.shared_task_queue.put(work) + + def acknowledge_tasks(self) -> None: + """Acknowledge processed tasks and publish updated results in main process.""" + while self.is_publishing: + try: + task = self.processed_task_queue.get_nowait() + except Empty: + time.sleep(0.01) + else: + self.backend.acknowledge(task) + self.processed_task_queue.task_done() + + def shutdown(self) -> None: + """Stop queue consumption and terminate all worker processes.""" + logger.info("Shutting down task executor") + self.is_acquiring = False + with suppress(ValueError): + self.shared_task_queue.join() + with suppress(ValueError): + self.processed_task_queue.join() + with ThreadPoolExecutor(max_workers=self.process_count) as executor: + executor.map(lambda worker: worker.shutdown(), self.worker_processes) + self.is_publishing = False + + def maintain_worker_pool(self) -> None: + """Restart worker processes that have exited.""" + while self.is_publishing: + for index, worker in enumerate(self.worker_processes): + if worker.is_alive(): + continue + worker.join(timeout=0) + self.worker_processes[index] = self.create_worker_process() + time.sleep(1) + + +class WorkerProcess(multiprocessing.Process): + """Single worker process running thread_count consumer threads.""" + + def __init__( + self, + task_queue: JoinableQueue[TaskResult], + processed_task_queue: JoinableQueue[TaskResult], + thread_count: int, + task_timeout: datetime.timedelta, + max_tasks: int | None = None, + ) -> None: + """Create process with dedicated thread pool for task execution.""" + self.shutdown_requested = multiprocessing.Event() + super().__init__(daemon=True) + self.task_queue = task_queue + self.processed_task_queue = processed_task_queue + self.thread_count = thread_count + self.task_timeout = task_timeout + self.max_tasks = max_tasks + self.task_count = 0 + self.lock: threading.Lock | None = None + self.expired: threading.Event | None = None + + def run(self) -> None: + """Start consumer execution inside this process.""" + logger.info("Starting worker process %s", self.name) + self.lock = threading.Lock() + self.expired = threading.Event() + consumer_threads = [ + WorkerThread(worker=self, index=index) for index in range(self.thread_count) + ] + for consumer_thread in consumer_threads: + consumer_thread.start() + for consumer_thread in consumer_threads: + consumer_thread.join(self.task_timeout.total_seconds()) + + def record_task(self) -> None: + """Record one processed task and stop when max_tasks is reached.""" + if self.max_tasks is None: + return + if self.lock is None or self.expired is None: + return + with self.lock: + self.task_count += 1 + if self.task_count >= self.max_tasks: + self.expired.set() + + def shutdown(self) -> None: + """Request graceful worker stop and wait for process exit.""" + logger.info("Stopping worker process %s", self.name) + self.shutdown_requested.set() + self.join() + + +class WorkerThread(threading.Thread): + """Single worker thread consuming tasks from the process queue.""" + + def __init__( + self, + *, + worker: WorkerProcess, + index: int, + ) -> None: + """Create worker thread bound to process worker state.""" + super().__init__(name=f"{socket.gethostname()}:{worker.pid}-{index}") + self.worker = worker + + def run(self) -> None: + """Start consuming tasks for this thread.""" + while self.worker.expired is None or not self.worker.expired.is_set(): + try: + task_result = self.worker.task_queue.get(timeout=1.0) + except Empty: + if self.worker.shutdown_requested.is_set(): + return + continue + try: + self.worker.processed_task_queue.put( + self.execute_task_result( + task_result, + ) + ) + finally: + self.worker.task_queue.task_done() + self.worker.record_task() + + def execute_task_result(self, task_result: TaskResult) -> TaskResult: + """Execute task from task result and update result lifecycle state.""" + logger.info("Executing task %r", task_result.id) + started_at = timezone.now() + task_result = dataclasses.replace( + task_result, + status=TaskResultStatus.RUNNING, + started_at=started_at, + last_attempted_at=started_at, + worker_ids=[*task_result.worker_ids, self.name], + ) + task_started.send(TaskExecutor, task_result=task_result) + + try: + return_value = WorkerThread.call_task(task_result) + except Exception as exception: + task_result = dataclasses.replace( + task_result, + status=TaskResultStatus.FAILED, + errors=[*task_result.errors, WorkerThread.create_task_error(exception)], + ) + logger.exception("Task failed %r", task_result.id) + else: + task_result = dataclasses.replace( + task_result, + status=TaskResultStatus.SUCCESSFUL, + ) + object.__setattr__( + task_result, "_return_value", normalize_json(return_value) + ) + logger.info("Task successful %r", task_result.id) + finally: + task_result = dataclasses.replace( + task_result, + finished_at=timezone.now(), + ) + task_finished.send(TaskExecutor, task_result=task_result) + + return task_result + + @staticmethod + def call_task(task_result: TaskResult) -> typing.Any: + """Call a task with context when required.""" + task = task_result.task + if task.takes_context: + args = [TaskContext(task_result=task_result), *task_result.args] + else: + args = task_result.args + if iscoroutinefunction(task.func): + return asyncio.run(task.func(*args, **task_result.kwargs)) + return task.func( + *args, + **task_result.kwargs, + ) + + @staticmethod + def create_task_error(exception: BaseException) -> TaskError: + """Build a task error payload for failed execution.""" + exception_type = type(exception) + return TaskError( + exception_class_path=f"{exception_type.__module__}.{exception_type.__qualname__}", + traceback="".join(format_exception(exception)), + ) diff --git a/threadmill/management/__init__.py b/threadmill/management/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/threadmill/management/commands/__init__.py b/threadmill/management/commands/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/threadmill/management/commands/threadmill.py b/threadmill/management/commands/threadmill.py new file mode 100644 index 0000000..0fa7249 --- /dev/null +++ b/threadmill/management/commands/threadmill.py @@ -0,0 +1,132 @@ +import datetime +import signal +import sys + +from django.core.management import BaseCommand, CommandError +from django.tasks import ( + DEFAULT_TASK_BACKEND_ALIAS, + DEFAULT_TASK_QUEUE_NAME, + InvalidTaskBackend, + task_backends, +) + +from ...executor import TaskExecutor + + +def kill_softly(signum, frame): + """Raise a KeyboardInterrupt to stop the worker gracefully.""" + signame = signal.Signals(signum).name + raise KeyboardInterrupt(f"Received {signame} ({signum}), shutting down…") + + +class Command(BaseCommand): + """Run task workers to process enqueued tasks from the specified backends and queues.""" + + help = __doc__ + + def add_arguments(self, parser): + parser.add_argument( + "-b", + "--backend", + default=DEFAULT_TASK_BACKEND_ALIAS, + help="Alias of the tasks backend to use.", + ) + parser.add_argument( + "-q", + "--queues", + nargs="+", + default=[DEFAULT_TASK_QUEUE_NAME], + help="Queue names to listen to and process tasks from.", + ) + parser.add_argument( + "-w", + "--workers", + type=int, + help="Number of worker processes to use. Defaults to the number of CPU cores minus one.", + ) + parser.add_argument( + "-t", + "--threads", + type=int, + default=1, + help="Number of threads to use. Defaults to 1. ", + ) + parser.add_argument( + "--max-tasks", + type=int, + default=0, + help=( + "Number of the maximum number of tasks to run until a worker is recycled." + " Defaults to 0, which means no limit." + ), + ) + parser.add_argument( + "--max-tasks-jitter", + type=int, + default=0, + help="Maximum random jitter to add to the max-tasks value by randint(0, max_tasks_jitter).", + ) + parser.add_argument( + "--task-timeout", + type=float, + default=3600.0, + help="Kill hung tasks after timeout seconds. Defaults to one hour.", + ) + parser.add_argument( + "--task-backlog-size", + type=int, + default=1, + help="The number of tasks to prefetch from the message queue while all workers are busy. Defaults to 1.", + ) + parser.add_argument( + "--exit-empty", + action="store_true", + help="Drain the task queue and exit with 0.", + ) + + def handle( + self, + *, + verbosity, + backend, + queues, + workers, + threads, + max_tasks, + max_tasks_jitter, + task_timeout, + exit_empty, + **options, + ): + match sys.platform: + case "win32": + signal.signal(signal.SIGBREAK, kill_softly) + case _: + signal.signal(signal.SIGHUP, kill_softly) + signal.signal(signal.SIGTERM, kill_softly) + signal.signal(signal.SIGINT, kill_softly) + self.stdout.write(self.style.SUCCESS("Starting workers…")) + try: + backend = task_backends[backend] + except InvalidTaskBackend: + raise CommandError(f"Invalid backend: {backend!r}") + if _non_queues := set(queues) - set(backend.queues): + raise CommandError( + f"Backend does not support all specified queues: {_non_queues!r}" + ) + exe = TaskExecutor( + backend=backend, + workers=workers, + threads=threads, + max_tasks=max_tasks, + max_tasks_jitter=max_tasks_jitter, + task_timeout=datetime.timedelta(seconds=task_timeout), + exit_empty=exit_empty, + queues=queues, + ) + try: + exe.run() + except KeyboardInterrupt as e: + self.stdout.write(self.style.WARNING(str(e))) + self.stdout.write(self.style.NOTICE("Shutting down workers…")) + exe.shutdown()