additional base commands for control server

This commit is contained in:
Daniil Fajnberg 2022-03-07 14:34:40 +01:00
parent 7c66604ad0
commit d05f84b2c3
2 changed files with 73 additions and 5 deletions

View File

@ -37,10 +37,21 @@ class CLIENT_INFO:
class CMD: class CMD:
__slots__ = () __slots__ = ()
# Base commands:
CMD = 'command' CMD = 'command'
NAME = 'name' NAME = 'name'
POOL_SIZE = 'pool-size' POOL_SIZE = 'pool-size'
IS_LOCKED = 'is-locked'
LOCK = 'lock'
UNLOCK = 'unlock'
NUM_RUNNING = 'num-running' NUM_RUNNING = 'num-running'
NUM_CANCELLATIONS = 'num-cancellations'
NUM_ENDED = 'num-ended'
NUM_FINISHED = 'num-finished'
IS_FULL = 'is-full'
GET_GROUP_IDS = 'get-group-ids'
# Simple commands:
START = 'start' START = 'start'
STOP = 'stop' STOP = 'stop'
STOP_ALL = 'stop-all' STOP_ALL = 'stop-all'

View File

@ -23,7 +23,7 @@ import logging
import json import json
from argparse import ArgumentError, HelpFormatter from argparse import ArgumentError, HelpFormatter
from asyncio.streams import StreamReader, StreamWriter from asyncio.streams import StreamReader, StreamWriter
from typing import Callable, Optional, Union, TYPE_CHECKING from typing import Callable, Optional, Type, Union, TYPE_CHECKING
from .constants import CMD, SESSION_WRITER, SESSION_MSG_BYTES, CLIENT_INFO from .constants import CMD, SESSION_WRITER, SESSION_MSG_BYTES, CLIENT_INFO
from .exceptions import HelpRequested, NotATaskPool, UnknownTaskPoolClass from .exceptions import HelpRequested, NotATaskPool, UnknownTaskPoolClass
@ -108,19 +108,36 @@ class ControlSession:
These include commands mapping to the following pool methods: These include commands mapping to the following pool methods:
- __str__ - __str__
- pool_size (get/set property) - pool_size (get/set property)
- is_locked
- lock & unlock
- num_running - num_running
""" """
self._add_command(CMD.NAME, short_help=get_first_doc_line(self._pool.__class__.__str__)) cls: Type[BaseTaskPool] = self._pool.__class__
self._add_command(CMD.NAME, short_help=get_first_doc_line(cls.__str__))
self._add_command( self._add_command(
CMD.POOL_SIZE, CMD.POOL_SIZE,
short_help="Get/set the maximum number of tasks in the pool.", short_help="Get/set the maximum number of tasks in the pool.",
formatter_class=HelpFormatter formatter_class=HelpFormatter
).add_optional_num_argument( ).add_optional_num_argument(
default=None, default=None,
help=f"If passed a number: {get_first_doc_line(self._pool.__class__.pool_size.fset)} " help=f"If passed a number: {get_first_doc_line(cls.pool_size.fset)} "
f"If omitted: {get_first_doc_line(self._pool.__class__.pool_size.fget)}" f"If omitted: {get_first_doc_line(cls.pool_size.fget)}"
)
self._add_command(CMD.IS_LOCKED, short_help=get_first_doc_line(cls.is_locked.fget))
self._add_command(CMD.LOCK, short_help=get_first_doc_line(cls.lock))
self._add_command(CMD.UNLOCK, short_help=get_first_doc_line(cls.unlock))
self._add_command(CMD.NUM_RUNNING, short_help=get_first_doc_line(cls.num_running.fget))
self._add_command(CMD.NUM_CANCELLATIONS, short_help=get_first_doc_line(cls.num_cancellations.fget))
self._add_command(CMD.NUM_ENDED, short_help=get_first_doc_line(cls.num_ended.fget))
self._add_command(CMD.NUM_FINISHED, short_help=get_first_doc_line(cls.num_finished.fget))
self._add_command(CMD.IS_FULL, short_help=get_first_doc_line(cls.is_full.fget))
self._add_command(
CMD.GET_GROUP_IDS, short_help=get_first_doc_line(cls.get_group_ids)
).add_argument(
'group_name',
nargs='*',
help="Must be a name of a task group that exists within the pool."
) )
self._add_command(CMD.NUM_RUNNING, short_help=get_first_doc_line(self._pool.__class__.num_running.fget))
def _add_simple_commands(self) -> None: def _add_simple_commands(self) -> None:
""" """
@ -227,11 +244,51 @@ class ControlSession:
log.debug("%s requests setting pool size to %s", self._client_class_name, num) log.debug("%s requests setting pool size to %s", self._client_class_name, num)
await self._write_function_output(self._pool.__class__.pool_size.fset, self._pool, num) await self._write_function_output(self._pool.__class__.pool_size.fset, self._pool, num)
async def _cmd_is_locked(self, **_kwargs) -> None:
"""Maps to the `is_locked` property of any task pool class."""
log.debug("%s checks locked status", self._client_class_name)
await self._write_function_output(self._pool.__class__.is_locked.fget, self._pool)
async def _cmd_lock(self, **_kwargs) -> None:
"""Maps to the `lock` method of any task pool class."""
log.debug("%s requests locking the pool", self._client_class_name)
await self._write_function_output(self._pool.lock)
async def _cmd_unlock(self, **_kwargs) -> None:
"""Maps to the `unlock` method of any task pool class."""
log.debug("%s requests unlocking the pool", self._client_class_name)
await self._write_function_output(self._pool.unlock)
async def _cmd_num_running(self, **_kwargs) -> None: async def _cmd_num_running(self, **_kwargs) -> None:
"""Maps to the `num_running` property of any task pool class.""" """Maps to the `num_running` property of any task pool class."""
log.debug("%s requests number of running tasks", self._client_class_name) log.debug("%s requests number of running tasks", self._client_class_name)
await self._write_function_output(self._pool.__class__.num_running.fget, self._pool) await self._write_function_output(self._pool.__class__.num_running.fget, self._pool)
async def _cmd_num_cancellations(self, **_kwargs) -> None:
"""Maps to the `num_cancellations` property of any task pool class."""
log.debug("%s requests number of cancelled tasks", self._client_class_name)
await self._write_function_output(self._pool.__class__.num_cancellations.fget, self._pool)
async def _cmd_num_ended(self, **_kwargs) -> None:
"""Maps to the `num_ended` property of any task pool class."""
log.debug("%s requests number of ended tasks", self._client_class_name)
await self._write_function_output(self._pool.__class__.num_ended.fget, self._pool)
async def _cmd_num_finished(self, **_kwargs) -> None:
"""Maps to the `num_finished` property of any task pool class."""
log.debug("%s requests number of finished tasks", self._client_class_name)
await self._write_function_output(self._pool.__class__.num_finished.fget, self._pool)
async def _cmd_is_full(self, **_kwargs) -> None:
"""Maps to the `is_full` property of any task pool class."""
log.debug("%s checks full status", self._client_class_name)
await self._write_function_output(self._pool.__class__.is_full.fget, self._pool)
async def _cmd_get_group_ids(self, **kwargs) -> None:
"""Maps to the `get_group_ids` method of any task pool class."""
log.debug("%s requests task ids for groups %s", self._client_class_name, kwargs['group_name'])
await self._write_function_output(self._pool.get_group_ids, *kwargs['group_name'])
async def _cmd_start(self, **kwargs) -> None: async def _cmd_start(self, **kwargs) -> None:
"""Maps to the `start` method of the `SimpleTaskPool` class.""" """Maps to the `start` method of the `SimpleTaskPool` class."""
num = kwargs[NUM] num = kwargs[NUM]