from collections.abc import Callable
from configparser import ConfigParser, MissingSectionHeaderError
from logging import getLogger
from pathlib import Path
from time import sleep
from typing import TYPE_CHECKING, Annotated
from xmlrpc.client import Fault
from typer import Argument, Exit, Option, Typer
from ..client import SupervisorMethodResult, SupervisorRemoteXMLRPCClient
from ..config import SupervisorConvenienceConfiguration
from .common import SupervisorTaskStep
if TYPE_CHECKING:
CommandConfig = Path | SupervisorConvenienceConfiguration
else:
CommandConfig = Path
log = getLogger(__name__)
__all__ = (
"check_programs",
"kill_supervisor",
"main",
"remove_supervisor_config",
"restart_programs",
"start_programs",
"start_supervisor",
"stop_programs",
"stop_supervisor",
"write_supervisor_config",
)
def _check_exists(cfg: SupervisorConvenienceConfiguration) -> bool:
log.info(f"Checking for working dir {cfg.working_dir} and config path {cfg.config_path}")
assert cfg.config_path is not None
assert cfg.working_dir is not None
if cfg.working_dir.exists() and cfg.config_path.exists():
# its probably already been written
log.info(f"Working dir {cfg.working_dir} and config path {cfg.config_path} exist")
return True
log.info(f"Working dir {cfg.working_dir} and/or config path {cfg.config_path} do not exist")
return False
def _check_same(cfg: SupervisorConvenienceConfiguration) -> bool:
log.info("Checking if config file matches")
assert cfg.config_path is not None
if _check_exists(cfg) and cfg.config_path.read_text().strip() == cfg.to_cfg().strip():
# same file contents
log.info(f"Config file {cfg.config_path} matches")
return True
elif not _check_exists(cfg):
# no file exists, so we can write it
log.info(f"Config file {cfg.config_path} does not exist, so we can write it")
return True
log.info(f"Config file {cfg.config_path} does not match")
return False
def _check_running(cfg: SupervisorConvenienceConfiguration) -> bool:
if _check_same(cfg):
# check if running
log.info("Checking if supervisor is running")
if cfg.running():
log.info("Supervisor is running")
return True
log.info("Supervisor is not running")
log.info("Supervisor config file does not match, so we can't check if it's running")
return False
def _wait_or_while(until: Callable, unless: Callable | None = None, timeout: int = 5) -> bool:
log.info(f"Waiting for {timeout} seconds")
for _ in range(timeout):
if until():
log.info("`until` condition met")
return True
if unless and unless():
log.info("`unless` condition met")
return False
sleep(1)
log.info("Timed out after {timeout} seconds")
return False
def _raise_or_exit(val: bool, exit: bool):
if exit:
raise Exit(int(not val))
return val
def _load_or_pass(cfg: str | Path | SupervisorConvenienceConfiguration) -> SupervisorConvenienceConfiguration:
if isinstance(cfg, Path):
cfg = SupervisorConvenienceConfiguration.model_validate_json(cfg.read_text())
if isinstance(cfg, str):
cfg = SupervisorConvenienceConfiguration.model_validate_json(cfg)
if not isinstance(cfg, SupervisorConvenienceConfiguration):
raise NotImplementedError
return cfg
def _write_config(cfg: SupervisorConvenienceConfiguration) -> bool:
assert cfg.config_path is not None
if not _check_same(cfg):
previous = cfg.model_copy(deep=True)
parser = ConfigParser(interpolation=None)
try:
parser.read(cfg.config_path)
except MissingSectionHeaderError:
pass
if parser.has_option("supervisord", "pidfile"):
pidfile = parser.get("supervisord", "pidfile").replace("%(here)s", str(cfg.config_path.parent))
previous.supervisord.pidfile = Path(pidfile)
if previous.running() and not stop_supervisor(previous, _exit=False):
log.critical("Supervisor did not stop; preserving existing configuration")
return False
log.info(f"Writing supervisor config to {cfg.config_path}")
cfg._write_self()
return True
[docs]
def write_supervisor_config(cfg_json: str, _exit: Annotated[bool, Argument(hidden=True)] = True):
"""Write a SupervisorConvenienceConfiguration JSON as a supervisor config file
Args:
cfg_json (str): JSON string of SupervisorConvenienceConfiguration
"""
# NOTE: typer does not support union types
log.info(f"Loading JSON config: {cfg_json}")
cfg_obj = _load_or_pass(cfg_json)
return _raise_or_exit(_write_config(cfg_obj), _exit)
[docs]
def start_supervisor(
cfg: Annotated[CommandConfig, Option(exists=True, file_okay=True, dir_okay=False, writable=False, readable=True, resolve_path=True)] = Path(
"pydantic.json"
),
_exit: Annotated[bool, Argument(hidden=True)] = True,
):
"""Start a supervisor instance using supervisord in background
Args:
cfg (Annotated[Path, Option, optional): Path to JSON file of SupervisorConvenienceConfiguration
"""
# NOTE: typer does not support union types
cfg_obj = _load_or_pass(cfg)
if not _check_same(cfg_obj) and not _write_config(cfg_obj):
return _raise_or_exit(False, _exit)
if cfg_obj.running():
log.info("Supervisor is already running")
return _raise_or_exit(True, _exit)
log.info("Starting supervisor")
cfg_obj.start(daemon=True)
running = _wait_or_while(until=lambda: cfg_obj.running(), timeout=cfg_obj.command_timeout)
if not running:
log.critical(f"Supervisor still not running {cfg_obj.command_timeout}s after start command!")
return _raise_or_exit(False, _exit)
log.info("Supervisor started")
return _raise_or_exit(True, _exit)
[docs]
def start_programs(
cfg: Annotated[CommandConfig, Option(exists=True, file_okay=True, dir_okay=False, writable=False, readable=True, resolve_path=True)] = Path(
"pydantic.json"
),
restart: bool = False,
_exit: Annotated[bool, Argument(hidden=True)] = True,
):
"""Start all programs in the supervisor instance
Args:
cfg (Annotated[Path, Option, optional): Path to JSON file of SupervisorConvenienceConfiguration
restart (bool, optional): if true, restart all programs. Defaults to False.
"""
# NOTE: typer does not support union types
cfg_obj = _load_or_pass(cfg)
log.info(f"Initializing client with config: {cfg_obj}")
client = SupervisorRemoteXMLRPCClient(cfg=cfg_obj)
if restart:
# No builtin restart, so stop everything then start again on the next line
log.info("Stopping all processes before restarting")
client.stopAllProcesses()
log.info("Starting all processes")
ret = client.startAllProcesses()
log.info(ret)
_wait_or_while(
until=lambda: all(_.running() for _ in client.getProgramProcessInfo()),
unless=lambda: any(_.stopped() for _ in client.getProgramProcessInfo()),
timeout=cfg_obj.command_timeout,
)
all_ok = _wait_or_while(
until=lambda: all(_.ok(ok_exitstatuses=cfg_obj.exitcodes) for _ in client.getProgramProcessInfo()),
unless=lambda: any(_.bad(ok_exitstatuses=cfg_obj.exitcodes) for _ in client.getProgramProcessInfo()),
timeout=cfg_obj.command_timeout,
)
if not all_ok:
log.critical("Not all processes started successfully")
for r in client.getProgramProcessInfo():
log.info(r.model_dump_json(exclude_unset=True))
return _raise_or_exit(False, _exit)
log.info("All processes started")
return _raise_or_exit(True, _exit)
[docs]
def check_programs(
cfg: Annotated[CommandConfig, Option(exists=True, file_okay=True, dir_okay=False, writable=False, readable=True, resolve_path=True)] = Path(
"pydantic.json"
),
check_running: bool = False,
check_done: bool = False,
_exit: Annotated[bool, Argument(hidden=True)] = True,
):
"""Check if programs are in a good state.
Args:
cfg (Annotated[Path, Option, optional): Path to JSON file of SupervisorConvenienceConfiguration
check_running (bool, optional): if true, only return true if they're running
check_done (bool, optional): if true, only return true if they're done (cleanly)
"""
# NOTE: typer does not support union types
cfg_obj = _load_or_pass(cfg)
log.info(f"Initializing client with config: {cfg_obj}")
client = SupervisorRemoteXMLRPCClient(cfg=cfg_obj)
log.info("Checking all processes")
ret = client.getProgramProcessInfo()
for r in ret:
log.info(r.model_dump_json(exclude_unset=True))
ok = False
if check_running:
if ret and all(p.running() for p in ret):
log.info("All processes running")
ok = True
else:
log.warning("Not all processes running")
elif check_done:
if ret and all(p.done(ok_exitstatuses=cfg_obj.exitcodes) for p in ret):
log.info("All processes done")
ok = True
else:
log.info("Not all processes done")
else:
if ret and all(p.ok(ok_exitstatuses=cfg_obj.exitcodes) for p in ret):
log.info("All processes ok")
ok = True
else:
log.info("Not all processes ok")
return _raise_or_exit(ok, _exit)
[docs]
def stop_programs(
cfg: Annotated[CommandConfig, Option(exists=True, file_okay=True, dir_okay=False, writable=False, readable=True, resolve_path=True)] = Path(
"pydantic.json"
),
_exit: Annotated[bool, Argument(hidden=True)] = True,
):
"""Stop all programs in the supervisor instance
Args:
cfg (Annotated[Path, Option, optional): Path to JSON file of SupervisorConvenienceConfiguration
"""
# NOTE: typer does not support union types
cfg_obj = _load_or_pass(cfg)
log.info(f"Initializing client with config: {cfg_obj}")
client = SupervisorRemoteXMLRPCClient(cfg=cfg_obj)
log.info("Stopping all processes")
ret = client.stopAllProcesses()
log.info(ret)
all_stopped = _wait_or_while(until=lambda: all(_.stopped() for _ in client.getProgramProcessInfo()), timeout=cfg_obj.command_timeout)
if not all_stopped:
log.critical("Not all processes stopped successfully")
for r in client.getProgramProcessInfo():
log.info(r.model_dump_json(exclude_unset=True))
return _raise_or_exit(False, _exit)
log.info("All processes stopped")
return _raise_or_exit(True, _exit)
[docs]
def restart_programs(
cfg: Annotated[CommandConfig, Option(exists=True, file_okay=True, dir_okay=False, writable=False, readable=True, resolve_path=True)] = Path(
"pydantic.json"
),
force: bool = False,
_exit: Annotated[bool, Argument(hidden=True)] = True,
):
"""Restart all programs in the supervisor instance
Args:
cfg (Annotated[Path, Option, optional): Path to JSON file of SupervisorConvenienceConfiguration
force (bool, optional): if true, force restart. Defaults to False.
"""
if force:
log.info("Force restarting all processes")
if not stop_programs(cfg, _exit=False):
log.warning("Could not stop programs")
return _raise_or_exit(False, _exit)
if not start_programs(cfg, _exit=False):
log.warning("Could not start programs")
return _raise_or_exit(False, _exit)
return _raise_or_exit(True, _exit)
[docs]
def stop_supervisor(
cfg: Annotated[CommandConfig, Option(exists=True, file_okay=True, dir_okay=False, writable=False, readable=True, resolve_path=True)] = Path(
"pydantic.json"
),
_exit: Annotated[bool, Argument(hidden=True)] = True,
):
"""Stop the supervisor instance
Args:
cfg (Annotated[Path, Option, optional): Path to JSON file of SupervisorConvenienceConfiguration
"""
# NOTE: typer does not support union types
cfg_obj = _load_or_pass(cfg)
log.info("Stopping supervisor")
cfg_obj.stop()
not_running = _wait_or_while(until=lambda: not cfg_obj.running(), timeout=cfg_obj.command_timeout)
if not not_running:
log.critical(f"Still running {cfg_obj.command_timeout}s after stop command!")
return _raise_or_exit(False, _exit)
return _raise_or_exit(True, _exit)
[docs]
def kill_supervisor(
cfg: Annotated[CommandConfig, Option(exists=True, file_okay=True, dir_okay=False, writable=False, readable=True, resolve_path=True)] = Path(
"pydantic.json"
),
_exit: Annotated[bool, Argument(hidden=True)] = True,
):
"""Kill the supervisor instance with os.kill
Args:
cfg (Annotated[Path, Option, optional): Path to JSON file of SupervisorConvenienceConfiguration
"""
try:
if not stop_programs(cfg, _exit=False):
log.warning("could not stop programs")
except ConnectionRefusedError:
# supervisor already down, continue
...
except Fault as error:
if error.faultCode != SupervisorMethodResult.SHUTDOWN_STATE.value:
raise
# NOTE: typer does not support union types
cfg_obj = _load_or_pass(cfg)
log.info("Killing supervisor")
cfg_obj.kill()
stopped = _wait_or_while(until=lambda: not cfg_obj.running(), timeout=cfg_obj.command_timeout)
if not stopped:
log.critical(f"Still running {cfg_obj.command_timeout}s after kill command!")
return _raise_or_exit(False, _exit)
return _raise_or_exit(True, _exit)
[docs]
def remove_supervisor_config(
cfg: Annotated[CommandConfig, Option(exists=True, file_okay=True, dir_okay=False, writable=False, readable=True, resolve_path=True)] = Path(
"pydantic.json"
),
_exit: Annotated[bool, Argument(hidden=True)] = True,
):
"""Remove the supervisor config file and working directory
Args:
cfg (Annotated[Path, Option, optional): Path to JSON file of SupervisorConvenienceConfiguration
"""
# NOTE: typer does not support union types
cfg_obj = _load_or_pass(cfg)
log.info("Removing supervisor config")
stopped = stop_supervisor(cfg_obj, _exit=False)
if not stopped:
log.critical("Supervisor still running after stop command!")
stopped = kill_supervisor(cfg_obj, _exit=False)
if not stopped:
log.critical("Supervisor still running after kill command!")
return _raise_or_exit(False, _exit)
log.info(f"Sleeping for {cfg_obj.command_timeout} seconds")
sleep(cfg_obj.command_timeout)
log.info("Removing supervisor config and folder")
cfg_obj.rmdir()
return _raise_or_exit(True, _exit)
def _add_to_typer(app, command: SupervisorTaskStep, foo):
"""Helper function to ensure correct command names"""
app.command(command)(foo)
def main():
app = Typer()
_add_to_typer(app, "configure-supervisor", write_supervisor_config)
_add_to_typer(app, "start-supervisor", start_supervisor)
_add_to_typer(app, "start-programs", start_programs)
_add_to_typer(app, "stop-programs", stop_programs)
_add_to_typer(app, "check-programs", check_programs)
_add_to_typer(app, "restart-programs", restart_programs)
_add_to_typer(app, "stop-supervisor", stop_supervisor)
_add_to_typer(app, "force-kill", kill_supervisor)
_add_to_typer(app, "unconfigure-supervisor", remove_supervisor_config)
app()