feat(BA-922): Add internal address config (#3989)

Author: HyeockJinKim

changes/3989.feature.md

@@ -0,0 +1 @@
+Separate internal api port

configs/account-manager/halfstack.toml

@@ -18,6 +18,7 @@ pool-pre-ping = false
 [account-manager]
 num-proc = 1
 service-addr = { host = "0.0.0.0", port = 8088 }
+internal-addr = { host = "0.0.0.0", port = 8087 }
 #user = "nobody"
 #group = "nobody"
 ssl-enabled = false

configs/account-manager/sample.toml

@@ -46,6 +46,8 @@ num-proc = 1
 
 # Set the service hostname/port to accept API requests.
 service-addr = { host = "0.0.0.0", port = 8088 }
+# Set the internal hostname/port to accept internal API requests.
+internal-addr = { host = "0.0.0.0", port = 8087 }
 
 # Specify the user/group used for the account-manager daemon,
 # to which the account-manager changes after reading the daemon configuration and SSL certifiactes.

configs/manager/halfstack.toml

@@ -18,6 +18,7 @@ pool-pre-ping = false
 [manager]
 num-proc = 4
 service-addr = { host = "0.0.0.0", port = 8081 }
+internal-addr = { host = "0.0.0.0", port = 8092 }
 #user = "nobody"
 #group = "nobody"
 ssl-enabled = false

configs/manager/sample.toml

@@ -73,6 +73,9 @@ group = "nobody"
 
 # Set the service hostname/port to accept API requests.
 service-addr = { host = "0.0.0.0", port = 8080 }
+
+# Set the internal hostname/port to accept internal API requests.
+internal-addr = { host = "0.0.0.0", port = 8082 }
 # env: BACKEND_SERVICE_IP, BACKEND_SERVICE_PORT
 
 # Set the SSL certificate chain and the private keys used for serving the API requests.

configs/prometheus/prometheus.yaml

@@ -21,19 +21,16 @@ scrape_configs:
       - targets: ['backendai-half-prometheus:9090']
   - job_name: backendai-half-manager
     static_configs:
-      - targets: ['host.docker.internal:8091']
+      - targets: ['host.docker.internal:8092']
   - job_name: backendai-half-account-manager
     static_configs:
-      - targets: ['host.docker.internal:8088']
+      - targets: ['host.docker.internal:8087']
   - job_name: backendai-half-storage-proxy
-    scheme: https
-    tls_config:
-      insecure_skip_verify: true
     static_configs:
-      - targets: ['host.docker.internal:6022']
+      - targets: ['host.docker.internal:6023']
   - job_name: backendai-half-wsproxy
     static_configs:
-      - targets: ['host.docker.internal:5050']
+      - targets: ['host.docker.internal:5051']
   - job_name: backendai-half-agent
     static_configs:
       - targets: ['host.docker.internal:6003']

configs/storage-proxy/halfstack.toml

@@ -56,6 +56,7 @@ ssl-enabled = false
 [api.manager]
 # Manager-facing API
 service-addr = { host = "0.0.0.0", port = 6022 }
+internal-addr = { host = "0.0.0.0", port = 6023 }
 ssl-enabled = false
 # ssl-cert = "configs/storage-proxy/ssl/manager-api-selfsigned.cert.pem"
 # ssl-privkey = "configs/storage-proxy/ssl/manager-api-selfsigned.key.pem"

configs/storage-proxy/sample.toml

@@ -68,6 +68,8 @@ ssl-enabled = false
 # Manager-facing API
 # Recommended to have SSL and bind on a private IP only accessible by managers
 service-addr = { host = "0.0.0.0", port = 6022 }
+# Set the internal hostname/port to accept internal API requests.
+internal-addr = { host = "0.0.0.0", port = 6023 }
 ssl-enabled = true
 ssl-cert = "configs/storage-proxy/ssl/manager-api-selfsigned.cert.pem"
 ssl-privkey = "configs/storage-proxy/ssl/manager-api-selfsigned.key.pem"

configs/wsproxy/halfstack.toml

@@ -3,6 +3,8 @@ bind_host = "0.0.0.0"
 advertised_host = "127.0.0.1"
 
 bind_api_port = 5050
+internal_api_port = 5051
+
 advertised_api_port = 5050
 
 # replace these values with your passphrase

configs/wsproxy/sample.toml

@@ -8,6 +8,8 @@ group = 501
 bind_host = "0.0.0.0"
 advertised_host = "example.com"
 bind_api_port = 5050
+# Set the internal hostname/port to accept internal API requests.
+internal_api_port = 5051
 advertised_api_port = 15050
 bind_proxy_port_range = [
     10200,

src/ai/backend/account_manager/config.py

@@ -294,6 +294,13 @@ class AccountManagerConfig(BaseSchema):
             examples=[HostPortPair(host="127.0.0.1", port=8099)],
         ),
     ]
