From c8b19409a551b85989730ef80b2111a4107ccc64 Mon Sep 17 00:00:00 2001 From: pierre Date: Mon, 7 Sep 2026 10:47:12 +0400 Subject: [PATCH] Worker pool restart --- README.md | 7 +++--- src/taskowl/main.py | 19 +++++++++++++++++ src/taskowl/mcp/tools.py | 20 +++++++++++++++++ src/taskowl/workers.py | 23 ++++++++++++++++++++ tests/test_api.py | 46 ++++++++++++++++++++++++++++++++++++++++ tests/test_workers.py | 46 ++++++++++++++++++++++++++++++++++++++++ 6 files changed, 158 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index e373738..34a7906 100644 --- a/README.md +++ b/README.md @@ -112,10 +112,10 @@ If authentication is enabled (see below), send the taskowl API key as |---|---| | **Tasks** | `list_tasks`, `get_task`, `get_task_timeline`, `get_task_chain`, `get_task_summary`, `list_task_types`, `list_orphaned_tasks` | | **Task actions** | `revoke_task`, `retry_task`, `execute_task` | -| **Workers** | `get_worker_status`, `list_workers`, `get_worker_stats`, `shutdown_worker`, `scale_worker_pool`, `get_active_tasks`, `get_scheduled_tasks`, `get_reserved_tasks` | +| **Workers** | `get_worker_status`, `list_workers`, `get_worker_stats`, `shutdown_worker`, `scale_worker_pool`, `restart_worker_pool`, `get_active_tasks`, `get_scheduled_tasks`, `get_reserved_tasks` | | **Queues** | `list_queues` | -**Total: 19 tools** +**Total: 20 tools** `list_tasks` supports exact filters (`state`, `name`, `worker`, `since`), a partial case-insensitive `search` on the task name, `offset` for pagination, and `sort_by` @@ -136,6 +136,7 @@ Questions you can ask your AI assistant when the MCP server is connected: | "Which workers are online?" | `get_worker_status`, `list_workers` | | "How many messages are in each queue?" | `list_queues` | | "Shutdown worker celery@worker1" | `shutdown_worker` | +| "Restart the pool on celery@worker1" | `restart_worker_pool` | | "What's scheduled to run next?" | `get_scheduled_tasks`, `get_reserved_tasks` | | "Retry task abc" | `retry_task` | | "Run myapp.tasks.process now" | `execute_task` | @@ -264,7 +265,7 @@ retries, and metrics. Interactive docs are available at: | **Tasks** | `GET /api/tasks`, `GET /api/tasks/{id}`, `GET /api/tasks/{id}/timeline`, `GET /api/tasks/{id}/chain`, `GET /api/tasks/summary`, `GET /api/tasks/types`, `GET /api/tasks/orphaned` | | **Task actions** | `POST /api/tasks/{id}/revoke`, `POST /api/tasks/{id}/retry`, `POST /api/tasks/execute` | | **Workers** | `GET /api/workers`, `GET /api/workers/list`, `GET /api/workers/{name}/stats`, `GET /api/workers/active-tasks`, `GET /api/workers/scheduled`, `GET /api/workers/reserved` | -| **Worker actions** | `POST /api/workers/{name}/shutdown`, `POST /api/workers/{name}/scale` | +| **Worker actions** | `POST /api/workers/{name}/shutdown`, `POST /api/workers/{name}/scale`, `POST /api/workers/{name}/restart` | | **Queues** | `GET /api/queues` | | **Ops** | `GET /health`, `GET /metrics` | diff --git a/src/taskowl/main.py b/src/taskowl/main.py index 6e841e7..9f68170 100644 --- a/src/taskowl/main.py +++ b/src/taskowl/main.py @@ -31,6 +31,7 @@ get_scheduled_tasks, get_worker_stats, list_workers, + restart_worker_pool, scale_worker_pool, shutdown_worker, ) @@ -342,6 +343,24 @@ async def api_scale_worker_pool( return result +@app.post("/api/workers/{worker_name}/restart") +async def api_restart_worker_pool( + worker_name: str, + reload: bool = False, + _: None = Depends(verify_api_key), +) -> dict: + """Restart a worker's execution pool. + + Args: + worker_name: Name of the worker + reload: If True, reload modules when restarting the pool + """ + result = await restart_worker_pool(worker_name, reload) + if "error" in result: + raise HTTPException(status_code=400, detail=result["error"]) + return result + + @app.get("/api/workers/active-tasks") async def api_get_active_tasks( worker_name: str | None = None, diff --git a/src/taskowl/mcp/tools.py b/src/taskowl/mcp/tools.py index fe9fcb9..2934d3c 100644 --- a/src/taskowl/mcp/tools.py +++ b/src/taskowl/mcp/tools.py @@ -356,6 +356,26 @@ async def scale_worker_pool(worker_name: str, delta: int) -> dict: response.raise_for_status() return response.json() + @server.tool( + name="restart_worker_pool", + description="Restart a worker's execution pool (optionally reloading modules)", + ) + async def restart_worker_pool(worker_name: str, reload: bool = False) -> dict: + """Restart a worker's execution pool. + + Args: + worker_name: Name of the worker + reload: If True, reload modules when restarting the pool + """ + async with httpx.AsyncClient() as client: + response = await client.post( + f"http://{settings.taskowl_host}:{settings.taskowl_port}/api/workers/{worker_name}/restart", + params={"reload": reload}, + headers=_get_headers(), + ) + response.raise_for_status() + return response.json() + @server.tool( name="get_active_tasks", description="Get currently executing tasks across all workers or a specific worker", diff --git a/src/taskowl/workers.py b/src/taskowl/workers.py index 1c1a4c7..14820b7 100644 --- a/src/taskowl/workers.py +++ b/src/taskowl/workers.py @@ -124,6 +124,29 @@ async def scale_worker_pool(worker_name: str, delta: int) -> dict: return {"error": f"Failed to scale worker pool: {str(e)}"} +async def restart_worker_pool(worker_name: str, reload: bool = False) -> dict: + """Restart a worker's execution pool. + + Args: + worker_name: Name of the worker + reload: If True, reload modules when restarting the pool + + Returns: + Dict with status and message + """ + try: + app = _get_celery_app() + app.control.pool_restart(destination=[worker_name], reload=reload) + + return { + "status": "success", + "message": f"Pool restart command sent to {worker_name}", + "worker": worker_name, + } + except Exception as e: + return {"error": f"Failed to restart worker pool: {str(e)}"} + + async def get_active_tasks(worker_name: str | None = None) -> dict: """Get currently executing tasks. diff --git a/tests/test_api.py b/tests/test_api.py index 1809e2e..de14f7d 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -858,6 +858,52 @@ async def test_api_scale_worker_pool_zero_delta(client: AsyncClient): assert "detail" in data +@pytest.mark.asyncio +async def test_api_restart_worker_pool(client: AsyncClient): + """Test POST /api/workers/{worker_name}/restart.""" + with patch("taskowl.workers._get_celery_app") as mock_get_app: + mock_app = MagicMock() + mock_get_app.return_value = mock_app + + response = await client.post("/api/workers/celery@worker1/restart") + assert response.status_code == 200 + data = response.json() + assert data["status"] == "success" + assert data["worker"] == "celery@worker1" + mock_app.control.pool_restart.assert_called_once_with( + destination=["celery@worker1"], reload=False + ) + + +@pytest.mark.asyncio +async def test_api_restart_worker_pool_reload(client: AsyncClient): + """Test POST /api/workers/{worker_name}/restart with reload=true.""" + with patch("taskowl.workers._get_celery_app") as mock_get_app: + mock_app = MagicMock() + mock_get_app.return_value = mock_app + + response = await client.post("/api/workers/celery@worker1/restart?reload=true") + assert response.status_code == 200 + data = response.json() + assert data["status"] == "success" + mock_app.control.pool_restart.assert_called_once_with( + destination=["celery@worker1"], reload=True + ) + + +@pytest.mark.asyncio +async def test_api_restart_worker_pool_error(client: AsyncClient): + """Test POST /api/workers/{worker_name}/restart on error.""" + with patch("taskowl.workers._get_celery_app") as mock_get_app: + mock_app = MagicMock() + mock_app.control.pool_restart.side_effect = Exception("Restart failed") + mock_get_app.return_value = mock_app + + response = await client.post("/api/workers/celery@worker1/restart") + assert response.status_code == 400 + assert "Restart failed" in response.json()["detail"] + + @pytest.mark.asyncio async def test_api_get_active_tasks(client: AsyncClient): """Test GET /api/workers/active-tasks.""" diff --git a/tests/test_workers.py b/tests/test_workers.py index c2af1ea..88c36e0 100644 --- a/tests/test_workers.py +++ b/tests/test_workers.py @@ -10,6 +10,7 @@ get_scheduled_tasks, get_worker_stats, list_workers, + restart_worker_pool, scale_worker_pool, shutdown_worker, ) @@ -178,6 +179,51 @@ async def test_scale_worker_pool_error(): assert "Failed to scale worker pool" in result["error"] +@pytest.mark.asyncio +async def test_restart_worker_pool_success(): + """Test restarting a worker pool.""" + with patch("taskowl.workers._get_celery_app") as mock_get_app: + mock_app = MagicMock() + mock_get_app.return_value = mock_app + + result = await restart_worker_pool("celery@worker1") + + assert result["status"] == "success" + assert result["worker"] == "celery@worker1" + mock_app.control.pool_restart.assert_called_once_with( + destination=["celery@worker1"], reload=False + ) + + +@pytest.mark.asyncio +async def test_restart_worker_pool_reload(): + """Test restarting a worker pool with reload enabled.""" + with patch("taskowl.workers._get_celery_app") as mock_get_app: + mock_app = MagicMock() + mock_get_app.return_value = mock_app + + result = await restart_worker_pool("celery@worker1", reload=True) + + assert result["status"] == "success" + mock_app.control.pool_restart.assert_called_once_with( + destination=["celery@worker1"], reload=True + ) + + +@pytest.mark.asyncio +async def test_restart_worker_pool_error(): + """Test restarting a worker pool when an error occurs.""" + with patch("taskowl.workers._get_celery_app") as mock_get_app: + mock_app = MagicMock() + mock_app.control.pool_restart.side_effect = Exception("Restart failed") + mock_get_app.return_value = mock_app + + result = await restart_worker_pool("celery@worker1") + + assert "error" in result + assert "Failed to restart worker pool" in result["error"] + + @pytest.mark.asyncio async def test_get_active_tasks_success(): """Test getting active tasks."""