Architecture and Scalability

This page describes the AgentSociety 2 runtime architecture and how it organizes agents, environments, LLM calls, trace data, and replay storage for larger-scale simulations.

AgentSociety 2 stores agent state in independent workspaces. During execution, Ray Tasks rebuild and advance agents in batches. Runtime components such as environment routing, LLM dispatch, trace, and replay run as shared services and are injected into each agent through ServiceProxy. The goal is to let the number of agents grow with experiment scale instead of keeping a long-lived process or actor for every agent.

Execution Model

Each AgentBase subclass stores its persistent state in its own workspace, including config.json, AGENT.json, state/*, and .runtime/logs/*. The agent object is only a runtime object rebuilt from those files for one task execution; after a step() completes, the new state is written back to the workspace.

This model matches the create / from_workspace / to_workspace contract described in Using Agents. To the scheduler, an agent can be represented as a lightweight execution spec: workspace path, agent type, agent id, and a shared ServiceProxy handle. After a Ray worker receives a task, it calls from_workspace locally, runs step, and then calls to_workspace to persist the result. The agent object itself is not passed back and forth through the Ray object store.

Environment modules use a symmetric persistence contract: EnvBase.to_workspace() writes dynamic state, and restore() restores it. Each module can choose its own state format, and usually stores the file at <workspace_root>/state/ENV_STATE.json. The router calls to_workspaces() at the end of each step. When --resume is used, the actor’s init() rebuilds modules through from_workspaces(). Restoration happens after the module init(), so checkpointed state overrides initialization-time resets. AgentSociety stores its own state in SOCIETY.json and SOCIETY_STEP.json. See Environment Modules and the resume section in Command Line Interface for details.

For each tick, AgentSociety performs the following flow:

  1. Split agent ids into batches.

  2. Submit one step_agent_batch Ray Task for each batch.

  3. Within each task, advance agents concurrently with asyncio.gather; across batches, Ray schedules work onto different workers.

  4. After all batches complete, call the environment step and then advance the simulation clock.

默认批大小为 AGENTSOCIETY_BATCH_SIZE=256``(CLI ``--batch-size 可覆盖)。每 tick 的 Ray Task 数 = ⌈N / 批大小⌉,受 AGENTSOCIETY_LLM_RAY_MAX_WORKERS``(默认机器 CPU 核数)封顶; 单个 task 内的并发 fan-out 由批大小决定,每进程的 LLM 并发则由本地 ``LLMClient 的 AIMD 自适应控制(无人工上下限)。

digraph exec_model {
    rankdir=LR;
    node [shape=box, style=rounded];

    Driver [label="AgentSociety (driver)"];
    Batch [label="step_agent_batch\nRay Task × 批", shape=note];
    Agent [label="agent.from_workspace()\n无状态重建 + step"];
    WS [label="workspace\n(磁盘状态)", shape=cylinder];

    Services [label="共享服务\nenv / llm / trace / replay", style=filled, color=lightyellow];

    Driver -> Batch [label="切批提交"];
    Batch -> Agent;
    Agent -> WS [label="读写状态"];
    Agent -> Services [label="经 ServiceProxy"];
}

ServiceProxy

Agents do not directly create or own the LLM dispatcher, environment router, trace actor, or replay actor. from_workspace receives a ServiceProxy and binds the relevant handles to the agent instance in _bind_services(service_proxy).

slot

Source

Purpose

_service_proxy

Input parameter

Raw service container with env / llm / trace / replay handles.

_env

service_proxy.env

Environment router (RouterBase / EnvRouterProxy).

_dispatcher

service_proxy.llm.default

Default LLM dispatcher (call(...)).

_trace

service_proxy.trace

Trace writer handle, usually TraceProxy.

ServiceProxy stores only serializable handles and configuration. It does not store locks, connection pools, or background threads. A Ray Task therefore only needs to receive one service container. The environment router, trace, and replay are created once by the driver or CLI; the LLM client only carries connection parameters, and workers create the litellm Router and concurrency controller on demand in their own event loop.

Environment Routing

Environment state is managed by the environment router. In production runtime, the CLI creates a dedicated Ray actor through get_env_router_actor_class and places an EnvRouterProxy in ServiceProxy. Agents call ask, step, init, get_world_description, and related interfaces through that proxy; all agents see the same environment state.

Routers (RouterBase subclasses) map an agent’s natural-language request to @tool methods exposed by environment modules:

Router

Characteristics

CodeGenRouter (default)

Extracts tool signatures from environment modules, generates call code, and executes it in a restricted environment; includes AST guards and cache statistics.

ReActRouter

ReAct-style tool selection.

PlanExecuteRouter

Plan first, then execute.

TwoTierReActRouter / TwoTierPlanExecuteRouter

Two-tier routing for large tool sets.

SearchToolRouter

Selects tools through search/retrieval.

For environment modules and the @tool decorator, see Environment Modules and Core Concepts.

LLM Dispatch

LLM calls are centrally wrapped by agentsociety2.config.llm_dispatcher. The current implementation uses a per-process dispatcher: LLMClient is a serializable connection config carrying only model_name, base_url, and api_key; it can enter Ray Tasks or the env router actor through ServiceProxy. When each consumer calls call() for the first time in its own event loop, it creates a local litellm Router and AdaptiveSemaphore. This avoids event-loop binding problems caused by reusing asyncio objects across Ray tasks.

Concurrency control happens inside each process / event loop. AdaptiveSemaphore adjusts with AIMD (additive increase, multiplicative decrease) based on recent latency and rate-limit errors, controlling LLM request fan-out within that process. Token usage is also accumulated locally per LLMClient instance and merged through Ray Task return values.

Defaults work out of the box; when tuning is needed, focus on these environment variables (see Installation):

Environment variable

Default

Meaning

AGENTSOCIETY_LLM_RAY_MAX_WORKERS

机器 CPU 核数

CPU budget for ray.init(num_cpus=...); caps how many step_agent_batch Ray Tasks can run concurrently per tick.

AGENTSOCIETY_LLM_RAY_CONCURRENCY

16

Initial AIMD concurrency for each local LLMClient. There is no manual upper or lower bound; it automatically approaches the API’s actual capacity.

AGENTSOCIETY_BATCH_SIZE

256

Number of agents handled by each step_agent_batch Ray Task; the CLI --batch-size option can override it.

AGENTSOCIETY_LLM_LATENCY_DEGRADE_FACTOR

4.0

Relative AIMD backoff factor. Concurrency is reduced when call latency is significantly higher than the recent healthy baseline.

AGENTSOCIETY_LLM_SLOW_LATENCY_MS

Not set

Optional absolute latency threshold in milliseconds. When unset, only the relative factor is used.

init_dispatchers() only initializes Ray for the env router actor and agent Ray Tasks. shutdown_dispatchers() is a no-op on the LLM side; local Routers are released when their process exits.

Trace

agentsociety2.trace provides distributed tracing:

  • ShardedTraceWriter shards writes by trace_id and flushes them in batches through a background thread.

  • JsonlTraceWriter provides start_span, trace_span, record_event, end_span, and TraceSpan.

  • TraceActor is a long-lived Ray actor; TraceProxy is the lightweight handle placed in ServiceProxy.

Agent execution, memory, environment calls, and LLM calls can all be recorded as spans. After a run completes, trace data can be used to analyze slow calls, tail-latency tasks, and environment tool-call bottlenecks.

Replay and Storage

Replay data is written by ReplaySink / ReplayWriter as append-only JSONL under run/replay/. The dataset catalog and column metadata are stored in _schema.json. ReplayProxy only carries the replay directory and enabled flag; the env router, society, and agent tasks each create a local ReplaySink in their own process and append concurrently, without a central SQLite actor. On the read side, ReplayReader uses DuckDB to register JSONL shards as views. For table structure and dataset registration, see Storage and Result Access.

Extension Points

The main AgentSociety 2 extension points are:

  • Add agent capabilities: write Agent Skills and place them under custom/skills/; see Agent Skills.

  • Add environment modules: inherit from EnvBase and expose callable methods with @tool; see Environment Modules.

  • Customize agent types: inherit from AgentBase and discover them through the registry; see Using Agents.

  • Analyze run results: read trace spans or replay datasets; see Storage and Result Access.

References