refactor: move standalone CLI into TaskiqCMD subcommand

Author: alexted

pyproject.toml

@@ -90,6 +90,7 @@ taskiq = "taskiq.__main__:main"
 [project.entry-points.taskiq_cli]
 worker = "taskiq.cli.worker.cmd:WorkerCMD"
 scheduler = "taskiq.cli.scheduler.cmd:SchedulerCMD"
+nng-hub = "taskiq.cli.nng.cmd:NNGHubCMD"
 
 [project.entry-points.opentelemetry_instrumentor]
 taskiq = "taskiq.instrumentation:TaskiqInstrumentor"

taskiq/brokers/nng/broker.py

@@ -42,7 +42,7 @@ def _ipc_addr(prefix: str = "taskiq-nng") -> str:
 
 class NNGBroker(AsyncBroker):
     """
-    Taskiq broker backed by a standalone :class:`~taskiq.brokers.nng_hub.NNGHub`.
+    Taskiq broker backed by a dedicated :class:`~taskiq.brokers.nng_hub.NNGHub`.
 
     The hub must be running before workers or clients start.  Launch it with::
 

taskiq/brokers/nng/hub.py

@@ -1,9 +1,9 @@
 """
 NNG hub: central control plane, task dispatcher, and lease manager.
 
-Run as a standalone process::
+Run as a taskiq sub-command::
 
-    taskiq-nng-hub --control-addr ipc:///tmp/taskiq-nng.ipc
+    taskiq nng-hub --control-addr ipc:///tmp/taskiq-nng.ipc
 
 Or embed it in an application for testing::
 
