chore: import upstream snapshot with attribution

python/sglang/srt/entrypoints/grpc_server.py

@@ -0,0 +1,263 @@
+"""
+Thin gRPC server wrapper — delegates to smg-grpc-servicer package.
+
+A lightweight HTTP sidecar is started alongside the gRPC server to expose:
+- /metrics (Prometheus, when --enable-metrics is set)
+- /start_profile, /stop_profile (profiling control)
+
+The sidecar is started on --smg-http-sidecar-port (default: --port + 1)
+once the gRPC request manager is ready, regardless of whether --enable-metrics
+is set.
+"""
+
+import inspect
+import json
+import logging
+import time
+
+from aiohttp import web
+
+from sglang.srt.managers.io_struct import ProfileReq, ProfileReqType
+from sglang.srt.utils.common import get_bool_env_var
+
+logger = logging.getLogger(__name__)
+
+
+async def _start_sidecar_server(host: str, port: int, app):
+    """Start the aiohttp sidecar and return the runner for cleanup."""
+    runner = web.AppRunner(app)
+    await runner.setup()
+    try:
+        site = web.TCPSite(runner, host, port)
+        await site.start()
+    except BaseException:
+        await runner.cleanup()
+        raise
+    logger.info("HTTP sidecar server started on http://%s:%d", host, port)
+    return runner
+
+
+def _add_metrics_routes(app):
+    """Add Prometheus /metrics endpoint to the aiohttp app."""
+    from prometheus_client import (
+        CollectorRegistry,
+        multiprocess,
+    )
+    from prometheus_client.openmetrics.exposition import (
+        CONTENT_TYPE_LATEST,
+        generate_latest,
+    )
+
+    async def metrics_handler(request):
+        try:
+            registry = CollectorRegistry()
+            multiprocess.MultiProcessCollector(registry)
+            data = generate_latest(registry)
+            return web.Response(
+                body=data,
+                headers={"Content-Type": CONTENT_TYPE_LATEST},
+            )
+        except Exception:
+            logger.exception("Failed to generate Prometheus metrics")
+            return web.Response(status=500, text="Failed to generate metrics")
+
+    app.router.add_get("/metrics", metrics_handler)
+
+
+def _check_communicator_results(results, action):
+    """Return a web.Response error if results indicate failure, else None."""
+    if not results:
+        return web.Response(status=500, text="No response from scheduler\n")
+    failures = [r for r in results if not r.success]
+    if failures:
+        msgs = " | ".join(r.message for r in failures)
+        return web.Response(status=500, text=f"{action} failed: {msgs}\n")
+    return None
+
+
+def _add_admin_routes(app, request_manager):
+    """Add admin endpoints to the aiohttp app.
+
+    Endpoints: /start_profile, /stop_profile.
+    Business logic (request construction, env var handling, response interpretation)
+    lives here; request_manager only provides the transport to the scheduler.
+    """
+
+    async def start_profile_handler(request):
+        try:
+            if request.content_length and request.content_length > 0:
+                try:
+                    body = await request.json()
+                except json.JSONDecodeError as e:
+                    return web.Response(
+                        status=400,
+                        text=f"Invalid JSON in request body: {e}",
+                    )
+            else:
+                body = {}
+
+            # Build ProfileReq with env var overrides (same as tokenizer_communicator_mixin)
+            with_stack = body.get("with_stack")
+            env_with_stack = get_bool_env_var("SGLANG_PROFILE_WITH_STACK", "true")
+            with_stack = (with_stack is not False) and env_with_stack
+            record_shapes = body.get("record_shapes")
+            env_record_shapes = get_bool_env_var("SGLANG_PROFILE_RECORD_SHAPES", "true")
+            record_shapes = (record_shapes is not False) and env_record_shapes
+
+            req = ProfileReq(
+                req_type=ProfileReqType.START_PROFILE,
+                output_dir=body.get("output_dir"),
+                start_step=body.get("start_step"),
+                num_steps=body.get("num_steps"),
+                activities=body.get("activities"),
+                with_stack=with_stack,
+                record_shapes=record_shapes,
+                profile_by_stage=body.get("profile_by_stage", False),
+                profile_id=str(time.time()),
+                merge_profiles=body.get("merge_profiles", False),
+                profile_prefix=body.get("profile_prefix"),
+                profile_stages=body.get("profile_stages"),
+            )
+            results = await request_manager.send_communicator_req(
+                req, "profile_communicator", timeout=600.0
+            )
+            err = _check_communicator_results(results, "Start Profile")
+            if err:
+                return err
+            return web.Response(text="Start profiling.\n")
+        except Exception as e:
+            logger.exception("Failed to start profile")
+            return web.Response(
+                status=500,
+                text=f"Internal error: {type(e).__name__}. Check server logs.\n",
+            )
+
+    async def stop_profile_handler(request):
+        try:
+            req = ProfileReq(req_type=ProfileReqType.STOP_PROFILE)
+            results = await request_manager.send_communicator_req(
+                req, "profile_communicator", timeout=600.0
+            )
+            err = _check_communicator_results(results, "Stop profile")
+            if err:
+                return err
+            return web.Response(text="Stop profiling. This will take some time.\n")
+        except Exception as e:
+            logger.exception("Failed to stop profile")
+            return web.Response(
+                status=500,
+                text=f"Internal error: {type(e).__name__}. Check server logs.\n",
+            )
+
+    app.router.add_post("/start_profile", start_profile_handler)
+    app.router.add_post("/stop_profile", stop_profile_handler)
+
+
+async def serve_grpc(server_args, model_info=None):
+    """Start the standalone gRPC server with integrated scheduler."""
+    try:
+        from smg_grpc_servicer.sglang.server import serve_grpc as _serve_grpc
+    except ImportError as e:
+        raise ImportError(
+            "gRPC mode requires the smg-grpc-servicer package. "
+            "If not installed, run: pip install smg-grpc-servicer[sglang]. "
+            "If already installed, there may be a broken import due to a "
+            "version mismatch — see the chained exception above for details."
+        ) from e
+
+    sidecar_app = web.Application()
+    sidecar_runner = None
+    sidecar_port = (
+        server_args.smg_http_sidecar_port
+        if server_args.smg_http_sidecar_port is not None
+        else server_args.port + 1
+    )
+
+    # Metrics setup: must set PROMETHEUS_MULTIPROC_DIR before scheduler
+    # processes import prometheus_client, since the env var is inherited
+    # at fork time.
+    if server_args.enable_metrics:
+        try:
+            from sglang.srt.observability.func_timer import enable_func_timer
+            from sglang.srt.utils import set_prometheus_multiproc_dir
+
+            set_prometheus_multiproc_dir()
+            enable_func_timer()
+            _add_metrics_routes(sidecar_app)
+        except Exception as e:
+            logger.error(
+                "Failed to set up metrics: %s. Continuing without metrics.",
+                e,
+                exc_info=True,
+            )
+
+    async def _on_request_manager_ready(request_manager, srv_args, sched_info):
+        nonlocal sidecar_runner
+        try:
+            _add_admin_routes(sidecar_app, request_manager)
+        except Exception as e:
+            logger.error(
+                "Failed to set up admin routes: %s. "
+                "Continuing without admin endpoints.",
+                e,
+                exc_info=True,
+            )
+        try:
+            sidecar_runner = await _start_sidecar_server(
+                server_args.host, sidecar_port, sidecar_app
+            )
+        except OSError as e:
+            logger.error(
+                "Failed to start HTTP sidecar server: %s. "
+                "Continuing without metrics/profile endpoints.",
+                e,
+                exc_info=True,
+            )
+        except Exception as e:
+            logger.error(
+                "Unexpected error starting HTTP sidecar server: %s. "
+                "Continuing without metrics/profile endpoints.",
+                e,
+                exc_info=True,
+            )
+
+    # Older smg-grpc-servicer releases (≤ 0.5.2) accept only (server_args,
+    # model_info) and reject the on_request_manager_ready hook. The hook is
+    # what calls _start_sidecar_server, so dropping the kwarg disables the
+    # entire HTTP sidecar (Prometheus /metrics and /start_profile +
+    # /stop_profile). Core gRPC serving still works without it.
+    serve_kwargs: dict = {}
+    sidecar_supported = (
+        "on_request_manager_ready" in inspect.signature(_serve_grpc).parameters
+    )
+    if sidecar_supported:
+        serve_kwargs["on_request_manager_ready"] = _on_request_manager_ready
+    elif server_args.enable_metrics:
+        # User explicitly asked for metrics but the installed servicer can't
+        # start the sidecar that serves them — fail loud rather than silently
+        # produce a server with no /metrics endpoint.
+        raise RuntimeError(
+            "--enable-metrics requires smg-grpc-servicer ≥ 0.5.3 (the version "
+            "that accepts 'on_request_manager_ready'); installed version "
+            "lacks the hook so the HTTP sidecar would never start. Upgrade "
+            "smg-grpc-servicer or remove --enable-metrics."
+        )
+    else:
+        logger.warning(
+            "Installed smg-grpc-servicer does not accept "
+            "'on_request_manager_ready'; HTTP sidecar disabled "
+            "(no /metrics, /start_profile, /stop_profile). "
+            "Upgrade smg-grpc-servicer to ≥ 0.5.3 to enable it."
+        )
+
+    try:
+        await _serve_grpc(server_args, model_info, **serve_kwargs)
+    finally:
+        if sidecar_runner is not None:
+            try:
+                await sidecar_runner.cleanup()
+            except Exception as e:
+                logger.exception(
+                    "Failed to cleanly shut down HTTP sidecar server: %s",
+                    e,
+                )