mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-06 08:16:00 +00:00
c85917a812
* feat(core): support live component replacement * fix(core): serialize replacement with app lifecycle * fix(core): preserve state during component replacement * fix(core): roll back partial component startup * chore(deps): add 'web' extra to reme-ai and remove reme_studio from core dependencies - Updated reme-ai dependency to include 'as' and 'web' extras - Removed reme_studio from core dependency list to avoid duplication or unnecessary install * fix(core): preserve persisted state on startup failure
394 lines
17 KiB
Python
394 lines
17 KiB
Python
"""Main application entry point."""
|
|
|
|
import asyncio
|
|
import heapq
|
|
from collections.abc import Mapping
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from pathlib import Path
|
|
from typing import Any, AsyncGenerator, TypeVar
|
|
|
|
from . import __version__
|
|
from .components import ApplicationContext, BaseComponent
|
|
from .components.job import BackgroundJob, BaseJob, CronJob, StreamJob
|
|
from .components.service import BaseService
|
|
from .enumeration import ComponentEnum, ComponentType, component_type_name
|
|
from .plugin import resolve_plugin_runtime
|
|
from .schema import ComponentConfig, Response, StreamChunk
|
|
from .utils import execute_stream_task, print_logo, get_logger
|
|
|
|
T = TypeVar("T", bound=BaseComponent)
|
|
_NodeKey = tuple[str, str]
|
|
|
|
|
|
class Application(BaseComponent):
|
|
"""Wires components from config and runs jobs against them."""
|
|
|
|
def __init__(self, **kwargs) -> None:
|
|
runtime = resolve_plugin_runtime(kwargs)
|
|
self.context = ApplicationContext(registry=runtime.registry, **runtime.config)
|
|
self._started_components: list[BaseComponent] = []
|
|
self._component_mutation_lock = asyncio.Lock()
|
|
|
|
self._setup_workspace_directories()
|
|
logger = get_logger(
|
|
log_to_console=self.config.log_to_console,
|
|
log_to_file=self.config.log_to_file,
|
|
force_init=True,
|
|
)
|
|
super().__init__()
|
|
self._init_service()
|
|
|
|
if self.config.enable_logo:
|
|
print_logo(self.config, self.context.service)
|
|
logger.info(f"Initializing {self.config.app_name} Application v{__version__}")
|
|
self._init_components()
|
|
self._init_jobs()
|
|
|
|
@property
|
|
def config(self):
|
|
"""Typed view onto the application config held by the context."""
|
|
return self.context.app_config
|
|
|
|
# ----- Wiring (called once during __init__) --------------------------
|
|
|
|
def _setup_workspace_directories(self) -> None:
|
|
"""Ensure the workspace root and configured subdirectories exist on disk."""
|
|
cfg = self.config
|
|
workspace_path = Path(cfg.workspace_dir).absolute()
|
|
workspace_path.mkdir(parents=True, exist_ok=True)
|
|
for subdir in [
|
|
cfg.metadata_dir,
|
|
cfg.session_dir,
|
|
cfg.mem_session_dir,
|
|
cfg.resource_dir,
|
|
cfg.daily_dir,
|
|
cfg.digest_dir,
|
|
]:
|
|
if subdir:
|
|
(workspace_path / subdir).mkdir(parents=True, exist_ok=True)
|
|
|
|
def _init_service(self) -> None:
|
|
"""Instantiate the single service backend declared in config.service."""
|
|
self.context.service = self._instantiate(
|
|
ComponentEnum.SERVICE,
|
|
self.config.service,
|
|
label="Service",
|
|
expected_type=BaseService,
|
|
)
|
|
|
|
def _init_components(self) -> None:
|
|
"""Instantiate every component declared under config.components."""
|
|
for ctype, group in self.config.components.items():
|
|
self.context.components[ctype] = {}
|
|
for name, cfg in group.items():
|
|
self.context.components[ctype][name] = self._instantiate(
|
|
ctype,
|
|
cfg,
|
|
label=f"Component '{name}'",
|
|
expected_type=BaseComponent,
|
|
name=name,
|
|
)
|
|
|
|
def _init_jobs(self) -> None:
|
|
"""Instantiate every job declared under config.jobs."""
|
|
for name, cfg in self.config.jobs.items():
|
|
self.context.jobs[name] = self._instantiate(
|
|
ComponentEnum.JOB,
|
|
cfg,
|
|
label=f"Job '{name}'",
|
|
expected_type=BaseJob,
|
|
name=name,
|
|
)
|
|
|
|
def _instantiate(
|
|
self,
|
|
ctype: ComponentType,
|
|
cfg: ComponentConfig,
|
|
*,
|
|
label: str,
|
|
expected_type: type[T],
|
|
name: str | None = None,
|
|
) -> T:
|
|
"""Resolve cfg.backend through the registry and construct the instance.
|
|
|
|
`label` is the human-readable identifier used only in error messages.
|
|
`expected_type` narrows the return type and guards against a backend
|
|
registered under the wrong component type.
|
|
`name` is forwarded to the constructor for named components/jobs;
|
|
leave it None for the service, which is keyed solely by type.
|
|
"""
|
|
if not cfg.backend:
|
|
raise ValueError(f"{label} is missing the required 'backend' field")
|
|
backend_cls = self.context.registry.get(ctype, cfg.backend)
|
|
if backend_cls is None:
|
|
raise ValueError(f"Unregistered backend '{cfg.backend}' for {label}")
|
|
|
|
params = cfg.model_dump()
|
|
params["app_context"] = self.context
|
|
if name is not None:
|
|
params.setdefault("name", name)
|
|
instance = backend_cls(**params)
|
|
if not isinstance(instance, expected_type):
|
|
got, want = type(instance).__name__, expected_type.__name__
|
|
raise TypeError(f"{label} backend '{cfg.backend}' produced {got}, expected {want} subclass")
|
|
return instance
|
|
|
|
# ----- Dependency ordering ------------------------------------------
|
|
|
|
def _topological_order(
|
|
self,
|
|
replacement: tuple[_NodeKey, BaseComponent] | None = None,
|
|
) -> list[BaseComponent]:
|
|
"""Return components in dependency order via Kahn's algorithm; raise on missing dep or cycle."""
|
|
nodes: dict[_NodeKey, BaseComponent] = {
|
|
(ctype, name): comp for ctype, group in self.context.components.items() for name, comp in group.items()
|
|
}
|
|
if replacement is not None:
|
|
nodes[replacement[0]] = replacement[1]
|
|
in_degree, dependents = self._build_dependency_graph(nodes)
|
|
|
|
ready = [k for k, d in in_degree.items() if d == 0]
|
|
heapq.heapify(ready)
|
|
ordered: list[BaseComponent] = []
|
|
while ready:
|
|
key = heapq.heappop(ready)
|
|
ordered.append(nodes[key])
|
|
for downstream in dependents[key]:
|
|
in_degree[downstream] -= 1
|
|
if in_degree[downstream] == 0:
|
|
heapq.heappush(ready, downstream)
|
|
|
|
if len(ordered) != len(nodes):
|
|
unresolved = [f"{k[0]}:{k[1]}" for k, d in in_degree.items() if d > 0]
|
|
raise ValueError(f"Circular dependency detected among: {unresolved}")
|
|
return ordered
|
|
|
|
@staticmethod
|
|
def _build_dependency_graph(
|
|
nodes: dict[_NodeKey, BaseComponent],
|
|
) -> tuple[dict[_NodeKey, int], dict[_NodeKey, list[_NodeKey]]]:
|
|
"""Compute in-degree and adjacency lists; raise if a required dep is missing."""
|
|
in_degree: dict[_NodeKey, int] = dict.fromkeys(nodes, 0)
|
|
dependents: dict[_NodeKey, list[_NodeKey]] = {k: [] for k in nodes}
|
|
for key, comp in nodes.items():
|
|
for dep in comp.dependencies:
|
|
dep_key = (dep.ctype, dep.name)
|
|
if dep_key in nodes:
|
|
dependents[dep_key].append(key)
|
|
in_degree[key] += 1
|
|
elif not dep.optional:
|
|
raise ValueError(
|
|
f"Component {key[0]}:{key[1]} depends on unregistered {dep.ctype}:{dep.name}",
|
|
)
|
|
return in_degree, dependents
|
|
|
|
# ----- Lifecycle -----------------------------------------------------
|
|
|
|
async def _start(self) -> None:
|
|
"""Start components, then jobs as base > stream > background > cron."""
|
|
async with self._component_mutation_lock:
|
|
pool_size = self.config.thread_pool_max_workers
|
|
if pool_size > 0:
|
|
self.context.thread_pool = ThreadPoolExecutor(max_workers=pool_size)
|
|
self.logger.info(f"Thread pool created with max_workers={pool_size}")
|
|
try:
|
|
components = self._topological_order()
|
|
jobs = list(self.context.jobs.values())
|
|
base_jobs = [j for j in jobs if not isinstance(j, (StreamJob, BackgroundJob))]
|
|
stream_jobs = [j for j in jobs if isinstance(j, StreamJob)]
|
|
background_jobs = [j for j in jobs if isinstance(j, BackgroundJob) and not isinstance(j, CronJob)]
|
|
cron_jobs = [j for j in jobs if isinstance(j, CronJob)]
|
|
for c in components + base_jobs + stream_jobs + background_jobs + cron_jobs:
|
|
await self._start_one(c)
|
|
except Exception:
|
|
await self._close_started_components()
|
|
raise
|
|
|
|
async def _start_one(self, c: BaseComponent) -> None:
|
|
"""Start one component and record it for ordered shutdown."""
|
|
try:
|
|
if isinstance(c, BackgroundJob):
|
|
self.logger.info(f"Starting background job: {c.name}")
|
|
await c.start()
|
|
self._started_components.append(c)
|
|
except Exception as e:
|
|
self.logger.exception(f"Failed to start {component_type_name(c.component_type)}:{c.name}: {e}")
|
|
raise
|
|
|
|
async def _close(self) -> None:
|
|
"""Close in reverse start order so every peer outlives its dependents."""
|
|
async with self._component_mutation_lock:
|
|
await self._close_started_components()
|
|
|
|
async def _close_started_components(self) -> None:
|
|
"""Close resources while the caller serializes component mutations."""
|
|
for c in reversed(self._started_components):
|
|
try:
|
|
await c.close()
|
|
except Exception as e:
|
|
self.logger.exception(f"Failed to close {component_type_name(c.component_type)}:{c.name}: {e}")
|
|
self._started_components.clear()
|
|
if self.context.thread_pool is not None:
|
|
self.context.thread_pool.shutdown(wait=True)
|
|
self.context.thread_pool = None
|
|
|
|
async def update_component(self, component_enum: ComponentType, name: str, /, **kwargs) -> BaseComponent:
|
|
"""Update an existing component by type/name; never creates missing components."""
|
|
async with self._component_mutation_lock:
|
|
component_type = component_type_name(component_enum)
|
|
group = self.context.components.get(component_type)
|
|
if not group or name not in group:
|
|
raise KeyError(f"Component '{name}' not found in {component_type}")
|
|
|
|
component = group[name]
|
|
for key in kwargs:
|
|
if not hasattr(component, key):
|
|
raise AttributeError(f"Component {component_type}:{name} has no attribute '{key}'")
|
|
for key, value in kwargs.items():
|
|
setattr(component, key, value)
|
|
return component
|
|
|
|
def _replacement_shutdown_order(
|
|
self,
|
|
old_component: BaseComponent,
|
|
replacement: BaseComponent,
|
|
replacement_order: list[BaseComponent],
|
|
) -> list[BaseComponent]:
|
|
"""Precompute shutdown tracking using identity, not component hashing."""
|
|
started_ids = {id(component) for component in self._started_components}
|
|
started_ids.discard(id(old_component))
|
|
started_ids.add(id(replacement))
|
|
component_ids = {
|
|
id(component) for components in self.context.components.values() for component in components.values()
|
|
}
|
|
non_components = [
|
|
component
|
|
for component in self._started_components
|
|
if component is not old_component and id(component) not in component_ids
|
|
]
|
|
return [component for component in replacement_order if id(component) in started_ids] + non_components
|
|
|
|
async def replace_component(
|
|
self,
|
|
component_enum: ComponentType,
|
|
name: str,
|
|
/,
|
|
*,
|
|
config: ComponentConfig | Mapping[str, Any],
|
|
runtime_updates: Mapping[str, Any] | None = None,
|
|
) -> BaseComponent:
|
|
"""Replace an existing component and synchronously commit its references.
|
|
|
|
``config`` is the complete declarative configuration for the new
|
|
component. ``runtime_updates`` injects non-serializable live objects,
|
|
such as an already verified model, before the replacement is started.
|
|
|
|
The old component is dumped before the replacement starts so compatible
|
|
``start()`` / ``load()`` implementations see its latest state. The
|
|
context, dependent bindings, in-memory application config, and shutdown
|
|
order are then switched without an await boundary. A dump or start
|
|
failure leaves the old generation authoritative. Hosts must quiesce
|
|
calls that may retain component references across this operation and
|
|
migrate state explicitly when changing between incompatible backends.
|
|
"""
|
|
async with self._component_mutation_lock:
|
|
component_type = component_type_name(component_enum)
|
|
group = self.context.components.get(component_type)
|
|
config_group = self.config.components.get(component_type)
|
|
if not group or name not in group or config_group is None:
|
|
raise KeyError(f"Component '{name}' not found in {component_type}")
|
|
|
|
old_component = group[name]
|
|
replacement_config = (
|
|
config.model_copy(deep=True)
|
|
if isinstance(config, ComponentConfig)
|
|
else ComponentConfig.model_validate(dict(config))
|
|
)
|
|
replacement = self._instantiate(
|
|
component_type,
|
|
replacement_config,
|
|
label=f"Component '{name}'",
|
|
expected_type=BaseComponent,
|
|
name=name,
|
|
)
|
|
for key, value in (runtime_updates or {}).items():
|
|
if not hasattr(replacement, key):
|
|
raise AttributeError(
|
|
f"Replacement {component_type}:{name} has no attribute '{key}'",
|
|
)
|
|
setattr(replacement, key, value)
|
|
|
|
node_key = (component_type, name)
|
|
replacement_order = Application._topological_order(self, replacement=(node_key, replacement))
|
|
was_started = old_component.is_started
|
|
|
|
consumers: list[tuple[BaseComponent, str]] = []
|
|
candidates: list[BaseComponent] = [
|
|
component for components in self.context.components.values() for component in components.values()
|
|
]
|
|
candidates.extend(self.context.jobs.values())
|
|
if self.context.service is not None:
|
|
candidates.append(self.context.service)
|
|
for consumer in candidates:
|
|
if consumer is old_component:
|
|
continue
|
|
for attr, dependency in consumer.dependency_bindings.items():
|
|
if (
|
|
dependency.ctype == component_type
|
|
and dependency.name == name
|
|
and consumer.__dict__.get(attr) is old_component
|
|
):
|
|
consumers.append((consumer, attr))
|
|
|
|
replacement_started_components = (
|
|
self._replacement_shutdown_order(old_component, replacement, replacement_order) if was_started else None
|
|
)
|
|
if was_started:
|
|
await old_component.dump()
|
|
await replacement.start()
|
|
|
|
# Commit the new generation synchronously so observers cannot see
|
|
# a context with only some dependency references updated.
|
|
group[name] = replacement
|
|
config_group[name] = replacement_config
|
|
for consumer, attr in consumers:
|
|
consumer.__dict__[attr] = replacement
|
|
if replacement_started_components is not None:
|
|
self._started_components = replacement_started_components
|
|
|
|
if was_started:
|
|
try:
|
|
await old_component.close()
|
|
except Exception as exc: # The committed replacement remains authoritative.
|
|
self.logger.exception(
|
|
f"Failed to close replaced component {component_type}:{name}: {exc}",
|
|
)
|
|
return replacement
|
|
|
|
# ----- Job execution -------------------------------------------------
|
|
|
|
async def run_job(self, name: str, /, **kwargs) -> Response:
|
|
"""Execute a registered job by name and return its final Response."""
|
|
if name not in self.context.jobs:
|
|
raise KeyError(f"Job '{name}' not found")
|
|
return await self.context.jobs[name](**kwargs)
|
|
|
|
async def run_stream_job(self, name: str, /, **kwargs) -> AsyncGenerator[StreamChunk, None]:
|
|
"""Execute a streaming job, yielding chunks as they are produced."""
|
|
if name not in self.context.jobs:
|
|
raise KeyError(f"Job '{name}' not found")
|
|
stream_queue: asyncio.Queue = asyncio.Queue()
|
|
task = asyncio.create_task(self.context.jobs[name](stream_queue=stream_queue, **kwargs))
|
|
async for chunk in execute_stream_task(
|
|
stream_queue=stream_queue,
|
|
task=task,
|
|
task_name=name,
|
|
output_format="chunk",
|
|
):
|
|
assert isinstance(chunk, StreamChunk)
|
|
yield chunk
|
|
|
|
def run_app(self):
|
|
"""Serve the application through the configured service backend."""
|
|
assert isinstance(self.context.service, BaseService)
|
|
self.context.service.run_app(app=self)
|