@@ -14,12 +14,9 @@
 """
 from __future__ import annotations
 
-import argparse
 import asyncio
 import base64
 import logging
-import os
-import signal
 import time
 import uuid
 from contextlib import suppress
@@ -369,91 +366,3 @@ async def _reaper_loop(self) -> None:
                 raise
             except Exception:
                 logger.exception("Reaper loop error")
-
-
-# ── standalone CLI entry point ────────────────────────────────────────────────
-
-def _build_config() -> HubConfig:
-    p = argparse.ArgumentParser(
-        description="taskiq-nng-hub — NNG task router, dispatcher, and lease manager",
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-    )
-    p.add_argument(
-        "--control-addr",
-        default=os.getenv("NNG_CONTROL_ADDR", "ipc:///tmp/taskiq-nng.ipc"),
-        help="NNG address the hub listens on.  Env: NNG_CONTROL_ADDR",
-    )
-    p.add_argument(
-        "--max-pending",
-        type=int,
-        default=int(os.getenv("NNG_MAX_PENDING", "10000")),
-    )
-    p.add_argument(
-        "--heartbeat-timeout",
-        type=float,
-        default=float(os.getenv("NNG_HEARTBEAT_TIMEOUT", "15.0")),
-        help="Seconds of silence before a worker is declared dead.",
-    )
-    p.add_argument(
-        "--lease-timeout",
-        type=float,
-        default=float(os.getenv("NNG_LEASE_TIMEOUT", "20.0")),
-        help="Seconds before an unacked task lease is reaped.",
-    )
-    p.add_argument(
-        "--routing-policy",
-        choices=["least_loaded", "p2c", "round_robin"],
-        default=os.getenv("NNG_ROUTING_POLICY", "least_loaded"),
-    )
-    p.add_argument(
-        "--control-concurrency",
-        type=int,
-        default=int(os.getenv("NNG_CONTROL_CONCURRENCY", "16")),
-        help="Number of concurrent Rep0 contexts.",
-    )
-    p.add_argument(
-        "--log-level",
-        default=os.getenv("NNG_LOG_LEVEL", "INFO"),
-        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
-    )
-    args = p.parse_args()
-    logging.basicConfig(
-        level=getattr(logging, args.log_level),
-        format="%(asctime)s %(name)-24s %(levelname)-8s %(message)s",
-    )
-    return HubConfig(
-        control_addr=args.control_addr,
-        max_pending=args.max_pending,
-        heartbeat_timeout=args.heartbeat_timeout,
-        lease_timeout=args.lease_timeout,
-        routing_policy=args.routing_policy,
-        control_concurrency=args.control_concurrency,
-    )
-
-
-async def _run(config: HubConfig) -> None:
-    hub = NNGHub(config)
-    loop = asyncio.get_running_loop()
-    stop_event = asyncio.Event()
-
-    def _on_signal() -> None:
-        logger.info("Shutdown signal received")
-        stop_event.set()
-
-    for sig in (signal.SIGTERM, signal.SIGINT):
-        loop.add_signal_handler(sig, _on_signal)
-
-    await hub.start()
-    try:
-        await stop_event.wait()
-    finally:
-        await hub.stop()
-
-
-def main() -> None:
-    """Entry point for the ``taskiq-nng-hub`` CLI command."""
-    config = _build_config()
-    try:
-        asyncio.run(_run(config))
-    except KeyboardInterrupt:
-        pass

taskiq/cli/nng/__init__.py

@@ -0,0 +1 @@
+"""NNG CLI command."""

taskiq/cli/nng/cmd.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import argparse
+import asyncio
+import logging
+import os
+import signal
+from collections.abc import Sequence
+from contextlib import suppress
+
+from taskiq.abc.cmd import TaskiqCMD
+
+from taskiq.brokers.nng.hub import HubConfig, NNGHub
+
+logger = logging.getLogger(__name__)
+
+
+class NNGHubCMD(TaskiqCMD):
+    """Command to run the NNG hub."""
+
+    short_help = "Run the NNG hub"
+
+    def exec(self, args: Sequence[str]) -> int | None:
+        parser = argparse.ArgumentParser(
+            prog="taskiq nng-hub",
+            description="NNG task router, dispatcher, and lease manager",
+            formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        )
+        parser.add_argument(
+            "--control-addr",
+            default=os.getenv("NNG_CONTROL_ADDR", "ipc:///tmp/taskiq-nng.ipc"),
+        )
+        parser.add_argument(
+            "--max-pending",
+            type=int,
+            default=int(os.getenv("NNG_MAX_PENDING", "10000")),
+        )
+        parser.add_argument(
+            "--heartbeat-timeout",
+            type=float,
+            default=float(os.getenv("NNG_HEARTBEAT_TIMEOUT", "15.0")),
+        )
+        parser.add_argument(
+            "--lease-timeout",
+            type=float,
+            default=float(os.getenv("NNG_LEASE_TIMEOUT", "20.0")),
+        )
+        parser.add_argument(
+            "--routing-policy",
+            choices=["least_loaded", "p2c", "round_robin"],
+            default=os.getenv("NNG_ROUTING_POLICY", "least_loaded"),
+        )
+        parser.add_argument(
+            "--control-concurrency",
+            type=int,
+            default=int(os.getenv("NNG_CONTROL_CONCURRENCY", "16")),
+        )
+        parser.add_argument(
+            "--log-level",
+            default=os.getenv("NNG_LOG_LEVEL", "INFO"),
+            choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+        )
+        ns = parser.parse_args(list(args))
+        logging.basicConfig(
+            level=getattr(logging, ns.log_level),
+            format="%(asctime)s %(name)-24s %(levelname)-8s %(message)s",
+        )
+        config = HubConfig(
+            control_addr=ns.control_addr,
+            max_pending=ns.max_pending,
+            heartbeat_timeout=ns.heartbeat_timeout,
+            lease_timeout=ns.lease_timeout,
+            routing_policy=ns.routing_policy,
+            control_concurrency=ns.control_concurrency,
+        )
+        asyncio.run(self._run(config))
+        return 0
+
+    async def _run(self, config: HubConfig) -> None:
+        hub = NNGHub(config)
+        loop = asyncio.get_running_loop()
+        stop_event = asyncio.Event()
+
+        def _on_signal() -> None:
+            logger.info("Shutdown signal received")
+            stop_event.set()
+
+        for sig in (signal.SIGTERM, signal.SIGINT):
+            with suppress(Exception):
+                loop.add_signal_handler(sig, _on_signal)
+
+        await hub.start()
+        try:
+            await stop_event.wait()
+        finally:
+            await hub.stop()