From 1c4ba740332a3b2a22f484c272bac2db6197ba0e Mon Sep 17 00:00:00 2001 From: Johannes Maron Date: Thu, 2 Jul 2026 14:15:36 +0200 Subject: [PATCH 1/4] Resolve #18 -- Add worker pool telemetry to inspector --- AGENTS.md | 14 +- CONTRIBUTING.md | 24 +- pyproject.toml | 3 + tests/backends/test_redis.py | 123 ++++++ tests/test_executor.py | 22 +- tests/test_inspector.py | 563 +++++++++++++++++++++++++++- tests/test_telemetry.py | 280 ++++++++++++++ threadmill/backends/base.py | 77 +++- threadmill/backends/redis.py | 75 +++- threadmill/executor.py | 27 +- threadmill/inspector/app.py | 461 +++++++---------------- threadmill/inspector/inspector.scss | 39 +- threadmill/inspector/queue_view.py | 309 +++++++++++++++ threadmill/inspector/utils.py | 34 ++ threadmill/inspector/worker_view.py | 215 +++++++++++ threadmill/telemetry.py | 150 ++++++++ 16 files changed, 2057 insertions(+), 359 deletions(-) create mode 100644 tests/test_telemetry.py create mode 100644 threadmill/inspector/queue_view.py create mode 100644 threadmill/inspector/utils.py create mode 100644 threadmill/inspector/worker_view.py create mode 100644 threadmill/telemetry.py diff --git a/AGENTS.md b/AGENTS.md index aa12e04..61ab711 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -18,7 +18,7 @@ uv run pytest -m "integration and benchmark" uv run pytest --benchmark-compare # compare vs main baseline (run main first) uvx prek run --all-files uv run manage.py threadmill worker # run the worker pool -uv run manage.py threadmill inspector # launch the textual TUI inspector +uv run manage.py threadmill inspector # launch the textual TUI inspector ``` CI additionally pins Django per matrix step: `uv run --with django~=6.1a1 pytest -m "not benchmark"`. @@ -28,8 +28,8 @@ Run a single test by node ID, e.g. `uv run pytest tests/test_command.py::TestCom ## Setup - Install pre-commit hooks before first commit: `uvx prek install` (not `pre-commit install`). -- `DJANGO_SETTINGS_MODULE=tests.testapp.settings` is already set in `.env` and in `pyproject.toml`. Pytest auto-loads it. -- CI's Linux job starts a Redis service and sets `REDIS_URL`; some integration tests may rely on it. Local runs of `-m integration` may need Redis if a backend test targets it. +- `DJANGO_SETTINGS_MODULE=tests.testapp.settings` is set in `.env` and `pyproject.toml`; pytest auto-loads it. +- CI's Linux job starts Redis and sets `REDIS_URL`; some integration tests rely on it. Local `-m integration` runs may need Redis. ## Code & style (repo-specific, beyond PEP 8) @@ -50,15 +50,15 @@ Codecov requires 100% patch coverage on PRs (`pyproject.toml`, `.codecov.yml`). ## Architecture notes -- Entry point for end users: `threadmill/management/commands/threadmill.py` (Django management command with `worker` and `inspector` subcommands). +- Entry point: `threadmill/management/commands/threadmill.py` (management command with `worker` and `inspector` subcommands). - Core runtime: `threadmill/executor.py` (`TaskExecutor`) — process/thread pool, graceful shutdown, worker recycling, task timeout/backlog. -- Integration point for queue authors: `threadmill/backends.py` (`AcknowledgeableTaskBackend`) — subclasses implement `acquire` (lock-without-remove) and `acknowledge` (remove + publish). Requires late-ack support from the underlying queue. +- Queue author integration point: `threadmill/backends/base.py` (`ThreadmillTaskBackend`) — subclasses implement `acquire` (lock-without-remove) and `acknowledge` (remove + publish); requires late-ack support from the underlying queue. - Test app backend `tests/testapp/backends.py` (`GeneratingTaskBackend`) generates tasks in-process for benchmarks; reset between runs via `default_task_backend.reset()`. ## Pre-commit -`.pre-commit-config.yaml` runs ruff (check + format), django-upgrade, pyupgrade, mdformat (excludes `.github/agents/`), yamlfmt, and `no-commit-to-branch` (protects `main`). Hooks auto-fix; ruff is configured `--exit-non-zero-on-fix`, so commit any fixes before pushing. +`.pre-commit-config.yaml` runs ruff (check + format), django-upgrade, pyupgrade, mdformat (excludes `.github/agents/`), yamlfmt, and `no-commit-to-branch` (protects `main`). Hooks auto-fix; ruff uses `--exit-non-zero-on-fix`, so commit fixes before pushing. ## PR / release -CI runs on `main` pushes and PRs. Releases are published to PyPI via `.github/workflows/release.yml` on GitHub release. Commits to `main` are blocked by `no-commit-to-branch`; work on a branch. +CI runs on `main` pushes and PRs. Releases publish to PyPI via `.github/workflows/release.yml` on GitHub release. `no-commit-to-branch` blocks commits to `main`; work on a branch. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index da729f9..db66148 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -15,35 +15,21 @@ curl -sSL https://raw.githubusercontent.com/codingjoe/naming-things/refs/heads/m ## 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: +The suite has unit tests, integration tests, and benchmarks. Avoid mocking where possible. ```bash +uv run pytest # full suite +uv run pytest -m integration # integration tests only 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: +Benchmark snapshots are created automatically. To compare a feature branch against main, run the suite on main first, then: ``` uv run pytest --benchmark-compare ``` -Before your first commit, ensure that the pre-commit hooks are installed by running: +Install pre-commit hooks before your first commit: ```bash uvx prek install diff --git a/pyproject.toml b/pyproject.toml index fee79cb..74cf408 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,9 @@ redis = ["redis>=5.0"] inspector = [ "textual>=8.2.7", ] +worker = [ + "psutil>=7.2.2", +] [project.urls] # https://packaging.python.org/en/latest/specifications/well-known-project-urls/#well-known-labels diff --git a/tests/backends/test_redis.py b/tests/backends/test_redis.py index cbb24c9..0395354 100644 --- a/tests/backends/test_redis.py +++ b/tests/backends/test_redis.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import dataclasses import datetime import logging @@ -7,6 +8,7 @@ from dataclasses import replace from unittest.mock import patch +import pytest from django.tasks import default_task_backend from django.tasks.base import TaskResultStatus from django.utils import timezone @@ -14,9 +16,12 @@ from tests.testapp.tasks import boom, compute_workload, echo from threadmill.backends.base import ( BackendTelemetry, + NodeTelemetry, QueueCounts, QueueRates, QueueStats, + ThreadmillTaskBackend, + WorkerTelemetry, ) from threadmill.backends.redis import RedisBroker, RedisTaskBackend # noqa: E402 @@ -687,3 +692,121 @@ def test_peek__empty_history_returns_nothing(self): ) assert successful == [] assert failed == [] + + +class TestWorkerTelemetrySerialization: + """Tests for worker telemetry JSON serialization/deserialization.""" + + def test_serialize_deserialize_roundtrip(self): + """A WorkerTelemetry snapshot survives a serialize→deserialize roundtrip.""" + from threadmill.backends.base import ( + NodeTelemetry, + WorkerTelemetry, + ) + from threadmill.backends.redis import ( + _deserialize_worker_telemetry, + _serialize_worker_telemetry, + ) + + sampled_at = datetime.datetime.now(tz=datetime.UTC) + node = NodeTelemetry( + hostname="host", + pid=100, + queues=("default", "high"), + process_count=1, + thread_count=4, + cpu_percent=45.0, + memory_bytes=8_000_000_000, + memory_total=16_000_000_000, + tasks_per_minute=12.5, + sampled_at=sampled_at, + ) + original = WorkerTelemetry( + nodes={"host": node}, + queues={"default": ("host",), "high": ("host",)}, + sampled_at=sampled_at, + ) + payload = _serialize_worker_telemetry(original) + restored = _deserialize_worker_telemetry(payload) + assert restored.sampled_at == original.sampled_at + assert set(restored.nodes) == {"host"} + assert set(restored.queues) == {"default", "high"} + r_node = restored.nodes["host"] + assert r_node.hostname == "host" + assert r_node.pid == 100 + assert r_node.process_count == 1 + assert r_node.thread_count == 4 + assert r_node.cpu_percent == 45.0 + assert r_node.memory_bytes == 8_000_000_000 + assert r_node.memory_total == 16_000_000_000 + assert r_node.queues == ("default", "high") + + +@pytest.mark.integration +class TestWorkerTelemetryRedis: + """Integration tests for worker telemetry storage over real Redis.""" + + async def test_publish_and_read_roundtrip(self): + """publish_worker_telemetry publishes, subscribe_worker_telemetry reads it back.""" + backend = RedisTaskBackend( + "worker_telemetry_store_test", + { + "QUEUES": ["default"], + "REDIS_URL": "redis://localhost:6379/0", + "OPTIONS": { + "result_ttl": datetime.timedelta(seconds=60), + }, + }, + ) + try: + sampled_at = datetime.datetime.now(tz=datetime.UTC) + node = NodeTelemetry( + hostname="testhost", + pid=200, + queues=("default",), + process_count=1, + thread_count=2, + cpu_percent=30.0, + memory_bytes=4_000_000_000, + memory_total=8_000_000_000, + tasks_per_minute=10.0, + sampled_at=sampled_at, + ) + snapshot = WorkerTelemetry( + nodes={"testhost": node}, + queues={"default": ("testhost",)}, + sampled_at=sampled_at, + ) + + async def _collect(): + async for received in backend.subscribe_worker_telemetry(): + return received + + task = asyncio.ensure_future(_collect()) + # Give the pub/sub subscription a moment to establish. + await asyncio.sleep(0.1) + backend.publish_worker_telemetry(snapshot) + + received = await asyncio.wait_for(task, timeout=5.0) + assert "testhost" in received.nodes + r_node = received.nodes["testhost"] + assert r_node.hostname == "testhost" + assert r_node.pid == 200 + assert r_node.cpu_percent == 30.0 + assert r_node.process_count == 1 + assert r_node.thread_count == 2 + assert "default" in received.queues + finally: + backend.close() + + async def test_worker_telemetry_empty_when_no_workers(self): + """The default subscribe_worker_telemetry yields an empty snapshot.""" + + class _StubBackend(ThreadmillTaskBackend): + def enqueue(self, *args, **kwargs): # pragma: no cover + raise NotImplementedError + + backend = _StubBackend(alias="test", params={}) + received = await backend.subscribe_worker_telemetry().__anext__() + assert received.nodes == {} + assert received.queues == {} diff --git a/tests/test_executor.py b/tests/test_executor.py index 19e1c09..9cfa8f1 100644 --- a/tests/test_executor.py +++ b/tests/test_executor.py @@ -55,12 +55,13 @@ def _task_result(task, *args, **kwargs) -> TaskResult: def _make_worker(*, max_tasks: int | None = None) -> WorkerProcess: - """Build an unstarted `WorkerProcess`.""" + """Build an unstarted `WorkerProcess` with a fresh shared task counter.""" return WorkerProcess( thread_count=1, max_tasks=max_tasks, backend_alias="default", queues=("default",), + total_counter=multiprocessing.Value("q", 0), ) @@ -190,6 +191,16 @@ def test_shutdown__shuts_down_worker_processes(self): executor.shutdown() assert not worker.is_alive() + def test_shutdown__stops_telemetry_sampler(self): + """Shutdown stops the telemetry sampler when one is running.""" + from unittest.mock import MagicMock + + sampler = MagicMock() + executor = TaskExecutor(backend=default_task_backend, queues=("default",)) + executor.telemetry_sampler = sampler + executor.shutdown() + sampler.stop.assert_called_once() + def test_maintain_worker_pool__restarts_dead_workers(self): """maintain_worker_pool replaces dead workers with new ones.""" executor = TaskExecutor( @@ -217,12 +228,13 @@ class TestWorkerProcess: """Tests for the WorkerProcess class.""" def test_record_task__increments_count(self): - """record_task increments task_count.""" + """record_task increments task_count and the shared counter.""" worker = _make_worker(max_tasks=5) worker.lock = threading.Lock() worker.expired = threading.Event() worker.record_task() assert worker.task_count == 1 + assert worker.total_counter.value == 1 def test_record_task__sets_expired_when_max_reached(self): """record_task sets expired event when max_tasks is reached.""" @@ -233,19 +245,21 @@ def test_record_task__sets_expired_when_max_reached(self): assert worker.expired.is_set() def test_record_task__noop_when_max_tasks_is_none(self): - """record_task is a no-op when max_tasks is None.""" + """record_task still counts throughput when max_tasks is None.""" worker = _make_worker(max_tasks=None) worker.lock = threading.Lock() worker.expired = threading.Event() worker.record_task() assert worker.task_count == 0 assert not worker.expired.is_set() + assert worker.total_counter.value == 1 def test_record_task__noop_before_run_sets_lock_and_expired(self): - """record_task is a safe no-op before run() initializes lock/expired.""" + """record_task counts but skips recycling before run() sets lock/expired.""" worker = _make_worker(max_tasks=5) worker.record_task() assert worker.task_count == 0 + assert worker.total_counter.value == 1 def test_shutdown_requested__is_settable(self): """shutdown_requested event can be set on an unstarted worker.""" diff --git a/tests/test_inspector.py b/tests/test_inspector.py index 11aaa62..a34926e 100644 --- a/tests/test_inspector.py +++ b/tests/test_inspector.py @@ -21,16 +21,21 @@ from tests.testapp.tasks import echo from threadmill.backends.base import ( BackendTelemetry, + NodeTelemetry, QueueCounts, QueueRates, QueueStats, ThreadmillTaskBackend, + WorkerTelemetry, ) from threadmill.inspector.app import ( InspectorApp, QueueList, + SelectionTree, TaskDetail, TaskList, + WorkerGraphs, + WorkerTreeNode, si_prefix, ) @@ -47,6 +52,9 @@ def peek(self, *args, **kwargs): def telemetry(self, *, interval=None): raise RuntimeError("telemetry failed") + def worker_telemetry(self): + raise RuntimeError("worker_telemetry failed") + def _failed_result() -> TaskResult: """Build a failed TaskResult with an error for detail-view tests.""" @@ -102,11 +110,16 @@ def _stats(**overrides: int | datetime.timedelta) -> QueueStats: (3_500_000, "3.5M"), ], ) -def test_si_prefix(count: int, expected: str) -> None: +def test_si_prefix(count: int, expected: str): """Large counts are compacted with k/M/G suffixes.""" assert si_prefix(count) == expected +def test_si_prefix__bytes(): + """Size in bytes can be abbreviated with base=1024.""" + assert si_prefix(2048 * 1024 * 1024, base=1024) == "2G" + + class TestInspectorApp: """Tests for the textual inspector TUI.""" @@ -511,3 +524,551 @@ async def test_action_refresh_refreshes_task_list(self): await pilot.pause() await pilot.pause() assert table.row_count == before - 1 + + +def _make_worker_telemetry( + *, + hostname: str = "node-1", + pid: int = 1234, + queue_name: str = "default", +) -> WorkerTelemetry: + """Build a WorkerTelemetry snapshot for inspector tests.""" + now = datetime.datetime.now(tz=datetime.UTC) + node = NodeTelemetry( + hostname=hostname, + pid=pid, + queues=(queue_name,), + process_count=1, + thread_count=2, + cpu_percent=45.0, + memory_bytes=4_000_000_000, + memory_total=8_000_000_000, + tasks_per_minute=30.0, + sampled_at=now, + ) + return WorkerTelemetry( + nodes={hostname: node}, + queues={queue_name: (hostname,)}, + sampled_at=now, + ) + + +def _queue_telemetry(*queues: str) -> BackendTelemetry: + """Build a BackendTelemetry listing the defined queues for the selection tree.""" + return BackendTelemetry(queues={name: _stats() for name in queues}) + + +class TestWorkerView: + """Tests for the worker view (selection tree, graphs, toggle).""" + + async def test_selection_tree_builds_from_telemetry(self): + """The selection tree renders Queue -> Node from telemetry.""" + app = InspectorApp(backend=default_task_backend, auto_refresh=False) + async with app.run_test() as pilot: + await pilot.pause() + tree = app.query_one("#selection-tree", SelectionTree) + tree.queue_telemetry = _queue_telemetry("default") + snapshot = _make_worker_telemetry() + tree.worker_telemetry = snapshot + await pilot.pause() + root = tree.root + assert len(root.children) == 1 + queue_node = root.children[0] + assert queue_node.data.kind == "queue" + assert queue_node.data.queue_name == "default" + assert queue_node.label.plain == "default (1)" + assert len(queue_node.children) == 1 + host_node = queue_node.children[0] + assert host_node.data.kind == "node" + assert host_node.data.hostname == "node-1" + + async def test_selection_tree_shows_unstaffed_queue_with_zero(self): + """A defined queue with no draining nodes renders as 'queue (0)'.""" + app = InspectorApp(backend=default_task_backend, auto_refresh=False) + async with app.run_test() as pilot: + await pilot.pause() + tree = app.query_one("#selection-tree", SelectionTree) + tree.queue_telemetry = _queue_telemetry("default", "io") + tree.worker_telemetry = _make_worker_telemetry() + await pilot.pause() + labels = [child.label.plain for child in tree.root.children] + assert "default (1)" in labels + assert "io (0)" in labels + io_node = next( + child for child in tree.root.children if child.data.queue_name == "io" + ) + assert len(io_node.children) == 0 + + def test_memory_percent_handles_zero_total(self): + """_memory_percent returns 0 when total RAM is unknown (<=0).""" + node = NodeTelemetry( + hostname="h", + pid=1, + queues=("default",), + process_count=1, + thread_count=1, + cpu_percent=0.0, + memory_bytes=100, + memory_total=0, + tasks_per_minute=0.0, + sampled_at=datetime.datetime.now(tz=datetime.UTC), + ) + assert WorkerGraphs._memory_percent(node) == 0.0 + + def test_memory_percent_computes_ratio(self): + """_memory_percent derives the usage ratio from used/total RAM.""" + node = NodeTelemetry( + hostname="h", + pid=1, + queues=("default",), + process_count=1, + thread_count=1, + cpu_percent=0.0, + memory_bytes=4_000_000_000, + memory_total=8_000_000_000, + tasks_per_minute=0.0, + sampled_at=datetime.datetime.now(tz=datetime.UTC), + ) + assert WorkerGraphs._memory_percent(node) == 50.0 + + async def test_worker_graphs_ram_border_title(self): + """The memory border title shows used/total RAM and the derived percent.""" + app = InspectorApp(backend=default_task_backend, auto_refresh=False) + async with app.run_test() as pilot: + await pilot.pause() + graphs = app.query_one("#worker-graphs", WorkerGraphs) + graphs.selection = WorkerTreeNode( + kind="node", label="node-1", hostname="node-1" + ) + graphs.worker_telemetry = _make_worker_telemetry() + await pilot.pause() + graphs._refresh_graphs() + await pilot.pause() + mem = graphs.query_one("#worker-memory-graph") + assert "RAM" in mem.border_title + assert "4GB/8GB" in mem.border_title + assert "50%" in mem.border_title + + async def test_toggle_view_shows_worker_widgets(self): + """Pressing 'v' shows the worker view and hides the queue view.""" + app = InspectorApp(backend=default_task_backend, auto_refresh=False) + async with app.run_test() as pilot: + await pilot.pause() + assert app.worker_view_enabled is False + assert app.query_one("#queue-view").display + assert not app.query_one("#worker-view").display + await pilot.press("v") + await pilot.pause() + assert app.worker_view_enabled is True + assert not app.query_one("#queue-view").display + assert app.query_one("#worker-view").display + + async def test_toggle_view_back_to_queue(self): + """Pressing 'v' again returns to the queue view.""" + app = InspectorApp(backend=default_task_backend, auto_refresh=False) + async with app.run_test() as pilot: + await pilot.pause() + await pilot.press("v") + await pilot.pause() + assert app.worker_view_enabled is True + await pilot.press("v") + await pilot.pause() + assert app.worker_view_enabled is False + assert app.query_one("#queue-list", QueueList).display + + async def test_toggle_binding_label_updates(self): + """The 'v' binding description changes with the current view.""" + from textual.binding import Binding + + app = InspectorApp(backend=default_task_backend, auto_refresh=False) + async with app.run_test() as pilot: + await pilot.pause() + assert app.worker_view_enabled is False + binding = app._bindings.key_to_bindings["v"][0] + assert isinstance(binding, Binding) + assert binding.description == "Toggle Worker View" + await pilot.press("v") + await pilot.pause() + assert app.worker_view_enabled is True + binding = app._bindings.key_to_bindings["v"][0] + assert binding.description == "Toggle Queue View" + + async def test_worker_graphs_update_on_selection(self): + """Selecting a node feeds the worker graphs.""" + app = InspectorApp(backend=default_task_backend, auto_refresh=False) + async with app.run_test() as pilot: + await pilot.pause() + tree = app.query_one("#selection-tree", SelectionTree) + tree.queue_telemetry = _queue_telemetry("default") + graphs = app.query_one("#worker-graphs", WorkerGraphs) + snapshot = _make_worker_telemetry() + tree.worker_telemetry = snapshot + await pilot.pause() + # Select the node + root = tree.root + queue_node = root.children[0] + host_node = queue_node.children[0] + tree.select_node(host_node) + tree.action_select_cursor() + await pilot.pause() + assert graphs.selection is not None + assert graphs.selection.kind == "node" + assert graphs.selection.hostname == "node-1" + + async def test_worker_graphs_append_history_on_telemetry(self): + """Worker graphs accumulate data points when a sample tick fires.""" + app = InspectorApp(backend=default_task_backend, auto_refresh=False) + async with app.run_test() as pilot: + await pilot.pause() + graphs = app.query_one("#worker-graphs", WorkerGraphs) + snapshot = _make_worker_telemetry() + # Set selection first + node_data = WorkerTreeNode( + kind="node", + label="node-1", + hostname="node-1", + ) + graphs.selection = node_data + await pilot.pause() + # History is pre-filled with zeros at the fixed window size. + assert len(graphs._cpu_history) == graphs.GRAPH_HISTORY_SIZE + graphs.worker_telemetry = snapshot + await pilot.pause() + # A telemetry update alone does not append a sample — only the + # fixed-interval timer does. Simulate a timer tick. + graphs._refresh_graphs() + await pilot.pause() + # The last entry is the new sample; length stays fixed. + assert len(graphs._cpu_history) == graphs.GRAPH_HISTORY_SIZE + assert graphs._cpu_history[-1] == 45.0 + + async def test_refresh_telemetry_fetches_worker_telemetry(self): + """_refresh_telemetry populates the worker_telemetry reactive.""" + app = InspectorApp(backend=default_task_backend, auto_refresh=False) + async with app.run_test() as pilot: + await pilot.pause() + app._refresh_telemetry() + await pilot.pause() + # The default backend returns an empty snapshot + assert app.worker_telemetry is not None + assert app.worker_telemetry.nodes == {} + + async def test_selection_tree_preserves_cursor_on_update(self): + """The tree restores the cursor to the same node after a telemetry update.""" + app = InspectorApp(backend=default_task_backend, auto_refresh=False) + async with app.run_test() as pilot: + await pilot.pause() + tree = app.query_one("#selection-tree", SelectionTree) + tree.queue_telemetry = _queue_telemetry("default") + tree.worker_telemetry = _make_worker_telemetry() + await pilot.pause() + # Move cursor to the node + root = tree.root + queue_node = root.children[0] + host_node = queue_node.children[0] + tree.select_node(host_node) + await pilot.pause() + # Send a new telemetry snapshot + tree.worker_telemetry = _make_worker_telemetry() + await pilot.pause() + assert tree.cursor_node is not None + assert tree.cursor_node.data.kind == "node" + + async def test_worker_graphs_show_node_selection(self): + """Selecting a node shows node-level metrics in the graphs.""" + app = InspectorApp(backend=default_task_backend, auto_refresh=False) + async with app.run_test() as pilot: + await pilot.pause() + tree = app.query_one("#selection-tree", SelectionTree) + tree.queue_telemetry = _queue_telemetry("default") + graphs = app.query_one("#worker-graphs", WorkerGraphs) + snapshot = _make_worker_telemetry() + tree.worker_telemetry = snapshot + await pilot.pause() + root = tree.root + queue_node = root.children[0] + host_node = queue_node.children[0] + tree.select_node(host_node) + await pilot.pause() + graphs.selection = host_node.data + graphs.worker_telemetry = snapshot + await pilot.pause() + graphs._refresh_graphs() + await pilot.pause() + assert len(graphs._cpu_history) == graphs.GRAPH_HISTORY_SIZE + assert graphs._cpu_history[-1] == 45.0 + + async def test_selection_tree_resets_cursor_when_node_gone(self): + """_find_node_by_data returns None when the previous node is gone.""" + app = InspectorApp(backend=default_task_backend, auto_refresh=False) + async with app.run_test() as pilot: + await pilot.pause() + tree = app.query_one("#selection-tree", SelectionTree) + tree.queue_telemetry = _queue_telemetry("default") + tree.worker_telemetry = _make_worker_telemetry() + await pilot.pause() + # Capture the node data, then build a different snapshot + root = tree.root + queue_node = root.children[0] + host_node = queue_node.children[0] + old_data = host_node.data + # New snapshot with a different hostname + tree.worker_telemetry = _make_worker_telemetry(hostname="node-2") + await pilot.pause() + # The old node data is not in the new tree + result = SelectionTree._find_node_by_data(tree.root, old_data) + assert result is None + + async def test_worker_graphs_reset_on_selection_change(self): + """Changing the selection resets the graph histories.""" + app = InspectorApp(backend=default_task_backend, auto_refresh=False) + async with app.run_test() as pilot: + await pilot.pause() + graphs = app.query_one("#worker-graphs", WorkerGraphs) + node_data = WorkerTreeNode(kind="node", label="node-1", hostname="node-1") + graphs.selection = node_data + await pilot.pause() + graphs.worker_telemetry = _make_worker_telemetry() + await pilot.pause() + assert len(graphs._cpu_history) == graphs.GRAPH_HISTORY_SIZE + # Change selection to a different node — watch_selection resets + # histories, then _refresh_graphs appends from the current telemetry. + graphs.selection = WorkerTreeNode( + kind="node", + label="node-2", + hostname="node-2", + ) + await pilot.pause() + # After reset, histories are pre-filled with zeros; node-2 is not + # in telemetry, so no new sample is appended. + assert len(graphs._cpu_history) == graphs.GRAPH_HISTORY_SIZE + assert graphs._cpu_history[-1] == 0.0 + + async def test_worker_graphs_ignore_missing_node(self): + """Graphs do nothing when the selected node is not in the telemetry.""" + app = InspectorApp(backend=default_task_backend, auto_refresh=False) + async with app.run_test() as pilot: + await pilot.pause() + graphs = app.query_one("#worker-graphs", WorkerGraphs) + graphs.selection = WorkerTreeNode( + kind="node", label="ghost", hostname="ghost" + ) + await pilot.pause() + graphs.worker_telemetry = _make_worker_telemetry() + await pilot.pause() + assert len(graphs._cpu_history) == graphs.GRAPH_HISTORY_SIZE + assert graphs._cpu_history[-1] == 0.0 + + async def test_refresh_telemetry_logs_on_worker_telemetry_error(self): + """_refresh_telemetry logs when worker_telemetry raises.""" + backend = FailingBackend(alias="default", params={}) + app = InspectorApp(backend=backend, auto_refresh=False) + async with app.run_test() as pilot: + await pilot.pause() + # The exception is caught and logged; the app does not crash. + assert app.worker_telemetry is not None + + async def test_selection_tree_skips_queue_with_missing_node(self): + """The tree skips a queue whose hostname is not in the nodes dict.""" + app = InspectorApp(backend=default_task_backend, auto_refresh=False) + async with app.run_test() as pilot: + await pilot.pause() + tree = app.query_one("#selection-tree", SelectionTree) + tree.queue_telemetry = _queue_telemetry("default") + now = datetime.datetime.now(tz=datetime.UTC) + snapshot = WorkerTelemetry( + nodes={}, + queues={"default": ("ghost-host",)}, + sampled_at=now, + ) + tree.worker_telemetry = snapshot + await pilot.pause() + root = tree.root + assert len(root.children) == 1 + queue_node = root.children[0] + assert len(queue_node.children) == 0 + + async def test_worker_graphs_ignore_non_node_selection(self): + """Graphs skip when the selection kind is not 'node'.""" + app = InspectorApp(backend=default_task_backend, auto_refresh=False) + async with app.run_test() as pilot: + await pilot.pause() + graphs = app.query_one("#worker-graphs", WorkerGraphs) + # Use a valid hostname but a queue kind so _refresh_graphs returns early + graphs.selection = WorkerTreeNode( + kind="queue", + label="default", + queue_name="default", + hostname="node-1", + ) + await pilot.pause() + graphs.worker_telemetry = _make_worker_telemetry() + await pilot.pause() + assert len(graphs._cpu_history) == graphs.GRAPH_HISTORY_SIZE + assert graphs._cpu_history[-1] == 0.0 + + async def test_worker_graphs_handles_unmounted_widgets(self): + """_refresh_graphs logs when Sparkline widgets are not yet mounted.""" + from unittest.mock import patch + + app = InspectorApp(backend=default_task_backend, auto_refresh=False) + async with app.run_test() as pilot: + await pilot.pause() + graphs = app.query_one("#worker-graphs", WorkerGraphs) + graphs.selection = WorkerTreeNode( + kind="node", label="node-1", hostname="node-1" + ) + await pilot.pause() + # Patch query_one to raise as if widgets are not mounted. + with patch.object(graphs, "query_one", side_effect=Exception("no widget")): + graphs.worker_telemetry = _make_worker_telemetry() + await pilot.pause() + # Histories are populated even though the graph redraw failed. + assert len(graphs._cpu_history) == graphs.GRAPH_HISTORY_SIZE + + +class _StubBackend(ThreadmillTaskBackend): + """Backend double that yields a fixed worker telemetry snapshot.""" + + def enqueue(self, task, args, kwargs): + raise NotImplementedError + + def peek(self, *args, **kwargs): + raise NotImplementedError + + def telemetry(self, *args, **kwargs): + raise NotImplementedError + + def worker_telemetry(self): + raise NotImplementedError + + async def subscribe_worker_telemetry(self): + yield _make_worker_telemetry() + + +class _ExplodingBackend(_StubBackend): + """Backend whose pub/sub subscription raises immediately.""" + + async def subscribe_worker_telemetry(self): + raise RuntimeError("subscribe failed") + yield # pragma: no cover + + +def _make_node( + *, + hostname: str = "node-1", + pid: int = 100, + queue_name: str = "default", + thread_count: int = 2, + sampled_at: datetime.datetime | None = None, +) -> NodeTelemetry: + """Build a per-process NodeTelemetry for cache-population tests.""" + now = sampled_at or datetime.datetime.now(tz=datetime.UTC) + return NodeTelemetry( + hostname=hostname, + pid=pid, + queues=(queue_name,), + process_count=1, + thread_count=thread_count, + cpu_percent=45.0, + memory_bytes=8_000_000_000, + memory_total=16_000_000_000, + tasks_per_minute=30.0, + sampled_at=now, + ) + + +class TestWorkerNodeCache: + """Tests for the inspector's latest-per-host node cache and rebuild.""" + + async def test_subscribe_caches_node_by_hostname(self): + """A subscribed snapshot is cached keyed by hostname.""" + app = InspectorApp( + backend=_StubBackend(alias="default", params={}), auto_refresh=False + ) + await app._subscribe_worker_telemetry() + assert "node-1" in app._node_cache + assert app._node_cache["node-1"].hostname == "node-1" + + def test_latest_snapshot_wins_for_same_host(self): + """Two snapshots from the same host keep only the most recent one.""" + app = InspectorApp( + backend=_StubBackend(alias="default", params={}), auto_refresh=False + ) + app._node_cache["node-1"] = _make_node(pid=100, queue_name="default") + app._node_cache["node-1"] = _make_node(pid=200, queue_name="celery") + assert len(app._node_cache) == 1 + assert app._node_cache["node-1"].pid == 200 + + def test_prune_drops_stale_nodes(self): + """_prune_nodes drops entries older than the TTL.""" + app = InspectorApp( + backend=_StubBackend(alias="default", params={}), auto_refresh=False + ) + old = datetime.datetime.now(tz=datetime.UTC) - datetime.timedelta(seconds=20) + app._node_cache["node-1"] = dataclasses.replace(_make_node(), sampled_at=old) + app._prune_nodes(10) + assert "node-1" not in app._node_cache + + def test_prune_keeps_fresh_nodes(self): + """_prune_nodes retains entries within the TTL window.""" + app = InspectorApp( + backend=_StubBackend(alias="default", params={}), auto_refresh=False + ) + app._node_cache["node-1"] = _make_node() + app._prune_nodes(10) + assert "node-1" in app._node_cache + + def test_build_includes_nodes_and_queue_index(self): + """_build_worker_telemetry returns cached nodes plus a queue->host index.""" + app = InspectorApp( + backend=_StubBackend(alias="default", params={}), auto_refresh=False + ) + app._node_cache["node-1"] = _make_node(pid=100, queue_name="default") + app._node_cache["node-2"] = _make_node( + hostname="node-2", pid=200, queue_name="celery" + ) + snapshot = app._build_worker_telemetry() + assert set(snapshot.nodes) == {"node-1", "node-2"} + assert set(snapshot.queues) == {"default", "celery"} + assert snapshot.queues["default"] == ("node-1",) + assert snapshot.queues["celery"] == ("node-2",) + + def test_build_empty_cache_returns_empty(self): + """_build_worker_telemetry returns empty nodes/queues from an empty cache.""" + app = InspectorApp( + backend=_StubBackend(alias="default", params={}), auto_refresh=False + ) + snapshot = app._build_worker_telemetry() + assert snapshot.nodes == {} + assert snapshot.queues == {} + + async def test_refresh_rebuilds_worker_telemetry_from_cache(self): + """_refresh_telemetry rebuilds worker_telemetry when the cache is populated.""" + app = InspectorApp(backend=default_task_backend, auto_refresh=False) + async with app.run_test() as pilot: + await pilot.pause() + app._node_cache["node-1"] = _make_node() + app._refresh_telemetry() + await pilot.pause() + assert app.worker_telemetry is not None + assert "node-1" in app.worker_telemetry.nodes + + +class TestWorkerTelemetrySubscription: + """Tests for the inspector's pub/sub subscription lifecycle.""" + + async def test_subscribe_worker_telemetry_caches_snapshot(self): + """_subscribe_worker_telemetry caches the snapshot per host.""" + backend = _StubBackend(alias="default", params={}) + app = InspectorApp(backend=backend, auto_refresh=False) + await app._subscribe_worker_telemetry() + assert app._node_cache + assert "node-1" in app._node_cache + + async def test_subscribe_worker_telemetry_logs_on_error(self): + """A failing subscription is caught and logged without crashing.""" + backend = _ExplodingBackend(alias="default", params={}) + app = InspectorApp(backend=backend, auto_refresh=False) + await app._subscribe_worker_telemetry() + assert not app._node_cache diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py new file mode 100644 index 0000000..4c20429 --- /dev/null +++ b/tests/test_telemetry.py @@ -0,0 +1,280 @@ +"""Tests for worker telemetry data types and the psutil sampler.""" + +from __future__ import annotations + +import datetime +import logging +import multiprocessing +from unittest.mock import MagicMock, patch + +import pytest + +from threadmill.backends.base import ( + NodeTelemetry, + ThreadmillTaskBackend, + WorkerTelemetry, +) +from threadmill.telemetry import TelemetrySampler + +utc = datetime.UTC + + +def _make_worker_telemetry( + *, + hostname: str = "node-1", + pid: int = 1234, + sampled_at: datetime.datetime | None = None, +) -> WorkerTelemetry: + """Build a minimal WorkerTelemetry snapshot for tests.""" + sampled_at = sampled_at or datetime.datetime.now(tz=utc) + node = NodeTelemetry( + hostname=hostname, + pid=pid, + queues=("default",), + process_count=1, + thread_count=2, + cpu_percent=45.0, + memory_bytes=4_000_000_000, + memory_total=8_000_000_000, + tasks_per_minute=30.0, + sampled_at=sampled_at, + ) + return WorkerTelemetry( + nodes={hostname: node}, + queues={"default": (hostname,)}, + sampled_at=sampled_at, + ) + + +class TestWorkerTelemetryDataTypes: + """Tests for the telemetry dataclasses defined in backends/base.py.""" + + def test_node_telemetry_fields(self): + """NodeTelemetry stores the publishing pid and node-level metrics.""" + now = datetime.datetime.now(tz=utc) + node = NodeTelemetry( + hostname="host", + pid=100, + queues=("default",), + process_count=1, + thread_count=1, + cpu_percent=20.0, + memory_bytes=4_000_000_000, + memory_total=8_000_000_000, + tasks_per_minute=0.0, + sampled_at=now, + ) + assert node.hostname == "host" + assert node.pid == 100 + assert node.process_count == 1 + assert node.thread_count == 1 + assert node.cpu_percent == 20.0 + assert node.memory_bytes == 4_000_000_000 + assert node.memory_total == 8_000_000_000 + + def test_worker_telemetry_queues_and_nodes(self): + """WorkerTelemetry maps queues to hostnames and nodes to NodeTelemetry.""" + now = datetime.datetime.now(tz=utc) + node = NodeTelemetry( + hostname="h", + pid=1, + queues=("default",), + process_count=1, + thread_count=1, + cpu_percent=0.0, + memory_bytes=0, + memory_total=0, + tasks_per_minute=0.0, + sampled_at=now, + ) + telemetry = WorkerTelemetry( + nodes={"h": node}, + queues={"default": ("h",)}, + sampled_at=now, + ) + assert "h" in telemetry.nodes + assert telemetry.queues["default"] == ("h",) + assert telemetry.sampled_at == now + + +class TestBaseBackendHooks: + """Tests for the default backend hooks.""" + + class _ConcreteBackend(ThreadmillTaskBackend): + """Minimal concrete backend for hook testing.""" + + def enqueue(self, task, args, kwargs): + raise NotImplementedError + + def test_publish_worker_telemetry_is_noop(self): + """The base publish hook does nothing and does not raise.""" + backend = self._ConcreteBackend(alias="test", params={}) + backend.publish_worker_telemetry(_make_worker_telemetry()) + + def test_worker_telemetry_returns_empty_snapshot(self): + """The base worker_telemetry hook returns an empty but valid snapshot.""" + backend = self._ConcreteBackend(alias="test", params={}) + snapshot = backend.worker_telemetry() + assert snapshot.nodes == {} + assert snapshot.queues == {} + assert snapshot.sampled_at is not None + + +class TestTelemetrySampler: + """Tests for the TelemetrySampler with stubbed psutil.""" + + def _make_executor(self): + """Build a fake executor with the attributes the sampler reads.""" + executor = MagicMock() + executor.queues = ("default",) + executor.process_count = 2 + executor.thread_count = 1 + executor.total_counter = multiprocessing.Value("q", 5) + return executor + + def test_is_available_when_psutil_imported(self): + """is_available is True when psutil is importable.""" + sampler = TelemetrySampler( + backend=MagicMock(), + executor=self._make_executor(), + ) + # psutil is installed in the test environment + assert sampler.is_available is True + + def test_sample_returns_worker_telemetry(self): + """sample() builds a WorkerTelemetry with node data and the executor pid.""" + fake_vmem = MagicMock() + fake_vmem.used = 4_400_000_000 + fake_vmem.total = 8_000_000_000 + + sampler = TelemetrySampler( + backend=MagicMock(), + executor=self._make_executor(), + ) + + with ( + patch("threadmill.telemetry.psutil") as fake_psutil, + patch("threadmill.telemetry.socket.gethostname", return_value="testhost"), + patch("threadmill.telemetry.os.getpid", return_value=100), + ): + fake_psutil.cpu_percent.return_value = 42.0 + fake_psutil.virtual_memory.return_value = fake_vmem + + snapshot = sampler.sample() + + assert "testhost" in snapshot.nodes + node = snapshot.nodes["testhost"] + assert node.hostname == "testhost" + assert node.pid == 100 + assert node.cpu_percent == 42.0 + assert node.memory_bytes == 4_400_000_000 + assert node.memory_total == 8_000_000_000 + assert node.process_count == 2 + assert node.thread_count == 2 # process_count * executor.thread_count + assert "default" in snapshot.queues + assert "testhost" in snapshot.queues["default"] + + def test_tasks_per_minute_first_call_returns_zero(self): + """The first throughput sample returns 0 (no previous baseline).""" + sampler = TelemetrySampler( + backend=MagicMock(), + executor=self._make_executor(), + clock=lambda: 100.0, + ) + assert sampler._tasks_per_minute(10, 100.0) == 0.0 + + def test_tasks_per_minute_computes_delta(self): + """Throughput is (delta_tasks / elapsed_seconds) * 60.""" + sampler = TelemetrySampler( + backend=MagicMock(), + executor=self._make_executor(), + clock=lambda: 100.0, + ) + sampler._tasks_per_minute(10, 100.0) + rate = sampler._tasks_per_minute(25, 130.0) + assert rate == pytest.approx(30.0) + + def test_tasks_per_minute_zero_elapsed_returns_zero(self): + """When elapsed time is zero, throughput returns 0 to avoid division by zero.""" + sampler = TelemetrySampler( + backend=MagicMock(), + executor=self._make_executor(), + clock=lambda: 100.0, + ) + sampler._tasks_per_minute(10, 100.0) + rate = sampler._tasks_per_minute(20, 100.0) + assert rate == 0.0 + + def test_start_without_psutil_is_noop(self): + """start() does nothing when psutil is not available.""" + sampler = TelemetrySampler( + backend=MagicMock(), + executor=self._make_executor(), + ) + with patch("threadmill.telemetry.psutil", None): + sampler.start() + assert sampler._thread is None + + def test_stop_without_start_is_noop(self): + """stop() does not raise when the sampler was never started.""" + sampler = TelemetrySampler( + backend=MagicMock(), + executor=self._make_executor(), + ) + sampler.stop() + + def test_publish_one_sample_calls_backend(self): + """_publish_one_sample calls backend.publish_worker_telemetry.""" + backend = MagicMock() + sampler = TelemetrySampler( + backend=backend, + executor=self._make_executor(), + ) + + fake_vmem = MagicMock() + fake_vmem.used = 1_600_000_000 + fake_vmem.total = 4_000_000_000 + + with ( + patch("threadmill.telemetry.psutil") as fake_psutil, + patch("threadmill.telemetry.socket.gethostname", return_value="h"), + patch("threadmill.telemetry.os.getpid", return_value=100), + ): + fake_psutil.cpu_percent.return_value = 10.0 + fake_psutil.virtual_memory.return_value = fake_vmem + sampler._publish_one_sample() + + backend.publish_worker_telemetry.assert_called_once() + snapshot = backend.publish_worker_telemetry.call_args.args[0] + assert isinstance(snapshot, WorkerTelemetry) + assert "h" in snapshot.nodes + + def test_run_logs_exception_on_sample_failure(self, caplog): + """_run logs an exception when _publish_one_sample raises.""" + backend = MagicMock() + backend.publish_worker_telemetry.side_effect = RuntimeError("boom") + sampler = TelemetrySampler( + backend=backend, + executor=self._make_executor(), + interval_seconds=0.01, + ) + + fake_vmem = MagicMock() + fake_vmem.used = 1_600_000_000 + fake_vmem.total = 4_000_000_000 + + with ( + patch("threadmill.telemetry.psutil") as fake_psutil, + patch("threadmill.telemetry.socket.gethostname", return_value="h"), + patch("threadmill.telemetry.os.getpid", return_value=100), + ): + fake_psutil.cpu_percent.return_value = 10.0 + fake_psutil.virtual_memory.return_value = fake_vmem + with caplog.at_level(logging.ERROR): + sampler.start() + import time as _time + + _time.sleep(0.1) + sampler.stop() + + assert "Failed to sample worker telemetry" in caplog.text diff --git a/threadmill/backends/base.py b/threadmill/backends/base.py index 14e9195..dea5e58 100644 --- a/threadmill/backends/base.py +++ b/threadmill/backends/base.py @@ -5,6 +5,7 @@ import datetime import json import threading +import typing from abc import ABC import django @@ -49,7 +50,7 @@ class QueueCounts: @dataclasses.dataclass(kw_only=True, slots=True) class QueueRates: - """Rolling ingress/egress throughput over a time window (time-series data).""" + """Rolling ingress/egress throughput over a time window.""" interval: datetime.timedelta ingress: int @@ -71,6 +72,45 @@ class BackendTelemetry: queues: dict[str, QueueStats] +@dataclasses.dataclass(kw_only=True, slots=True) +class NodeTelemetry: + """Telemetry for a single node (host) running a worker pool. + + The executor's telemetry sampler publishes one snapshot per node with + system-level CPU/memory sampled via `psutil` and the executor's + process/thread counts. `memory_bytes` is physical RAM in use and + `memory_total` is the machine's total physical RAM; swap is tracked + separately by `psutil.swap_memory()` and is deliberately excluded. + The inspector derives the usage percentage from these two values. + """ + + hostname: str + pid: int + queues: tuple[str, ...] + process_count: int + thread_count: int + cpu_percent: float + memory_bytes: int + memory_total: int + tasks_per_minute: float + sampled_at: datetime.datetime + + +@dataclasses.dataclass(kw_only=True, slots=True) +class WorkerTelemetry: + """Snapshot of worker pool health across a backend's queues and nodes. + + `nodes` is keyed by hostname and carries system-level CPU/mem plus + per-process counters. `queues` maps each queue name to the hostnames + listening on it, so the inspector can render a Queue -> Node selection + tree. + """ + + nodes: dict[str, NodeTelemetry] + queues: dict[str, tuple[str, ...]] + sampled_at: datetime.datetime + + class Broker(threading.Thread): """Backend maintenance thread launched by the task executor.""" @@ -200,7 +240,7 @@ def peek( count: int = 1, ) -> collections.abc.Generator[TaskResult, None, None]: """ - Yield up to ``count`` tasks from a queue in the given status segment. + Yield up to `count` tasks from a queue in the given status segment. Args: queue_name: The name of the queue to peek into. @@ -218,3 +258,36 @@ def telemetry( interval: The time window for rolling rates. """ raise NotImplementedError + + def publish_worker_telemetry(self, telemetry: WorkerTelemetry) -> None: + """Publish a worker-pool telemetry snapshot to a shared store. + + Backends without a shared store implement this as a no-op so the + worker command still runs without worker telemetry. + """ + + def worker_telemetry(self) -> WorkerTelemetry: + """Return the latest worker-pool telemetry snapshot, or an empty one. + + Backends without a shared store return an empty snapshot so the + inspector can render the worker view without raising. + """ + return WorkerTelemetry( + nodes={}, + queues={}, + sampled_at=datetime.datetime.now(tz=datetime.UTC), + ) + + async def subscribe_worker_telemetry(self) -> typing.AsyncIterator[WorkerTelemetry]: + """Yield worker telemetry snapshots as they arrive. + + Backends with pub/sub capability override this to maintain a + persistent subscription. The default implementation yields a + single empty snapshot and stops, so non-Redis backends still + compile and the inspector can fall back to polling. + """ + yield WorkerTelemetry( + nodes={}, + queues={}, + sampled_at=datetime.datetime.now(tz=datetime.UTC), + ) diff --git a/threadmill/backends/redis.py b/threadmill/backends/redis.py index 23fb0c9..b902b71 100644 --- a/threadmill/backends/redis.py +++ b/threadmill/backends/redis.py @@ -2,10 +2,13 @@ from __future__ import annotations +import dataclasses import datetime +import json import logging import queue import time +import typing import uuid from collections.abc import Generator, Sequence from pathlib import Path @@ -19,10 +22,12 @@ from threadmill.backends.base import ( BackendTelemetry, Broker, + NodeTelemetry, QueueCounts, QueueRates, QueueStats, ThreadmillTaskBackend, + WorkerTelemetry, ) logger = logging.getLogger(__name__) @@ -31,7 +36,7 @@ def _load_lua(name: str) -> str: - """Load a Lua script from the lua directory.""" + """Load a Lua script by name.""" return (_LUA_DIR / f"{name}.lua").read_text() @@ -111,6 +116,10 @@ def main(self) -> None: logger.exception("Telemetry trim error for queue %r", queue_name) +WORKER_TELEMETRY_CHANNEL = "{prefix}:worker-telemetry" +WORKER_TELEMETRY_TTL = 10 # seconds before stale entries are pruned + + class RedisTaskBackend(ThreadmillTaskBackend): """Redis-backed durable priority queue backend. @@ -149,6 +158,7 @@ def __init__(self, alias: str, params: dict) -> None: raise ValueError( f"REDIS_URL must be specified in your settings for the {type(self).__name__}." ) from e + self.redis_url = redis_url self.client = redis.from_url(redis_url) self.key_prefix = f"threadmill:{{{alias}}}" self.lease_ttl = self.options.get("lease_ttl", datetime.timedelta(hours=1)) @@ -464,6 +474,69 @@ def telemetry( ) return BackendTelemetry(queues=queues) + def publish_worker_telemetry(self, telemetry: WorkerTelemetry) -> None: + channel = WORKER_TELEMETRY_CHANNEL.format(prefix=self.key_prefix) + self.client.publish(channel, _serialize_worker_telemetry(telemetry)) + + async def subscribe_worker_telemetry(self) -> typing.AsyncIterator[WorkerTelemetry]: + import redis.asyncio as aioredis + + redis_url = self.redis_url + channel = WORKER_TELEMETRY_CHANNEL.format(prefix=self.key_prefix) + async with aioredis.from_url(redis_url) as client: + pubsub = client.pubsub() + await pubsub.subscribe(channel) + try: + while True: + message = await pubsub.get_message( + timeout=1.0, ignore_subscribe_messages=True + ) + if message is None: + continue + payload = message["data"] + if isinstance(payload, bytes): + payload = payload.decode() + yield _deserialize_worker_telemetry(payload) + finally: + await pubsub.unsubscribe(channel) + await pubsub.aclose() + def close(self) -> None: """Close the Redis connection.""" self.client.close() + + +def _serialize_worker_telemetry(telemetry: WorkerTelemetry) -> str: + """Serialize a worker telemetry snapshot to JSON.""" + return json.dumps(dataclasses.asdict(telemetry), default=lambda o: o.isoformat()) + + +def _deserialize_worker_telemetry(payload: str) -> WorkerTelemetry: + """Deserialize a JSON payload into a worker telemetry snapshot.""" + data = json.loads(payload) + nodes = { + hostname: _deserialize_node(node_data) + for hostname, node_data in data.get("nodes", {}).items() + } + queues = { + queue_name: tuple(hostnames) + for queue_name, hostnames in data.get("queues", {}).items() + } + sampled_at = datetime.datetime.fromisoformat(data["sampled_at"]) + return WorkerTelemetry(nodes=nodes, queues=queues, sampled_at=sampled_at) + + +def _deserialize_node(data: dict) -> NodeTelemetry: + """Deserialize a JSON dict into a node telemetry snapshot.""" + return NodeTelemetry( + hostname=data["hostname"], + pid=data["pid"], + queues=tuple(data.get("queues", [])), + process_count=data["process_count"], + thread_count=data["thread_count"], + cpu_percent=data["cpu_percent"], + memory_bytes=data["memory_bytes"], + memory_total=data["memory_total"], + tasks_per_minute=data["tasks_per_minute"], + sampled_at=datetime.datetime.fromisoformat(data["sampled_at"]), + ) diff --git a/threadmill/executor.py b/threadmill/executor.py index 3531e22..828ca91 100644 --- a/threadmill/executor.py +++ b/threadmill/executor.py @@ -23,6 +23,8 @@ from django.utils import timezone from django.utils.json import normalize_json +from .telemetry import TelemetrySampler + if typing.TYPE_CHECKING: from .backends.base import Broker, ThreadmillTaskBackend @@ -38,7 +40,7 @@ @dataclasses.dataclass(kw_only=True, slots=True) class TaskExecutor: - """Tasks consumed from shared joinable queues via process and thread pools.""" + """Pool of processes and threads consuming tasks from shared joinable queues.""" backend: ThreadmillTaskBackend workers: int | None = None @@ -54,11 +56,17 @@ class TaskExecutor: queues: tuple[str] broker: Broker | None = dataclasses.field(default=None, init=False) exit_empty: bool = False + total_counter: multiprocessing.Value = dataclasses.field(init=False) + telemetry_sampler: TelemetrySampler | None = dataclasses.field( + default=None, init=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) + # Shared aggregate task counter; created before workers fork so each + # worker process inherits the same value for throughput telemetry. + self.total_counter = multiprocessing.Value("q", 0) def get_maximum_tasks_per_child(self) -> int | None: """Return worker recycling limit based on config and thread count.""" @@ -75,6 +83,7 @@ def create_worker_process(self) -> WorkerProcess: self.backend.alias, self.queues, self.exit_empty, + total_counter=self.total_counter, ) worker.start() return worker @@ -90,6 +99,8 @@ def run(self) -> None: if self.backend.broker_class: self.broker = self.backend.broker_class(self.backend) threads.append(self.broker) + self.telemetry_sampler = TelemetrySampler(self.backend, executor=self) + self.telemetry_sampler.start() for thread in threads: thread.start() @@ -101,6 +112,8 @@ def shutdown(self) -> None: logger.info("Shutting down task executor") if self.broker is not None: self.broker.shutdown() + if self.telemetry_sampler is not None: + self.telemetry_sampler.stop() with ThreadPoolExecutor(max_workers=self.process_count) as executor: executor.map(lambda worker: worker.shutdown(), self.worker_processes) self.is_publishing = False @@ -133,6 +146,8 @@ def __init__( backend_alias: str = "", queues: tuple[str, ...] = (), exit_empty: bool = False, + *, + total_counter: multiprocessing.Value, ) -> None: """Create process with dedicated thread pool for task execution.""" self.shutdown_requested = multiprocessing.Event() @@ -143,6 +158,7 @@ def __init__( self.queues = queues self.exit_empty = exit_empty self.task_count = 0 + self.total_counter = total_counter self.lock: threading.Lock | None = None self.expired: threading.Event | None = None @@ -164,7 +180,10 @@ def run(self) -> None: ) def record_task(self) -> None: - """Record one processed task and stop when max_tasks is reached.""" + """Record one processed task and recycle the process at max_tasks.""" + # Always count towards the executor's aggregate throughput counter. + with self.total_counter.get_lock(): + self.total_counter.value += 1 if self.max_tasks is None: return if self.lock is None or self.expired is None: @@ -217,7 +236,7 @@ def run(self) -> None: self.worker.record_task() def execute_task_result(self, task_result: TaskResult) -> TaskResult: - """Execute task from task result and update result lifecycle state.""" + """Execute the wrapped task and update result lifecycle state.""" logger.info("Executing task %r", task_result.id) started_at = timezone.now() task_result = dataclasses.replace( diff --git a/threadmill/inspector/app.py b/threadmill/inspector/app.py index eb6193c..5bc2717 100644 --- a/threadmill/inspector/app.py +++ b/threadmill/inspector/app.py @@ -4,342 +4,62 @@ import datetime import logging -import math -import typing -from typing import Any - -from django.contrib.humanize.templatetags.humanize import naturaltime -from django.tasks import ( - DEFAULT_TASK_QUEUE_NAME, - TaskResult, - TaskResultStatus, - task_backends, -) + +from django.tasks import task_backends from textual.app import App, ComposeResult from textual.binding import Binding from textual.containers import Horizontal, Vertical from textual.reactive import reactive -from textual.theme import BUILTIN_THEMES from textual.widgets import ( - DataTable, Footer, ListItem, ListView, Select, Static, - TabbedContent, - TabPane, + Tree, ) -from ..backends.base import BackendTelemetry, ThreadmillTaskBackend +from ..backends.base import ( + BackendTelemetry, + NodeTelemetry, + ThreadmillTaskBackend, + WorkerTelemetry, +) +from .queue_view import ( + TAB_KEYS, + QueueItem, + QueueList, + TaskDetail, + TaskList, +) +from .utils import ( + si_prefix, + supported_aliases as _supported_aliases, +) +from .worker_view import ( + SelectionTree, + WorkerGraphs, + WorkerTreeNode, +) logger = logging.getLogger(__name__) TELEMETRY_INTERVAL_SECONDS = 2.0 """Seconds between automatic queue-stat refreshes; the task list stays manual.""" -TAB_STATUSES: list[tuple[str, TaskResultStatus | None]] = [ - ("Running", TaskResultStatus.RUNNING), - ("Ready", TaskResultStatus.READY), - ("Successful", TaskResultStatus.SUCCESSFUL), - ("Failed", TaskResultStatus.FAILED), +__all__ = [ + "InspectorApp", + "QueueItem", + "QueueList", + "SelectionTree", + "TaskDetail", + "TaskList", + "WorkerGraphs", + "WorkerTreeNode", + "si_prefix", ] -class QueueItem(ListItem): - """A list item that exposes its queue name.""" - - def __init__(self, queue_name: str, label: str, **kwargs: Any) -> None: - super().__init__(Static(label), id=f"queue-{queue_name}", **kwargs) - self.queue_name = queue_name - - -def _supported_aliases() -> typing.Generator[tuple[str, str]]: - for alias in task_backends: - if isinstance(task_backends[alias], ThreadmillTaskBackend): - yield alias, alias - - -def _format_dt(value: datetime.datetime | None) -> str: - return value.isoformat() if value else "" - - -def si_prefix(n: int) -> str: - """Round and shorten a number with an SI prefix (k, M, G, etc.).""" - if n < 1000: - return str(n) - prefixes = ["", "k", "M", "G", "T", "P", "E", "Z", "Y"] - m = min(int(math.log10(abs(n)) // 3), len(prefixes) - 1) - if (value := n / 1000**m) and value >= 100: - return f"{int(value)}{prefixes[m]}" - return f"{value:.1f}".rstrip("0").rstrip(".") + prefixes[m] - - -COLUMN_LABELS: dict[str, str] = { - "id": "ID", - "priority": "Priority", - "function": "Function", - "enqueued": "Enqueued", - "started": "Started", - "finished": "Finished", - "workers": "Workers", -} - -TAB_COLUMNS: dict[str, tuple[str, ...]] = { - "running": ("id", "function", "enqueued", "started", "workers"), - "ready": ("id", "function", "priority", "enqueued"), - "successful": ("id", "function", "enqueued", "started", "finished", "workers"), - "failed": ("id", "function", "enqueued", "started", "finished", "workers"), -} - -TAB_KEYS = { - label.lower(): str(index + 1) for index, (label, _) in enumerate(TAB_STATUSES) -} - - -def _cell(result: TaskResult, column: str) -> str: - match column: - case "id": - return result.id[:8] - case "function": - return result.task.module_path - case "priority": - return str(result.task.priority) - case "enqueued": - return naturaltime(result.enqueued_at) - case "started": - return naturaltime(result.started_at) - case "finished": - return naturaltime(result.finished_at) - case "workers": - return ", ".join(result.worker_ids) or "-" - - -class TaskDetail(Static): - """Read-only detail view for the selected task result.""" - - task_result: reactive[TaskResult | None] = reactive(None) - - def compose(self) -> ComposeResult: - yield from super().compose() - self.border_title = "Detail" - - def watch_task_result(self) -> None: - self.update(self._detail_content()) - - def _detail_content(self) -> str: - """Build the detail text for the current task.""" - result = self.task_result - if result is None: - return "Select a task to view details." - lines = [ - f"[b]Task:[/b] {result.task.module_path}", - f"[b]ID:[/b] {result.id}", - f"[b]Status:[/b] {result.status.name}", - f"[b]Queue:[/b] {result.task.queue_name}", - f"[b]Priority:[/b] {result.task.priority}", - f"[b]Enqueued:[/b] {_format_dt(result.enqueued_at)}", - f"[b]Started:[/b] {_format_dt(result.started_at)}", - f"[b]Finished:[/b] {_format_dt(result.finished_at)}", - f"[b]Last attempted:[/b] {_format_dt(result.last_attempted_at)}", - f"[b]Worker IDs:[/b] {', '.join(result.worker_ids) or '-'}", - f"[b]Args:[/b] {result.args!r}", - f"[b]Kwargs:[/b] {result.kwargs!r}", - ] - if result.errors: - lines.append("") - lines.append("[b]Errors:[/b]") - for error in result.errors: - lines.append(f" {error.exception_class_path}") - for line in error.traceback.splitlines(): - lines.append(f" {line}") - return "\n".join(lines) - - -class TaskList(Vertical): - """Tabbed list of tasks for the selected queue.""" - - backend: reactive[ThreadmillTaskBackend | None] = reactive(None) - queue_name: reactive[str] = reactive("") - telemetry: reactive[BackendTelemetry | None] = reactive(None) - counts: reactive[dict[str, int]] = reactive({}) - selected_task: reactive[TaskResult | None] = reactive(None) - - def compose(self) -> ComposeResult: - with TabbedContent(initial="tab-ready") as tabs: - for label, _ in TAB_STATUSES: - with TabPane(label, id=f"tab-{label.lower()}"): - yield DataTable(id=f"table-{label.lower()}", cursor_type="row") - self._tabs = tabs - self.border_title = "Tasks" - - def watch_backend(self) -> None: - self._refresh_data() - - def watch_queue_name(self) -> None: - self._refresh_data() - - def refresh_tasks(self) -> None: - """Re-fetch and render tasks for the current queue and tab.""" - self._refresh_data() - - def compute_counts(self) -> dict[str, int]: - if self.telemetry is None or not self.queue_name: - return {label.lower(): 0 for label, _ in TAB_STATUSES} - counts = self.telemetry.queues[self.queue_name].counts - return { - label.lower(): getattr(counts, label.lower()) for label, _ in TAB_STATUSES - } - - def watch_counts(self, counts: dict[str, int]) -> None: - for label, _status in TAB_STATUSES: - tab = self._tabs.get_tab(f"tab-{label.lower()}") - tab.label = f"{label} ({si_prefix(counts.get(label.lower(), 0))})" - - def watch_selected_task(self, task: TaskResult | None) -> None: - if self.app._task_detail is not None: - self.app._task_detail.task_result = task - - def on_mount(self) -> None: - """Configure per-status columns on first mount.""" - for label, _ in TAB_STATUSES: - tab_id = label.lower() - table = self.query_one(f"#table-{tab_id}", DataTable) - table.add_columns( - *(COLUMN_LABELS[column] for column in TAB_COLUMNS[tab_id]) - ) - table.disabled = tab_id != "ready" - - def switch_tab(self, tab_id: str) -> None: - """Activate the tab with the given id.""" - self._tabs.active = tab_id - - def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None: - """Notify the app when a task row is selected.""" - if isinstance(event.row_key.value, str): - self._select_task_by_id(event.row_key.value) - - def on_data_table_row_highlighted(self, event: DataTable.RowHighlighted) -> None: - """Update the detail preview when a row is highlighted.""" - if isinstance(event.row_key.value, str): - self._select_task_by_id(event.row_key.value) - - def on_tabbed_content_tab_activated( - self, _event: TabbedContent.TabActivated - ) -> None: - """Refresh the visible table when the user switches tabs.""" - self._refresh_data() - - def _select_task_by_id(self, task_id: str) -> None: - """Find the task result matching the row key in the current results.""" - self.selected_task = next( - (result for result in self._current_results if result.id == task_id), - None, - ) - - def _refresh_data(self) -> None: - """Fetch and display tasks for the current queue and active tab.""" - backend = self.backend - queue_name = self.queue_name - if backend is None or not queue_name: - self._current_results = [] - return - tab_id = self._tabs.active.removeprefix("tab-") - status = next( - (status for label, status in TAB_STATUSES if label.lower() == tab_id), - None, - ) - for label, _ in TAB_STATUSES: - table = self.query_one(f"#table-{label.lower()}", DataTable) - table.clear() - table.disabled = label.lower() != tab_id - table = self.query_one(f"#table-{tab_id}", DataTable) - table.disabled = False - try: - self._current_results = list( - backend.peek(queue_name=queue_name, status=status, count=100) - ) - except Exception: # noqa: BLE001 - logger.exception("Failed to peek tasks for %r", queue_name) - self._current_results = [] - columns = TAB_COLUMNS[tab_id] - for result in self._current_results: - table.add_row( - *(_cell(result, column) for column in columns), - key=result.id, - ) - self._restore_cursor(table) - - def _restore_cursor(self, table: DataTable) -> None: - """Keep the highlighted task selected across a refresh, falling back to the first row.""" - previous_id = self.selected_task.id if self.selected_task else None - row = next( - ( - index - for index, result in enumerate(self._current_results) - if result.id == previous_id - ), - None, - ) - if row is None: - row = 0 if self._current_results else None - if row is not None: - table.move_cursor(row=row) - self.selected_task = self._current_results[row] if row is not None else None - - -class QueueList(ListView): - """List of queues for the selected backend with ingress/egress deltas.""" - - telemetry: reactive[BackendTelemetry | None] = reactive(None) - - def __init__(self, **kwargs: Any) -> None: - super().__init__(**kwargs) - self._items: dict[str, QueueItem] = {} - - def compose(self) -> ComposeResult: - yield from super().compose() - self.border_title = "Queues" - - def watch_telemetry(self, telemetry: BackendTelemetry | None) -> None: - """Refresh queue labels when a new telemetry snapshot arrives.""" - if telemetry is not None: - self.update_telemetry(telemetry) - - def update_telemetry(self, telemetry: BackendTelemetry) -> None: - """Refresh queue labels from a new telemetry snapshot.""" - queues = telemetry.queues - for queue_name, stats in sorted(queues.items()): - theme = BUILTIN_THEMES[self.app.theme] - rates = stats.rates - label = ( - f"{queue_name:24} " - f"[{theme.success}]+{si_prefix(rates.ingress):3}[/] " - f"[{theme.error}]-{si_prefix(rates.egress):3}[/]" - ) - if queue_name in self._items: - self._items[queue_name].children[0].update(label) - else: - item = QueueItem(queue_name, label) - self._items[queue_name] = item - self.append(item) - stale = [name for name in self._items if name not in queues] - for queue_name in stale: - self._items.pop(queue_name).remove() - if self.index is None and self._items: - target = ( - DEFAULT_TASK_QUEUE_NAME - if DEFAULT_TASK_QUEUE_NAME in self._items - else next(iter(self._items)) - ) - self.index = list(self._items.keys()).index(target) - self._notify_selection(target) - - def _notify_selection(self, queue_name: str) -> None: - """Tell the app which queue to display.""" - self.app._task_list.queue_name = queue_name - - class InspectorApp(App): """Threadmill TUI inspector with backend/queue/task panes.""" @@ -347,6 +67,7 @@ class InspectorApp(App): BINDINGS = [ Binding("q", "quit", "Quit"), Binding("f5", "refresh", "Refresh"), + Binding("v", "toggle_view", "Toggle Worker View"), *( Binding(key, f"switch_tab('tab-{tab_id}')", tab_id.capitalize()) for tab_id, key in TAB_KEYS.items() @@ -355,6 +76,10 @@ class InspectorApp(App): backend: reactive[ThreadmillTaskBackend] = reactive(None) telemetry: reactive[BackendTelemetry] = reactive(None, always_update=True) + worker_telemetry: reactive[WorkerTelemetry | None] = reactive( + None, always_update=True + ) + worker_view_enabled: reactive[bool] = reactive(False) def __init__( self, @@ -367,8 +92,11 @@ def __init__( self._task_list: TaskList | None = None self._task_detail: TaskDetail | None = None self._options_static: Static | None = None + self._selection_tree: SelectionTree | None = None + self._worker_graphs: WorkerGraphs | None = None self._telemetry_timer = None self._auto_refresh = auto_refresh + self._node_cache: dict[str, NodeTelemetry] = {} self.set_reactive(InspectorApp.backend, backend) def compose(self) -> ComposeResult: @@ -377,23 +105,35 @@ def compose(self) -> ComposeResult: with Vertical(id="split-view"): with Horizontal(id="backend-row"): yield Select( - _supported_aliases(), + list(_supported_aliases()), id="backend-select", value=self.backend.alias, allow_blank=False, ) yield Static(id="backend-options") - with Horizontal(): - with Vertical(id="left-pane"): + with Horizontal(id="queue-view"): + with Vertical(classes="left-pane"): yield QueueList(id="queue-list").data_bind( telemetry=InspectorApp.telemetry ) - with Vertical(id="right-pane"): + with Vertical(classes="right-pane"): yield TaskList(id="task-list").data_bind( backend=InspectorApp.backend, telemetry=InspectorApp.telemetry, ) yield TaskDetail(id="task-detail", name="Task Detail") + with Horizontal(id="worker-view"): + with Vertical(classes="left-pane"): + yield SelectionTree( + "Worker Telemetry", id="selection-tree" + ).data_bind( + worker_telemetry=InspectorApp.worker_telemetry, + queue_telemetry=InspectorApp.telemetry, + ) + with Vertical(classes="right-pane"): + yield WorkerGraphs(id="worker-graphs").data_bind( + worker_telemetry=InspectorApp.worker_telemetry, + ) yield Footer(show_command_palette=True) def on_mount(self) -> None: @@ -403,14 +143,23 @@ def on_mount(self) -> None: self._task_list = self.query_one("#task-list", TaskList) self._task_detail = self.query_one("#task-detail", TaskDetail) self._options_static = self.query_one("#backend-options", Static) + self._selection_tree = self.query_one("#selection-tree", SelectionTree) + self._worker_graphs = self.query_one("#worker-graphs", WorkerGraphs) self._refresh_options() self._refresh_telemetry() + self.worker_telemetry = WorkerTelemetry( + nodes={}, + queues={}, + sampled_at=datetime.datetime.now(tz=datetime.UTC), + ) if self._auto_refresh: self._telemetry_timer = self.set_interval( TELEMETRY_INTERVAL_SECONDS, self._refresh_telemetry, name="telemetry-refresh", ) + self.run_worker(self._subscribe_worker_telemetry, group="telemetry") + self._apply_worker_view_visibility() self._queue_list.focus() def action_quit(self) -> None: @@ -426,6 +175,36 @@ def action_switch_tab(self, tab_id: str) -> None: """Activate the task status tab matching the given id.""" self._task_list.switch_tab(tab_id) + def action_toggle_view(self) -> None: + """Switch between queue view and worker view.""" + self.worker_view_enabled = not self.worker_view_enabled + + def watch_worker_view_enabled(self, enabled: bool) -> None: + """Show/hide widgets and update the toggle binding label.""" + self._apply_worker_view_visibility() + self._bindings.key_to_bindings["v"] = [ + Binding( + "v", + "toggle_view", + "Toggle Queue View" if enabled else "Toggle Worker View", + ) + ] + self.refresh_bindings() + + def _apply_worker_view_visibility(self) -> None: + """Toggle between queue view and worker view.""" + self.query_one("#queue-view").display = not self.worker_view_enabled + self.query_one("#worker-view").display = self.worker_view_enabled + if self.worker_view_enabled: + self._selection_tree.focus() + else: + self._queue_list.focus() + + def on_tree_node_selected(self, event: Tree.NodeSelected) -> None: + """Update worker graphs when a tree node is selected.""" + if event.node.data is not None and isinstance(event.node.data, WorkerTreeNode): + self._worker_graphs.selection = event.node.data + def watch_backend(self) -> None: self._refresh_options() self._refresh_telemetry() @@ -456,8 +235,50 @@ def _refresh_options(self) -> None: self._options_static.update(" ".join(parts) or "No options") def _refresh_telemetry(self) -> None: - """Poll the backend for fresh telemetry.""" + """Poll backend for queue telemetry and rebuild the worker view.""" try: self.telemetry = self.backend.telemetry() except Exception: # noqa: BLE001 logger.exception("Failed to refresh telemetry") + if self._node_cache: + self.worker_telemetry = self._build_worker_telemetry() + + def _build_worker_telemetry(self) -> WorkerTelemetry: + """Build a snapshot from the latest per-host node cache.""" + queues: dict[str, set[str]] = {} + for node in self._node_cache.values(): + for queue_name in node.queues: + queues.setdefault(queue_name, set()).add(node.hostname) + return WorkerTelemetry( + nodes=dict(self._node_cache), + queues={name: tuple(sorted(hosts)) for name, hosts in queues.items()}, + sampled_at=datetime.datetime.now(tz=datetime.UTC), + ) + + async def _subscribe_worker_telemetry(self) -> None: + """Maintain a pub/sub subscription; cache the latest snapshot per host. + + Each message updates the per-host cache silently; the reactive + `worker_telemetry` is only rebuilt on the 2-second timer tick + (`_refresh_telemetry`) to avoid flooding the UI with redraws. + """ + from ..backends.redis import WORKER_TELEMETRY_TTL as _TTL + + try: + async for snapshot in self.backend.subscribe_worker_telemetry(): + for hostname, node in snapshot.nodes.items(): + self._node_cache[hostname] = node + self._prune_nodes(_TTL) + except Exception: # noqa: BLE001 + logger.exception("Worker telemetry subscription failed") + + def _prune_nodes(self, ttl_seconds: int) -> None: + """Drop cached nodes whose latest sample is older than the TTL.""" + cutoff = datetime.datetime.now(tz=datetime.UTC) - datetime.timedelta( + seconds=ttl_seconds + ) + self._node_cache = { + hostname: node + for hostname, node in self._node_cache.items() + if node.sampled_at > cutoff + } diff --git a/threadmill/inspector/inspector.scss b/threadmill/inspector/inspector.scss index eab55ba..813c774 100644 --- a/threadmill/inspector/inspector.scss +++ b/threadmill/inspector/inspector.scss @@ -33,10 +33,14 @@ Select { } } -#left-pane { +.left-pane { width: 42; } +#queue-view, #worker-view { + height: 1fr; +} + #backend-row { height: 3; } @@ -104,3 +108,36 @@ TaskDetail { overflow: scroll; padding: 0 1; } + +#selection-tree { + height: 1fr; + width: 100%; + border: round $accent 40%; + border-title-color: $primary; + border-title-align: right; + background: $background; + padding: 0 1; + + &:focus-within { + border: round $accent 100%; + } +} + +#worker-graphs Sparkline { + height: 1fr; + border: round $accent 40%; + border-title-color: $primary; + border-title-align: right; +} + +#worker-throughput-graph { + color: $primary; +} + +#worker-cpu-graph { + color: $warning; +} + +#worker-memory-graph { + color: $success; +} diff --git a/threadmill/inspector/queue_view.py b/threadmill/inspector/queue_view.py new file mode 100644 index 0000000..a962df0 --- /dev/null +++ b/threadmill/inspector/queue_view.py @@ -0,0 +1,309 @@ +"""Queue and task widgets for the inspector TUI.""" + +from __future__ import annotations + +import logging +from typing import Any + +from django.contrib.humanize.templatetags.humanize import naturaltime +from django.tasks import ( + DEFAULT_TASK_QUEUE_NAME, + TaskResult, + TaskResultStatus, +) +from textual.app import ComposeResult +from textual.containers import Vertical +from textual.reactive import reactive +from textual.theme import BUILTIN_THEMES +from textual.widgets import ( + DataTable, + ListItem, + ListView, + Static, + TabbedContent, + TabPane, +) + +from ..backends.base import BackendTelemetry, ThreadmillTaskBackend +from .utils import format_dt, si_prefix + +logger = logging.getLogger(__name__) + +TAB_STATUSES: list[tuple[str, TaskResultStatus | None]] = [ + ("Running", TaskResultStatus.RUNNING), + ("Ready", TaskResultStatus.READY), + ("Successful", TaskResultStatus.SUCCESSFUL), + ("Failed", TaskResultStatus.FAILED), +] + +COLUMN_LABELS: dict[str, str] = { + "id": "ID", + "priority": "Priority", + "function": "Function", + "enqueued": "Enqueued", + "started": "Started", + "finished": "Finished", + "workers": "Workers", +} + +TAB_COLUMNS: dict[str, tuple[str, ...]] = { + "running": ("id", "function", "enqueued", "started", "workers"), + "ready": ("id", "function", "priority", "enqueued"), + "successful": ("id", "function", "enqueued", "started", "finished", "workers"), + "failed": ("id", "function", "enqueued", "started", "finished", "workers"), +} + +TAB_KEYS = { + label.lower(): str(index + 1) for index, (label, _) in enumerate(TAB_STATUSES) +} + + +def _cell(result: TaskResult, column: str) -> str: + match column: + case "id": + return result.id[:8] + case "function": + return result.task.module_path + case "priority": + return str(result.task.priority) + case "enqueued": + return naturaltime(result.enqueued_at) + case "started": + return naturaltime(result.started_at) + case "finished": + return naturaltime(result.finished_at) + case "workers": + return ", ".join(result.worker_ids) or "-" + + +class QueueItem(ListItem): + """A list item that exposes its queue name.""" + + def __init__(self, queue_name: str, label: str, **kwargs: Any) -> None: + super().__init__(Static(label), id=f"queue-{queue_name}", **kwargs) + self.queue_name = queue_name + + +class TaskDetail(Static): + """Read-only detail view for the selected task result.""" + + task_result: reactive[TaskResult | None] = reactive(None) + + def compose(self) -> ComposeResult: + yield from super().compose() + self.border_title = "Detail" + + def watch_task_result(self) -> None: + self.update(self._detail_content()) + + def _detail_content(self) -> str: + """Build the detail text for the current task.""" + result = self.task_result + if result is None: + return "Select a task to view details." + lines = [ + f"[b]Task:[/b] {result.task.module_path}", + f"[b]ID:[/b] {result.id}", + f"[b]Status:[/b] {result.status.name}", + f"[b]Queue:[/b] {result.task.queue_name}", + f"[b]Priority:[/b] {result.task.priority}", + f"[b]Enqueued:[/b] {format_dt(result.enqueued_at)}", + f"[b]Started:[/b] {format_dt(result.started_at)}", + f"[b]Finished:[/b] {format_dt(result.finished_at)}", + f"[b]Last attempted:[/b] {format_dt(result.last_attempted_at)}", + f"[b]Worker IDs:[/b] {', '.join(result.worker_ids) or '-'}", + f"[b]Args:[/b] {result.args!r}", + f"[b]Kwargs:[/b] {result.kwargs!r}", + ] + if result.errors: + lines.append("") + lines.append("[b]Errors:[/b]") + for error in result.errors: + lines.append(f" {error.exception_class_path}") + for line in error.traceback.splitlines(): + lines.append(f" {line}") + return "\n".join(lines) + + +class TaskList(Vertical): + """Tabbed list of tasks for the selected queue.""" + + backend: reactive[ThreadmillTaskBackend | None] = reactive(None) + queue_name: reactive[str] = reactive("") + telemetry: reactive[BackendTelemetry | None] = reactive(None) + counts: reactive[dict[str, int]] = reactive({}) + selected_task: reactive[TaskResult | None] = reactive(None) + + def compose(self) -> ComposeResult: + with TabbedContent(initial="tab-ready") as tabs: + for label, _ in TAB_STATUSES: + with TabPane(label, id=f"tab-{label.lower()}"): + yield DataTable(id=f"table-{label.lower()}", cursor_type="row") + self._tabs = tabs + self.border_title = "Tasks" + + def watch_backend(self) -> None: + self._refresh_data() + + def watch_queue_name(self) -> None: + self._refresh_data() + + def refresh_tasks(self) -> None: + """Re-fetch and render tasks for the current queue and tab.""" + self._refresh_data() + + def compute_counts(self) -> dict[str, int]: + if self.telemetry is None or not self.queue_name: + return {label.lower(): 0 for label, _ in TAB_STATUSES} + counts = self.telemetry.queues[self.queue_name].counts + return { + label.lower(): getattr(counts, label.lower()) for label, _ in TAB_STATUSES + } + + def watch_counts(self, counts: dict[str, int]) -> None: + for label, _status in TAB_STATUSES: + tab = self._tabs.get_tab(f"tab-{label.lower()}") + tab.label = f"{label} ({si_prefix(counts.get(label.lower(), 0))})" + + def watch_selected_task(self, task: TaskResult | None) -> None: + if self.app._task_detail is not None: + self.app._task_detail.task_result = task + + def on_mount(self) -> None: + """Configure per-status columns on first mount.""" + for label, _ in TAB_STATUSES: + tab_id = label.lower() + table = self.query_one(f"#table-{tab_id}", DataTable) + table.add_columns( + *(COLUMN_LABELS[column] for column in TAB_COLUMNS[tab_id]) + ) + table.disabled = tab_id != "ready" + + def switch_tab(self, tab_id: str) -> None: + """Activate the tab with the given id.""" + self._tabs.active = tab_id + + def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None: + """Notify the app when a task row is selected.""" + if isinstance(event.row_key.value, str): + self._select_task_by_id(event.row_key.value) + + def on_data_table_row_highlighted(self, event: DataTable.RowHighlighted) -> None: + """Update the detail preview when a row is highlighted.""" + if isinstance(event.row_key.value, str): + self._select_task_by_id(event.row_key.value) + + def on_tabbed_content_tab_activated( + self, _event: TabbedContent.TabActivated + ) -> None: + """Refresh the visible table when the user switches tabs.""" + self._refresh_data() + + def _select_task_by_id(self, task_id: str) -> None: + """Find the task result matching the row key in the current results.""" + self.selected_task = next( + (result for result in self._current_results if result.id == task_id), + None, + ) + + def _refresh_data(self) -> None: + """Fetch and display tasks for the current queue and active tab.""" + backend = self.backend + queue_name = self.queue_name + if backend is None or not queue_name: + self._current_results = [] + return + tab_id = self._tabs.active.removeprefix("tab-") + status = next( + (status for label, status in TAB_STATUSES if label.lower() == tab_id), + None, + ) + for label, _ in TAB_STATUSES: + table = self.query_one(f"#table-{label.lower()}", DataTable) + table.clear() + table.disabled = label.lower() != tab_id + table = self.query_one(f"#table-{tab_id}", DataTable) + table.disabled = False + try: + self._current_results = list( + backend.peek(queue_name=queue_name, status=status, count=100) + ) + except Exception: # noqa: BLE001 + logger.exception("Failed to peek tasks for %r", queue_name) + self._current_results = [] + columns = TAB_COLUMNS[tab_id] + for result in self._current_results: + table.add_row( + *(_cell(result, column) for column in columns), + key=result.id, + ) + self._restore_cursor(table) + + def _restore_cursor(self, table: DataTable) -> None: + """Keep the highlighted task selected across a refresh, falling back to the first row.""" + previous_id = self.selected_task.id if self.selected_task else None + row = next( + ( + index + for index, result in enumerate(self._current_results) + if result.id == previous_id + ), + None, + ) + if row is None: + row = 0 if self._current_results else None + if row is not None: + table.move_cursor(row=row) + self.selected_task = self._current_results[row] if row is not None else None + + +class QueueList(ListView): + """List of queues for the selected backend with ingress/egress deltas.""" + + telemetry: reactive[BackendTelemetry | None] = reactive(None) + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._items: dict[str, QueueItem] = {} + + def compose(self) -> ComposeResult: + yield from super().compose() + self.border_title = "Queues" + + def watch_telemetry(self, telemetry: BackendTelemetry | None) -> None: + """Refresh queue labels when a new telemetry snapshot arrives.""" + if telemetry is not None: + self.update_telemetry(telemetry) + + def update_telemetry(self, telemetry: BackendTelemetry) -> None: + """Refresh queue labels from a new telemetry snapshot.""" + queues = telemetry.queues + for queue_name, stats in sorted(queues.items()): + theme = BUILTIN_THEMES[self.app.theme] + rates = stats.rates + label = ( + f"{queue_name:24} " + f"[{theme.success}]+{si_prefix(rates.ingress):3}[/] " + f"[{theme.error}]-{si_prefix(rates.egress):3}[/]" + ) + if queue_name in self._items: + self._items[queue_name].children[0].update(label) + else: + item = QueueItem(queue_name, label) + self._items[queue_name] = item + self.append(item) + stale = [name for name in self._items if name not in queues] + for queue_name in stale: + self._items.pop(queue_name).remove() + if self.index is None and self._items: + target = ( + DEFAULT_TASK_QUEUE_NAME + if DEFAULT_TASK_QUEUE_NAME in self._items + else next(iter(self._items)) + ) + self.index = list(self._items.keys()).index(target) + self._notify_selection(target) + + def _notify_selection(self, queue_name: str) -> None: + """Tell the app which queue to display.""" + self.app._task_list.queue_name = queue_name diff --git a/threadmill/inspector/utils.py b/threadmill/inspector/utils.py new file mode 100644 index 0000000..995e82f --- /dev/null +++ b/threadmill/inspector/utils.py @@ -0,0 +1,34 @@ +"""Shared helpers for the inspector TUI.""" + +from __future__ import annotations + +import datetime +import math +import typing + +from django.tasks import task_backends + +from ..backends.base import ThreadmillTaskBackend + + +def si_prefix(n: int | float, base: int = 1_000) -> str: + """Round and shorten a number with an SI prefix (k, M, G, etc.).""" + if n < 1000: + return str(n) + prefixes = ["", "k", "M", "G", "T", "P", "E", "Z", "Y"] + m = min(int(math.log10(abs(n)) // 3), len(prefixes) - 1) + if (value := n / base**m) and value >= 100: + return f"{int(value)}{prefixes[m]}" + return f"{value:.1f}".rstrip("0").rstrip(".") + prefixes[m] + + +def format_dt(value: datetime.datetime | None) -> str: + """Format a datetime for display, or return empty string.""" + return value.isoformat() if value else "" + + +def supported_aliases() -> typing.Generator[tuple[str, str]]: + """Yield (alias, alias) pairs for all ThreadmillTaskBackend instances.""" + for alias in task_backends: + if isinstance(task_backends[alias], ThreadmillTaskBackend): + yield alias, alias diff --git a/threadmill/inspector/worker_view.py b/threadmill/inspector/worker_view.py new file mode 100644 index 0000000..f0dada3 --- /dev/null +++ b/threadmill/inspector/worker_view.py @@ -0,0 +1,215 @@ +"""Worker telemetry widgets for the inspector TUI.""" + +from __future__ import annotations + +import dataclasses +import logging +from collections import deque +from typing import Any + +from textual.app import ComposeResult +from textual.reactive import reactive +from textual.widgets import Sparkline, Static, Tree +from textual.widgets.tree import TreeNode + +from ..backends.base import BackendTelemetry, NodeTelemetry, WorkerTelemetry +from .utils import si_prefix + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class WorkerTreeNode: + """A node in the Queue -> Node selection tree.""" + + kind: str + label: str + queue_name: str = "" + hostname: str = "" + + +class SelectionTree(Tree[WorkerTreeNode]): + """Queue -> Node hierarchy built from worker telemetry.""" + + worker_telemetry: reactive[WorkerTelemetry | None] = reactive(None) + queue_telemetry: reactive[BackendTelemetry | None] = reactive(None) + _last_structure: frozenset[tuple[str, str]] | None = None + + def compose(self) -> ComposeResult: + yield from super().compose() + self.border_title = "Selection" + self.show_root = False + + def watch_worker_telemetry(self, telemetry: WorkerTelemetry | None) -> None: + """Rebuild the tree when a new worker telemetry snapshot arrives.""" + if telemetry is not None: + self.update_telemetry() + + def watch_queue_telemetry(self, telemetry: BackendTelemetry | None) -> None: + """Rebuild the tree when the set of defined queues changes.""" + if telemetry is not None: + self.update_telemetry() + + def update_telemetry(self) -> None: + """Rebuild the tree from the latest telemetry, preserving the cursor. + + Every defined queue is shown (labelled with the number of nodes + draining it), even when no node is listening on it. The tree is only + rebuilt when the set of (queue, hostname) pairs changes. + """ + worker = self.worker_telemetry + queues = self.queue_telemetry + if worker is None or queues is None: + return + structure: frozenset[tuple[str, str]] = frozenset( + (queue, host) + for queue in queues.queues + for host in worker.queues.get(queue, ()) + ) + if structure == self._last_structure: + return + self._last_structure = structure + + previous_data = self.cursor_node.data if self.cursor_node else None + self.clear() + for queue_name in sorted(queues.queues): + draining = worker.queues.get(queue_name, ()) + queue_node = self.root.add( + f"{queue_name} ({len(draining)})", + WorkerTreeNode( + kind="queue", + label=queue_name, + queue_name=queue_name, + ), + expand=True, + ) + for hostname in draining: + if worker.nodes.get(hostname) is None: + continue + queue_node.add_leaf( + hostname, + WorkerTreeNode( + kind="node", + label=hostname, + queue_name=queue_name, + hostname=hostname, + ), + ) + self._restore_cursor(previous_data) + + def _restore_cursor(self, previous_data: WorkerTreeNode | None) -> None: + """Move the cursor to the previously selected node, if still present.""" + if previous_data is None or not self.root.children: + return + node = self._find_node_by_data(self.root, previous_data) + if node is not None: + self.call_after_refresh(self.select_node, node) + + @staticmethod + def _find_node_by_data( + root: TreeNode[WorkerTreeNode], target: WorkerTreeNode + ) -> TreeNode[WorkerTreeNode] | None: + """Depth-first search for a tree node whose data matches *target*.""" + for child in root.children: + if child.data == target: + return child + result = SelectionTree._find_node_by_data(child, target) + if result is not None: + return result + return None + + +class WorkerGraphs(Static): + """Throughput/CPU/memory graphs for the worker view.""" + + # 60 s of data at a 2 s sample interval = 30 data points. + GRAPH_HISTORY_SIZE = 30 + worker_telemetry: reactive[WorkerTelemetry | None] = reactive(None) + selection: reactive[WorkerTreeNode | None] = reactive(None) + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._throughput_history: deque[float] = deque( + [0.0] * self.GRAPH_HISTORY_SIZE, maxlen=self.GRAPH_HISTORY_SIZE + ) + self._cpu_history: deque[float] = deque( + [0.0] * self.GRAPH_HISTORY_SIZE, maxlen=self.GRAPH_HISTORY_SIZE + ) + self._memory_history: deque[float] = deque( + [0.0] * self.GRAPH_HISTORY_SIZE, maxlen=self.GRAPH_HISTORY_SIZE + ) + + def compose(self) -> ComposeResult: + yield from super().compose() + yield Sparkline( + id="worker-throughput-graph", data=list(self._throughput_history) + ) + yield Sparkline(id="worker-cpu-graph", data=list(self._cpu_history)) + yield Sparkline(id="worker-memory-graph", data=list(self._memory_history)) + + def watch_selection(self) -> None: + """Reset histories when the selection changes.""" + self._reset_histories() + self._refresh_graphs() + + def watch_worker_telemetry(self) -> None: + """Append the latest telemetry sample to the graphs and redraw.""" + self._refresh_graphs() + + def _reset_histories(self) -> None: + """Clear all histories back to pre-filled zeros.""" + for history in ( + self._throughput_history, + self._cpu_history, + self._memory_history, + ): + history.clear() + history.extend([0.0] * self.GRAPH_HISTORY_SIZE) + + @staticmethod + def _memory_percent(node: NodeTelemetry) -> float: + """Derive memory usage percentage from used/total physical RAM.""" + if node.memory_total <= 0: + return 0.0 + return node.memory_bytes / node.memory_total * 100.0 + + def _refresh_graphs(self) -> None: + """Append the current sample to each graph and redraw.""" + telemetry = self.worker_telemetry + selection = self.selection + if telemetry is None or selection is None: + return + node = telemetry.nodes.get(selection.hostname) + if node is None or selection.kind != "node": + return + self._throughput_history.append(node.tasks_per_minute) + self._cpu_history.append(node.cpu_percent) + self._memory_history.append(self._memory_percent(node)) + self._update_border_titles(node) + try: + throughput = self.query_one("#worker-throughput-graph", Sparkline) + cpu = self.query_one("#worker-cpu-graph", Sparkline) + mem = self.query_one("#worker-memory-graph", Sparkline) + except Exception: # noqa: BLE001 + logger.debug("Worker graph widgets not yet mounted") + return + throughput.data = list(self._throughput_history) + cpu.data = list(self._cpu_history) + mem.data = list(self._memory_history) + + def _update_border_titles(self, node: NodeTelemetry) -> None: + """Show current values in the sparkline border titles.""" + try: + throughput = self.query_one("#worker-throughput-graph", Sparkline) + cpu = self.query_one("#worker-cpu-graph", Sparkline) + mem = self.query_one("#worker-memory-graph", Sparkline) + except Exception: # noqa: BLE001 + return + memory_percent = self._memory_percent(node) + throughput.border_title = f"Throughput {node.tasks_per_minute:.0f} tasks/min" + cpu.border_title = f"CPU {node.cpu_percent:.0f}%" + mem.border_title = ( + f"RAM {si_prefix(node.memory_bytes, base=1024)}B/{si_prefix(node.memory_total, base=1024)}B " + f"({memory_percent:.0f}%) " + f"procs {node.process_count} threads {node.thread_count}" + ) diff --git a/threadmill/telemetry.py b/threadmill/telemetry.py new file mode 100644 index 0000000..59dfc4c --- /dev/null +++ b/threadmill/telemetry.py @@ -0,0 +1,150 @@ +"""Worker-pool telemetry sampling via psutil and a Redis key store. + +The :class:`TelemetrySampler` runs as a single thread on the +:class:`~threadmill.executor.TaskExecutor` (one per node/host), periodically +sampling system CPU/memory plus rolling task throughput and publishing a +:class:`~threadmill.backends.base.WorkerTelemetry` snapshot through the +backend's :meth:`publish_worker_telemetry` hook. `psutil` is optional; when +missing, the sampler is a silent no-op so the worker command still functions. +""" + +from __future__ import annotations + +import collections +import datetime +import logging +import os +import socket +import threading +import time +import typing + +from .backends.base import ( + NodeTelemetry, + WorkerTelemetry, +) + +logger = logging.getLogger(__name__) + +try: + import psutil +except ImportError: # pragma: no cover - exercised via the guarded path + psutil = None # type: ignore[assignment] + + +SAMPLE_INTERVAL_SECONDS = 2.0 +"""Default seconds between worker telemetry samples.""" + + +class TelemetrySampler: + """Periodically sample and publish per-node worker-pool telemetry. + + Sampling only runs when `psutil` is importable; without it the sampler + publishes nothing and the worker keeps working. + """ + + def __init__( + self, + backend, + *, + executor, + interval_seconds: float = SAMPLE_INTERVAL_SECONDS, + clock: typing.Callable[[], float] | None = None, + ) -> None: + self.backend = backend + self.executor = executor + self.interval_seconds = interval_seconds + self._clock = clock or time.monotonic + self._stop_requested = threading.Event() + self._thread: threading.Thread | None = None + self._task_count_previous = 0 + self._sampled_at_previous: float | None = None + + @property + def is_available(self) -> bool: + """Whether psutil is importable and sampling is active.""" + return psutil is not None + + def start(self) -> None: + """Start the background sampling thread, if psutil is available.""" + if not self.is_available or self._thread is not None: + return + self._thread = threading.Thread( + target=self._run, name="worker-telemetry-sampler", daemon=True + ) + self._thread.start() + + def stop(self) -> None: + """Request the sampling thread to stop and wait for it to exit.""" + self._stop_requested.set() + if self._thread is not None: + self._thread.join(timeout=self.interval_seconds + 1) + self._thread = None + + def _run(self) -> None: + # Prime the CPU percent counter so the first sample is meaningful. + psutil.cpu_percent(interval=None) + while not self._stop_requested.wait(self.interval_seconds): + try: + self._publish_one_sample() + except Exception: # noqa: BLE001 + logger.exception("Failed to sample worker telemetry") + + def _publish_one_sample(self) -> None: + snapshot = self.sample() + self.backend.publish_worker_telemetry(snapshot) + + def sample(self) -> WorkerTelemetry: + """Build and return one :class:`WorkerTelemetry` snapshot.""" + now = datetime.datetime.now(tz=datetime.UTC) + now_monotonic = self._clock() + + queues = tuple(self.executor.queues) + process_count = self.executor.process_count + thread_count = process_count * self.executor.thread_count + task_count = self.executor.total_counter.value + + tasks_per_minute = self._tasks_per_minute(task_count, now_monotonic) + + # virtual_memory() reports physical RAM only; swap is tracked + # separately by psutil.swap_memory() and is deliberately excluded. + virtual_memory = psutil.virtual_memory() + node = NodeTelemetry( + hostname=socket.gethostname(), + pid=os.getpid(), + queues=queues, + process_count=process_count, + thread_count=thread_count, + cpu_percent=psutil.cpu_percent(interval=None), + memory_bytes=virtual_memory.used, + memory_total=virtual_memory.total, + tasks_per_minute=tasks_per_minute, + sampled_at=now, + ) + + queue_index: dict[str, list[str]] = collections.defaultdict(list) + for queue_name in queues: + queue_index[queue_name].append(node.hostname) + + return WorkerTelemetry( + nodes={node.hostname: node}, + queues={ + queue_name: tuple(hostnames) + for queue_name, hostnames in queue_index.items() + }, + sampled_at=now, + ) + + def _tasks_per_minute(self, task_count: int, now_monotonic: float) -> float: + """Compute tasks/minute from the delta since the previous sample.""" + if self._sampled_at_previous is None: + self._task_count_previous = task_count + self._sampled_at_previous = now_monotonic + return 0.0 + elapsed_seconds = now_monotonic - self._sampled_at_previous + if elapsed_seconds <= 0: + return 0.0 + delta = task_count - self._task_count_previous + self._task_count_previous = task_count + self._sampled_at_previous = now_monotonic + return delta / elapsed_seconds * 60.0 From 33fc4e8e710ddf31cf7f635b396729e528ce2b5e Mon Sep 17 00:00:00 2001 From: Johannes Maron Date: Sat, 4 Jul 2026 22:12:41 +0200 Subject: [PATCH 2/4] Simplify --- threadmill/backends/base.py | 18 ++---------------- threadmill/backends/redis.py | 9 +++------ threadmill/inspector/app.py | 4 +--- threadmill/inspector/worker_view.py | 7 +------ 4 files changed, 7 insertions(+), 31 deletions(-) diff --git a/threadmill/backends/base.py b/threadmill/backends/base.py index dea5e58..9dc5bd8 100644 --- a/threadmill/backends/base.py +++ b/threadmill/backends/base.py @@ -74,15 +74,7 @@ class BackendTelemetry: @dataclasses.dataclass(kw_only=True, slots=True) class NodeTelemetry: - """Telemetry for a single node (host) running a worker pool. - - The executor's telemetry sampler publishes one snapshot per node with - system-level CPU/memory sampled via `psutil` and the executor's - process/thread counts. `memory_bytes` is physical RAM in use and - `memory_total` is the machine's total physical RAM; swap is tracked - separately by `psutil.swap_memory()` and is deliberately excluded. - The inspector derives the usage percentage from these two values. - """ + """Telemetry for a single node (host) running a worker pool.""" hostname: str pid: int @@ -98,13 +90,7 @@ class NodeTelemetry: @dataclasses.dataclass(kw_only=True, slots=True) class WorkerTelemetry: - """Snapshot of worker pool health across a backend's queues and nodes. - - `nodes` is keyed by hostname and carries system-level CPU/mem plus - per-process counters. `queues` maps each queue name to the hostnames - listening on it, so the inspector can render a Queue -> Node selection - tree. - """ + """Snapshot of worker pool health across a backend's queues and nodes.""" nodes: dict[str, NodeTelemetry] queues: dict[str, tuple[str, ...]] diff --git a/threadmill/backends/redis.py b/threadmill/backends/redis.py index b902b71..49d9065 100644 --- a/threadmill/backends/redis.py +++ b/threadmill/backends/redis.py @@ -116,10 +116,6 @@ def main(self) -> None: logger.exception("Telemetry trim error for queue %r", queue_name) -WORKER_TELEMETRY_CHANNEL = "{prefix}:worker-telemetry" -WORKER_TELEMETRY_TTL = 10 # seconds before stale entries are pruned - - class RedisTaskBackend(ThreadmillTaskBackend): """Redis-backed durable priority queue backend. @@ -143,6 +139,7 @@ class RedisTaskBackend(ThreadmillTaskBackend): SUCCESSFUL_RESULTS_KEY = "{prefix}:results:{queue_name}:successful" FAILED_RESULTS_KEY = "{prefix}:results:{queue_name}:failed" INGRESS_KEY = "{prefix}:ingress:{queue_name}:events" + WORKER_TELEMETRY_CHANNEL = "{prefix}:worker-telemetry" ACQUIRE_SCRIPT = _load_lua("acquire") """Pop the next task from a priority queue and move it directly to the running set.""" @@ -475,14 +472,14 @@ def telemetry( return BackendTelemetry(queues=queues) def publish_worker_telemetry(self, telemetry: WorkerTelemetry) -> None: - channel = WORKER_TELEMETRY_CHANNEL.format(prefix=self.key_prefix) + channel = self.WORKER_TELEMETRY_CHANNEL.format(prefix=self.key_prefix) self.client.publish(channel, _serialize_worker_telemetry(telemetry)) async def subscribe_worker_telemetry(self) -> typing.AsyncIterator[WorkerTelemetry]: import redis.asyncio as aioredis redis_url = self.redis_url - channel = WORKER_TELEMETRY_CHANNEL.format(prefix=self.key_prefix) + channel = self.WORKER_TELEMETRY_CHANNEL.format(prefix=self.key_prefix) async with aioredis.from_url(redis_url) as client: pubsub = client.pubsub() await pubsub.subscribe(channel) diff --git a/threadmill/inspector/app.py b/threadmill/inspector/app.py index 5bc2717..0fbf8d5 100644 --- a/threadmill/inspector/app.py +++ b/threadmill/inspector/app.py @@ -262,13 +262,11 @@ async def _subscribe_worker_telemetry(self) -> None: `worker_telemetry` is only rebuilt on the 2-second timer tick (`_refresh_telemetry`) to avoid flooding the UI with redraws. """ - from ..backends.redis import WORKER_TELEMETRY_TTL as _TTL - try: async for snapshot in self.backend.subscribe_worker_telemetry(): for hostname, node in snapshot.nodes.items(): self._node_cache[hostname] = node - self._prune_nodes(_TTL) + self._prune_nodes(10) except Exception: # noqa: BLE001 logger.exception("Worker telemetry subscription failed") diff --git a/threadmill/inspector/worker_view.py b/threadmill/inspector/worker_view.py index f0dada3..2f98676 100644 --- a/threadmill/inspector/worker_view.py +++ b/threadmill/inspector/worker_view.py @@ -51,12 +51,7 @@ def watch_queue_telemetry(self, telemetry: BackendTelemetry | None) -> None: self.update_telemetry() def update_telemetry(self) -> None: - """Rebuild the tree from the latest telemetry, preserving the cursor. - - Every defined queue is shown (labelled with the number of nodes - draining it), even when no node is listening on it. The tree is only - rebuilt when the set of (queue, hostname) pairs changes. - """ + """Rebuild the tree from the latest telemetry, preserving the cursor.""" worker = self.worker_telemetry queues = self.queue_telemetry if worker is None or queues is None: From 1f62bca587c4aa875714863ac62914dc77cde8cb Mon Sep 17 00:00:00 2001 From: Johannes Maron Date: Sun, 5 Jul 2026 11:10:40 +0200 Subject: [PATCH 3/4] Use LRU --- tests/test_inspector.py | 29 +++++++++++++-------- threadmill/inspector/app.py | 40 ++++++++++++++--------------- threadmill/inspector/worker_view.py | 7 +---- threadmill/telemetry.py | 12 ++------- 4 files changed, 42 insertions(+), 46 deletions(-) diff --git a/tests/test_inspector.py b/tests/test_inspector.py index a34926e..332ac7c 100644 --- a/tests/test_inspector.py +++ b/tests/test_inspector.py @@ -1000,24 +1000,33 @@ def test_latest_snapshot_wins_for_same_host(self): assert len(app._node_cache) == 1 assert app._node_cache["node-1"].pid == 200 - def test_prune_drops_stale_nodes(self): - """_prune_nodes drops entries older than the TTL.""" + def test_push_evicts_oldest_host_when_cache_is_full(self): + """Pushing beyond the cache cap evicts the least-recently-updated host.""" app = InspectorApp( backend=_StubBackend(alias="default", params={}), auto_refresh=False ) - old = datetime.datetime.now(tz=datetime.UTC) - datetime.timedelta(seconds=20) - app._node_cache["node-1"] = dataclasses.replace(_make_node(), sampled_at=old) - app._prune_nodes(10) + app._node_cache.maxlen = 2 + app._node_cache["node-1"] = _make_node(hostname="node-1") + app._node_cache["node-2"] = _make_node(hostname="node-2") + app._node_cache["node-3"] = _make_node(hostname="node-3") + assert len(app._node_cache) == 2 assert "node-1" not in app._node_cache + assert "node-2" in app._node_cache + assert "node-3" in app._node_cache - def test_prune_keeps_fresh_nodes(self): - """_prune_nodes retains entries within the TTL window.""" + def test_push_for_existing_host_marks_it_most_recent(self): + """Re-pushing an existing host keeps it and evicts the oldest other host.""" app = InspectorApp( backend=_StubBackend(alias="default", params={}), auto_refresh=False ) - app._node_cache["node-1"] = _make_node() - app._prune_nodes(10) - assert "node-1" in app._node_cache + app._node_cache.maxlen = 2 + app._node_cache["node-1"] = _make_node(hostname="node-1") + app._node_cache["node-2"] = _make_node(hostname="node-2") + # Refresh node-1 so node-2 becomes the oldest entry. + app._node_cache["node-1"] = _make_node(hostname="node-1", pid=999) + app._node_cache["node-3"] = _make_node(hostname="node-3") + assert set(app._node_cache) == {"node-1", "node-3"} + assert app._node_cache["node-1"].pid == 999 def test_build_includes_nodes_and_queue_index(self): """_build_worker_telemetry returns cached nodes plus a queue->host index.""" diff --git a/threadmill/inspector/app.py b/threadmill/inspector/app.py index 0fbf8d5..7aab37e 100644 --- a/threadmill/inspector/app.py +++ b/threadmill/inspector/app.py @@ -2,6 +2,7 @@ from __future__ import annotations +import collections import datetime import logging @@ -44,7 +45,7 @@ logger = logging.getLogger(__name__) -TELEMETRY_INTERVAL_SECONDS = 2.0 +TELEMETRY_INTERVAL_SECONDS = 1.0 """Seconds between automatic queue-stat refreshes; the task list stays manual.""" __all__ = [ @@ -60,6 +61,21 @@ ] +class NodeTelemetryCache(collections.OrderedDict[str, NodeTelemetry]): + """Bounded LRU of the latest per-host worker telemetry.""" + + def __init__(self, maxlen: int) -> None: + super().__init__() + self.maxlen = maxlen + + def __setitem__(self, key, value) -> None: # type: ignore[override] + if key in self: + self.move_to_end(key) + super().__setitem__(key, value) + while len(self) > self.maxlen: + self.popitem(last=False) + + class InspectorApp(App): """Threadmill TUI inspector with backend/queue/task panes.""" @@ -86,6 +102,7 @@ def __init__( backend: ThreadmillTaskBackend, *, auto_refresh: bool = True, + node_cache_max: int = 100, ) -> None: super().__init__() self._queue_list: QueueList | None = None @@ -96,7 +113,7 @@ def __init__( self._worker_graphs: WorkerGraphs | None = None self._telemetry_timer = None self._auto_refresh = auto_refresh - self._node_cache: dict[str, NodeTelemetry] = {} + self._node_cache: NodeTelemetryCache = NodeTelemetryCache(maxlen=node_cache_max) self.set_reactive(InspectorApp.backend, backend) def compose(self) -> ComposeResult: @@ -256,27 +273,10 @@ def _build_worker_telemetry(self) -> WorkerTelemetry: ) async def _subscribe_worker_telemetry(self) -> None: - """Maintain a pub/sub subscription; cache the latest snapshot per host. - - Each message updates the per-host cache silently; the reactive - `worker_telemetry` is only rebuilt on the 2-second timer tick - (`_refresh_telemetry`) to avoid flooding the UI with redraws. - """ + """Maintain a pub/sub subscription; cache the latest snapshot per host.""" try: async for snapshot in self.backend.subscribe_worker_telemetry(): for hostname, node in snapshot.nodes.items(): self._node_cache[hostname] = node - self._prune_nodes(10) except Exception: # noqa: BLE001 logger.exception("Worker telemetry subscription failed") - - def _prune_nodes(self, ttl_seconds: int) -> None: - """Drop cached nodes whose latest sample is older than the TTL.""" - cutoff = datetime.datetime.now(tz=datetime.UTC) - datetime.timedelta( - seconds=ttl_seconds - ) - self._node_cache = { - hostname: node - for hostname, node in self._node_cache.items() - if node.sampled_at > cutoff - } diff --git a/threadmill/inspector/worker_view.py b/threadmill/inspector/worker_view.py index 2f98676..6be38e3 100644 --- a/threadmill/inspector/worker_view.py +++ b/threadmill/inspector/worker_view.py @@ -41,12 +41,10 @@ def compose(self) -> ComposeResult: self.show_root = False def watch_worker_telemetry(self, telemetry: WorkerTelemetry | None) -> None: - """Rebuild the tree when a new worker telemetry snapshot arrives.""" if telemetry is not None: self.update_telemetry() def watch_queue_telemetry(self, telemetry: BackendTelemetry | None) -> None: - """Rebuild the tree when the set of defined queues changes.""" if telemetry is not None: self.update_telemetry() @@ -117,8 +115,7 @@ def _find_node_by_data( class WorkerGraphs(Static): """Throughput/CPU/memory graphs for the worker view.""" - # 60 s of data at a 2 s sample interval = 30 data points. - GRAPH_HISTORY_SIZE = 30 + GRAPH_HISTORY_SIZE = 60 # one minute of 1-second samples worker_telemetry: reactive[WorkerTelemetry | None] = reactive(None) selection: reactive[WorkerTreeNode | None] = reactive(None) @@ -143,12 +140,10 @@ def compose(self) -> ComposeResult: yield Sparkline(id="worker-memory-graph", data=list(self._memory_history)) def watch_selection(self) -> None: - """Reset histories when the selection changes.""" self._reset_histories() self._refresh_graphs() def watch_worker_telemetry(self) -> None: - """Append the latest telemetry sample to the graphs and redraw.""" self._refresh_graphs() def _reset_histories(self) -> None: diff --git a/threadmill/telemetry.py b/threadmill/telemetry.py index 59dfc4c..d78b369 100644 --- a/threadmill/telemetry.py +++ b/threadmill/telemetry.py @@ -1,12 +1,4 @@ -"""Worker-pool telemetry sampling via psutil and a Redis key store. - -The :class:`TelemetrySampler` runs as a single thread on the -:class:`~threadmill.executor.TaskExecutor` (one per node/host), periodically -sampling system CPU/memory plus rolling task throughput and publishing a -:class:`~threadmill.backends.base.WorkerTelemetry` snapshot through the -backend's :meth:`publish_worker_telemetry` hook. `psutil` is optional; when -missing, the sampler is a silent no-op so the worker command still functions. -""" +"""Worker-pool telemetry sampling via psutil and a Redis pub/sub.""" from __future__ import annotations @@ -32,7 +24,7 @@ psutil = None # type: ignore[assignment] -SAMPLE_INTERVAL_SECONDS = 2.0 +SAMPLE_INTERVAL_SECONDS = 1.0 """Default seconds between worker telemetry samples.""" From 42fafff7c5f7350e900da90cc6743fe655ed2f7e Mon Sep 17 00:00:00 2001 From: Johannes Maron Date: Sun, 5 Jul 2026 20:45:42 +0200 Subject: [PATCH 4/4] Add plot --- pyproject.toml | 1 + tests/test_inspector.py | 2 +- threadmill/inspector/app.py | 9 +++++ threadmill/inspector/inspector.scss | 26 +++++++++++-- threadmill/inspector/queue_view.py | 60 ++++++++++++++++++++++++++++- threadmill/inspector/worker_view.py | 32 +++++++++++---- 6 files changed, 116 insertions(+), 14 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 74cf408..1116fe9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,7 @@ dependencies = ["django>=6.0"] redis = ["redis>=5.0"] inspector = [ "textual>=8.2.7", + "textual-plotext>=1.0.1", ] worker = [ "psutil>=7.2.2", diff --git a/tests/test_inspector.py b/tests/test_inspector.py index 332ac7c..019c8ed 100644 --- a/tests/test_inspector.py +++ b/tests/test_inspector.py @@ -646,7 +646,7 @@ async def test_worker_graphs_ram_border_title(self): await pilot.pause() mem = graphs.query_one("#worker-memory-graph") assert "RAM" in mem.border_title - assert "4GB/8GB" in mem.border_title + assert "4.37GB/7.5GB" in mem.border_title assert "50%" in mem.border_title async def test_toggle_view_shows_worker_widgets(self): diff --git a/threadmill/inspector/app.py b/threadmill/inspector/app.py index 7aab37e..4b110f5 100644 --- a/threadmill/inspector/app.py +++ b/threadmill/inspector/app.py @@ -30,6 +30,7 @@ TAB_KEYS, QueueItem, QueueList, + QueueSparklines, TaskDetail, TaskList, ) @@ -52,6 +53,7 @@ "InspectorApp", "QueueItem", "QueueList", + "QueueSparklines", "SelectionTree", "TaskDetail", "TaskList", @@ -108,6 +110,7 @@ def __init__( self._queue_list: QueueList | None = None self._task_list: TaskList | None = None self._task_detail: TaskDetail | None = None + self._queue_sparklines: QueueSparklines | None = None self._options_static: Static | None = None self._selection_tree: SelectionTree | None = None self._worker_graphs: WorkerGraphs | None = None @@ -134,6 +137,9 @@ def compose(self) -> ComposeResult: telemetry=InspectorApp.telemetry ) with Vertical(classes="right-pane"): + yield QueueSparklines(id="queue-sparklines").data_bind( + telemetry=InspectorApp.telemetry, + ) yield TaskList(id="task-list").data_bind( backend=InspectorApp.backend, telemetry=InspectorApp.telemetry, @@ -159,6 +165,7 @@ def on_mount(self) -> None: self._queue_list = self.query_one("#queue-list", QueueList) self._task_list = self.query_one("#task-list", TaskList) self._task_detail = self.query_one("#task-detail", TaskDetail) + self._queue_sparklines = self.query_one("#queue-sparklines", QueueSparklines) self._options_static = self.query_one("#backend-options", Static) self._selection_tree = self.query_one("#selection-tree", SelectionTree) self._worker_graphs = self.query_one("#worker-graphs", WorkerGraphs) @@ -238,11 +245,13 @@ def on_list_view_selected(self, event: ListView.Selected) -> None: """Update the task list when a queue is selected.""" if queue_name := self._queue_from_item(event.item): self._task_list.queue_name = queue_name + self._queue_sparklines.queue_name = queue_name def on_list_view_highlighted(self, event: ListView.Highlighted) -> None: """Preview the selected queue without requiring Enter.""" if queue_name := self._queue_from_item(event.item): self._task_list.queue_name = queue_name + self._queue_sparklines.queue_name = queue_name def _refresh_options(self) -> None: """Show the selected backend's constructor options.""" diff --git a/threadmill/inspector/inspector.scss b/threadmill/inspector/inspector.scss index 813c774..f09970c 100644 --- a/threadmill/inspector/inspector.scss +++ b/threadmill/inspector/inspector.scss @@ -123,7 +123,8 @@ TaskDetail { } } -#worker-graphs Sparkline { +#worker-graphs Sparkline, +#worker-graphs PlotextPlot { height: 1fr; border: round $accent 40%; border-title-color: $primary; @@ -134,10 +135,27 @@ TaskDetail { color: $primary; } -#worker-cpu-graph { - color: $warning; + + +#queue-sparklines { + height: auto; + padding: 0 1; } -#worker-memory-graph { +#ingress-sparkline { + width: 1fr; + height: 3; color: $success; + border: round $accent 40%; + border-title-color: $success; + border-title-align: right; +} + +#egress-sparkline { + width: 1fr; + height: 3; + color: $warning; + border: round $accent 40%; + border-title-color: $warning; + border-title-align: right; } diff --git a/threadmill/inspector/queue_view.py b/threadmill/inspector/queue_view.py index a962df0..596b6a8 100644 --- a/threadmill/inspector/queue_view.py +++ b/threadmill/inspector/queue_view.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +from collections import deque from typing import Any from django.contrib.humanize.templatetags.humanize import naturaltime @@ -12,13 +13,14 @@ TaskResultStatus, ) from textual.app import ComposeResult -from textual.containers import Vertical +from textual.containers import Horizontal, Vertical from textual.reactive import reactive from textual.theme import BUILTIN_THEMES from textual.widgets import ( DataTable, ListItem, ListView, + Sparkline, Static, TabbedContent, TabPane, @@ -257,6 +259,62 @@ def _restore_cursor(self, table: DataTable) -> None: self.selected_task = self._current_results[row] if row is not None else None +class QueueSparklines(Horizontal): + """Slim continuous ingress/egress sparklines for the selected queue.""" + + GRAPH_HISTORY_SIZE = 60 # one minute of 1-second samples + telemetry: reactive[BackendTelemetry | None] = reactive(None) + queue_name: reactive[str] = reactive("") + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._ingress_history: deque[float] = deque( + [0.0] * self.GRAPH_HISTORY_SIZE, maxlen=self.GRAPH_HISTORY_SIZE + ) + self._egress_history: deque[float] = deque( + [0.0] * self.GRAPH_HISTORY_SIZE, maxlen=self.GRAPH_HISTORY_SIZE + ) + + def compose(self) -> ComposeResult: + yield Sparkline(id="ingress-sparkline", data=list(self._ingress_history)) + yield Sparkline(id="egress-sparkline", data=list(self._egress_history)) + + def watch_queue_name(self) -> None: + self._reset_histories() + self._refresh() + + def watch_telemetry(self) -> None: + self._refresh() + + def _reset_histories(self) -> None: + """Clear the sparkline histories back to pre-filled zeros.""" + for history in (self._ingress_history, self._egress_history): + history.clear() + history.extend([0.0] * self.GRAPH_HISTORY_SIZE) + + def _refresh(self) -> None: + """Append the latest rate sample and redraw both sparklines.""" + telemetry = self.telemetry + if telemetry is None or not self.queue_name: + return + stats = telemetry.queues.get(self.queue_name) + if stats is None: + return + rates = stats.rates + self._ingress_history.append(rates.ingress) + self._egress_history.append(rates.egress) + try: + ingress = self.query_one("#ingress-sparkline", Sparkline) + egress = self.query_one("#egress-sparkline", Sparkline) + except Exception: # noqa: BLE001 + logger.debug("Queue sparkline widgets not yet mounted") + return + ingress.data = list(self._ingress_history) + egress.data = list(self._egress_history) + ingress.border_title = f"Ingress +{si_prefix(rates.ingress)}" + egress.border_title = f"Egress -{si_prefix(rates.egress)}" + + class QueueList(ListView): """List of queues for the selected backend with ingress/egress deltas.""" diff --git a/threadmill/inspector/worker_view.py b/threadmill/inspector/worker_view.py index 6be38e3..b2b082c 100644 --- a/threadmill/inspector/worker_view.py +++ b/threadmill/inspector/worker_view.py @@ -11,6 +11,7 @@ from textual.reactive import reactive from textual.widgets import Sparkline, Static, Tree from textual.widgets.tree import TreeNode +from textual_plotext import PlotextPlot from ..backends.base import BackendTelemetry, NodeTelemetry, WorkerTelemetry from .utils import si_prefix @@ -136,8 +137,8 @@ def compose(self) -> ComposeResult: yield Sparkline( id="worker-throughput-graph", data=list(self._throughput_history) ) - yield Sparkline(id="worker-cpu-graph", data=list(self._cpu_history)) - yield Sparkline(id="worker-memory-graph", data=list(self._memory_history)) + yield PlotextPlot(id="worker-cpu-graph") + yield PlotextPlot(id="worker-memory-graph") def watch_selection(self) -> None: self._reset_histories() @@ -178,21 +179,36 @@ def _refresh_graphs(self) -> None: self._update_border_titles(node) try: throughput = self.query_one("#worker-throughput-graph", Sparkline) - cpu = self.query_one("#worker-cpu-graph", Sparkline) - mem = self.query_one("#worker-memory-graph", Sparkline) + cpu = self.query_one("#worker-cpu-graph", PlotextPlot) + mem = self.query_one("#worker-memory-graph", PlotextPlot) except Exception: # noqa: BLE001 logger.debug("Worker graph widgets not yet mounted") return throughput.data = list(self._throughput_history) - cpu.data = list(self._cpu_history) - mem.data = list(self._memory_history) + self._draw_plot(cpu, self._cpu_history, ylim=(0, 100), color="yellow") + self._draw_plot(mem, self._memory_history, ylim=(0, 100), color="green") + + @staticmethod + def _draw_plot( + plot: PlotextPlot, data: deque[float], *, ylim: tuple[float, float], color: str + ) -> None: + """Redraw a PlotextPlot line chart from the given samples.""" + plt = plot.plt + plt.clear_data() + plt.bar(list(data), color=color) + plt.ylim(*ylim) + plt.xlabel(None) + plt.xticks([]) + plt.yticks([]) + plt.frame(False) + plot.refresh() def _update_border_titles(self, node: NodeTelemetry) -> None: """Show current values in the sparkline border titles.""" try: throughput = self.query_one("#worker-throughput-graph", Sparkline) - cpu = self.query_one("#worker-cpu-graph", Sparkline) - mem = self.query_one("#worker-memory-graph", Sparkline) + cpu = self.query_one("#worker-cpu-graph", PlotextPlot) + mem = self.query_one("#worker-memory-graph", PlotextPlot) except Exception: # noqa: BLE001 return memory_percent = self._memory_percent(node)