mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-07 03:46:38 +00:00
feat(cloud): harden multi-tenant runtime resources
This commit is contained in:
Junyan Qin
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,466 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Exercise long-lived Core registries and verify that they reach a plateau.
|
||||
|
||||
This probe is intentionally separate from the default test suite because the
|
||||
audit profile creates tens of thousands of historical identities. It uses the
|
||||
real admission, eviction, and cleanup code while replacing external platform
|
||||
objects that are irrelevant to registry retention.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import gc
|
||||
import json
|
||||
import time
|
||||
import tracemalloc
|
||||
from dataclasses import asdict, dataclass
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import psutil
|
||||
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
|
||||
# Import the Application graph before taskmgr. The production boot path has
|
||||
# this same ordering; importing taskmgr first exposes its historical cycle
|
||||
# through HTTP route annotations.
|
||||
from langbot.pkg.core import app as _core_app # noqa: F401
|
||||
from langbot.pkg.core.taskmgr import AsyncTaskManager
|
||||
from langbot.pkg.pipeline.pool import QueryPool
|
||||
from langbot.pkg.pipeline.ratelimit.algos.fixedwin import FixedWindowAlgo
|
||||
from langbot.pkg.plugin.connector import PluginRuntimeConnector
|
||||
from langbot.pkg.platform.sources.websocket_adapter import (
|
||||
WebSocketMessage,
|
||||
WebSocketSession,
|
||||
)
|
||||
from langbot.pkg.provider.modelmgr.modelmgr import ModelManager
|
||||
from langbot.pkg.provider.session.sessionmgr import SessionManager
|
||||
from langbot_plugin.api.entities.builtin.provider.session import LauncherTypes
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ProbeScale:
|
||||
query_churn_per_phase: int
|
||||
session_churn_per_phase: int
|
||||
rate_limit_churn_per_phase: int
|
||||
task_churn_per_phase: int
|
||||
websocket_churn_per_phase: int
|
||||
empty_workspace_churn_per_phase: int
|
||||
|
||||
|
||||
SCALES = {
|
||||
'quick': ProbeScale(
|
||||
query_churn_per_phase=2_500,
|
||||
session_churn_per_phase=500,
|
||||
rate_limit_churn_per_phase=10_000,
|
||||
task_churn_per_phase=1_000,
|
||||
websocket_churn_per_phase=500,
|
||||
empty_workspace_churn_per_phase=1_000,
|
||||
),
|
||||
'audit': ProbeScale(
|
||||
query_churn_per_phase=25_000,
|
||||
session_churn_per_phase=2_500,
|
||||
rate_limit_churn_per_phase=10_000,
|
||||
task_churn_per_phase=5_000,
|
||||
websocket_churn_per_phase=2_500,
|
||||
empty_workspace_churn_per_phase=10_000,
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
class _ProbeQuery:
|
||||
"""Small weak-referenceable stand-in for SDK Query construction."""
|
||||
|
||||
def __init__(self, **values):
|
||||
self.__dict__.update(values)
|
||||
|
||||
|
||||
class _EmptyResult:
|
||||
def all(self) -> list:
|
||||
return []
|
||||
|
||||
|
||||
class _EmptyPluginRuntimeHandler:
|
||||
async def reconcile_plugin_installations(self, _states: tuple) -> dict:
|
||||
return {
|
||||
'applied': [],
|
||||
'removed': [],
|
||||
'missing_artifacts': [],
|
||||
'failed_installations': [],
|
||||
}
|
||||
|
||||
def unregister_installation_binding(self, _binding) -> None:
|
||||
raise AssertionError('An empty Workspace exposed an installation binding')
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ProcessSample:
|
||||
rss_bytes: int
|
||||
traced_current_bytes: int
|
||||
traced_peak_bytes: int
|
||||
asyncio_tasks: int
|
||||
threads: int
|
||||
open_fds: int | None
|
||||
|
||||
|
||||
def _sample_process() -> ProcessSample:
|
||||
gc.collect()
|
||||
process = psutil.Process()
|
||||
try:
|
||||
open_fds = process.num_fds()
|
||||
except (AttributeError, psutil.Error):
|
||||
open_fds = None
|
||||
traced_current, traced_peak = tracemalloc.get_traced_memory()
|
||||
return ProcessSample(
|
||||
rss_bytes=process.memory_info().rss,
|
||||
traced_current_bytes=traced_current,
|
||||
traced_peak_bytes=traced_peak,
|
||||
asyncio_tasks=len(asyncio.all_tasks()),
|
||||
threads=process.num_threads(),
|
||||
open_fds=open_fds,
|
||||
)
|
||||
|
||||
|
||||
def _execution_context(index: int, *, query_uuid: str | None = None) -> ExecutionContext:
|
||||
return ExecutionContext(
|
||||
instance_uuid='runtime-resource-probe',
|
||||
workspace_uuid=f'workspace-{index}',
|
||||
placement_generation=1,
|
||||
bot_uuid='probe-bot',
|
||||
pipeline_uuid='probe-pipeline',
|
||||
query_uuid=query_uuid,
|
||||
)
|
||||
|
||||
|
||||
class CoreRuntimeProbe:
|
||||
"""Own the same manager instances across two equal churn phases."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.query_pool = QueryPool(max_queries=100, max_queries_per_workspace=1)
|
||||
app = SimpleNamespace(
|
||||
event_loop=asyncio.get_running_loop(),
|
||||
persistence_mgr=None,
|
||||
instance_config=SimpleNamespace(
|
||||
data={
|
||||
'concurrency': {'session': 1},
|
||||
'system': {
|
||||
'session_retention': {
|
||||
'idle_ttl_seconds': 86_400,
|
||||
'max_entries': 200,
|
||||
'max_entries_per_workspace': 200,
|
||||
'max_conversations_per_session': 20,
|
||||
'max_messages_per_conversation': 100,
|
||||
},
|
||||
'task_retention': {
|
||||
'completed_limit': 200,
|
||||
'max_log_chars': 4_096,
|
||||
'max_active_user_tasks': 256,
|
||||
'max_active_user_tasks_per_workspace': 8,
|
||||
},
|
||||
},
|
||||
}
|
||||
),
|
||||
)
|
||||
self.session_manager = SessionManager(app)
|
||||
self.task_manager = AsyncTaskManager(app)
|
||||
self.rate_limit = FixedWindowAlgo(SimpleNamespace())
|
||||
self.websocket_session = WebSocketSession(
|
||||
'resource-probe',
|
||||
max_conversations=200,
|
||||
max_messages=100,
|
||||
)
|
||||
logger = SimpleNamespace(
|
||||
debug=lambda *_args, **_kwargs: None,
|
||||
info=lambda *_args, **_kwargs: None,
|
||||
warning=lambda *_args, **_kwargs: None,
|
||||
error=lambda *_args, **_kwargs: None,
|
||||
)
|
||||
self.empty_model_queries = 0
|
||||
|
||||
async def execute_empty(_statement):
|
||||
self.empty_model_queries += 1
|
||||
return _EmptyResult()
|
||||
|
||||
model_app = SimpleNamespace(
|
||||
logger=logger,
|
||||
persistence_mgr=SimpleNamespace(execute_async=execute_empty),
|
||||
)
|
||||
self.empty_model_manager = ModelManager(model_app)
|
||||
|
||||
async def runtime_disconnect_callback(_connector) -> None:
|
||||
return None
|
||||
|
||||
plugin_app = SimpleNamespace(
|
||||
instance_config=SimpleNamespace(data={'plugin': {'enable': True}}),
|
||||
deployment=SimpleNamespace(mode='cloud'),
|
||||
logger=logger,
|
||||
)
|
||||
self.empty_plugin_connector = PluginRuntimeConnector(
|
||||
plugin_app,
|
||||
runtime_disconnect_callback,
|
||||
)
|
||||
self.empty_plugin_connector.handler = _EmptyPluginRuntimeHandler()
|
||||
|
||||
async def validate_context(context):
|
||||
return context
|
||||
|
||||
async def load_desired_states(_context):
|
||||
return []
|
||||
|
||||
self.empty_plugin_connector._validate_execution_context = validate_context
|
||||
self.empty_plugin_connector._load_workspace_desired_states = load_desired_states
|
||||
|
||||
async def initialize(self) -> None:
|
||||
await self.rate_limit.initialize()
|
||||
|
||||
async def run_phase(self, scale: ProbeScale, phase: int) -> None:
|
||||
offsets = {
|
||||
'query': (phase - 1) * scale.query_churn_per_phase,
|
||||
'session': (phase - 1) * scale.session_churn_per_phase,
|
||||
'rate': (phase - 1) * scale.rate_limit_churn_per_phase,
|
||||
'task': (phase - 1) * scale.task_churn_per_phase,
|
||||
'websocket': (phase - 1) * scale.websocket_churn_per_phase,
|
||||
'empty_workspace': ((phase - 1) * scale.empty_workspace_churn_per_phase),
|
||||
}
|
||||
await self._churn_queries(offsets['query'], scale.query_churn_per_phase)
|
||||
await self._churn_sessions(offsets['session'], scale.session_churn_per_phase)
|
||||
await self._churn_rate_limits(offsets['rate'], scale.rate_limit_churn_per_phase)
|
||||
await self._churn_tasks(offsets['task'], scale.task_churn_per_phase)
|
||||
self._churn_websocket_history(
|
||||
offsets['websocket'],
|
||||
scale.websocket_churn_per_phase,
|
||||
)
|
||||
await self._churn_empty_workspaces(
|
||||
offsets['empty_workspace'],
|
||||
scale.empty_workspace_churn_per_phase,
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
async def _churn_queries(self, start: int, count: int) -> None:
|
||||
def make_query(**values):
|
||||
return _ProbeQuery(**values)
|
||||
|
||||
with patch(
|
||||
'langbot.pkg.pipeline.pool.pipeline_query.Query',
|
||||
side_effect=make_query,
|
||||
):
|
||||
for index in range(start, start + count):
|
||||
context = _execution_context(index)
|
||||
query = await self.query_pool.add_query(
|
||||
bot_uuid='probe-bot',
|
||||
launcher_type=LauncherTypes.PERSON,
|
||||
launcher_id=f'launcher-{index}',
|
||||
sender_id=f'sender-{index}',
|
||||
message_event=SimpleNamespace(),
|
||||
message_chain=SimpleNamespace(),
|
||||
adapter=None,
|
||||
pipeline_uuid='probe-pipeline',
|
||||
execution_context=context,
|
||||
)
|
||||
removed = await self.query_pool.remove_query(query)
|
||||
if not removed:
|
||||
raise AssertionError('Query cleanup failed')
|
||||
|
||||
async def _churn_sessions(self, start: int, count: int) -> None:
|
||||
for index in range(start, start + count):
|
||||
workspace_index = index % 100
|
||||
context = _execution_context(
|
||||
workspace_index,
|
||||
query_uuid=f'session-query-{index}',
|
||||
)
|
||||
query = SimpleNamespace(
|
||||
launcher_type=LauncherTypes.PERSON,
|
||||
launcher_id=f'launcher-{index}',
|
||||
sender_id=f'sender-{index}',
|
||||
bot_uuid='probe-bot',
|
||||
pipeline_uuid='probe-pipeline',
|
||||
query_uuid=context.query_uuid,
|
||||
_execution_context=context,
|
||||
)
|
||||
await self.session_manager.get_session(query)
|
||||
|
||||
async def _churn_rate_limits(self, start: int, count: int) -> None:
|
||||
for index in range(start, start + count):
|
||||
context = _execution_context(
|
||||
index % 1_000,
|
||||
query_uuid=f'rate-query-{index}',
|
||||
)
|
||||
query = SimpleNamespace(
|
||||
bot_uuid='probe-bot',
|
||||
pipeline_uuid='probe-pipeline',
|
||||
_execution_context=context,
|
||||
pipeline_config={
|
||||
'safety': {
|
||||
'rate-limit': {
|
||||
'window-length': 60,
|
||||
'limitation': 100_000,
|
||||
'strategy': 'drop',
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
admitted = await self.rate_limit.require_access(
|
||||
query,
|
||||
LauncherTypes.PERSON,
|
||||
f'rate-identity-{index}',
|
||||
)
|
||||
if not admitted:
|
||||
raise AssertionError('Rate-limit registry rejected bounded churn')
|
||||
|
||||
async def _churn_tasks(self, start: int, count: int) -> None:
|
||||
async def complete_immediately() -> None:
|
||||
return None
|
||||
|
||||
for batch_start in range(start, start + count, 256):
|
||||
batch_size = min(256, start + count - batch_start)
|
||||
wrappers = [
|
||||
self.task_manager.create_task(
|
||||
complete_immediately(),
|
||||
name=f'resource-probe-{batch_start + offset}',
|
||||
)
|
||||
for offset in range(batch_size)
|
||||
]
|
||||
await asyncio.gather(*(wrapper.task for wrapper in wrappers))
|
||||
await asyncio.sleep(0)
|
||||
|
||||
def _churn_websocket_history(self, start: int, count: int) -> None:
|
||||
for index in range(start, start + count):
|
||||
conversation_key = f'conversation-{index}'
|
||||
response_id = f'response-{index}'
|
||||
indexes = self.websocket_session.get_stream_message_indexes(conversation_key)
|
||||
indexes[response_id] = 0
|
||||
self.websocket_session.append_message(
|
||||
conversation_key,
|
||||
WebSocketMessage(
|
||||
id=self.websocket_session.next_message_id(conversation_key),
|
||||
role='assistant',
|
||||
content='probe',
|
||||
message_chain=[],
|
||||
timestamp='1970-01-01T00:00:00+00:00',
|
||||
is_final=True,
|
||||
),
|
||||
)
|
||||
|
||||
async def _churn_empty_workspaces(self, start: int, count: int) -> None:
|
||||
for index in range(start, start + count):
|
||||
await self.empty_model_manager._load_workspace_models(_execution_context(index))
|
||||
await self.empty_plugin_connector.reconcile_projected_workspaces(
|
||||
_execution_context(index) for index in range(start, start + count)
|
||||
)
|
||||
|
||||
def retained_state(self) -> dict[str, int]:
|
||||
return {
|
||||
'query_cached': len(self.query_pool.cached_queries),
|
||||
'query_queued': len(self.query_pool.queries),
|
||||
'query_active_workspaces': len(self.query_pool.active_query_count_by_workspace),
|
||||
'query_scope_counters': len(self.query_pool.query_count_by_scope),
|
||||
'sessions': len(self.session_manager.session_list),
|
||||
'session_index': len(self.session_manager._session_index),
|
||||
'rate_limit_containers': len(self.rate_limit.containers),
|
||||
'task_records': len(self.task_manager.tasks),
|
||||
'websocket_conversations': len(self.websocket_session.message_lists),
|
||||
'websocket_stream_indexes': len(self.websocket_session.stream_message_indexes),
|
||||
'empty_model_scopes': len(self.empty_model_manager._scope_generations),
|
||||
'empty_model_providers': len(self.empty_model_manager.provider_dict),
|
||||
'empty_model_llms': len(self.empty_model_manager.llm_model_dict),
|
||||
'empty_plugin_workspace_sets': len(self.empty_plugin_connector._workspace_installations),
|
||||
'empty_plugin_installations': len(self.empty_plugin_connector._known_desired_states),
|
||||
}
|
||||
|
||||
def assert_bounded(self) -> None:
|
||||
state = self.retained_state()
|
||||
expected_maximums = {
|
||||
'query_cached': 0,
|
||||
'query_queued': 0,
|
||||
'query_active_workspaces': 0,
|
||||
'query_scope_counters': 100,
|
||||
'sessions': 200,
|
||||
'session_index': 200,
|
||||
'rate_limit_containers': 10_000,
|
||||
'task_records': 200,
|
||||
'websocket_conversations': 200,
|
||||
'websocket_stream_indexes': 200,
|
||||
'empty_model_scopes': 0,
|
||||
'empty_model_providers': 0,
|
||||
'empty_model_llms': 0,
|
||||
'empty_plugin_workspace_sets': 0,
|
||||
'empty_plugin_installations': 0,
|
||||
}
|
||||
violations = {key: (state[key], maximum) for key, maximum in expected_maximums.items() if state[key] > maximum}
|
||||
if violations:
|
||||
raise AssertionError(f'Core retained-state limits failed: {violations}')
|
||||
|
||||
|
||||
async def _run(args: argparse.Namespace) -> dict:
|
||||
scale = SCALES[args.scale]
|
||||
tracemalloc.start()
|
||||
started_at = time.monotonic()
|
||||
probe = CoreRuntimeProbe()
|
||||
await probe.initialize()
|
||||
|
||||
baseline = _sample_process()
|
||||
await probe.run_phase(scale, 1)
|
||||
probe.assert_bounded()
|
||||
phase_one = _sample_process()
|
||||
state_one = probe.retained_state()
|
||||
|
||||
await probe.run_phase(scale, 2)
|
||||
probe.assert_bounded()
|
||||
phase_two = _sample_process()
|
||||
state_two = probe.retained_state()
|
||||
|
||||
if state_two != state_one:
|
||||
raise AssertionError(f'Core retained state did not plateau: phase_one={state_one}, phase_two={state_two}')
|
||||
traced_growth = phase_two.traced_current_bytes - phase_one.traced_current_bytes
|
||||
rss_growth = phase_two.rss_bytes - phase_one.rss_bytes
|
||||
max_traced_growth = int(args.max_traced_growth_mib * 1024 * 1024)
|
||||
max_rss_growth = int(args.max_rss_growth_mib * 1024 * 1024)
|
||||
if traced_growth > max_traced_growth:
|
||||
raise AssertionError(f'Second-phase traced memory grew by {traced_growth} bytes (limit {max_traced_growth})')
|
||||
if rss_growth > max_rss_growth:
|
||||
raise AssertionError(f'Second-phase RSS grew by {rss_growth} bytes (limit {max_rss_growth})')
|
||||
|
||||
return {
|
||||
'component': 'langbot-core',
|
||||
'scale': args.scale,
|
||||
'work_per_phase': asdict(scale),
|
||||
'elapsed_seconds': round(time.monotonic() - started_at, 3),
|
||||
'samples': {
|
||||
'baseline': asdict(baseline),
|
||||
'phase_one': asdict(phase_one),
|
||||
'phase_two': asdict(phase_two),
|
||||
},
|
||||
'second_phase_growth': {
|
||||
'rss_bytes': rss_growth,
|
||||
'traced_current_bytes': traced_growth,
|
||||
},
|
||||
'retained_state': {
|
||||
'phase_one': state_one,
|
||||
'phase_two': state_two,
|
||||
},
|
||||
'passed': True,
|
||||
}
|
||||
|
||||
|
||||
def _parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument('--scale', choices=tuple(SCALES), default='quick')
|
||||
parser.add_argument('--max-traced-growth-mib', type=float, default=8.0)
|
||||
parser.add_argument('--max-rss-growth-mib', type=float, default=64.0)
|
||||
parser.add_argument('--json', action='store_true', help='Print compact JSON')
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = _parse_args()
|
||||
result = asyncio.run(_run(args))
|
||||
if args.json:
|
||||
print(json.dumps(result, sort_keys=True))
|
||||
else:
|
||||
print(json.dumps(result, indent=2, sort_keys=True))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,561 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Measure populated Workspace runtime replacement cost and retention.
|
||||
|
||||
Unlike ``runtime_resource_probe.py``, which stresses historical request keys
|
||||
and empty tenants, this probe keeps one representative Provider, LLM,
|
||||
Embedding model, Rerank model, Pipeline, Bot, and Knowledge Base per Workspace.
|
||||
It then advances every Workspace to a new placement generation and verifies
|
||||
that old runtime objects are closed and collectible while active registry
|
||||
cardinality remains constant.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import gc
|
||||
import json
|
||||
import time
|
||||
import tracemalloc
|
||||
import weakref
|
||||
from dataclasses import asdict, dataclass
|
||||
from types import SimpleNamespace
|
||||
|
||||
import psutil
|
||||
|
||||
from langbot.pkg.api.http.context import ExecutionContext
|
||||
|
||||
# Match the production import order; importing a leaf manager first exposes a
|
||||
# historical annotation cycle that the application graph resolves.
|
||||
from langbot.pkg.core import app as _core_app # noqa: F401
|
||||
from langbot.pkg.entity.persistence import bot as persistence_bot
|
||||
from langbot.pkg.entity.persistence import model as persistence_model
|
||||
from langbot.pkg.entity.persistence import pipeline as persistence_pipeline
|
||||
from langbot.pkg.entity.persistence import rag as persistence_rag
|
||||
from langbot.pkg.pipeline.pipelinemgr import PipelineManager
|
||||
from langbot.pkg.platform.botmgr import PlatformManager
|
||||
from langbot.pkg.provider.modelmgr import requester
|
||||
from langbot.pkg.provider.modelmgr.modelmgr import ModelManager
|
||||
from langbot.pkg.provider.tools.loaders.mcp import MCPLoader
|
||||
from langbot.pkg.rag.knowledge.kbmgr import RAGManager
|
||||
from langbot.pkg.workspace.entities import WorkspaceExecutionBinding
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ProbeScale:
|
||||
workspaces: int
|
||||
|
||||
|
||||
SCALES = {
|
||||
'quick': ProbeScale(workspaces=250),
|
||||
'audit': ProbeScale(workspaces=5_000),
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ProcessSample:
|
||||
rss_bytes: int
|
||||
traced_current_bytes: int
|
||||
traced_peak_bytes: int
|
||||
asyncio_tasks: int
|
||||
threads: int
|
||||
open_fds: int | None
|
||||
|
||||
|
||||
class _ProbeLogger:
|
||||
def debug(self, *_args, **_kwargs) -> None:
|
||||
return None
|
||||
|
||||
def info(self, *_args, **_kwargs) -> None:
|
||||
return None
|
||||
|
||||
def warning(self, *_args, **_kwargs) -> None:
|
||||
return None
|
||||
|
||||
def error(self, *_args, **_kwargs) -> None:
|
||||
return None
|
||||
|
||||
|
||||
class _ProbeWorkspaceService:
|
||||
instance_uuid = 'runtime-capacity-probe'
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.generations: dict[str, int] = {}
|
||||
self.binding_lookups = 0
|
||||
|
||||
async def get_execution_binding(
|
||||
self,
|
||||
workspace_uuid: str,
|
||||
*,
|
||||
expected_generation: int | None = None,
|
||||
) -> WorkspaceExecutionBinding:
|
||||
self.binding_lookups += 1
|
||||
generation = self.generations[workspace_uuid]
|
||||
if expected_generation is not None and expected_generation != generation:
|
||||
raise AssertionError(f'stale probe generation {expected_generation} != {generation}')
|
||||
return WorkspaceExecutionBinding(
|
||||
instance_uuid=self.instance_uuid,
|
||||
workspace_uuid=workspace_uuid,
|
||||
placement_generation=generation,
|
||||
write_fenced=False,
|
||||
state='active',
|
||||
)
|
||||
|
||||
|
||||
class _ProbeRequester(requester.ProviderAPIRequester):
|
||||
name = 'capacity-probe'
|
||||
closed = 0
|
||||
|
||||
async def invoke_llm(
|
||||
self,
|
||||
query,
|
||||
model,
|
||||
messages,
|
||||
funcs=None,
|
||||
extra_args=None,
|
||||
remove_think=False,
|
||||
):
|
||||
return None
|
||||
|
||||
async def aclose(self) -> None:
|
||||
type(self).closed += 1
|
||||
|
||||
|
||||
class _ProbeAdapter:
|
||||
killed = 0
|
||||
|
||||
def __init__(self, _config, _logger) -> None:
|
||||
self.listeners = []
|
||||
|
||||
def register_listener(self, event_type, listener) -> None:
|
||||
self.listeners.append((event_type, listener))
|
||||
|
||||
async def kill(self) -> None:
|
||||
type(self).killed += 1
|
||||
|
||||
|
||||
class _ProbeMCPSession:
|
||||
closed = 0
|
||||
|
||||
def __init__(self, server_name: str) -> None:
|
||||
self.server_name = server_name
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
type(self).closed += 1
|
||||
|
||||
|
||||
def _sample_process() -> ProcessSample:
|
||||
gc.collect()
|
||||
process = psutil.Process()
|
||||
try:
|
||||
open_fds = process.num_fds()
|
||||
except (AttributeError, psutil.Error):
|
||||
open_fds = None
|
||||
traced_current, traced_peak = tracemalloc.get_traced_memory()
|
||||
return ProcessSample(
|
||||
rss_bytes=process.memory_info().rss,
|
||||
traced_current_bytes=traced_current,
|
||||
traced_peak_bytes=traced_peak,
|
||||
asyncio_tasks=len(asyncio.all_tasks()),
|
||||
threads=process.num_threads(),
|
||||
open_fds=open_fds,
|
||||
)
|
||||
|
||||
|
||||
class PopulatedWorkspaceProbe:
|
||||
def __init__(self) -> None:
|
||||
_ProbeRequester.closed = 0
|
||||
_ProbeAdapter.killed = 0
|
||||
_ProbeMCPSession.closed = 0
|
||||
self.workspace_service = _ProbeWorkspaceService()
|
||||
self.logger = _ProbeLogger()
|
||||
self.app = SimpleNamespace(
|
||||
logger=self.logger,
|
||||
workspace_service=self.workspace_service,
|
||||
persistence_mgr=SimpleNamespace(
|
||||
mode=SimpleNamespace(value='cloud_runtime'),
|
||||
),
|
||||
pipeline_config_meta_trigger={'name': 'trigger', 'stages': []},
|
||||
pipeline_config_meta_safety={'name': 'safety', 'stages': []},
|
||||
pipeline_config_meta_ai={'name': 'ai', 'stages': []},
|
||||
pipeline_config_meta_output={'name': 'output', 'stages': []},
|
||||
task_mgr=SimpleNamespace(
|
||||
cancel_by_scope=lambda *_args, **_kwargs: None,
|
||||
cancel_task=lambda *_args, **_kwargs: None,
|
||||
),
|
||||
)
|
||||
self.model_manager = ModelManager(self.app)
|
||||
self.model_manager.requester_dict = {
|
||||
_ProbeRequester.name: _ProbeRequester,
|
||||
}
|
||||
self.pipeline_manager = PipelineManager(self.app)
|
||||
self.pipeline_manager.stage_dict = {}
|
||||
self.rag_manager = RAGManager(self.app)
|
||||
self.mcp_loader = MCPLoader(self.app)
|
||||
self.platform_manager = PlatformManager(self.app)
|
||||
self.platform_manager.adapter_dict = {
|
||||
'capacity-probe': _ProbeAdapter,
|
||||
}
|
||||
self.generation_refs: dict[
|
||||
int,
|
||||
list[weakref.ReferenceType],
|
||||
] = {}
|
||||
|
||||
def _context(
|
||||
self,
|
||||
workspace_uuid: str,
|
||||
generation: int,
|
||||
*,
|
||||
bot_uuid: str | None = None,
|
||||
pipeline_uuid: str | None = None,
|
||||
) -> ExecutionContext:
|
||||
return ExecutionContext(
|
||||
instance_uuid=self.workspace_service.instance_uuid,
|
||||
workspace_uuid=workspace_uuid,
|
||||
placement_generation=generation,
|
||||
bot_uuid=bot_uuid,
|
||||
pipeline_uuid=pipeline_uuid,
|
||||
)
|
||||
|
||||
async def load_generation(self, workspaces: int, generation: int) -> None:
|
||||
for index in range(workspaces):
|
||||
workspace_uuid = f'workspace-{index}'
|
||||
provider_uuid = f'provider-{index}'
|
||||
llm_uuid = f'llm-{index}'
|
||||
embedding_uuid = f'embedding-{index}'
|
||||
rerank_uuid = f'rerank-{index}'
|
||||
pipeline_uuid = f'pipeline-{index}'
|
||||
bot_uuid = f'bot-{index}'
|
||||
kb_uuid = f'knowledge-{index}'
|
||||
mcp_server_name = f'mcp-{index}'
|
||||
self.workspace_service.generations[workspace_uuid] = generation
|
||||
context = self._context(workspace_uuid, generation)
|
||||
|
||||
runtime_provider = await self.model_manager.load_provider(
|
||||
context,
|
||||
persistence_model.ModelProvider(
|
||||
uuid=provider_uuid,
|
||||
workspace_uuid=workspace_uuid,
|
||||
name='Capacity Provider',
|
||||
requester=_ProbeRequester.name,
|
||||
base_url='https://capacity.invalid',
|
||||
api_keys=['probe'],
|
||||
),
|
||||
)
|
||||
await self.model_manager.cache_provider(context, runtime_provider)
|
||||
|
||||
runtime_llm = await self.model_manager.load_llm_model_with_provider(
|
||||
context,
|
||||
persistence_model.LLMModel(
|
||||
uuid=llm_uuid,
|
||||
workspace_uuid=workspace_uuid,
|
||||
name='Capacity LLM',
|
||||
provider_uuid=provider_uuid,
|
||||
abilities=['func_call'],
|
||||
extra_args={'temperature': 0.1},
|
||||
),
|
||||
runtime_provider,
|
||||
)
|
||||
await self.model_manager.cache_llm_model(context, runtime_llm)
|
||||
runtime_embedding = await self.model_manager.load_embedding_model_with_provider(
|
||||
context,
|
||||
persistence_model.EmbeddingModel(
|
||||
uuid=embedding_uuid,
|
||||
workspace_uuid=workspace_uuid,
|
||||
name='Capacity Embedding',
|
||||
provider_uuid=provider_uuid,
|
||||
extra_args={'dimensions': 1_024},
|
||||
),
|
||||
runtime_provider,
|
||||
)
|
||||
await self.model_manager.cache_embedding_model(
|
||||
context,
|
||||
runtime_embedding,
|
||||
)
|
||||
runtime_rerank = await self.model_manager.load_rerank_model_with_provider(
|
||||
context,
|
||||
persistence_model.RerankModel(
|
||||
uuid=rerank_uuid,
|
||||
workspace_uuid=workspace_uuid,
|
||||
name='Capacity Rerank',
|
||||
provider_uuid=provider_uuid,
|
||||
extra_args={},
|
||||
),
|
||||
runtime_provider,
|
||||
)
|
||||
await self.model_manager.cache_rerank_model(
|
||||
context,
|
||||
runtime_rerank,
|
||||
)
|
||||
|
||||
pipeline_context = self._context(
|
||||
workspace_uuid,
|
||||
generation,
|
||||
pipeline_uuid=pipeline_uuid,
|
||||
)
|
||||
await self.pipeline_manager.load_pipeline(
|
||||
pipeline_context,
|
||||
persistence_pipeline.LegacyPipeline(
|
||||
uuid=pipeline_uuid,
|
||||
workspace_uuid=workspace_uuid,
|
||||
name='Capacity Pipeline',
|
||||
description='',
|
||||
for_version='probe',
|
||||
is_default=True,
|
||||
stages=[],
|
||||
config={},
|
||||
extensions_preferences={},
|
||||
),
|
||||
_binding_validated=True,
|
||||
)
|
||||
runtime_pipeline = self.pipeline_manager._pipelines_by_key[
|
||||
(
|
||||
self.workspace_service.instance_uuid,
|
||||
workspace_uuid,
|
||||
pipeline_uuid,
|
||||
)
|
||||
]
|
||||
|
||||
runtime_kb = await self.rag_manager.load_knowledge_base(
|
||||
context,
|
||||
persistence_rag.KnowledgeBase(
|
||||
uuid=kb_uuid,
|
||||
workspace_uuid=workspace_uuid,
|
||||
name='Capacity Knowledge',
|
||||
description='',
|
||||
knowledge_engine_plugin_id=None,
|
||||
collection_id=kb_uuid,
|
||||
creation_settings={},
|
||||
retrieval_settings={},
|
||||
),
|
||||
_binding_validated=True,
|
||||
)
|
||||
|
||||
await self.mcp_loader._assert_execution_active(context)
|
||||
runtime_mcp = _ProbeMCPSession(mcp_server_name)
|
||||
self.mcp_loader._register_session(
|
||||
context,
|
||||
mcp_server_name,
|
||||
runtime_mcp,
|
||||
)
|
||||
|
||||
bot_context = self._context(
|
||||
workspace_uuid,
|
||||
generation,
|
||||
bot_uuid=bot_uuid,
|
||||
)
|
||||
runtime_bot = await self.platform_manager.load_bot(
|
||||
bot_context,
|
||||
persistence_bot.Bot(
|
||||
uuid=bot_uuid,
|
||||
workspace_uuid=workspace_uuid,
|
||||
name='Capacity Bot',
|
||||
description='',
|
||||
adapter='capacity-probe',
|
||||
adapter_config={},
|
||||
enable=True,
|
||||
use_pipeline_uuid=pipeline_uuid,
|
||||
pipeline_routing_rules=[],
|
||||
),
|
||||
_binding_validated=True,
|
||||
)
|
||||
|
||||
self.generation_refs.setdefault(generation, []).extend(
|
||||
(
|
||||
weakref.ref(runtime_provider),
|
||||
weakref.ref(runtime_llm),
|
||||
weakref.ref(runtime_embedding),
|
||||
weakref.ref(runtime_rerank),
|
||||
weakref.ref(runtime_pipeline),
|
||||
weakref.ref(runtime_kb),
|
||||
weakref.ref(runtime_mcp),
|
||||
weakref.ref(runtime_bot),
|
||||
)
|
||||
)
|
||||
|
||||
await asyncio.sleep(0)
|
||||
|
||||
def retained_state(self) -> dict[str, int]:
|
||||
return {
|
||||
'model_providers': len(self.model_manager.provider_dict),
|
||||
'llm_models': len(self.model_manager.llm_model_dict),
|
||||
'embedding_models': len(self.model_manager.embedding_model_dict),
|
||||
'rerank_models': len(self.model_manager.rerank_model_dict),
|
||||
'model_scopes': len(self.model_manager._scope_generations),
|
||||
'pipelines': len(self.pipeline_manager._pipelines_by_key),
|
||||
'pipeline_scopes': len(self.pipeline_manager._scope_generations),
|
||||
'knowledge_bases': len(self.rag_manager.knowledge_bases),
|
||||
'knowledge_scopes': len(self.rag_manager._scope_generations),
|
||||
'mcp_sessions': len(self.mcp_loader.sessions),
|
||||
'mcp_scopes': len(self.mcp_loader._scope_generations),
|
||||
'bots': len(self.platform_manager._bots_by_key),
|
||||
'bot_scopes': len(self.platform_manager._scope_generations),
|
||||
'requesters_closed': _ProbeRequester.closed,
|
||||
'adapters_killed': _ProbeAdapter.killed,
|
||||
'mcp_sessions_closed': _ProbeMCPSession.closed,
|
||||
'binding_lookups': self.workspace_service.binding_lookups,
|
||||
}
|
||||
|
||||
def assert_generation_state(
|
||||
self,
|
||||
workspaces: int,
|
||||
generation: int,
|
||||
) -> None:
|
||||
state = self.retained_state()
|
||||
cardinality_keys = (
|
||||
'model_providers',
|
||||
'llm_models',
|
||||
'embedding_models',
|
||||
'rerank_models',
|
||||
'model_scopes',
|
||||
'pipelines',
|
||||
'pipeline_scopes',
|
||||
'knowledge_bases',
|
||||
'knowledge_scopes',
|
||||
'mcp_sessions',
|
||||
'mcp_scopes',
|
||||
'bots',
|
||||
'bot_scopes',
|
||||
)
|
||||
invalid = {key: value for key in cardinality_keys if (value := state[key]) != workspaces}
|
||||
if invalid:
|
||||
raise AssertionError(f'populated Workspace cardinality mismatch: {invalid}')
|
||||
expected_retired = (generation - 1) * workspaces
|
||||
if state['requesters_closed'] != expected_retired:
|
||||
raise AssertionError(f'retired requester count {state["requesters_closed"]} != {expected_retired}')
|
||||
if state['adapters_killed'] != expected_retired:
|
||||
raise AssertionError(f'retired adapter count {state["adapters_killed"]} != {expected_retired}')
|
||||
if state['mcp_sessions_closed'] != expected_retired:
|
||||
raise AssertionError(f'retired MCP session count {state["mcp_sessions_closed"]} != {expected_retired}')
|
||||
|
||||
def assert_generation_collected(self, generation: int) -> None:
|
||||
gc.collect()
|
||||
references = self.generation_refs.pop(generation)
|
||||
retained = sum(reference() is not None for reference in references)
|
||||
if retained:
|
||||
raise AssertionError(f'{retained} generation-{generation} runtime objects remain reachable')
|
||||
|
||||
|
||||
async def _run(args: argparse.Namespace) -> dict:
|
||||
scale = SCALES[args.scale]
|
||||
tracemalloc.start()
|
||||
probe = PopulatedWorkspaceProbe()
|
||||
baseline = _sample_process()
|
||||
|
||||
phase_one_started = time.monotonic()
|
||||
await probe.load_generation(scale.workspaces, 1)
|
||||
phase_one_seconds = time.monotonic() - phase_one_started
|
||||
probe.assert_generation_state(scale.workspaces, 1)
|
||||
phase_one = _sample_process()
|
||||
phase_one_state = probe.retained_state()
|
||||
|
||||
phase_two_started = time.monotonic()
|
||||
await probe.load_generation(scale.workspaces, 2)
|
||||
phase_two_seconds = time.monotonic() - phase_two_started
|
||||
probe.assert_generation_state(scale.workspaces, 2)
|
||||
probe.assert_generation_collected(1)
|
||||
phase_two = _sample_process()
|
||||
phase_two_state = probe.retained_state()
|
||||
|
||||
phase_three_started = time.monotonic()
|
||||
await probe.load_generation(scale.workspaces, 3)
|
||||
phase_three_seconds = time.monotonic() - phase_three_started
|
||||
probe.assert_generation_state(scale.workspaces, 3)
|
||||
probe.assert_generation_collected(2)
|
||||
phase_three = _sample_process()
|
||||
phase_three_state = probe.retained_state()
|
||||
|
||||
cardinality_keys = (
|
||||
'model_providers',
|
||||
'llm_models',
|
||||
'embedding_models',
|
||||
'rerank_models',
|
||||
'model_scopes',
|
||||
'pipelines',
|
||||
'pipeline_scopes',
|
||||
'knowledge_bases',
|
||||
'knowledge_scopes',
|
||||
'mcp_sessions',
|
||||
'mcp_scopes',
|
||||
'bots',
|
||||
'bot_scopes',
|
||||
)
|
||||
if any(
|
||||
phase_two_state[key] != phase_one_state[key] or phase_three_state[key] != phase_one_state[key]
|
||||
for key in cardinality_keys
|
||||
):
|
||||
raise AssertionError(
|
||||
'populated Workspace registries did not plateau: '
|
||||
f'phase_one={phase_one_state}, phase_two={phase_two_state}, '
|
||||
f'phase_three={phase_three_state}'
|
||||
)
|
||||
|
||||
traced_growth = phase_three.traced_current_bytes - phase_two.traced_current_bytes
|
||||
rss_growth = phase_three.rss_bytes - phase_two.rss_bytes
|
||||
max_traced_growth = int(args.max_traced_growth_mib * 1024 * 1024)
|
||||
max_rss_growth = int(args.max_rss_growth_mib * 1024 * 1024)
|
||||
if traced_growth > max_traced_growth:
|
||||
raise AssertionError(f'replacement traced memory grew by {traced_growth} bytes (limit {max_traced_growth})')
|
||||
if rss_growth > max_rss_growth:
|
||||
raise AssertionError(f'replacement RSS grew by {rss_growth} bytes (limit {max_rss_growth})')
|
||||
phase_ratio = max(
|
||||
phase_two_seconds,
|
||||
phase_three_seconds,
|
||||
) / max(phase_one_seconds, 0.000_001)
|
||||
if phase_ratio > args.max_replacement_time_ratio:
|
||||
raise AssertionError(f'replacement phase ratio {phase_ratio:.3f} exceeds {args.max_replacement_time_ratio:.3f}')
|
||||
|
||||
return {
|
||||
'component': 'langbot-populated-workspaces',
|
||||
'scale': args.scale,
|
||||
'workspaces': scale.workspaces,
|
||||
'passed': True,
|
||||
'phase_seconds': {
|
||||
'initial': round(phase_one_seconds, 3),
|
||||
'replacement_one': round(phase_two_seconds, 3),
|
||||
'replacement_two': round(phase_three_seconds, 3),
|
||||
'maximum_replacement_ratio': round(phase_ratio, 3),
|
||||
},
|
||||
'samples': {
|
||||
'baseline': asdict(baseline),
|
||||
'phase_one': asdict(phase_one),
|
||||
'phase_two': asdict(phase_two),
|
||||
'phase_three': asdict(phase_three),
|
||||
},
|
||||
'replacement_growth': {
|
||||
'rss_bytes': rss_growth,
|
||||
'traced_current_bytes': traced_growth,
|
||||
},
|
||||
'retained_state': {
|
||||
'phase_one': phase_one_state,
|
||||
'phase_two': phase_two_state,
|
||||
'phase_three': phase_three_state,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument('--scale', choices=tuple(SCALES), default='quick')
|
||||
parser.add_argument('--max-traced-growth-mib', type=float, default=16.0)
|
||||
parser.add_argument('--max-rss-growth-mib', type=float, default=64.0)
|
||||
parser.add_argument(
|
||||
'--max-replacement-time-ratio',
|
||||
type=float,
|
||||
default=3.0,
|
||||
)
|
||||
parser.add_argument('--json', action='store_true')
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = _parse_args()
|
||||
result = asyncio.run(_run(args))
|
||||
if args.json:
|
||||
print(json.dumps(result, sort_keys=True))
|
||||
else:
|
||||
print(json.dumps(result, indent=2, sort_keys=True))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user