+    internal_addr: Annotated[
+        HostPortPair,
+        Field(
+            description="Address of account-manager internal service for internal infra communication.",
+            examples=[HostPortPair(host="127.0.0.1", port=8098)],
+        ),
+    ]
     ipc_base_path: Annotated[
         Path,
         Field(

src/ai/backend/account_manager/server.py

@@ -265,9 +265,13 @@ async def _call_cleanup_context_shutdown_handlers(app: web.Application) -> None:
     # should be done in create_app() in other modules.
     cors.add(app.router.add_route("GET", r"", hello))
     cors.add(app.router.add_route("GET", r"/", hello))
-    cors.add(
-        app.router.add_route("GET", r"/metrics", build_prometheus_metrics_handler(metric_registry))
-    )
+    return app
+
+
+def build_internal_app() -> web.Application:
+    app = web.Application()
+    metric_registry = CommonMetricRegistry.instance()
+    app.router.add_route("GET", r"/metrics", build_prometheus_metrics_handler(metric_registry))
     return app
 
 
@@ -278,6 +282,7 @@ async def server_main(
     _args: list[Any],
 ) -> AsyncIterator[None]:
     root_app = build_root_app(pidx, _args[0], subapp_pkgs=global_subapp_pkgs)
+    internal_app = build_internal_app()
     root_ctx: RootContext = root_app["_root.context"]
 
     local_cfg = cast(ServerConfig, root_ctx.local_config)
@@ -327,7 +332,9 @@ async def server_main(
             )
 
         runner = web.AppRunner(root_app, keepalive_timeout=30.0)
+        internal_runner = web.AppRunner(internal_app, keepalive_timeout=30.0)
         await runner.setup()
+        await internal_runner.setup()
         service_addr = am_cfg.service_addr
         site = web.TCPSite(
             runner,
@@ -337,7 +344,15 @@ async def server_main(
             reuse_port=True,
             ssl_context=ssl_ctx,
         )
+        internal_site = web.TCPSite(
+            internal_runner,
+            str(am_cfg.internal_addr.host),
+            am_cfg.internal_addr.port,
+            backlog=1024,
+            reuse_port=True,
+        )
         await site.start()
+        await internal_site.start()
 
         if os.geteuid() == 0:
             uid = am_cfg.user

src/ai/backend/manager/config.py

@@ -265,6 +265,7 @@
             t.Key("user", default=None): tx.UserID(default_uid=_file_perm.st_uid),
             t.Key("group", default=None): tx.GroupID(default_gid=_file_perm.st_gid),
             t.Key("service-addr", default=("0.0.0.0", 8080)): tx.HostPortPair,
+            t.Key("internal-addr", default=("0.0.0.0", 8081)): tx.HostPortPair,
             t.Key(
                 "rpc-auth-manager-keypair", default="fixtures/manager/manager.key_secret"
             ): tx.Path(type="file"),

src/ai/backend/manager/server.py

@@ -855,7 +855,7 @@ def build_root_app(
             sample_rate=local_config["pyroscope"]["sample-rate"],
         )
     )
-    root_ctx = RootContext(metrics=CommonMetricRegistry())
+    root_ctx = RootContext(metrics=CommonMetricRegistry.instance())
     app = web.Application(
         middlewares=[
             exception_middleware,
@@ -942,9 +942,6 @@ async def _call_cleanup_context_shutdown_handlers(app: web.Application) -> None:
     # should be done in create_app() in other modules.
     cors.add(app.router.add_route("GET", r"", hello))
     cors.add(app.router.add_route("GET", r"/", hello))
-    cors.add(
-        app.router.add_route("GET", r"/metrics", build_prometheus_metrics_handler(root_ctx.metrics))
-    )
     if subapp_pkgs is None:
         subapp_pkgs = []
     for pkg_name in subapp_pkgs:
@@ -959,6 +956,13 @@ async def _call_cleanup_context_shutdown_handlers(app: web.Application) -> None:
     return app
 
 
+def build_internal_app() -> web.Application:
+    app = web.Application()
+    metric_registry = CommonMetricRegistry.instance()
+    app.router.add_route("GET", r"/metrics", build_prometheus_metrics_handler(metric_registry))
+    return app
+
+
 def build_public_app(
     root_ctx: RootContext,
     subapp_pkgs: Iterable[str] | None = None,
@@ -982,6 +986,7 @@ async def server_main(
     _args: List[Any],
 ) -> AsyncIterator[None]:
     root_app = build_root_app(pidx, _args[0], subapp_pkgs=global_subapp_pkgs)
+    internal_app = build_internal_app()
     root_ctx: RootContext = root_app["_root.context"]
 
     # Start aiomonitor.
@@ -1021,8 +1026,11 @@ async def server_main(
                 )
 
             runner = web.AppRunner(root_app, keepalive_timeout=30.0)
+            internal_runner = web.AppRunner(internal_app, keepalive_timeout=30.0)
             await runner.setup()
+            await internal_runner.setup()
             service_addr = cast(HostPortPair, root_ctx.local_config["manager"]["service-addr"])
+            internal_addr = cast(HostPortPair, root_ctx.local_config["manager"]["internal-addr"])
             site = web.TCPSite(
                 runner,
                 str(service_addr.host),
@@ -1031,7 +1039,15 @@ async def server_main(
                 reuse_port=True,
                 ssl_context=ssl_ctx,
             )
+            internal_site = web.TCPSite(
+                internal_runner,
+                str(internal_addr.host),
+                internal_addr.port,
+                backlog=1024,
+                reuse_port=True,
+            )
             await site.start()
+            await internal_site.start()
             public_metrics_port = cast(
                 Optional[int], root_ctx.local_config["manager"]["public-metrics-port"]
             )

src/ai/backend/storage/api/manager.py

@@ -42,7 +42,11 @@
     VolumeMounted,
     VolumeUnmounted,
 )
-from ai.backend.common.metrics.http import build_api_metric_middleware
+from ai.backend.common.metrics.http import (
+    build_api_metric_middleware,
+    build_prometheus_metrics_handler,
+)
+from ai.backend.common.metrics.metric import CommonMetricRegistry
 from ai.backend.common.types import AgentId, BinarySize, ItemResult, QuotaScopeID, ResultSet
 from ai.backend.logging import BraceStyleAdapter
 
@@ -105,13 +109,6 @@ async def check_status(request: web.Request) -> web.Response:
         )
 
 
-@skip_token_auth
-async def prometheus_metrics_handler(request: web.Request) -> web.Response:
-    root_ctx: RootContext = request.app["ctx"]
-    metrics = root_ctx.metric_registry.to_prometheus()
-    return web.Response(text=metrics, content_type="text/plain")
-
-
 @ctxmgr
 def handle_fs_errors(
     volume: AbstractVolume,
@@ -1214,7 +1211,6 @@ async def init_manager_app(ctx: RootContext) -> web.Application:
     app["app_ctx"] = app_ctx
     app.on_shutdown.append(_shutdown)
     app.router.add_route("GET", "/", check_status)
-    app.router.add_route("GET", "/metrics", prometheus_metrics_handler)
     app.router.add_route("GET", "/status", check_status)
     app.router.add_route("GET", "/volumes", get_volumes)
     app.router.add_route("GET", "/volume/hwinfo", get_hwinfo)
@@ -1251,6 +1247,13 @@ async def init_manager_app(ctx: RootContext) -> web.Application:
     return app
 
 
+def init_internal_app() -> web.Application:
+    app = web.Application()
+    metric_registry = CommonMetricRegistry.instance()
+    app.router.add_route("GET", "/metrics", build_prometheus_metrics_handler(metric_registry))
+    return app
+
+
 async def handle_volume_mount(
     context: RootContext,
     source: AgentId,

src/ai/backend/storage/config.py

@@ -100,6 +100,9 @@
                             t.Key("service-addr"): tx.HostPortPair(
                                 allow_blank_host=True,
                             ),
+                            t.Key("internal-addr", default=("127.0.0.1", 6023)): tx.HostPortPair(
+                                allow_blank_host=True,
+                            ),
                             t.Key("ssl-enabled"): t.ToBool,
                             t.Key("ssl-cert", default=None): t.Null | tx.Path(type="file"),
                             t.Key("ssl-privkey", default=None): t.Null | tx.Path(type="file"),

src/ai/backend/storage/context.py

@@ -26,7 +26,7 @@
 from ai.backend.logging import BraceStyleAdapter
 
 from .api.client import init_client_app
-from .api.manager import init_manager_app
+from .api.manager import init_internal_app, init_manager_app
 from .exception import InvalidVolumeError
 from .plugin import (
     BasePluginContext,
@@ -173,6 +173,7 @@ async def __aenter__(self) -> None:
         # TODO: Setup the apps outside of the context.
         self.client_api_app = await init_client_app(self)
         self.manager_api_app = await init_manager_app(self)
+        self.internal_api_app = init_internal_app()
         self.backends = {
             **DEFAULT_BACKENDS,
         }

src/ai/backend/storage/server.py

@@ -206,10 +206,13 @@ async def server_main(
                 )
             client_api_runner = web.AppRunner(ctx.client_api_app)
             manager_api_runner = web.AppRunner(ctx.manager_api_app)
+            internal_api_runner = web.AppRunner(ctx.internal_api_app)
             await client_api_runner.setup()
             await manager_api_runner.setup()
+            await internal_api_runner.setup()
             client_service_addr = local_config["api"]["client"]["service-addr"]
             manager_service_addr = local_config["api"]["manager"]["service-addr"]
+            internal_addr = local_config["api"]["manager"]["internal-addr"]
             client_api_site = web.TCPSite(
                 client_api_runner,
                 str(client_service_addr.host),
@@ -226,8 +229,16 @@ async def server_main(
                 reuse_port=True,
                 ssl_context=manager_ssl_ctx,
             )
+            internal_api_site = web.TCPSite(
+                internal_api_runner,
+                str(internal_addr.host),
+                internal_addr.port,
+                backlog=1024,
+                reuse_port=True,
+            )
             await client_api_site.start()
             await manager_api_site.start()
+            await internal_api_site.start()
             if _is_root():
                 uid = local_config["storage-proxy"]["user"]
                 gid = local_config["storage-proxy"]["group"]

src/ai/backend/wsproxy/config.py

@@ -390,6 +390,9 @@ class WSProxyConfig(BaseSchema):
     bind_api_port: Annotated[
         int, Field(default=5050, description="Port number to bind for API server")
     ]
+    internal_api_port: Annotated[
+        int, Field(default=5051, description="Port number to bind for internal API server")
+    ]
     advertised_api_port: Annotated[
         int | None,
         Field(default=None, examples=[15050], description="API port number reachable from client"),