Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 28 additions & 12 deletions aiida_restapi/routers/daemon.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@

from aiida.cmdline.utils.decorators import with_dbenv
from aiida.engine.daemon.client import DaemonException, get_daemon_client
from fastapi import APIRouter
from fastapi import APIRouter, Query

from aiida_restapi.jsonapi.models import errors
from aiida_restapi.models.daemon import DaemonStatus, DaemonWorker
Expand Down Expand Up @@ -41,15 +41,15 @@ async def get_daemon_status() -> DaemonStatus:


@read_router.get(
'/worker',
'/workers',
response_model=dict[str, DaemonWorker],
responses={
500: {'model': errors.DaemonError},
},
)
@with_dbenv()
async def get_daemon_worker() -> dict[str, dict[str, t.Any]]:
"""Return daemon worker metadata."""
async def get_daemon_workers() -> dict[str, dict[str, t.Any]]:
"""Return daemon workers metadata."""
client = get_daemon_client()

if not client.is_daemon_running:
Expand Down Expand Up @@ -158,16 +158,24 @@ async def _wait_for_num_workers(
},
)
@with_dbenv()
async def increase_daemon_worker() -> dict[str, t.Any]:
"""Increase the number of daemon workers by one."""
async def increase_daemon_workers(
amount: t.Annotated[
int,
Query(
description='The number of workers to increase by',
ge=1,
),
] = 1,
) -> dict[str, t.Any]:
"""Increase the number of daemon workers by a specified amount (default=1)."""
client = get_daemon_client()

if not client.is_daemon_running:
raise DaemonException('The daemon is not running.')

initial = client.get_numprocesses()['numprocesses']
client.increase_workers(1)
num_workers = await _wait_for_num_workers(initial + 1)
client.increase_workers(amount)
num_workers = await _wait_for_num_workers(initial + amount)

return {
'running': True,
Expand All @@ -183,16 +191,24 @@ async def increase_daemon_worker() -> dict[str, t.Any]:
},
)
@with_dbenv()
async def decrease_daemon_worker() -> dict[str, t.Any]:
"""Decrease the number of daemon workers by one."""
async def decrease_daemon_workers(
amount: t.Annotated[
int,
Query(
description='The number of workers to decrease by',
ge=1,
),
] = 1,
) -> dict[str, t.Any]:
"""Decrease the number of daemon workers by a specified amount (default=1)."""
client = get_daemon_client()

if not client.is_daemon_running:
raise DaemonException('The daemon is not running.')

initial = client.get_numprocesses()['numprocesses']
client.decrease_workers(1)
num_workers = await _wait_for_num_workers(max(initial - 1, 0))
client.decrease_workers(amount)
num_workers = await _wait_for_num_workers(max(initial - amount, 0))

return {
'running': True,
Expand Down
34 changes: 28 additions & 6 deletions tests/test_daemon.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,17 +26,17 @@ def test_status_and_start(client: TestClient):


@pytest.mark.usefixtures('stopped_daemon_client')
def test_worker_when_stopped(client: TestClient):
"""Test ``/daemon/worker`` when the daemon is not running."""
response = client.get('/daemon/worker')
def test_workers_when_stopped(client: TestClient):
"""Test ``/daemon/workers`` when the daemon is not running."""
response = client.get('/daemon/workers')
assert response.status_code == 200, response.content
assert response.json() == {}


@pytest.mark.usefixtures('started_daemon_client')
def test_worker_when_running(client: TestClient):
"""Test ``/daemon/worker`` when the daemon is running."""
response = client.get('/daemon/worker')
def test_workers_when_running(client: TestClient):
"""Test ``/daemon/workers`` when the daemon is running."""
response = client.get('/daemon/workers')
assert response.status_code == 200, response.content

results = response.json()
Expand Down Expand Up @@ -116,3 +116,25 @@ def test_increase_and_decrease_when_running(client: TestClient):
results = response.json()
assert results['running'] is True
assert results['num_workers'] == initial_workers


@pytest.mark.usefixtures('started_daemon_client')
def test_increase_and_decrease_by_amount(client: TestClient):
"""Test ``/daemon/increase`` and ``/daemon/decrease`` with a specified amount."""
response = client.get('/daemon/status')
assert response.status_code == 200, response.content
initial_workers = response.json()['num_workers']

response = client.post('/daemon/increase', params={'amount': 3})
assert response.status_code == 200, response.content

results = response.json()
assert results['running'] is True
assert results['num_workers'] == initial_workers + 3

response = client.post('/daemon/decrease', params={'amount': 2})
assert response.status_code == 200, response.content

results = response.json()
assert results['running'] is True
assert results['num_workers'] == initial_workers + 1
Loading