wehub-resource-sync
416 lines
14 KiB
Python
416 lines
14 KiB
Python
"""Tests for fair queueing per-agent round-robin.
|
||
|
||
Covers per-agent_id RR scheduling without starvation, DM-not-starved
|
||
mixed-load scenario, round-robin order within a single agent_id,
|
||
cross-agent_id non-interference, terminal RR-state cleanup, strict
|
||
ABCABC ordering at max_concurrency=1, and no-underutilization soak.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import random
|
||
from typing import Any
|
||
from unittest.mock import MagicMock
|
||
|
||
import pytest
|
||
|
||
from opensquilla.gateway.routing import RouteEnvelope, SourceKind
|
||
from opensquilla.gateway.task_runtime import TaskRuntime
|
||
from opensquilla.session.models import AgentTaskRecord
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Shared helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _make_envelope(agent_id: str, session_key: str) -> RouteEnvelope:
|
||
return RouteEnvelope(
|
||
source_kind=SourceKind.WEB,
|
||
source_name="test",
|
||
agent_id=agent_id,
|
||
session_key=session_key,
|
||
input_provenance={"kind": "test"},
|
||
)
|
||
|
||
|
||
def _make_storage() -> Any:
|
||
storage = MagicMock()
|
||
task_db: dict[str, AgentTaskRecord] = {}
|
||
|
||
async def create(record: AgentTaskRecord) -> None:
|
||
task_db[record.task_id] = record
|
||
|
||
async def update(task_id: str, **kwargs: Any) -> None:
|
||
rec = task_db.get(task_id)
|
||
if rec is None:
|
||
return
|
||
for k, v in kwargs.items():
|
||
if hasattr(rec, k):
|
||
object.__setattr__(rec, k, v)
|
||
|
||
async def get(task_id: str) -> AgentTaskRecord | None:
|
||
return task_db.get(task_id)
|
||
|
||
async def list_tasks(**kwargs: Any) -> list[AgentTaskRecord]:
|
||
return list(task_db.values())
|
||
|
||
storage.create_agent_task = create
|
||
storage.update_agent_task = update
|
||
storage.get_agent_task = get
|
||
storage.list_agent_tasks = list_tasks
|
||
return storage
|
||
|
||
|
||
def _make_runtime(
|
||
max_concurrency: int = 4,
|
||
max_pending_per_session: int | None = None,
|
||
) -> TaskRuntime:
|
||
async def turn_handler(run: Any) -> None:
|
||
# Tiny random delay to simulate real work without LLM calls.
|
||
await asyncio.sleep(random.uniform(0.001, 0.005))
|
||
|
||
return TaskRuntime(
|
||
storage=_make_storage(),
|
||
turn_handler=turn_handler,
|
||
max_concurrency=max_concurrency,
|
||
max_pending_per_session=max_pending_per_session,
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# DM not starved
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_dm_not_starved() -> None:
|
||
"""Multiple channel sessions push 100 tasks total; one DM session pushes 10.
|
||
|
||
All sessions share the same agent_id. With fair queuing the DM session
|
||
must complete all 10 tasks before the channel sessions have drained more
|
||
than 80 of their 100 tasks (i.e. DM is not pushed to the end of the queue).
|
||
|
||
We use 10 channel sessions × 10 tasks each (total 100) so that multiple
|
||
channel tasks are simultaneously contending for global-sem slots alongside
|
||
the DM tasks, making the per-agent fairness gate meaningful.
|
||
"""
|
||
random.seed(42)
|
||
# max_pending_per_session=None: no per-session queue cap so we can enqueue
|
||
# 100 tasks to a single session without hitting TaskQueueFullError.
|
||
runtime = _make_runtime(max_concurrency=4, max_pending_per_session=None)
|
||
|
||
agent_id = "agent-dm-starve"
|
||
# Single channel session with 100 tasks vs single DM session with 10 tasks.
|
||
# Both share the same agent_id. Without fair queuing, channel grabs every
|
||
# slot as soon as it becomes free (FIFO) and DM waits until all 100 channel
|
||
# tasks drain. With per-agent_id session round-robin, once channel has more
|
||
# completions than DM, it must yield until DM catches up — so DM completes
|
||
# its 10 tasks while channel has completed at most ~20 tasks (10 lead + 10
|
||
# interleaved rounds before DM is done).
|
||
channel_env = _make_envelope(agent_id, f"{agent_id}::channel-1")
|
||
dm_env = _make_envelope(agent_id, f"{agent_id}::dm-1")
|
||
|
||
channel_completed: list[str] = []
|
||
dm_completed: list[str] = []
|
||
|
||
original_turn = runtime._turn_handler
|
||
|
||
async def instrumented_turn(run: Any) -> None:
|
||
await original_turn(run)
|
||
if "channel" in run.session_key:
|
||
channel_completed.append(run.task_id)
|
||
else:
|
||
dm_completed.append(run.task_id)
|
||
|
||
runtime._turn_handler = instrumented_turn # type: ignore[method-assign]
|
||
|
||
# Enqueue channel tasks first (100), then DM tasks (10).
|
||
channel_handles = []
|
||
for i in range(100):
|
||
h = await runtime.enqueue(channel_env, f"channel msg {i}")
|
||
channel_handles.append(h)
|
||
|
||
dm_handles = []
|
||
for i in range(10):
|
||
h = await runtime.enqueue(dm_env, f"dm msg {i}")
|
||
dm_handles.append(h)
|
||
|
||
# Wait for all DM tasks to finish.
|
||
for h in dm_handles:
|
||
await runtime.wait(h.task_id, timeout=30.0)
|
||
|
||
channel_done_when_dm_finished = len(channel_completed)
|
||
|
||
# Drain remaining channel tasks.
|
||
for h in channel_handles:
|
||
try:
|
||
await runtime.wait(h.task_id, timeout=30.0)
|
||
except KeyError:
|
||
pass
|
||
|
||
assert len(dm_completed) == 10, f"DM only completed {len(dm_completed)}/10"
|
||
assert channel_done_when_dm_finished <= 80, (
|
||
f"DM starved: channel had already completed {channel_done_when_dm_finished} "
|
||
f"tasks when DM finished (expected <=80)"
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Round-robin within same agent_id
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_round_robin_within_agent() -> None:
|
||
"""Same agent_id, 3 sessions A/B/C each push 9 tasks.
|
||
|
||
Execution order should approximate ABCABCABC... with <=20% deviation
|
||
from ideal round-robin (i.e. each session gets between 7 and 11 slots
|
||
in the first 27 completions out of 27 total).
|
||
"""
|
||
random.seed(7)
|
||
runtime = _make_runtime(max_concurrency=3, max_pending_per_session=None)
|
||
|
||
agent_id = "agent-rr"
|
||
envs = {
|
||
"A": _make_envelope(agent_id, f"{agent_id}::sess-a"),
|
||
"B": _make_envelope(agent_id, f"{agent_id}::sess-b"),
|
||
"C": _make_envelope(agent_id, f"{agent_id}::sess-c"),
|
||
}
|
||
|
||
completion_order: list[str] = []
|
||
original_turn = runtime._turn_handler
|
||
|
||
async def instrumented_turn(run: Any) -> None:
|
||
await original_turn(run)
|
||
for label, env in envs.items():
|
||
if run.session_key == env.session_key:
|
||
completion_order.append(label)
|
||
break
|
||
|
||
runtime._turn_handler = instrumented_turn # type: ignore[method-assign]
|
||
|
||
handles = []
|
||
for label in ("A", "B", "C"):
|
||
for i in range(9):
|
||
h = await runtime.enqueue(envs[label], f"msg {i} from {label}")
|
||
handles.append(h)
|
||
|
||
for h in handles:
|
||
try:
|
||
await runtime.wait(h.task_id, timeout=30.0)
|
||
except KeyError:
|
||
pass
|
||
|
||
counts = {label: completion_order.count(label) for label in ("A", "B", "C")}
|
||
total = len(completion_order)
|
||
assert total == 27, f"Expected 27 completions, got {total}: {counts}"
|
||
|
||
# Each session should complete all 9 tasks; check no session was starved
|
||
# relative to the others (tolerance: each session completes between 6 and 12
|
||
# of the first 27, i.e. 9 ± 20% of 27/3 = 9).
|
||
tolerance = int(9 * 0.20) + 1 # floor(1.8)+1 = 2 extra buffer
|
||
for label, count in counts.items():
|
||
assert abs(count - 9) <= tolerance, (
|
||
f"Session {label} completed {count}/9 tasks — outside tolerance {tolerance}: {counts}"
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Cross-agent_id unaffected
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_cross_agent_unaffected() -> None:
|
||
"""Two different agent_ids each have 1 session.
|
||
|
||
Both should complete all tasks without interference; total completions
|
||
must equal the sum of what was submitted.
|
||
"""
|
||
random.seed(99)
|
||
runtime = _make_runtime(max_concurrency=2)
|
||
|
||
env_a = _make_envelope("agent-alpha", "agent-alpha::sess-1")
|
||
env_b = _make_envelope("agent-beta", "agent-beta::sess-1")
|
||
|
||
completed: list[str] = []
|
||
original_turn = runtime._turn_handler
|
||
|
||
async def instrumented_turn(run: Any) -> None:
|
||
await original_turn(run)
|
||
completed.append(run.session_key)
|
||
|
||
runtime._turn_handler = instrumented_turn # type: ignore[method-assign]
|
||
|
||
handles = []
|
||
for i in range(10):
|
||
handles.append(await runtime.enqueue(env_a, f"alpha {i}"))
|
||
handles.append(await runtime.enqueue(env_b, f"beta {i}"))
|
||
|
||
for h in handles:
|
||
try:
|
||
await runtime.wait(h.task_id, timeout=30.0)
|
||
except KeyError:
|
||
pass
|
||
|
||
alpha_count = completed.count("agent-alpha::sess-1")
|
||
beta_count = completed.count("agent-beta::sess-1")
|
||
assert alpha_count == 10, f"agent-alpha completed {alpha_count}/10"
|
||
assert beta_count == 10, f"agent-beta completed {beta_count}/10"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Terminal cleanup includes RR state
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_terminal_cleanup_includes_rr() -> None:
|
||
"""After all tasks complete, _agent_in_flight must be empty (no leaks)."""
|
||
random.seed(1)
|
||
runtime = _make_runtime(max_concurrency=2)
|
||
|
||
agent_id = "agent-cleanup"
|
||
env = _make_envelope(agent_id, f"{agent_id}::sess-1")
|
||
|
||
handles = []
|
||
for i in range(5):
|
||
h = await runtime.enqueue(env, f"msg {i}")
|
||
handles.append(h)
|
||
|
||
for h in handles:
|
||
try:
|
||
await runtime.wait(h.task_id, timeout=10.0)
|
||
except KeyError:
|
||
pass
|
||
|
||
# Give the event loop a tick to finish any post-terminal cleanup.
|
||
await asyncio.sleep(0)
|
||
|
||
assert runtime._agent_in_flight == {}, (
|
||
f"_agent_in_flight not cleaned up after run: {runtime._agent_in_flight}"
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Strict ABCABC round-robin order (max_concurrency=1 for determinism)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_round_robin_strict_abcabc_order() -> None:
|
||
"""With max_concurrency=1, 3 sessions A/B/C enqueued first-to-last must
|
||
execute in strict ABCABCABC... order (true RR, not count-minimum).
|
||
|
||
max_concurrency=1 ensures only one task runs at a time, making the RR
|
||
ordering deterministic and verifiable without timing sensitivity.
|
||
"""
|
||
random.seed(0)
|
||
# max_concurrency=1: strictly sequential; RR order must be exact.
|
||
runtime = _make_runtime(max_concurrency=1, max_pending_per_session=None)
|
||
|
||
agent_id = "agent-strict-rr"
|
||
# Enqueue sessions in A, B, C order so the deque starts [A, B, C].
|
||
envs = {
|
||
"A": _make_envelope(agent_id, f"{agent_id}::sess-a"),
|
||
"B": _make_envelope(agent_id, f"{agent_id}::sess-b"),
|
||
"C": _make_envelope(agent_id, f"{agent_id}::sess-c"),
|
||
}
|
||
|
||
completion_order: list[str] = []
|
||
original_turn = runtime._turn_handler
|
||
|
||
async def instrumented_turn(run: Any) -> None:
|
||
await original_turn(run)
|
||
for label, env in envs.items():
|
||
if run.session_key == env.session_key:
|
||
completion_order.append(label)
|
||
break
|
||
|
||
runtime._turn_handler = instrumented_turn # type: ignore[method-assign]
|
||
|
||
# Enqueue 3 tasks per session, one session at a time, in A→B→C order.
|
||
# This puts all A tasks in session A's pending list, etc.
|
||
tasks_per_session = 3
|
||
handles = []
|
||
for label in ("A", "B", "C"):
|
||
for i in range(tasks_per_session):
|
||
h = await runtime.enqueue(envs[label], f"msg {i} from {label}")
|
||
handles.append(h)
|
||
|
||
for h in handles:
|
||
try:
|
||
await runtime.wait(h.task_id, timeout=30.0)
|
||
except KeyError:
|
||
pass
|
||
|
||
total = len(completion_order)
|
||
assert total == tasks_per_session * 3, (
|
||
f"Expected {tasks_per_session * 3} completions, got {total}: {completion_order}"
|
||
)
|
||
|
||
# With max_concurrency=1 and true RR, the order must be A,B,C,A,B,C,A,B,C.
|
||
expected = ["A", "B", "C"] * tasks_per_session
|
||
assert completion_order == expected, (
|
||
f"Expected strict ABCABC order, got: {completion_order}"
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# No underutilization: 5 sessions, max_concurrency=4, slot always full
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_no_underutilization() -> None:
|
||
"""1 agent, 5 sessions, 10 messages each, max_concurrency=4.
|
||
|
||
This asserts the RR scheduler fills all available slots while work is
|
||
available without relying on wall-clock thresholds that include CI logging
|
||
and enqueue overhead.
|
||
"""
|
||
random.seed(5)
|
||
tasks_per_session = 10
|
||
num_sessions = 5
|
||
max_concurrency = 4
|
||
peak_in_flight = 0
|
||
current_in_flight = 0
|
||
peak_reached = asyncio.Event()
|
||
release_handlers = asyncio.Event()
|
||
counter_lock = asyncio.Lock()
|
||
|
||
agent_id = "agent-util"
|
||
|
||
async def timed_handler(_run: Any) -> None:
|
||
nonlocal current_in_flight, peak_in_flight
|
||
async with counter_lock:
|
||
current_in_flight += 1
|
||
peak_in_flight = max(peak_in_flight, current_in_flight)
|
||
if peak_in_flight >= max_concurrency:
|
||
peak_reached.set()
|
||
try:
|
||
await release_handlers.wait()
|
||
finally:
|
||
async with counter_lock:
|
||
current_in_flight -= 1
|
||
|
||
storage = _make_storage()
|
||
runtime = TaskRuntime(
|
||
storage=storage,
|
||
turn_handler=timed_handler,
|
||
max_concurrency=max_concurrency,
|
||
max_pending_per_session=None,
|
||
)
|
||
|
||
envs = [
|
||
_make_envelope(agent_id, f"{agent_id}::sess-{i}")
|
||
for i in range(num_sessions)
|
||
]
|
||
|
||
handles = []
|
||
for env in envs:
|
||
for i in range(tasks_per_session):
|
||
h = await runtime.enqueue(env, f"msg {i}")
|
||
handles.append(h)
|
||
|
||
await asyncio.wait_for(peak_reached.wait(), timeout=5.0)
|
||
release_handlers.set()
|
||
await asyncio.gather(*(runtime.wait(h.task_id, timeout=60.0) for h in handles))
|
||
|
||
assert peak_in_flight == max_concurrency
|