mirror of https://github.com/aliasrobotics/cai.git
512 lines
18 KiB
Python
512 lines
18 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Performance tests for TUI streaming.
|
|
|
|
This test suite ensures that streaming:
|
|
- Doesn't introduce significant latency
|
|
- Handles high-frequency updates efficiently
|
|
- Doesn't leak memory
|
|
- Scales with multiple concurrent streams
|
|
- Maintains responsiveness under load
|
|
"""
|
|
|
|
import asyncio
|
|
import gc
|
|
import os
|
|
import sys
|
|
import time
|
|
import unittest
|
|
import tracemalloc
|
|
from typing import Dict, List, Optional, Any
|
|
from unittest.mock import MagicMock, patch
|
|
import psutil
|
|
import threading
|
|
|
|
# Add src to path
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "src"))
|
|
|
|
from cai.tui.display.streaming_display import StreamingDisplay
|
|
from cai.tui.display.tool_display import ToolDisplay
|
|
from cai.tui.display.agent_display import AgentDisplay
|
|
from cai.tui.display.manager import DisplayManager
|
|
from cai.tui.display.base import DisplayContext
|
|
|
|
|
|
class PerformanceMetrics:
|
|
"""Track performance metrics."""
|
|
|
|
def __init__(self):
|
|
self.latencies = []
|
|
self.memory_snapshots = []
|
|
self.cpu_usage = []
|
|
self.update_counts = {}
|
|
self.start_time = time.time()
|
|
|
|
def record_latency(self, operation: str, duration: float):
|
|
"""Record operation latency."""
|
|
self.latencies.append({
|
|
'operation': operation,
|
|
'duration': duration,
|
|
'timestamp': time.time() - self.start_time
|
|
})
|
|
|
|
def record_memory(self):
|
|
"""Record current memory usage."""
|
|
process = psutil.Process()
|
|
self.memory_snapshots.append({
|
|
'rss': process.memory_info().rss,
|
|
'vms': process.memory_info().vms,
|
|
'timestamp': time.time() - self.start_time
|
|
})
|
|
|
|
def record_cpu(self):
|
|
"""Record CPU usage."""
|
|
process = psutil.Process()
|
|
self.cpu_usage.append({
|
|
'percent': process.cpu_percent(interval=0.1),
|
|
'timestamp': time.time() - self.start_time
|
|
})
|
|
|
|
def increment_update(self, stream_id: str):
|
|
"""Increment update count for a stream."""
|
|
self.update_counts[stream_id] = self.update_counts.get(stream_id, 0) + 1
|
|
|
|
def get_summary(self) -> Dict[str, Any]:
|
|
"""Get performance summary."""
|
|
if not self.latencies:
|
|
return {'error': 'No metrics recorded'}
|
|
|
|
latency_by_op = {}
|
|
for metric in self.latencies:
|
|
op = metric['operation']
|
|
if op not in latency_by_op:
|
|
latency_by_op[op] = []
|
|
latency_by_op[op].append(metric['duration'])
|
|
|
|
summary = {
|
|
'total_duration': time.time() - self.start_time,
|
|
'total_updates': sum(self.update_counts.values()),
|
|
'latency_stats': {}
|
|
}
|
|
|
|
for op, durations in latency_by_op.items():
|
|
summary['latency_stats'][op] = {
|
|
'count': len(durations),
|
|
'avg': sum(durations) / len(durations),
|
|
'min': min(durations),
|
|
'max': max(durations),
|
|
'p95': sorted(durations)[int(len(durations) * 0.95)] if len(durations) > 1 else durations[0]
|
|
}
|
|
|
|
if self.memory_snapshots:
|
|
initial_mem = self.memory_snapshots[0]['rss']
|
|
final_mem = self.memory_snapshots[-1]['rss']
|
|
summary['memory'] = {
|
|
'initial_mb': initial_mem / 1024 / 1024,
|
|
'final_mb': final_mem / 1024 / 1024,
|
|
'growth_mb': (final_mem - initial_mem) / 1024 / 1024
|
|
}
|
|
|
|
if self.cpu_usage:
|
|
cpu_values = [c['percent'] for c in self.cpu_usage]
|
|
summary['cpu'] = {
|
|
'avg_percent': sum(cpu_values) / len(cpu_values),
|
|
'max_percent': max(cpu_values)
|
|
}
|
|
|
|
return summary
|
|
|
|
|
|
class TestStreamingPerformance(unittest.TestCase):
|
|
"""Test streaming performance characteristics."""
|
|
|
|
def setUp(self):
|
|
"""Set up test environment."""
|
|
os.environ["CAI_TELEMETRY"] = "false"
|
|
os.environ["CAI_TRACING"] = "false"
|
|
os.environ["CAI_STREAM"] = "true"
|
|
|
|
self.metrics = PerformanceMetrics()
|
|
self.display_manager = DisplayManager()
|
|
|
|
# Mock terminal with performance tracking
|
|
self.mock_terminal = self._create_performance_terminal()
|
|
|
|
# Patch get_terminal_output
|
|
self.patcher = patch('cai.tui.core.terminal_console.get_terminal_output')
|
|
self.mock_get_terminal = self.patcher.start()
|
|
self.mock_get_terminal.return_value = self.mock_terminal
|
|
|
|
def tearDown(self):
|
|
"""Clean up."""
|
|
self.patcher.stop()
|
|
|
|
def _create_performance_terminal(self):
|
|
"""Create mock terminal that tracks performance."""
|
|
terminal = MagicMock()
|
|
|
|
def write_with_metrics(content):
|
|
start = time.time()
|
|
# Simulate write operation
|
|
time.sleep(0.0001) # 0.1ms write time
|
|
self.metrics.record_latency('terminal_write', time.time() - start)
|
|
|
|
def start_streaming_with_metrics(line_id, header):
|
|
start = time.time()
|
|
self.metrics.record_latency('stream_start', time.time() - start)
|
|
|
|
def update_streaming_with_metrics(line_id, content):
|
|
start = time.time()
|
|
self.metrics.increment_update(line_id)
|
|
# Simulate update processing
|
|
time.sleep(0.0001)
|
|
self.metrics.record_latency('stream_update', time.time() - start)
|
|
|
|
def finish_streaming_with_metrics(line_id, content, stats=None):
|
|
start = time.time()
|
|
self.metrics.record_latency('stream_finish', time.time() - start)
|
|
|
|
terminal.write = MagicMock(side_effect=write_with_metrics)
|
|
terminal.start_streaming_line = MagicMock(side_effect=start_streaming_with_metrics)
|
|
terminal.update_streaming_line = MagicMock(side_effect=update_streaming_with_metrics)
|
|
terminal.finish_streaming_line = MagicMock(side_effect=finish_streaming_with_metrics)
|
|
|
|
return terminal
|
|
|
|
def test_single_stream_latency(self):
|
|
"""Test latency of single stream operations."""
|
|
display = StreamingDisplay()
|
|
context = DisplayContext(
|
|
terminal_id="perf-1",
|
|
terminal_number=1,
|
|
agent_name="Latency Test",
|
|
interaction_counter=1
|
|
)
|
|
|
|
# Record initial memory
|
|
self.metrics.record_memory()
|
|
|
|
# Measure streaming operations
|
|
stream_id = "latency-stream"
|
|
|
|
# Start streaming
|
|
start = time.time()
|
|
display.start_streaming(context, stream_id, {
|
|
"content_type": "text",
|
|
"model": "gpt-4"
|
|
})
|
|
self.metrics.record_latency('display_start', time.time() - start)
|
|
|
|
# Multiple updates
|
|
content = "x" * 1000 # 1KB of content
|
|
for i in range(100): # 100 updates
|
|
start = time.time()
|
|
display.update_streaming(stream_id, {
|
|
"content": content[:i*10]
|
|
})
|
|
self.metrics.record_latency('display_update', time.time() - start)
|
|
|
|
# Finish streaming
|
|
start = time.time()
|
|
display.finish_streaming(stream_id, {
|
|
"final_stats": {"output_tokens": 100}
|
|
})
|
|
self.metrics.record_latency('display_finish', time.time() - start)
|
|
|
|
# Record final memory
|
|
self.metrics.record_memory()
|
|
|
|
# Analyze results
|
|
summary = self.metrics.get_summary()
|
|
|
|
# Assert performance requirements
|
|
self.assertLess(summary['latency_stats']['display_update']['avg'], 0.005) # < 5ms avg update
|
|
self.assertLess(summary['latency_stats']['display_update']['p95'], 0.010) # < 10ms p95
|
|
self.assertLess(summary['memory']['growth_mb'], 10) # < 10MB memory growth
|
|
|
|
async def test_concurrent_streams_performance(self):
|
|
"""Test performance with multiple concurrent streams."""
|
|
display = StreamingDisplay()
|
|
|
|
# Record initial state
|
|
self.metrics.record_memory()
|
|
self.metrics.record_cpu()
|
|
|
|
# Create multiple concurrent streams
|
|
num_streams = 10
|
|
contexts = []
|
|
for i in range(num_streams):
|
|
contexts.append(DisplayContext(
|
|
terminal_id=f"concurrent-{i}",
|
|
terminal_number=i+1,
|
|
agent_name=f"Agent {i}",
|
|
interaction_counter=1
|
|
))
|
|
|
|
async def stream_agent(context, stream_num):
|
|
"""Simulate streaming for one agent."""
|
|
stream_id = f"concurrent-stream-{stream_num}"
|
|
|
|
# Start
|
|
display.start_streaming(context, stream_id, {
|
|
"content_type": "text"
|
|
})
|
|
|
|
# Stream updates
|
|
for j in range(50): # 50 updates per stream
|
|
display.update_streaming(stream_id, {
|
|
"content": f"Update {j} from agent {stream_num}"
|
|
})
|
|
await asyncio.sleep(0.001) # 1ms between updates
|
|
|
|
# Finish
|
|
display.finish_streaming(stream_id, {})
|
|
|
|
# Run all streams concurrently
|
|
start_time = time.time()
|
|
tasks = []
|
|
for i, context in enumerate(contexts):
|
|
task = asyncio.create_task(stream_agent(context, i))
|
|
tasks.append(task)
|
|
|
|
await asyncio.gather(*tasks)
|
|
total_time = time.time() - start_time
|
|
|
|
# Record final state
|
|
self.metrics.record_memory()
|
|
self.metrics.record_cpu()
|
|
|
|
# Analyze performance
|
|
summary = self.metrics.get_summary()
|
|
total_updates = num_streams * 50
|
|
updates_per_second = total_updates / total_time
|
|
|
|
# Performance assertions
|
|
self.assertGreater(updates_per_second, 100) # > 100 updates/sec
|
|
self.assertLess(summary['memory']['growth_mb'], 50) # < 50MB for 10 streams
|
|
|
|
def test_rapid_update_handling(self):
|
|
"""Test handling of very rapid updates."""
|
|
display = StreamingDisplay()
|
|
context = DisplayContext(
|
|
terminal_id="rapid-1",
|
|
terminal_number=1,
|
|
agent_name="Rapid Test",
|
|
interaction_counter=1
|
|
)
|
|
|
|
stream_id = "rapid-stream"
|
|
display.start_streaming(context, stream_id, {"content_type": "text"})
|
|
|
|
# Blast updates as fast as possible
|
|
start_time = time.time()
|
|
update_count = 1000
|
|
|
|
for i in range(update_count):
|
|
display.update_streaming(stream_id, {
|
|
"content": f"Rapid update {i}"
|
|
})
|
|
# No sleep - maximum speed
|
|
|
|
display.finish_streaming(stream_id, {})
|
|
|
|
elapsed = time.time() - start_time
|
|
updates_per_second = update_count / elapsed
|
|
|
|
# Should handle at least 1000 updates/sec
|
|
self.assertGreater(updates_per_second, 1000)
|
|
|
|
# Check that terminal wasn't overwhelmed
|
|
actual_terminal_updates = self.metrics.update_counts.get(stream_id, 0)
|
|
# Should batch updates (not every update hits terminal)
|
|
self.assertLess(actual_terminal_updates, update_count)
|
|
|
|
def test_memory_leak_prevention(self):
|
|
"""Test that streaming doesn't leak memory."""
|
|
display = StreamingDisplay()
|
|
|
|
# Enable memory tracking
|
|
tracemalloc.start()
|
|
|
|
# Take initial snapshot
|
|
snapshot1 = tracemalloc.take_snapshot()
|
|
|
|
# Create and destroy many streams
|
|
for cycle in range(10):
|
|
contexts = []
|
|
for i in range(10): # 10 streams per cycle
|
|
context = DisplayContext(
|
|
terminal_id=f"leak-{cycle}-{i}",
|
|
terminal_number=1,
|
|
agent_name=f"Leak Test {cycle}-{i}",
|
|
interaction_counter=1
|
|
)
|
|
contexts.append(context)
|
|
|
|
stream_id = f"leak-stream-{cycle}-{i}"
|
|
|
|
# Full streaming lifecycle
|
|
display.start_streaming(context, stream_id, {"content_type": "text"})
|
|
|
|
# Multiple updates
|
|
for j in range(100):
|
|
display.update_streaming(stream_id, {
|
|
"content": f"Memory test update {j}"
|
|
})
|
|
|
|
display.finish_streaming(stream_id, {})
|
|
|
|
# Force garbage collection
|
|
gc.collect()
|
|
|
|
# Take final snapshot
|
|
snapshot2 = tracemalloc.take_snapshot()
|
|
|
|
# Compare snapshots
|
|
top_stats = snapshot2.compare_to(snapshot1, 'lineno')
|
|
|
|
# Calculate total growth
|
|
total_growth = sum(stat.size_diff for stat in top_stats if stat.size_diff > 0)
|
|
|
|
# Should not grow more than 10MB after 100 streams
|
|
self.assertLess(total_growth / 1024 / 1024, 10)
|
|
|
|
tracemalloc.stop()
|
|
|
|
def test_display_responsiveness_under_load(self):
|
|
"""Test that display remains responsive under heavy load."""
|
|
display = StreamingDisplay()
|
|
|
|
# Create background load
|
|
load_streams = []
|
|
for i in range(5):
|
|
context = DisplayContext(
|
|
terminal_id=f"load-{i}",
|
|
terminal_number=i+1,
|
|
agent_name=f"Load Agent {i}",
|
|
interaction_counter=1
|
|
)
|
|
stream_id = f"load-stream-{i}"
|
|
display.start_streaming(context, stream_id, {"content_type": "text"})
|
|
load_streams.append((stream_id, context))
|
|
|
|
# Start background updates
|
|
stop_load = threading.Event()
|
|
|
|
def generate_load():
|
|
"""Generate continuous updates."""
|
|
counter = 0
|
|
while not stop_load.is_set():
|
|
for stream_id, _ in load_streams:
|
|
display.update_streaming(stream_id, {
|
|
"content": f"Background update {counter}"
|
|
})
|
|
counter += 1
|
|
time.sleep(0.001)
|
|
|
|
load_thread = threading.Thread(target=generate_load)
|
|
load_thread.start()
|
|
|
|
# Test responsiveness of new stream
|
|
test_context = DisplayContext(
|
|
terminal_id="responsive-test",
|
|
terminal_number=10,
|
|
agent_name="Responsiveness Test",
|
|
interaction_counter=1
|
|
)
|
|
|
|
test_stream_id = "responsive-stream"
|
|
|
|
# Measure latency of operations under load
|
|
latencies = []
|
|
|
|
# Start
|
|
start = time.time()
|
|
display.start_streaming(test_context, test_stream_id, {"content_type": "text"})
|
|
latencies.append(('start', time.time() - start))
|
|
|
|
# Updates
|
|
for i in range(10):
|
|
start = time.time()
|
|
display.update_streaming(test_stream_id, {
|
|
"content": f"Responsive test {i}"
|
|
})
|
|
latencies.append(('update', time.time() - start))
|
|
time.sleep(0.01)
|
|
|
|
# Finish
|
|
start = time.time()
|
|
display.finish_streaming(test_stream_id, {})
|
|
latencies.append(('finish', time.time() - start))
|
|
|
|
# Stop background load
|
|
stop_load.set()
|
|
load_thread.join()
|
|
|
|
# Clean up load streams
|
|
for stream_id, _ in load_streams:
|
|
display.finish_streaming(stream_id, {})
|
|
|
|
# Check responsiveness
|
|
max_latency = max(lat[1] for lat in latencies)
|
|
avg_latency = sum(lat[1] for lat in latencies) / len(latencies)
|
|
|
|
# Even under load, operations should be fast
|
|
self.assertLess(max_latency, 0.050) # < 50ms max
|
|
self.assertLess(avg_latency, 0.010) # < 10ms average
|
|
|
|
def test_large_content_streaming(self):
|
|
"""Test streaming performance with large content."""
|
|
display = StreamingDisplay()
|
|
context = DisplayContext(
|
|
terminal_id="large-1",
|
|
terminal_number=1,
|
|
agent_name="Large Content",
|
|
interaction_counter=1
|
|
)
|
|
|
|
# Generate large content (10KB)
|
|
large_content = "x" * 10000
|
|
|
|
stream_id = "large-stream"
|
|
display.start_streaming(context, stream_id, {"content_type": "text"})
|
|
|
|
# Stream in chunks
|
|
chunk_size = 100
|
|
start_time = time.time()
|
|
|
|
for i in range(0, len(large_content), chunk_size):
|
|
display.update_streaming(stream_id, {
|
|
"content": large_content[:i+chunk_size]
|
|
})
|
|
|
|
display.finish_streaming(stream_id, {})
|
|
|
|
elapsed = time.time() - start_time
|
|
throughput_mbps = (len(large_content) / 1024 / 1024) / elapsed
|
|
|
|
# Should handle at least 1MB/s throughput
|
|
self.assertGreater(throughput_mbps, 1.0)
|
|
|
|
|
|
def run_async_test(coro):
|
|
"""Helper to run async test."""
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
return loop.run_until_complete(coro)
|
|
finally:
|
|
loop.close()
|
|
|
|
|
|
class TestStreamingPerformanceSync(TestStreamingPerformance):
|
|
"""Synchronous wrapper for async tests."""
|
|
|
|
def test_concurrent_streams_performance_sync(self):
|
|
run_async_test(self.test_concurrent_streams_performance())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("🚀 Running TUI Streaming Performance Tests")
|
|
print("=" * 60)
|
|
unittest.main(verbosity=2) |