docker/space_runtime.py

"""Resolve Space-specific runtime defaults for Hermes containers."""

from __future__ import annotations

import os
import shlex
import sys
from collections.abc import Mapping, Sequence
from pathlib import Path

SPACE_WEB_ARGS = ["dashboard", "--host", "0.0.0.0", "--port", "7860", "--no-open"]
DEFAULT_HERMES_HOME = "/opt/data"
SPACE_HERMES_HOME = "/data/hermes"
DEFAULT_API_SERVER_PORT = "8642"
DEFAULT_API_SERVER_UPSTREAM_HOST = "127.0.0.1"
API_SERVER_PORT_PATHS = (
    ("platforms", "api_server", "extra", "port"),
    ("platforms", "api_server", "port"),
    ("api_server", "extra", "port"),
    ("api_server", "port"),
)


def is_hf_space(env: Mapping[str, str]) -> bool:
    return bool(env.get("SPACE_ID") or env.get("SPACE_HOST"))


def resolve_hermes_home(env: Mapping[str, str]) -> str:
    current = env.get("HERMES_HOME") or DEFAULT_HERMES_HOME
    if is_hf_space(env) and current == DEFAULT_HERMES_HOME:
        return SPACE_HERMES_HOME
    return current


def resolve_default_command(env: Mapping[str, str], argv: Sequence[str]) -> list[str]:
    if argv:
        return list(argv)
    if is_hf_space(env):
        return SPACE_WEB_ARGS.copy()
    return []


def resolve_hermes_command(
    env: Mapping[str, str],
    argv: Sequence[str],
    executable: str = "hermes",
) -> list[str]:
    return [executable, *resolve_default_command(env, argv)]


def _resolve_config_path(env: Mapping[str, str]) -> Path:
    return Path(resolve_hermes_home(env)) / "config.yaml"


def _strip_inline_comment(text: str) -> str:
    quote: str | None = None
    result: list[str] = []
    for char in text:
        if char in {"'", '"'}:
            if quote == char:
                quote = None
            elif quote is None:
                quote = char
        if char == "#" and quote is None:
            break
        result.append(char)
    return "".join(result).rstrip()


def _parse_scalar(text: str) -> str | None:
    value = _strip_inline_comment(text).strip()
    if not value:
        return None
    if len(value) >= 2 and value[0] == value[-1] and value[0] in {"'", '"'}:
        return value[1:-1]
    return value


def _read_yaml_path_value(text: str, target_paths: Sequence[tuple[str, ...]]) -> str | None:
    stack: list[tuple[int, str]] = []
    for raw_line in text.splitlines():
        line = raw_line.rstrip()
        stripped = line.lstrip(" ")
        if not stripped or stripped.startswith("#") or stripped.startswith("-"):
            continue

        indent = len(line) - len(stripped)
        key, separator, remainder = stripped.partition(":")
        if not separator:
            continue

        while stack and indent <= stack[-1][0]:
            stack.pop()

        path = tuple(part for _, part in stack) + (key.strip(),)
        value = _parse_scalar(remainder)
        if path in target_paths and value is not None:
            return value
        if value is None:
            stack.append((indent, key.strip()))

    return None


def _config_api_server_port(env: Mapping[str, str]) -> str | None:
    config_path = _resolve_config_path(env)
    try:
        config_text = config_path.read_text(encoding="utf-8")
    except OSError:
        return None

    value = _read_yaml_path_value(config_text, API_SERVER_PORT_PATHS)
    if value is None:
        return None

    try:
        port = int(value, 10)
    except ValueError:
        return None
    if 0 < port <= 65535:
        return str(port)
    return None


def resolve_api_server_port(env: Mapping[str, str]) -> str:
    explicit = (env.get("API_SERVER_PORT") or "").strip()
    if explicit:
        return explicit
    return _config_api_server_port(env) or DEFAULT_API_SERVER_PORT


def resolve_api_server_upstream(env: Mapping[str, str]) -> str:
    explicit = (env.get("UPSTREAM") or "").strip()
    if explicit:
        return explicit
    return f"http://{DEFAULT_API_SERVER_UPSTREAM_HOST}:{resolve_api_server_port(env)}"


def _command_string(env: Mapping[str, str], argv: Sequence[str]) -> str:
    executable = env.get("HERMES_EXECUTABLE", "hermes")
    return shlex.join(resolve_hermes_command(env, argv, executable=executable))


def main(argv: Sequence[str] | None = None) -> int:
    argv = list(sys.argv[1:] if argv is None else argv)
    env = os.environ

    if not argv or argv[0] == "home":
        print(resolve_hermes_home(env))
        return 0

    if argv[0] == "command":
        print(_command_string(env, argv[1:]))
        return 0

    if argv[0] == "api-server-port":
        print(resolve_api_server_port(env))
        return 0

    if argv[0] == "api-server-upstream":
        print(resolve_api_server_upstream(env))
        return 0

    if argv[0] == "exec":
        print_only = False
        args = argv[1:]
        if args and args[0] == "--print-only":
            print_only = True
            args = args[1:]
        command = resolve_hermes_command(
            env,
            args,
            executable=env.get("HERMES_EXECUTABLE", "hermes"),
        )
        if print_only:
            print(shlex.join(command))
            return 0
        os.execvpe(command[0], command, dict(env))

    prog = os.path.basename(sys.argv[0] or "space_runtime.py")
    raise SystemExit(
        f"Usage: {prog} [home|command|api-server-port|api-server-upstream|exec [--print-only] [args...]]"
    )


if __name__ == "__main__":
    raise SystemExit(main())