mirror of
https://github.com/FunnyWolf/agentic-soc-platform.git
synced 2026-08-22 13:12:56 +02:00
128 lines
4.4 KiB
Python
128 lines
4.4 KiB
Python
import logging
|
|
import time
|
|
from dataclasses import dataclass
|
|
|
|
from django.core.management.base import CommandError
|
|
|
|
from apps.common.logging import configure_process_file_logging
|
|
from apps.common.worker_health import WorkerHealthReporter
|
|
|
|
SLEEP_ALWAYS = "always"
|
|
SLEEP_WHEN_IDLE = "when_idle"
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class WorkerIterationResult:
|
|
processed: bool = False
|
|
message: str = ""
|
|
|
|
|
|
def add_worker_arguments(parser, *, interval_help="Seconds between worker iterations."):
|
|
parser.add_argument("--once", action="store_true", help="Run one worker iteration and exit.")
|
|
parser.add_argument("--interval", type=float, help=interval_help)
|
|
|
|
|
|
def _coerce_result(result):
|
|
if isinstance(result, WorkerIterationResult):
|
|
return result
|
|
if isinstance(result, bool):
|
|
return WorkerIterationResult(processed=result)
|
|
return WorkerIterationResult(processed=bool(result))
|
|
|
|
|
|
def _resolve_interval(options, default_interval):
|
|
interval = options.get("interval")
|
|
if interval is None:
|
|
interval = default_interval() if callable(default_interval) else default_interval
|
|
if interval <= 0:
|
|
raise CommandError("--interval must be greater than 0.")
|
|
return interval
|
|
|
|
|
|
def _styled(command, style_name, message):
|
|
styler = getattr(command.style, style_name, None)
|
|
return styler(message) if styler else message
|
|
|
|
|
|
def _worker_label(worker_name):
|
|
return worker_name if worker_name.endswith("worker") else f"{worker_name} worker"
|
|
|
|
|
|
def _refresh_runtime_config_cache():
|
|
from apps.settings.runtime_config import invalidate
|
|
|
|
invalidate()
|
|
|
|
|
|
def _run_once_or_raise(worker_label, run_once):
|
|
try:
|
|
_refresh_runtime_config_cache()
|
|
return _coerce_result(run_once())
|
|
except CommandError:
|
|
raise
|
|
except Exception as exc:
|
|
raise CommandError(f"{worker_label} iteration failed: {type(exc).__name__}: {exc}") from exc
|
|
|
|
|
|
def _should_sleep(result, sleep_policy):
|
|
if sleep_policy == SLEEP_ALWAYS:
|
|
return True
|
|
if sleep_policy == SLEEP_WHEN_IDLE:
|
|
return not result.processed
|
|
raise ValueError(f"unsupported worker sleep policy: {sleep_policy}")
|
|
|
|
|
|
def run_worker(
|
|
command,
|
|
*,
|
|
options,
|
|
worker_name,
|
|
worker_type,
|
|
run_once,
|
|
default_interval,
|
|
sleep_policy=SLEEP_WHEN_IDLE,
|
|
started_message=None,
|
|
stopped_message=None,
|
|
sleep_seconds=None,
|
|
log_role=None,
|
|
):
|
|
worker_label = _worker_label(worker_name)
|
|
if sleep_policy not in {SLEEP_ALWAYS, SLEEP_WHEN_IDLE}:
|
|
raise ValueError(f"unsupported worker sleep policy: {sleep_policy}")
|
|
if log_role and not configure_process_file_logging(log_role):
|
|
raise CommandError(f"unsupported worker log role: {log_role}")
|
|
interval = _resolve_interval(options, default_interval)
|
|
|
|
if options["once"]:
|
|
result = _run_once_or_raise(worker_label, run_once)
|
|
command.stdout.write(result.message or f"{worker_label} completed one iteration")
|
|
return
|
|
|
|
command.stdout.write(_styled(command, "SUCCESS", started_message or f"{worker_label} started"))
|
|
health = WorkerHealthReporter(worker_type)
|
|
health.start()
|
|
try:
|
|
while True:
|
|
iteration_started = time.perf_counter()
|
|
health.iteration_started()
|
|
try:
|
|
result = _run_once_or_raise(worker_label, run_once)
|
|
health.iteration_succeeded(result, (time.perf_counter() - iteration_started) * 1000)
|
|
if result.message:
|
|
command.stdout.write(result.message)
|
|
if _should_sleep(result, sleep_policy):
|
|
current_sleep_seconds = sleep_seconds() if sleep_seconds else interval
|
|
if current_sleep_seconds <= 0:
|
|
raise CommandError("worker sleep interval must be greater than 0.")
|
|
time.sleep(current_sleep_seconds)
|
|
except Exception as exc:
|
|
health.iteration_failed(exc.__cause__ or exc, (time.perf_counter() - iteration_started) * 1000)
|
|
logger.exception("%s iteration failed", worker_label)
|
|
command.stderr.write(_styled(command, "ERROR", f"{worker_label} failed: {type(exc).__name__}: {exc}"))
|
|
time.sleep(interval)
|
|
except KeyboardInterrupt:
|
|
health.stop()
|
|
command.stdout.write(_styled(command, "WARNING", stopped_message or f"{worker_label} stopped."))
|