cleveragents-core/benchmarks/subplan_spawn_bench.py

"""Airspeed Velocity benchmarks for subplan spawn service throughput.

Measures spawn creation, validation, and entry building overhead for
SubplanService across varying numbers of spawn entries.
"""

from __future__ import annotations

from unittest.mock import MagicMock

from cleveragents.application.services.subplan_service import (
    SpawnEntry,
    SubplanService,
)
from cleveragents.domain.models.core.decision import (
    ContextSnapshot,
    Decision,
    DecisionType,
    ResourceRef,
)
from cleveragents.domain.models.core.plan import (
    ExecutionMode,
    NamespacedName,
    Plan,
    PlanIdentity,
    SubplanConfig,
)

_PLAN_ID = "01HGZ6FE0AQDYTR4BXVQZ6PN00"


def _mock_decision_service() -> MagicMock:
    svc: MagicMock = MagicMock()
    svc.list_by_type = MagicMock(return_value=[])
    return svc


def _make_decision(
    decision_id: str,
    sequence: int = 0,
    resource_ids: list[str] | None = None,
) -> Decision:
    resources: list[ResourceRef] = []
    if resource_ids:
        resources = [ResourceRef(resource_id=rid) for rid in resource_ids]
    return Decision(
        decision_id=decision_id,
        plan_id=_PLAN_ID,
        decision_type=DecisionType.SUBPLAN_SPAWN,
        sequence_number=sequence,
        question="Spawn child plan?",
        chosen_option="local/sub-action",
        context_snapshot=ContextSnapshot(relevant_resources=resources),
    )


def _make_plan() -> Plan:
    return Plan(
        identity=PlanIdentity(plan_id=_PLAN_ID),
        namespaced_name=NamespacedName(namespace="local", name="bench-plan"),
        description="Benchmark plan",
        action_name="local/bench-action",
    )


class SubplanSpawnThroughputSuite:
    """Benchmark SubplanService spawn throughput."""

    def setup(self) -> None:
        """Prepare fixtures for spawn benchmarks."""
        self.service = SubplanService(
            decision_service=_mock_decision_service(),
        )
        self.plan = _make_plan()
        self.seq_config = SubplanConfig(
            execution_mode=ExecutionMode.SEQUENTIAL,
        )
        self.par_config = SubplanConfig(
            execution_mode=ExecutionMode.PARALLEL,
            max_parallel=50,
        )
        # Build entries for varying sizes
        self.entries_1 = self._build_entries(1)
        self.entries_5 = self._build_entries(5)
        self.entries_20 = self._build_entries(20)

    def _build_entries(self, count: int) -> list[SpawnEntry]:
        entries: list[SpawnEntry] = []
        for i in range(count):
            did = f"01HGZ6FE0AQDYTR4BXVQZ6D{i:02d}0"
            dec = _make_decision(did, sequence=i)
            entries.append(SpawnEntry(decision=dec, action_name=f"local/sub-{i}"))
        return entries

    def time_spawn_single_sequential(self) -> None:
        """Time spawning a single child plan in sequential mode."""
        self.service.spawn(
            parent_plan=self.plan,
            config=self.seq_config,
            spawn_entries=self.entries_1,
        )

    def time_spawn_5_sequential(self) -> None:
        """Time spawning 5 child plans in sequential mode."""
        self.service.spawn(
            parent_plan=self.plan,
            config=self.seq_config,
            spawn_entries=self.entries_5,
        )

    def time_spawn_20_parallel(self) -> None:
        """Time spawning 20 child plans in parallel mode."""
        self.service.spawn(
            parent_plan=self.plan,
            config=self.par_config,
            spawn_entries=self.entries_20,
        )

    def time_validate_5_entries(self) -> None:
        """Time validating 5 spawn entries."""
        self.service.validate_spawn(
            config=self.seq_config,
            spawn_entries=self.entries_5,
        )

    def time_validate_20_entries_with_resources(self) -> None:
        """Time validating 20 entries with resource scope check."""
        self.service.validate_spawn(
            config=self.par_config,
            spawn_entries=self.entries_20,
            available_resources=frozenset(f"res-{i}" for i in range(100)),
        )

    def time_build_entries_from_decisions(self) -> None:
        """Time building spawn entries from 5 decisions."""
        decisions = [
            _make_decision(
                f"01HGZ6FE0AQDYTR4BXVQZ6D{i:02d}0",
                sequence=i,
                resource_ids=[f"01HGZ6FE0AQDYTR4BXVQZ6R{i:02d}0"],
            )
            for i in range(5)
        ]
        self.service.build_spawn_entries(decisions)

    def time_service_construction(self) -> None:
        """Time constructing SubplanService."""
        SubplanService(decision_service=_mock_decision_service())