fengqun/examples/run_large_scale_stress_acceptance.py

from __future__ import annotations

from concurrent.futures import ProcessPoolExecutor, as_completed
from dataclasses import asdict, dataclass
from pathlib import Path
import hashlib
import json
import os
import sys
from time import perf_counter


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from swarm_minimal.core import Agent, InMemorySwarmStore, SwarmCoordinator, Task, TaskStatus  # noqa: E402


@dataclass(frozen=True)
class StressConfig:
    logical_cpus: int
    worker_processes: int
    agents_per_process: int
    tasks_per_process: int
    cpu_cycles_per_task: int
    timeout_seconds: int

    @property
    def total_agents(self) -> int:
        return self.worker_processes * self.agents_per_process

    @property
    def total_tasks(self) -> int:
        return self.worker_processes * self.tasks_per_process


def main() -> None:
    config = build_config()
    started = perf_counter()
    shard_results: list[dict[str, object]] = []
    failures: list[str] = []

    with ProcessPoolExecutor(max_workers=config.worker_processes) as executor:
        futures = [
            executor.submit(
                run_stress_shard,
                shard_index,
                config.agents_per_process,
                config.tasks_per_process,
                config.cpu_cycles_per_task,
            )
            for shard_index in range(config.worker_processes)
        ]
        for future in as_completed(futures, timeout=config.timeout_seconds):
            try:
                shard_results.append(future.result())
            except Exception as exc:  # pragma: no cover - defensive failure reporting.
                failures.append(str(exc))

    duration = perf_counter() - started
    completed_tasks = sum(int(item["completed_tasks"]) for item in shard_results)
    failed_tasks = sum(int(item["failed_tasks"]) for item in shard_results)
    duplicate_claims = sum(int(item["duplicate_claim_count"]) for item in shard_results)
    participating_agents = sum(int(item["participating_agents"]) for item in shard_results)
    converged_shards = sum(1 for item in shard_results if item["converged"])
    status = (
        "PASS"
        if not failures
        and len(shard_results) == config.worker_processes
        and completed_tasks == config.total_tasks
        and failed_tasks == 0
        and duplicate_claims == 0
        and participating_agents == config.total_agents
        and converged_shards == config.worker_processes
        else "FAIL"
    )
    report = {
        "standard": "local-large-scale-stress-v1",
        "status": status,
        "scope_note": (
            "This is a bounded local maximum-performance stress run. It uses all detected logical CPUs by default "
            "and validates in-memory autonomous claim, convergence and duplicate-claim safety. It is not a "
            "production Kubernetes runtime benchmark."
        ),
        "config": asdict(config),
        "success_thresholds": {
            "worker_processes": "equals detected logical_cpus unless overridden",
            "completed_tasks": config.total_tasks,
            "failed_tasks": 0,
            "duplicate_claims": 0,
            "participating_agents": config.total_agents,
            "converged_shards": config.worker_processes,
        },
        "summary": {
            "duration_seconds": round(duration, 4),
            "tasks_per_second": round(completed_tasks / duration, 2) if duration else completed_tasks,
            "completed_tasks": completed_tasks,
            "failed_tasks": failed_tasks,
            "duplicate_claims": duplicate_claims,
            "participating_agents": participating_agents,
            "converged_shards": converged_shards,
            "failures": failures,
        },
        "shards": sorted(shard_results, key=lambda item: int(item["shard_index"])),
    }
    print(json.dumps(report, ensure_ascii=False, indent=2))
    if status != "PASS":
        raise SystemExit(1)


def build_config() -> StressConfig:
    logical_cpus = max(1, os.cpu_count() or 1)
    worker_processes = read_positive_int("SWARM_STRESS_PROCESSES", logical_cpus)
    agents_per_process = read_positive_int("SWARM_STRESS_AGENTS_PER_PROCESS", max(4, logical_cpus))
    tasks_per_process = read_positive_int("SWARM_STRESS_TASKS_PER_PROCESS", max(1024, logical_cpus * 256))
    cpu_cycles_per_task = read_positive_int("SWARM_STRESS_CPU_CYCLES", 64)
    timeout_seconds = read_positive_int("SWARM_STRESS_TIMEOUT_SECONDS", 300)
    return StressConfig(
        logical_cpus=logical_cpus,
        worker_processes=worker_processes,
        agents_per_process=agents_per_process,
        tasks_per_process=tasks_per_process,
        cpu_cycles_per_task=cpu_cycles_per_task,
        timeout_seconds=timeout_seconds,
    )


def read_positive_int(name: str, default: int) -> int:
    raw = os.environ.get(name)
    if raw is None:
        return default
    try:
        value = int(raw)
    except ValueError as exc:
        raise ValueError(f"{name} must be an integer") from exc
    if value <= 0:
        raise ValueError(f"{name} must be positive")
    return value


def run_stress_shard(
    shard_index: int,
    agents_per_process: int,
    tasks_per_process: int,
    cpu_cycles_per_task: int,
) -> dict[str, object]:
    store = InMemorySwarmStore()
    run_id = f"large-stress-{os.getpid()}-{shard_index}"
    store.shared_state[f"run:{run_id}:goal"] = "large scale autonomous claim stress"
    store.shared_state[f"run:{run_id}:status"] = "running"
    for task_index in range(tasks_per_process):
        store.add_task(Task(kind="stress", input=f"shard={shard_index}; task={task_index}"))

    agents = [
        Agent(
            id=f"stress-agent-{shard_index}-{agent_index}",
            capability="stress",
            run=lambda task, shared_state, agent_index=agent_index: run_cpu_bound_task(
                task.input,
                agent_index,
                cpu_cycles_per_task,
            ),
        )
        for agent_index in range(agents_per_process)
    ]

    started = perf_counter()
    report = SwarmCoordinator(store=store, agents=agents).run_autonomous_until_converged(
        run_id,
        max_workers=agents_per_process,
    )
    duration = perf_counter() - started
    participating_agents = {event.agent_id for event in report.claim_events}
    all_done = all(task.status == TaskStatus.DONE for task in store.tasks.values())

    return {
        "shard_index": shard_index,
        "process_id": os.getpid(),
        "agent_count": agents_per_process,
        "task_count": tasks_per_process,
        "completed_tasks": report.completed_tasks,
        "failed_tasks": report.failed_tasks,
        "duplicate_claim_count": len(report.duplicate_claims),
        "participating_agents": len(participating_agents),
        "converged": report.converged,
        "all_tasks_done": all_done,
        "duration_seconds": round(duration, 4),
        "tasks_per_second": round(report.completed_tasks / duration, 2) if duration else report.completed_tasks,
        "accepted_score": report.result.accepted_score,
        "observation_count": len(report.result.observations),
    }


def run_cpu_bound_task(task_input: str, agent_index: int, cpu_cycles_per_task: int) -> tuple[str, float]:
    payload = f"{task_input}; agent={agent_index}".encode()
    for _ in range(cpu_cycles_per_task):
        payload = hashlib.blake2b(payload, digest_size=16).digest()
    checksum = int.from_bytes(payload[:4], "big")
    score = 0.5 + (checksum % 5000) / 10000
    return f"agent={agent_index}; checksum={checksum}; stress=ok", score


if __name__ == "__main__":
    main()