mantisdk 0.1.0__py3-none-any.whl

mantisdk/cli/__init__.py

@@ -0,0 +1,56 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Mantisdk command line interface entry point."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import sys
+from typing import Dict, Iterable, Tuple
+
+_SUBCOMMANDS: Dict[str, Tuple[str, str]] = {
+    "vllm": ("mantisdk.cli.vllm", "Run the vLLM CLI with Mantisdk instrumentation."),
+    "store": ("mantisdk.cli.store", "Run a LightningStore server."),
+    "prometheus": ("mantisdk.cli.prometheus", "Serve Prometheus metrics from the multiprocess registry."),
+    "agentops": ("mantisdk.cli.agentops_server", "Start the AgentOps server manager."),
+}
+
+_DESCRIPTION = "Mantisdk CLI entry point.\n\nAvailable subcommands:\n" + "\n".join(
+    f"  {name:<10}{desc}" for name, (_, desc) in _SUBCOMMANDS.items()
+)
+
+
+def main(argv: Iterable[str] | None = None) -> int:
+    """Dispatch to the requested Mantisdk subcommand."""
+    parser = argparse.ArgumentParser(
+        prog="msk",
+        description=_DESCRIPTION,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("subcommand", choices=_SUBCOMMANDS.keys(), help="Subcommand to run.")
+    parser.add_argument("args", nargs=argparse.REMAINDER, help=argparse.SUPPRESS)
+
+    parsed = parser.parse_args(list(argv) if argv is not None else None)
+    module_name, _ = _SUBCOMMANDS[parsed.subcommand]
+    module = importlib.import_module(module_name)
+
+    entry_point = getattr(module, "main", None)
+    if entry_point is None:
+        parser.error(f"Subcommand '{parsed.subcommand}' does not define a callable 'main'")
+
+    dispatch_args = parsed.args
+    original_argv = sys.argv
+    sys.argv = [f"{parser.prog} {parsed.subcommand}", *dispatch_args]
+    try:
+        result = entry_point(dispatch_args or None)
+    finally:
+        sys.argv = original_argv
+
+    if isinstance(result, int):
+        return result
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

mantisdk/cli/prometheus.py

@@ -0,0 +1,115 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Serve Prometheus metrics from the Mantisdk multiprocess registry."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import logging
+import os
+from pathlib import Path
+from typing import Iterable
+
+from fastapi import FastAPI
+from prometheus_client import make_asgi_app  # pyright: ignore[reportUnknownVariableType]
+
+from mantisdk.logging import setup as setup_logging
+from mantisdk.utils.metrics import get_prometheus_registry
+from mantisdk.utils.server_launcher import PythonServerLauncher, PythonServerLauncherArgs
+
+logger = logging.getLogger(__name__)
+
+
+def ensure_prometheus_dir() -> str:
+    """Ensure PROMETHEUS_MULTIPROC_DIR is set and the directory exists."""
+
+    directory = os.getenv("PROMETHEUS_MULTIPROC_DIR")
+    if directory is None:
+        raise ValueError("PROMETHEUS_MULTIPROC_DIR is not set.")
+
+    Path(directory).mkdir(parents=True, exist_ok=True)
+    logger.info("Serving Prometheus multiprocess metrics from %s", directory)
+    return directory
+
+
+def create_prometheus_app(metrics_path: str = "/v1/prometheus") -> FastAPI:
+    """Create a FastAPI app that exposes Prometheus metrics and a health endpoint.
+
+    Args:
+        metrics_path: URL path to expose the Prometheus metrics endpoint on.
+
+    Returns:
+        A FastAPI application ready to serve metrics.
+    """
+
+    if not metrics_path.startswith("/"):
+        raise ValueError("metrics_path must start with '/'.")
+
+    normalized_path = metrics_path.rstrip("/")
+    if normalized_path in ("", "/"):
+        raise ValueError("metrics_path must not be '/'. Choose a sub-path such as /v1/prometheus.")
+
+    app = FastAPI(title="Mantisdk Prometheus exporter", docs_url=None, redoc_url=None)
+    metrics_app = make_asgi_app(registry=get_prometheus_registry())  # pyright: ignore[reportUnknownVariableType]
+    app.mount(normalized_path, metrics_app)  # pyright: ignore[reportUnknownArgumentType]
+
+    @app.get("/health")
+    async def healthcheck() -> dict[str, str]:  # pyright: ignore[reportUnusedFunction]
+        return {"status": "ok"}
+
+    return app
+
+
+def main(argv: Iterable[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Serve Prometheus metrics outside the LightningStore server.")
+    parser.add_argument("--host", default="0.0.0.0", help="Host to bind the metrics server to.")
+    parser.add_argument("--port", type=int, default=4748, help="Port to expose the Prometheus metrics on.")
+    parser.add_argument(
+        "--metrics-path",
+        default="/v1/prometheus",
+        help="HTTP path used to expose metrics. Must start with '/' and not be the root path.",
+    )
+    parser.add_argument(
+        "--log-level",
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+        help="Configure the logging level for the metrics server.",
+    )
+    parser.add_argument(
+        "--access-log",
+        action="store_true",
+        help="Enable uvicorn access logs. Disabled by default to reduce noise.",
+    )
+    args = parser.parse_args(list(argv) if argv is not None else None)
+
+    setup_logging(args.log_level)
+    ensure_prometheus_dir()
+
+    try:
+        app = create_prometheus_app(args.metrics_path)
+    except ValueError as exc:
+        logger.error("Failed to configure prometheus app: %s", exc)
+        return 1
+
+    launcher_args = PythonServerLauncherArgs(
+        host=args.host,
+        port=args.port,
+        log_level=getattr(logging, args.log_level),
+        access_log=args.access_log,
+        healthcheck_url="/health",
+    )
+    launcher = PythonServerLauncher(app, launcher_args)
+
+    try:
+        asyncio.run(launcher.run_forever())
+    except KeyboardInterrupt:
+        logger.info("Received shutdown signal. Stopping Prometheus server.")
+    except RuntimeError as exc:
+        logger.error("Prometheus server failed to start: %s", exc, exc_info=True)
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

mantisdk/cli/store.py

@@ -0,0 +1,131 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Run a LightningStore server for persistent access from multiple processes."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import logging
+from typing import Iterable, List
+
+from mantisdk import setup_logging
+from mantisdk.store.client_server import LightningStoreServer
+from mantisdk.store.memory import InMemoryLightningStore
+from mantisdk.utils.metrics import (
+    ConsoleMetricsBackend,
+    MetricsBackend,
+    MultiMetricsBackend,
+    PrometheusMetricsBackend,
+    setup_multiprocess_prometheus,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def main(argv: Iterable[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Run a LightningStore server")
+    parser.add_argument("--host", default="0.0.0.0", help="Host to bind the server to")
+    parser.add_argument("--port", type=int, default=4747, help="Port to run the server on")
+    parser.add_argument(
+        "--cors-origin",
+        dest="cors_origins",
+        action="append",
+        help="Allowed CORS origin. Repeat for multiple origins. Use '*' to allow all origins.",
+    )
+    parser.add_argument(
+        "--log-level",
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+        help="Configure the logging level for the store.",
+    )
+    parser.add_argument(
+        "--tracker",
+        nargs="+",
+        choices=["prometheus", "console"],
+        help="Enable metrics tracking. Repeat for multiple trackers.",
+    )
+    parser.add_argument(
+        "--n-workers",
+        default=1,
+        type=int,
+        help=(
+            "Number of workers to run in the server. When it's greater than 1, the server will be run using `mp` launch mode. "
+            "Only applicable for zero-copy stores such as MongoDB backend."
+        ),
+    )
+
+    parser.add_argument(
+        "--backend",
+        choices=["memory", "mongo"],
+        default="memory",
+        help="Backend to use for the store.",
+    )
+    parser.add_argument(
+        "--mongo-uri",
+        default="mongodb://localhost:27017/?replicaSet=rs0",
+        help="MongoDB URI to use for the store. Applicable only if --backend is 'mongo'.",
+    )
+
+    args = parser.parse_args(list(argv) if argv is not None else None)
+
+    setup_logging(args.log_level)
+
+    trackers: List[MetricsBackend] = []
+    if args.tracker:
+        if "prometheus" in args.tracker:
+            logger.info("Enabling Prometheus metrics tracking.")
+            if args.n_workers > 1:
+                # This has to be done before prometheus_client is imported
+                setup_multiprocess_prometheus()
+                logger.info("Setting up Prometheus multiprocess directory for metrics tracking.")
+            trackers.append(PrometheusMetricsBackend())
+
+        if "console" in args.tracker:
+            logger.info("Enabling console metrics tracking.")
+            trackers.append(ConsoleMetricsBackend())
+
+    if len(trackers) == 0:
+        tracker: MetricsBackend | None = None
+    elif len(trackers) == 1:
+        tracker = trackers[0]
+    else:
+        tracker = MultiMetricsBackend(trackers)
+
+    if args.backend == "memory":
+        store = InMemoryLightningStore(
+            thread_safe=True,  # Using thread_safe store for server
+            tracker=tracker,
+        )
+    elif args.backend == "mongo":
+        from mantisdk.store.mongo import MongoLightningStore
+
+        store = MongoLightningStore(mongo_uri=args.mongo_uri, tracker=tracker)
+    else:
+        raise ValueError(f"Invalid backend: {args.backend}")
+
+    if args.n_workers > 1:
+        logger.info(f"Running the server using `mp` launch mode with {args.n_workers} workers.")
+        launch_mode = "mp"
+    else:
+        logger.info("Running the server using `asyncio` launch mode.")
+        launch_mode = "asyncio"
+    server = LightningStoreServer(
+        store,
+        host=args.host,
+        port=args.port,
+        cors_allow_origins=args.cors_origins,
+        launch_mode=launch_mode,
+        tracker=tracker,
+        n_workers=args.n_workers,
+    )
+    try:
+        asyncio.run(server.run_forever())
+    except RuntimeError as exc:
+        logger.error("LightningStore server failed to start: %s", exc, exc_info=True)
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

mantisdk/cli/vllm.py

@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import Iterable
+
+
+def main(argv: Iterable[str] | None = None) -> int:
+    import sys
+
+    from vllm.entrypoints.cli.main import main as vllm_main
+
+    from mantisdk.instrumentation.vllm import instrument_vllm
+
+    instrument_vllm()
+    if argv is not None:
+        original_argv = sys.argv
+        sys.argv = [original_argv[0], *list(argv)]
+        try:
+            vllm_main()
+        finally:
+            sys.argv = original_argv
+    else:
+        vllm_main()
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())