diff --git a/README.md b/README.md index 3582bb05..7151d65a 100644 --- a/README.md +++ b/README.md @@ -288,6 +288,7 @@ are mainly for maintenance, debugging, or advanced integration. Run `reme help` |-------------------------------------------|----------------------------------------------------------------------------------------| | `reme start` | Start the local ReMe service. | | `reme version` / `reme health_check` | Check package and component status. | +| `reme status` | Show stateful data-component memory estimates and process RSS. | | [`reme search`](docs/en/memory_search.md) | Retrieve memory with BM25 and wikilinks by default, plus vectors when enabled. | | `reme read` / `reme write` / `reme edit` | Inspect and maintain Markdown memory files. | | `reme auto_memory` | Turn conversation messages into daily memory cards. Requires LLM credentials. | diff --git a/README_ZH.md b/README_ZH.md index 6a01a79b..13c5a68c 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -277,6 +277,7 @@ frontmatter 和文件操作接口主要用于维护、调试或高级集成。 |-------------------------------------------|---------------------------------------------| | `reme start` | 启动本地 ReMe 服务。 | | `reme version` / `reme health_check` | 检查包版本和组件状态。 | +| `reme status` | 查看有状态数据组件的内存估算及进程 RSS。 | | [`reme search`](docs/zh/memory_search.md) | 默认使用 BM25 和 wikilink 检索，启用后增加向量检索。 | | `reme read` / `reme write` / `reme edit` | 检查和维护 Markdown 记忆文件。 | | `reme auto_memory` | 将对话 messages 转为 daily 记忆卡片；需要 LLM 凭证。 | diff --git a/reme/config/config_parser.py b/reme/config/config_parser.py index c97a84ec..3f6ec1b4 100644 --- a/reme/config/config_parser.py +++ b/reme/config/config_parser.py @@ -201,9 +201,12 @@ def parse_args(*args) -> tuple[str, dict]: return first, parsed -def resolve_app_config(**kwargs) -> dict: +def resolve_app_config(*, log_config: bool = True, **kwargs) -> dict: """Resolve full app-start config: load `config=path` file, fall back to `default`, then deep-merge with the remaining kwargs as overrides. + + Set ``log_config=False`` for user-facing client calls that should print only + the requested job's output. """ from ..utils import get_logger @@ -215,10 +218,12 @@ def resolve_app_config(**kwargs) -> dict: config_value = kwargs.get("config") if isinstance(config_value, str): kwargs.pop("config") - logger.info(f"Loading config: {config_value}") + if log_config: + logger.info(f"Loading config: {config_value}") configs.append(_load_config(config_value)) elif "default" in _CONFIG_REGISTRY: - logger.info("No config specified, loading 'default'") + if log_config: + logger.info("No config specified, loading 'default'") configs.append(_load_config("default")) configs.append(kwargs) diff --git a/reme/config/default.yaml b/reme/config/default.yaml index 2276ca23..48df45fd 100644 --- a/reme/config/default.yaml +++ b/reme/config/default.yaml @@ -216,6 +216,15 @@ jobs: steps: - backend: health_check_step + status: + backend: base + description: "report memory estimates for stateful data components and process RSS" + parameters: + type: object + properties: { } + steps: + - backend: status_step + help: backend: base description: "list all registered jobs with their metadata" @@ -225,6 +234,24 @@ jobs: steps: - backend: help_step + shell: + backend: base + description: "execute a shell command asynchronously in the workspace" + parameters: + type: object + properties: + cmd: + type: string + description: "shell command to execute" + shell_timeout: + type: number + description: "maximum execution time in seconds" + default: 86400 + required: + - cmd + steps: + - backend: shell_step + traverse: backend: base description: "Walk the wikilink graph from a path." diff --git a/reme/reme.py b/reme/reme.py index 728ca437..6af9da89 100644 --- a/reme/reme.py +++ b/reme/reme.py @@ -39,7 +39,7 @@ async def call_server(action: str, **kwargs): # Prefer the running server's real config; fall back to the local config file. service = running_service_config() if service is None: - service = resolve_app_config(**resolve_kwargs).get("service") + service = resolve_app_config(log_config=False, **resolve_kwargs).get("service") service = service if isinstance(service, dict) else {} backend: str = kwargs.pop("backend", None) or service.get("backend", "http") @@ -61,9 +61,9 @@ async def call_server(action: str, **kwargs): def main(): """Parse CLI arguments and launch the appropriate mode.""" + load_env() action, kwargs = parse_args(*sys.argv[1:]) if action == "start": - load_env() kwargs = prepare_start_config(kwargs) if should_precheck_start(kwargs) and not precheck_start(kwargs.get("service")): return diff --git a/reme/steps/common/__init__.py b/reme/steps/common/__init__.py index bdc07d58..61c55167 100644 --- a/reme/steps/common/__init__.py +++ b/reme/steps/common/__init__.py @@ -6,8 +6,10 @@ from .health_check import HealthCheckStep from .help import HelpStep from .llm_demo import LLMDemoStep from .python_execute import PythonExecuteStep +from .shell import ShellStep from .stream_demo import StreamDemoStep1, StreamDemoStep2 from .stream_llm_demo import StreamLLMDemoStep +from .status import StatusStep from .version import VersionStep __all__ = [ @@ -18,8 +20,10 @@ __all__ = [ "HelpStep", "LLMDemoStep", "PythonExecuteStep", + "ShellStep", "StreamDemoStep1", "StreamDemoStep2", "StreamLLMDemoStep", + "StatusStep", "VersionStep", ] diff --git a/reme/steps/common/shell.py b/reme/steps/common/shell.py new file mode 100644 index 00000000..a80c600b --- /dev/null +++ b/reme/steps/common/shell.py @@ -0,0 +1,119 @@ +"""Execute a shell command and return its stdout.""" + +import asyncio +import os +import signal +from dataclasses import dataclass +from typing import Any + +from ..base_step import BaseStep +from ...components import R + +DEFAULT_TIMEOUT = 60.0 * 60 * 24 + + +@dataclass(frozen=True) +class _ShellResult: + stdout: str + stderr: str + returncode: int | None + timed_out: bool = False + + +@R.register("shell_step") +class ShellStep(BaseStep): + """Run a command in a shell and return stdout as the response answer.""" + + async def execute(self): + assert self.context is not None + + command = self.context.get("cmd", "") + timeout, timeout_error = self._parse_timeout(self.context.get("shell_timeout", DEFAULT_TIMEOUT)) + if not isinstance(command, str) or not command.strip(): + self.context.response.success = False + self.context.response.answer = "cmd is required" + return self.context.response + if timeout_error: + self.context.response.success = False + self.context.response.answer = timeout_error + return self.context.response + + result = await self._run_shell(command, timeout) + if result.timed_out: + self.context.response.success = False + self.context.response.answer = f"Shell command timed out after {timeout:g}s" + else: + self.context.response.success = result.returncode == 0 + self.context.response.answer = result.stdout if result.stdout or result.returncode == 0 else result.stderr + + self.context.response.metadata.update( + { + "returncode": result.returncode, + "stderr": result.stderr, + "shell_timeout": timeout, + }, + ) + return self.context.response + + async def _run_shell(self, command: str, timeout: float) -> _ShellResult: + process_kwargs = {"start_new_session": True} if os.name == "posix" else {} + process = await asyncio.create_subprocess_shell( + command, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + cwd=self.workspace_path, + **process_kwargs, + ) + try: + stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=timeout) + return _ShellResult( + stdout=stdout.decode(errors="replace"), + stderr=stderr.decode(errors="replace"), + returncode=process.returncode, + ) + except TimeoutError: + self._kill_process_tree(process) + stdout, stderr = await process.communicate() + return _ShellResult( + stdout=stdout.decode(errors="replace"), + stderr=stderr.decode(errors="replace"), + returncode=process.returncode, + timed_out=True, + ) + + @staticmethod + def _kill_process_tree(process: asyncio.subprocess.Process) -> None: + if os.name == "posix": + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + return + + # Windows has no process groups with POSIX kill semantics. Walk the + # descendants explicitly so commands do not survive their shell. + import psutil # pylint: disable=import-outside-toplevel + + try: + descendants = psutil.Process(process.pid).children(recursive=True) + except psutil.Error: + descendants = [] + for child in reversed(descendants): + try: + child.kill() + except psutil.Error: + pass + try: + process.kill() + except ProcessLookupError: + pass + + @staticmethod + def _parse_timeout(raw: Any) -> tuple[float, str]: + try: + timeout = float(raw) + except (TypeError, ValueError): + return DEFAULT_TIMEOUT, "shell_timeout must be a positive number" + if timeout <= 0: + return DEFAULT_TIMEOUT, "shell_timeout must be a positive number" + return timeout, "" diff --git a/reme/steps/common/status.py b/reme/steps/common/status.py new file mode 100644 index 00000000..2b04f1ad --- /dev/null +++ b/reme/steps/common/status.py @@ -0,0 +1,143 @@ +"""Report ReMe runtime memory usage.""" + +import sys +from collections.abc import Mapping +from types import ModuleType + +import numpy as np +import psutil + +from ..base_step import BaseStep +from ...components import BaseComponent, R +from ...enumeration import ComponentEnum + +_SKIPPED_COMPONENT_ATTRIBUTES = {"app_context", "logger"} +_TRACKED_COMPONENT_TYPES = { + ComponentEnum.EMBEDDING_STORE, + ComponentEnum.FILE_GRAPH, + ComponentEnum.FILE_STORE, + ComponentEnum.KEYWORD_INDEX, +} + + +def _format_bytes(size: int) -> str: + """Format a byte count using binary units.""" + value = float(size) + units = ("B", "KiB", "MiB", "GiB", "TiB") + for unit in units[:-1]: + if abs(value) < 1024: + return f"{int(value)} B" if unit == "B" else f"{value:.2f} {unit}" + value /= 1024 + return f"{value:.2f} {units[-1]}" + + +def _component_size(obj: object) -> int: + """Estimate the memory owned by one component's Python object graph. + + References to the application context, logger, and other components are + excluded so shared application state is not charged repeatedly. + """ + seen: set[int] = set() + root_id = id(obj) + + def walk(value: object) -> int: # pylint: disable=too-many-return-statements + value_id = id(value) + if value_id in seen: + return 0 + seen.add(value_id) + + if isinstance(value, BaseComponent) and value_id != root_id: + return 0 + if isinstance(value, (type, ModuleType)): + return 0 + if isinstance(value, np.ndarray): + # ndarray.__sizeof__ normally includes its owned data buffer. A + # view may be smaller, so retain at least the visible buffer size. + return max(sys.getsizeof(value), int(value.nbytes)) + + size = sys.getsizeof(value) + if isinstance( + value, + (str, bytes, bytearray, memoryview, int, float, bool, complex, type(None)), + ): + return size + if isinstance(value, Mapping): + return size + sum(walk(key) + walk(item) for key, item in value.items()) + if isinstance(value, (list, tuple, set, frozenset)): + return size + sum(walk(item) for item in value) + if hasattr(value, "__dict__"): + attributes = vars(value) + if id(attributes) in seen: + return size + seen.add(id(attributes)) + return ( + size + + sys.getsizeof(attributes) + + sum( + walk(key) + walk(item) + for key, item in attributes.items() + if key not in _SKIPPED_COMPONENT_ATTRIBUTES + ) + ) + slots = getattr(value, "__slots__", ()) + if isinstance(slots, str): + slots = (slots,) + return size + sum(walk(getattr(value, slot)) for slot in slots if hasattr(value, slot)) + + return walk(obj) + + +def _collect_memory(app_context) -> dict: + """Collect per-component estimates, their sum, and process RSS.""" + components: dict[str, dict[str, dict[str, int | str]]] = {} + total = 0 + if app_context is not None: + for component_type in sorted(_TRACKED_COMPONENT_TYPES, key=lambda item: item.value): + group = {} + for name, component in sorted(app_context.components.get(component_type, {}).items()): + # _component_size keeps an independent seen set so each component + # remains understandable in isolation. Consequently, a non-component + # object shared by multiple components may be included more than once + # in components_total_bytes; this is an estimate, not unique RSS. + size = _component_size(component) + group[name] = {"bytes": size, "human": _format_bytes(size)} + total += size + if group: + components[component_type.value] = group + + rss = psutil.Process().memory_info().rss + return { + "components": components, + "components_total_bytes": total, + "components_total": _format_bytes(total), + "process_rss_bytes": rss, + "process_rss": _format_bytes(rss), + } + + +def _format_status(memory: dict) -> str: + """Build the human-readable CLI response.""" + lines = ["Memory (estimated component object size)"] + for component_type, group in memory["components"].items(): + for name, usage in group.items(): + lines.append(f" {component_type}:{name} {usage['human']}") + lines.extend( + [ + f" Components total {memory['components_total']}", + f" Process RSS {memory['process_rss']}", + ], + ) + return "\n".join(lines) + + +@R.register("status_step") +class StatusStep(BaseStep): + """Report per-component memory estimates and process RSS.""" + + async def execute(self): + assert self.context is not None + + memory = _collect_memory(self.app_context) + self.context.response.answer = _format_status(memory) + self.context.response.metadata["status"] = {"memory": memory} + return self.context.response diff --git a/reme/utils/service_utils.py b/reme/utils/service_utils.py index 912d0b26..d7edb8fe 100644 --- a/reme/utils/service_utils.py +++ b/reme/utils/service_utils.py @@ -104,7 +104,7 @@ def running_service_config() -> dict | None: _, kwargs = parse_args("start", *argv) except ValueError: continue - service = resolve_app_config(**kwargs).get("service") + service = resolve_app_config(log_config=False, **kwargs).get("service") if isinstance(service, dict): return service return None diff --git a/tests/unit/test_config_parser.py b/tests/unit/test_config_parser.py index effe99b3..7e0ed024 100644 --- a/tests/unit/test_config_parser.py +++ b/tests/unit/test_config_parser.py @@ -10,6 +10,7 @@ from reme.config.config_parser import ( _read_config_file, parse_args, parse_dot_notation, + resolve_app_config, ) @@ -20,6 +21,24 @@ def test_load_builtin_config_by_filename_with_suffix(): assert cfg["service"]["backend"] == "http" +def test_resolve_app_config_can_suppress_config_log(monkeypatch): + """Client-side config resolution can avoid polluting command output.""" + messages = [] + + class FakeLogger: + """Capture config log messages.""" + + def info(self, message): + """Record one INFO message.""" + messages.append(message) + + monkeypatch.setattr("reme.utils.get_logger", lambda **_kwargs: FakeLogger()) + + resolve_app_config(log_config=False) + + assert not messages + + def test_default_config_registers_daily_write_job(): """``daily_write`` is exposed as a base job backed by ``daily_write_step``.""" cfg = _load_config("default.yaml") @@ -30,6 +49,17 @@ def test_default_config_registers_daily_write_job(): assert job["parameters"]["required"] == ["name", "description", "session_id", "content"] +def test_default_config_registers_shell_job(): + """``shell`` exposes command execution through ``shell_step``.""" + cfg = _load_config("default.yaml") + + job = cfg["jobs"]["shell"] + assert job["backend"] == "base" + assert job["steps"] == [{"backend": "shell_step"}] + assert job["parameters"]["required"] == ["cmd"] + assert job["parameters"]["properties"]["shell_timeout"]["default"] == 86400 + + def test_default_config_keeps_frontmatter_chunk_metadata_opt_in(): """Markdown frontmatter-to-chunk metadata is disabled by default for compatibility.""" cfg = _load_config("default.yaml") diff --git a/tests/unit/test_reme_cli.py b/tests/unit/test_reme_cli.py index 7becea3a..03f3dd1a 100644 --- a/tests/unit/test_reme_cli.py +++ b/tests/unit/test_reme_cli.py @@ -10,6 +10,24 @@ from reme.components.service.cli_service import CliService from reme import reme as reme_module +def test_main_loads_env_before_calling_server(monkeypatch): + """Client actions can resolve connection settings from the local .env.""" + events = [] + + main_globals = reme_module.main.__globals__ + monkeypatch.setitem(main_globals, "load_env", lambda: events.append("load_env")) + monkeypatch.setitem(main_globals, "parse_args", lambda *_args: ("shell", {"cmd": "pwd"})) + + async def fake_call_server(action, **kwargs): + events.append(("call_server", action, kwargs)) + + monkeypatch.setitem(main_globals, "call_server", fake_call_server) + + reme_module.main() + + assert events == ["load_env", ("call_server", "shell", {"cmd": "pwd"})] + + def test_prepare_start_config_moves_unknown_start_args_to_job_args(monkeypatch): """``reme start job=...`` is translated into a one-shot cli service config.""" @@ -216,3 +234,37 @@ def test_call_server_treats_show_metadata_as_client_kwarg(monkeypatch, capsys): assert seen["action"] == "version" assert seen["payload"] == {} assert capsys.readouterr().out == "ok\n" + + +def test_call_server_passes_shell_parameters_as_payload(monkeypatch, capsys): + """Shell-specific parameter names do not collide with client options.""" + seen = {} + + class FakeClient: + """Async client stub that records shell request arguments.""" + + def __init__(self, **kwargs): + seen["client_kwargs"] = kwargs + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + return None + + async def __call__(self, action: str, **kwargs): + seen["action"] = action + seen["payload"] = kwargs + yield "ok" + + monkeypatch.setattr(reme_module.R, "get", lambda component_type, backend: FakeClient) + monkeypatch.setattr(reme_module, "running_service_config", lambda: None) + + async def run(): + await reme_module.call_server("shell", backend="http", cmd="ls", shell_timeout=5) + + asyncio.run(run()) + + assert seen["action"] == "shell" + assert seen["payload"] == {"cmd": "ls", "shell_timeout": 5} + assert capsys.readouterr().out == "ok\n" diff --git a/tests/unit/test_shell_step.py b/tests/unit/test_shell_step.py new file mode 100644 index 00000000..db04c393 --- /dev/null +++ b/tests/unit/test_shell_step.py @@ -0,0 +1,100 @@ +"""Tests for asynchronous shell command execution.""" + +import asyncio +import time + +from reme.components.application_context import ApplicationContext +from reme.enumeration import ComponentEnum +from reme.steps.common.shell import DEFAULT_TIMEOUT, ShellStep +from reme.components import R + + +def _run(coro): + """Run a coroutine on a fresh event loop.""" + asyncio.run(coro) + + +def test_shell_step_is_registered(): + """Importing common steps makes shell_step discoverable.""" + assert R.get(ComponentEnum.STEP, "shell_step") is ShellStep + + +def test_shell_step_default_timeout_is_one_day(): + """Shell commands may run for one day when no timeout is supplied.""" + assert DEFAULT_TIMEOUT == 86400 + + +def test_shell_step_returns_stdout_from_workspace(tmp_path): + """Successful commands return stdout and run in the configured workspace.""" + + async def run(): + step = ShellStep(app_context=ApplicationContext(workspace_dir=str(tmp_path))) + response = await step(cmd="pwd") + + assert response.success is True + assert response.answer.strip() == str(tmp_path) + assert response.metadata["returncode"] == 0 + assert response.metadata["stderr"] == "" + + _run(run()) + + +def test_shell_step_reports_stderr_on_failure(): + """Failed commands expose stderr and their non-zero exit status.""" + + async def run(): + step = ShellStep() + response = await step(cmd="echo boom >&2; exit 7") + + assert response.success is False + assert response.answer == "boom\n" + assert response.metadata["returncode"] == 7 + assert response.metadata["stderr"] == "boom\n" + + _run(run()) + + +def test_shell_step_times_out(): + """Timeout terminates child processes as well as their parent shell.""" + + async def run(): + step = ShellStep() + started = time.monotonic() + response = await step(cmd="sleep 3 & wait", shell_timeout=0.01) + + assert response.success is False + assert response.answer == "Shell command timed out after 0.01s" + assert response.metadata["shell_timeout"] == 0.01 + assert time.monotonic() - started < 1 + + _run(run()) + + +def test_shell_step_requires_a_command(): + """Blank commands are rejected without creating a subprocess.""" + + async def run(): + response = await ShellStep()(cmd=" ") + + assert response.success is False + assert response.answer == "cmd is required" + + _run(run()) + + +def test_shell_step_does_not_accept_legacy_parameter_names(): + """Only cmd and shell_timeout configure shell execution.""" + + async def run(): + response = await ShellStep()(command="printf legacy", timeout=1) + + assert response.success is False + assert response.answer == "cmd is required" + + response = await ShellStep()(cmd="printf current", timeout=1) + + assert response.success is True + assert response.answer == "current" + assert response.metadata["shell_timeout"] == DEFAULT_TIMEOUT + + _run(run()) diff --git a/tests/unit/test_status_step.py b/tests/unit/test_status_step.py new file mode 100644 index 00000000..1ff97e8d --- /dev/null +++ b/tests/unit/test_status_step.py @@ -0,0 +1,98 @@ +"""Tests for the built-in runtime memory status report.""" + +import asyncio + +from reme.components.application_context import ApplicationContext +from reme.components.base_component import BaseComponent +from reme.enumeration import ComponentEnum +from reme.steps.common.status import ( + StatusStep, + _collect_memory, + _component_size, +) + + +class _SizedComponent(BaseComponent): + """Small component with predictable owned payload for accounting tests.""" + + component_type = ComponentEnum.FILE_STORE + + def __init__(self, payload: bytes, **kwargs): + super().__init__(**kwargs) + self.payload = payload + self.peer = None + + +def test_component_size_does_not_charge_referenced_components_twice(): + """A dependency component is accounted under its own status entry.""" + dependency = _SizedComponent(b"x" * 4096) + owner = _SizedComponent(b"y") + owner.peer = dependency + + owner_size = _component_size(owner) + dependency_size = _component_size(dependency) + + assert dependency_size > owner_size + + +def test_collect_memory_reports_only_stateful_data_components_and_sum(tmp_path): + """Status includes only the data components whose state can grow.""" + context = ApplicationContext(workspace_dir=str(tmp_path)) + context.components = { + ComponentEnum.FILE_STORE: { + "default": _SizedComponent(b"abc", app_context=context), + }, + ComponentEnum.AGENT_WRAPPER: { + "default": _SizedComponent(b"agent", app_context=context), + }, + ComponentEnum.AS_LLM: { + "default": _SizedComponent(b"llm", app_context=context), + }, + ComponentEnum.FILE_CATALOG: { + "default": _SizedComponent(b"catalog", app_context=context), + }, + ComponentEnum.FILE_CHUNKER: { + "default": _SizedComponent(b"chunker", app_context=context), + }, + ComponentEnum.TOKENIZER: { + "words": _SizedComponent(b"defgh", app_context=context), + }, + ComponentEnum.FILE_GRAPH: { + "default": _SizedComponent(b"graph", app_context=context), + }, + } + + memory = _collect_memory(context) + + assert set(memory["components"]) == {"file_graph", "file_store"} + assert ( + not { + "agent_wrapper", + "as_llm", + "file_catalog", + "file_chunker", + "tokenizer", + } + & memory["components"].keys() + ) + sizes = [usage["bytes"] for group in memory["components"].values() for usage in group.values()] + assert memory["components_total_bytes"] == sum(sizes) + assert memory["process_rss_bytes"] > 0 + + +def test_status_step_returns_human_summary_and_exact_metadata(tmp_path): + """The public step response serves CLI users and programmatic clients.""" + context = ApplicationContext(workspace_dir=str(tmp_path)) + context.components = { + ComponentEnum.FILE_STORE: { + "default": _SizedComponent(b"abc", app_context=context), + }, + } + + response = asyncio.run(StatusStep(app_context=context)()) + + assert not response.answer.startswith("ReMe status") + assert "Memory (estimated component object size)" in response.answer + assert "file_store:default" in response.answer + assert "Storage" not in response.answer + assert set(response.metadata["status"]) == {"memory"}