backend.ai-appproxy-worker 25.19.1__tar.gz

backend_ai_appproxy_worker-25.19.1/MANIFEST.in

@@ -0,0 +1 @@
+include *.py

backend_ai_appproxy_worker-25.19.1/PKG-INFO

@@ -0,0 +1,164 @@
+Metadata-Version: 2.4
+Name: backend.ai-appproxy-worker
+Version: 25.19.1
+Summary: Backend.AI AppProxy Worker
+Home-page: https://github.com/lablup/backend.ai
+Author: Lablup Inc. and contributors
+License: LGPLv3
+Project-URL: Documentation, https://docs.backend.ai/
+Project-URL: Source, https://github.com/lablup/backend.ai
+Classifier: Intended Audience :: Developers
+Classifier: Operating System :: MacOS :: MacOS X
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Environment :: No Input/Output (Daemon)
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Software Development
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Programming Language :: Python :: 3.13
+Classifier: License :: OSI Approved :: GNU Lesser General Public License v3 or later (LGPLv3+)
+Requires-Python: >=3.13,<3.14
+Description-Content-Type: text/markdown
+Requires-Dist: Jinja2~=3.1.6
+Requires-Dist: PyJWT~=2.10.1
+Requires-Dist: aiohttp_cors~=0.8.1
+Requires-Dist: aiohttp_jinja2~=1.6
+Requires-Dist: aiohttp~=3.13.0
+Requires-Dist: aiomonitor~=0.7.0
+Requires-Dist: aiotools~=2.2.3
+Requires-Dist: attrs>=25.3
+Requires-Dist: backend.ai-appproxy-common==25.19.1
+Requires-Dist: backend.ai-common==25.19.1
+Requires-Dist: backend.ai-logging==25.19.1
+Requires-Dist: backend.ai-plugin==25.19.1
+Requires-Dist: click~=8.1.7
+Requires-Dist: memray~=1.17.2
+Requires-Dist: multidict~=6.6.4
+Requires-Dist: prometheus-client~=0.21.1
+Requires-Dist: pydantic[email]~=2.11.3
+Requires-Dist: pyroscope-io~=0.8.8
+Requires-Dist: setproctitle~=1.3.5
+Requires-Dist: tenacity>=9.0
+Requires-Dist: tomli-w~=1.2.0
+Requires-Dist: types-Jinja2
+Requires-Dist: uvloop~=0.22.1; sys_platform != "Windows"
+Requires-Dist: yarl~=1.19.0
+Dynamic: author
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: project-url
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# Backend.AI App Proxy Worker
+
+## Purpose
+
+The App Proxy Worker is a high-performance reverse proxy that routes user traffic to compute session services (Jupyter, SSH, TensorBoard, etc.) running on agents. It receives routing information from the Coordinator and handles SSL/TLS termination, load balancing, and traffic forwarding.
+
+## Key Responsibilities
+
+### 1. Traffic Proxying
+- Proxy HTTP/HTTPS requests to session services
+- Proxy WebSocket connections for interactive services
+- Handle SSL/TLS termination
+- Stream responses efficiently
+
+### 2. Route Resolution
+- Receive routing tables from Coordinator
+- Resolve session services from URLs
+- Cache routing information locally
+- Update routes dynamically
+
+### 3. Health Checking
+- Monitor backend service health
+- Detect failed services
+- Report health status to Coordinator
+- Handle service failover
+
+## Architecture
+
+### 1. Traffic Proxy (Main)
+
+**Framework**: aiohttp + custom reverse proxy
+
+**Port**: 5050 (default, HTTPS)
+
+**Protocol**: HTTP/HTTPS, WebSocket
+
+**Key Features**:
+
+#### HTTP/HTTPS Proxy
+- Route user requests to session services
+- URL Pattern: `https://<worker-domain>/<session-id>/<service-name>/...`
+
+#### WebSocket Proxy
+- Interactive service communication (Jupyter Kernel, SSH, etc.)
+- Real-time log streaming
+
+**Key Characteristics**:
+- SSL/TLS termination (Let's Encrypt auto-certificate)
+- High-performance async proxy
+- Connection pooling and reuse
+- Streaming support (large file downloads)
+- Sticky session support
+- Auto-retry and failover
+
+**Processing Flow**:
+
+#### HTTP Proxy Flow
+```
+User → HTTPS Request → Worker (SSL termination)
+                           ↓
+                       Parse URL (extract session_id, service_name)
+                           ↓
+                       Lookup route from local cache
+                           ↓
+                       Resolve backend address (agent:port)
+                           ↓
+                       Proxy request to agent
+                           ↓
+                       Stream response back to user
+```
+
+#### WebSocket Proxy Flow
+```
+User → WS Upgrade Request → Worker
+                               ↓
+                           Establish WS connection to agent
+                               ↓
+                           Bidirectional message forwarding
+```
+
+### 2. REST API (Management)
+
+**Framework**: aiohttp (async HTTP server)
+
+**Port**: 6040 (default, separate management port)
+
+**Key Features**:
+- Communication with Coordinator
+- Health check endpoints
+- Metrics exposure (Prometheus)
+- Internal management (no external access)
+
+### Component Interaction
+
+**Traffic Proxy Flow**:
+```
+User (Browser) → Worker (Port 5050) → Kernel (on Agent)
+                    │
+                    ├─ SSL/TLS termination
+                    ├─ Route resolution
+                    └─ Traffic proxying
+```
+
+**Management Flow**:
+```
+Coordinator → Worker REST API (Port 6040) → Route updates
+```

backend_ai_appproxy_worker-25.19.1/ai/backend/appproxy/worker/VERSION

@@ -0,0 +1 @@
+25.19.1

backend_ai_appproxy_worker-25.19.1/ai/backend/appproxy/worker/__init__.py

@@ -0,0 +1,3 @@
+from pathlib import Path
+
+__version__ = (Path(__file__).parent / "VERSION").read_text().strip()

backend_ai_appproxy_worker-25.19.1/ai/backend/appproxy/worker/api/health.py

@@ -0,0 +1,78 @@
+from typing import Iterable
+
+import aiohttp_cors
+from aiohttp import web
+
+from ai.backend.appproxy.common.types import CORSOptions, FrontendMode, WebMiddleware
+
+from .. import __version__
+from ..errors import MissingPortProxyConfigError
+from ..types import RootContext
+
+
+async def hello(request: web.Request) -> web.Response:
+    """Health check endpoint with dependency connectivity status"""
+    from ai.backend.common.dto.internal.health import HealthResponse, HealthStatus
+
+    request["do_not_print_access_log"] = True
+
+    root_ctx: RootContext = request.app["_root.context"]
+    connectivity = await root_ctx.health_probe.get_connectivity_status()
+    response = HealthResponse(
+        status=HealthStatus.OK if connectivity.overall_healthy else HealthStatus.DEGRADED,
+        version=__version__,
+        component="appproxy-worker",
+        connectivity=connectivity,
+    )
+    return web.json_response(response.model_dump_json())
+
+
+async def status(request: web.Request) -> web.Response:
+    """
+    Returns health status of worker.
+    """
+    request["do_not_print_access_log"] = True
+
+    root_ctx: RootContext = request.app["_root.context"]
+    worker_config = root_ctx.local_config.proxy_worker
+    if worker_config.frontend_mode == FrontendMode.WILDCARD_DOMAIN:
+        available_slots = 0
+    else:
+        if worker_config.port_proxy is None:
+            raise MissingPortProxyConfigError("Port proxy configuration is required for PORT mode")
+        available_slots = (
+            worker_config.port_proxy.bind_port_range[1]
+            - worker_config.port_proxy.bind_port_range[0]
+            + 1
+        )
+    return web.json_response({
+        "version": __version__,
+        "authority": worker_config.authority,
+        "app_mode": worker_config.frontend_mode,
+        "protocol": worker_config.protocol,
+        "occupied_slots": len(root_ctx.proxy_frontend.circuits),
+        "available_slots": available_slots,
+    })
+
+
+async def init(app: web.Application) -> None:
+    pass
+
+
+async def shutdown(app: web.Application) -> None:
+    pass
+
+
+def create_app(
+    default_cors_options: CORSOptions,
+) -> tuple[web.Application, Iterable[WebMiddleware]]:
+    app = web.Application()
+    app["prefix"] = "health"
+    app.on_startup.append(init)
+    app.on_shutdown.append(shutdown)
+    cors = aiohttp_cors.setup(app, defaults=default_cors_options)
+    add_route = app.router.add_route
+    root_resource = cors.add(app.router.add_resource(r""))
+    cors.add(root_resource.add_route("GET", hello))
+    cors.add(add_route("GET", "/status", status))
+    return app, []

backend_ai_appproxy_worker-25.19.1/ai/backend/appproxy/worker/api/setup.py

@@ -0,0 +1,229 @@
+import urllib.parse
+from typing import Iterable
+from uuid import UUID
+
+import aiohttp
+import jwt
+import yarl
+from aiohttp import web
+from pydantic import AnyUrl, BaseModel
+from tenacity import AsyncRetrying, TryAgain, retry_if_exception_type, wait_exponential
+from tenacity.stop import stop_after_attempt
+
+from ai.backend.appproxy.common.defs import PERMIT_COOKIE_NAME
+from ai.backend.appproxy.common.errors import (
+    InvalidAPIParameters,
+    ServerMisconfiguredError,
+)
+from ai.backend.appproxy.common.types import (
+    CORSOptions,
+    FrontendMode,
+    ProxyProtocol,
+    PydanticResponse,
+    WebMiddleware,
+)
+from ai.backend.appproxy.common.types import SerializableCircuit as Circuit
+from ai.backend.appproxy.common.utils import calculate_permit_hash, pydantic_api_handler
+
+from ..config import (
+    PortProxyConfig,
+    TraefikPortProxyConfig,
+    TraefikWildcardDomainConfig,
+    WildcardDomainConfig,
+)
+from ..coordinator_client import get_circuit_info
+from ..errors import MissingPortConfigError
+from ..types import FrontendServerMode, InteractiveAppInfo, RootContext
+
+
+def generate_proxy_url(
+    config: PortProxyConfig
+    | WildcardDomainConfig
+    | TraefikPortProxyConfig
+    | TraefikWildcardDomainConfig,
+    protocol: str,
+    circuit: Circuit,
+    redirect_path: str | None = None,
+) -> str:
+    # Generate base URL based on config type
+    match config:
+        case PortProxyConfig():
+            base_url = f"{protocol}://{config.advertised_host or config.bind_host}:{circuit.port}"
+        case TraefikPortProxyConfig():
+            base_url = f"{protocol}://{config.advertised_host}:{circuit.port}"
+        case WildcardDomainConfig():
+            base_url = f"{protocol}://{circuit.subdomain}{config.domain}:{config.advertised_port or config.bind_addr.port}"
+        case TraefikWildcardDomainConfig():
+            base_url = f"{protocol}://{circuit.subdomain}{config.domain}:{config.advertised_port}"
+
+    # Append redirect path if provided
+    if redirect_path:
+        # Ensure redirect path starts with /
+        if not redirect_path.startswith("/"):
+            redirect_path = f"/{redirect_path}"
+        return f"{base_url}{redirect_path}"
+
+    return base_url
+
+
+async def ensure_traefik_route_set_up(traefik_api_port: int, circuit: Circuit) -> None:
+    match circuit.protocol:
+        case ProxyProtocol.TCP:
+            proto = "tcp"
+        case _:
+            proto = "http"
+    path = f"/api/{proto}/routers/{circuit.traefik_router_name}"
+
+    base_url = yarl.URL("http://127.0.0.1").with_port(traefik_api_port)
+    async for attempt in AsyncRetrying(
+        wait=wait_exponential(multiplier=0.02, min=0.02, max=5.0),
+        stop=stop_after_attempt(20),
+        retry=retry_if_exception_type(TryAgain),
+    ):
+        with attempt:
+            async with aiohttp.ClientSession(base_url=base_url) as sess:
+                async with sess.get(path) as resp:
+                    if resp.status == 404:
+                        raise TryAgain
+
+
+class ProxySetupRequestModel(BaseModel):
+    token: str
+
+
+class ProxySetupResponseModel(BaseModel):
+    redirect: AnyUrl
+    redirectURI: AnyUrl
+
+
+@pydantic_api_handler(ProxySetupRequestModel)
+async def setup(
+    request: web.Request, params: ProxySetupRequestModel
+) -> web.StreamResponse | PydanticResponse[ProxySetupResponseModel]:
+    root_ctx: RootContext = request.app["_root.context"]
+    jwt_body = jwt.decode(
+        params.token, root_ctx.local_config.secrets.jwt_secret, algorithms=["HS256"]
+    )
+    requested_circuit_id = UUID(jwt_body["circuit"])
+
+    config = root_ctx.local_config.proxy_worker
+    port_config = config.port_proxy  # As a default fallback
+    circuit = await get_circuit_info(root_ctx, request["request_id"], str(requested_circuit_id))
+
+    match config.frontend_mode:
+        case FrontendServerMode.TRAEFIK:
+            if config.traefik is None:
+                raise ServerMisconfiguredError("proxy_worker: Missing 'traefik' config section")
+            match config.traefik.frontend_mode:
+                case FrontendMode.PORT:
+                    port_config = config.traefik.port_proxy
+                    if port_config is None:
+                        raise ServerMisconfiguredError(
+                            "proxy_worker.traefik: Missing 'port_proxy' config section"
+                        )
+                case FrontendMode.WILDCARD_DOMAIN:
+                    port_config = config.traefik.wildcard_domain
+                    if port_config is None:
+                        raise ServerMisconfiguredError(
+                            "proxy_worker.traefik: Missing 'wildcard_domain' config section"
+                        )
+            await ensure_traefik_route_set_up(config.traefik.api_port, circuit)
+        case FrontendServerMode.PORT:
+            port_config = config.port_proxy
+            if port_config is None:
+                raise ServerMisconfiguredError(
+                    "proxy_worker: Missing root-level 'port_proxy' config section"
+                )
+        case FrontendServerMode.WILDCARD_DOMAIN:
+            port_config = config.wildcard_domain
+            if port_config is None:
+                raise ServerMisconfiguredError(
+                    "proxy_worker: Missing root-level 'wildcard_domain' config section"
+                )
+        case _:
+            raise ServerMisconfiguredError(
+                f"proxy_worker: Invalid root-level 'frontend_mode': {config.frontend_mode}"
+            )
+    if port_config is None:
+        raise MissingPortConfigError("Port configuration is required")
+
+    use_tls = config.tls_advertised or config.tls_listen
+    if not isinstance(circuit.app_info, InteractiveAppInfo):
+        raise InvalidAPIParameters("E20011: Not supported for inference apps")
+
+    # Web browsers block redirect between cross-origins if Access-Control-Allow-Origin value is set to a concrete Origin instead of wildcard;
+    # Hence we need to send "*" as allowed origin manually, instead of benefiting from aiohttp-cors
+    cors_headers = {
+        "Access-Control-Allow-Origin": "*",
+        "Access-Control-Allow-Headers": "*",
+        "Access-Control-Expose-Headers": "*",
+    }
+    match circuit.protocol:
+        case ProxyProtocol.HTTP:
+            protocol = "https" if use_tls else "http"
+            redirect_path = jwt_body.get("redirect", "")
+            proxy_url = generate_proxy_url(port_config, protocol, circuit, redirect_path)
+            response = web.HTTPPermanentRedirect(proxy_url, headers=cors_headers)
+            cookie_domain = None
+            if circuit.frontend_mode == FrontendMode.WILDCARD_DOMAIN:
+                wildcard_info = config.wildcard_domain
+                if not wildcard_info:
+                    raise ServerMisconfiguredError("worker:proxy-worker.wildcard-domain")
+                cookie_domain = wildcard_info.domain
+            response.set_cookie(
+                PERMIT_COOKIE_NAME,
+                calculate_permit_hash(root_ctx.local_config.permit_hash, circuit.app_info.user_id),
+                domain=cookie_domain,
+                httponly=True,
+                secure=use_tls,
+                samesite="Lax",
+                max_age=604800,  # 7 days
+            )
+            return response
+        case ProxyProtocol.TCP:
+            protocol = "tcp"
+            queryparams = {
+                "directTCP": "true",
+                "auth": params.token,
+                "proto": protocol,
+                "gateway": generate_proxy_url(port_config, protocol, circuit, redirect_path=None),
+            }
+            if jwt_body["redirect"]:
+                return web.HTTPPermanentRedirect(
+                    f"http://localhost:45678/start?{urllib.parse.urlencode(queryparams)}",
+                    headers=cors_headers,
+                )
+            else:
+                return PydanticResponse(
+                    ProxySetupResponseModel(
+                        redirect=AnyUrl(
+                            f"http://localhost:45678/start?{urllib.parse.urlencode(queryparams)}"
+                        ),
+                        redirectURI=AnyUrl(
+                            f"http://localhost:45678/start?{urllib.parse.urlencode(queryparams)}"
+                        ),
+                    ),
+                    headers=cors_headers,
+                )
+        case _:
+            raise InvalidAPIParameters("E20002: Protocol not available as interactive app")
+
+
+async def init(app: web.Application) -> None:
+    pass
+
+
+async def shutdown(app: web.Application) -> None:
+    pass
+
+
+def create_app(
+    default_cors_options: CORSOptions,
+) -> tuple[web.Application, Iterable[WebMiddleware]]:
+    app = web.Application()
+    app["prefix"] = "setup"
+    app.on_startup.append(init)
+    app.on_shutdown.append(shutdown)
+    add_route = app.router.add_route
+    add_route("GET", "", setup)
+    return app, []

backend_ai_appproxy_worker-25.19.1/ai/backend/appproxy/worker/cli/__main__.py

@@ -0,0 +1,144 @@
+from pathlib import Path
+from typing import Any, Optional
+
+import click
+
+from ai.backend.common.cli import LazyGroup
+
+from .context import CLIContext
+
+# LogLevel values for click.Choice - avoid importing ai.backend.logging at module level
+_LOG_LEVELS = ["CRITICAL", "ERROR", "WARNING", "INFO", "DEBUG", "TRACE", "NOTSET"]
+
+
+@click.group(invoke_without_command=False, context_settings={"help_option_names": ["-h", "--help"]})
+@click.option(
+    "-f",
+    "--config-path",
+    "--config",
+    type=click.Path(
+        file_okay=True,
+        dir_okay=False,
+        exists=True,
+        path_type=Path,
+    ),
+    default=None,
+    help="The config file path. (default: ./app-proxy-worker.toml)",
+)
+@click.option(
+    "--debug",
+    is_flag=True,
+    help="Set the logging level to DEBUG",
+)
+@click.option(
+    "--log-level",
+    type=click.Choice(_LOG_LEVELS, case_sensitive=False),
+    default="INFO",
+    help="Set the logging verbosity level",
+)
+@click.pass_context
+def main(
+    ctx: click.Context,
+    debug: bool,
+    log_level: str,
+    config_path: Optional[Path] = None,
+) -> None:
+    """
+    Proxy Worker Administration CLI
+    """
+    from setproctitle import setproctitle
+
+    from ai.backend.logging.types import LogLevel
+
+    setproctitle("backend.ai: proxy-worker.cli")
+    if debug:
+        log_level = "DEBUG"
+    ctx.obj = ctx.with_resource(CLIContext(config_path=config_path, log_level=LogLevel(log_level)))
+
+
+@main.command()
+@click.option(
+    "--output",
+    "-o",
+    default="-",
+    type=click.Path(dir_okay=False, writable=True),
+    help="Output file path (default: stdout)",
+)
+def generate_example_configuration(output: Path) -> None:
+    """
+    Generates example TOML configuration file for Backend.AI Proxy Worker.
+    """
+    import tomli_w
+
+    from ai.backend.appproxy.common.config import generate_example_json
+    from ai.backend.appproxy.common.utils import ensure_json_serializable
+
+    from ..config import ServerConfig
+
+    generated_example = generate_example_json(ServerConfig)
+    if output == "-" or output is None:
+        print(tomli_w.dumps(ensure_json_serializable(generated_example)))
+    else:
+        with open(output, mode="w") as fw:
+            fw.write(tomli_w.dumps(ensure_json_serializable(generated_example)))
+
+
+async def _generate() -> dict[str, Any]:
+    import importlib
+
+    import aiohttp_cors
+    from aiohttp import web
+
+    from ai.backend.appproxy.common.openapi import generate_openapi
+
+    from ..server import global_subapp_pkgs
+
+    cors_options = {
+        "*": aiohttp_cors.ResourceOptions(
+            allow_credentials=False, expose_headers="*", allow_headers="*"
+        ),
+    }
+
+    subapps: list[web.Application] = []
+    for subapp in global_subapp_pkgs:
+        pkg = importlib.import_module("ai.backend.appproxy.worker.api" + subapp)
+        app, _ = pkg.create_app(cors_options)
+        subapps.append(app)
+    return generate_openapi("Proxy Worker", subapps, verbose=True)
+
+
+@main.command()
+@click.option(
+    "--output",
+    "-o",
+    default="-",
+    type=click.Path(dir_okay=False, writable=True),
+    help="Output file path (default: stdout)",
+)
+def generate_openapi_spec(output: Path) -> None:
+    """
+    Generates OpenAPI specification of Backend.AI API.
+    """
+    import asyncio
+    import json
+
+    openapi = asyncio.run(_generate())
+    if output == "-" or output is None:
+        print(json.dumps(openapi, ensure_ascii=False, indent=2))
+    else:
+        with open(output, mode="w") as fw:
+            fw.write(json.dumps(openapi, ensure_ascii=False, indent=2))
+
+
+@main.group(cls=LazyGroup, import_name="ai.backend.appproxy.worker.cli.dependencies:cli")
+def dependencies():
+    """Command set for dependency verification and validation."""
+
+
+@main.group(cls=LazyGroup, import_name="ai.backend.appproxy.worker.cli.health:cli")
+def health():
+    """Command set for health checking."""
+
+
+if __name__ == "__main__":
+    main()

backend_ai_appproxy_worker-25.19.1/ai/backend/appproxy/worker/cli/context.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from contextlib import AbstractContextManager
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Optional, Self
+
+import click
+
+if TYPE_CHECKING:
+    from ai.backend.logging import AbstractLogger
+    from ai.backend.logging.types import LogLevel
+
+
+class CLIContext(AbstractContextManager):
+    _logger: AbstractLogger
+
+    def __init__(self, log_level: LogLevel, config_path: Optional[Path] = None) -> None:
+        self.config_path = config_path
+        self.log_level = log_level
+
+    def __enter__(self) -> Self:
+        from ai.backend.logging import LocalLogger
+
+        # The "start-server" command is injected by ai.backend.cli from the entrypoint
+        # and it has its own multi-process-aware logging initialization.
+        # If we duplicate the local logging with it, the process termination may hang.
+        click_ctx = click.get_current_context()
+        if click_ctx.invoked_subcommand != "start-server":
+            self._logger = LocalLogger(log_level=self.log_level)
+            self._logger.__enter__()
+        return self
+
+    def __exit__(self, *exc_info: Any) -> None:
+        click_ctx = click.get_current_context()
+        if click_ctx.invoked_subcommand != "start-server":
+            self._logger.__exit__()