""" Data aggregator for report generation. Aggregates data from test_tasks and test_results tables into report structures. """ import statistics from typing import Dict, List, Optional from src.clients.tiworkload_client import get_task_metrics from src.database.db import get_db from src.models.models import TestResult, TestTask from .metrics_parser import ( compute_spu_identifier, extract_workload_params, extract_workload_type_from_template, parse_performance_metrics, ) from .models import ( ClusterConfig, ComparisonData, MetricStats, SPUComparisonReport, SPUReportData, TaskComparisonRun, ThreadComparisonData, WorkflowData, WorkloadRunData, ) def aggregate_spu_data(spu_identifier: str) -> Optional[SPUReportData]: """Aggregate all test data for a specific SPU configuration. Args: spu_identifier: SPU identifier string Returns: SPUReportData or None if no data found """ with get_db() as db: # Query all tasks matching this SPU tasks = query_tasks_by_spu_components(db, spu_identifier) if not tasks: return None # Extract cluster config from first task (they should all be the same) cluster_config = extract_cluster_config_from_task(tasks[0]) # Group tasks by workflow type workflow_groups = {} for task in tasks: wf_type = task.workload_template if wf_type not in workflow_groups: workflow_groups[wf_type] = [] workflow_groups[wf_type].append(task) # Process each workflow type workflows = [] for wf_type, task_list in workflow_groups.items(): for task in task_list: # Get result for this task result = ( db.query(TestResult) .filter_by(task_id=task.id) .filter(TestResult.deleted == False) .order_by(TestResult.id.desc()) .first() ) if not result or not result.result_data: continue # Extract runs from result runs = extract_runs_from_result(result, task) # Create workflow data workflow = WorkflowData( workflow_type=wf_type, workflow_execution_id=task.workflow_execution_id or "", task_id=task.id, status=task.status, runs=runs, ) workflows.append(workflow) return SPUReportData( spu_identifier=spu_identifier, cluster_config=cluster_config, workflows=workflows, ) def query_tasks_by_spu_components(db, spu_identifier: str) -> List[TestTask]: """Query tasks matching SPU identifier components. Since we don't have spu_identifier stored in the DB, we need to compute it from task fields and filter. """ # Get all non-deleted tasks all_tasks = db.query(TestTask).filter(TestTask.deleted == False).all() # Filter by computed SPU identifier matching_tasks = [] for task in all_tasks: task_spu_id = compute_spu_identifier( tidb_spu=task.tidb_spu, tidb_count=task.tidb_count, tikv_spu=task.tikv_spu, tikv_count=task.tikv_count, cloud_provider=task.cloud_provider, region=task.region, cluster_version=task.cluster_version, workload_config=task.workload_config, ) if task_spu_id == spu_identifier: matching_tasks.append(task) return matching_tasks def extract_cluster_config_from_task(task: TestTask) -> ClusterConfig: """Extract cluster configuration from a task.""" return ClusterConfig( tidb_spu=task.tidb_spu, tidb_count=task.tidb_count, tidb_config=task.tidb_config, tikv_spu=task.tikv_spu, tikv_count=task.tikv_count, tikv_storage_type=task.tikv_storage_type, tikv_storage_size_gi=task.tikv_storage_size_gi, tikv_config=task.tikv_config, pd_spu=task.pd_spu, pd_count=task.pd_count, pd_config=task.pd_config, tiflash_spu=task.tiflash_spu, tiflash_count=task.tiflash_count, tiflash_storage_type=task.tiflash_storage_type, tiflash_storage_size_gi=task.tiflash_storage_size_gi, tiflash_config=task.tiflash_config, tiproxy_spu=task.tiproxy_spu, tiproxy_count=task.tiproxy_count, tiproxy_config=task.tiproxy_config, cloud_provider=task.cloud_provider, region=task.region, cluster_version=task.cluster_version, workload_config=task.workload_config, ) def extract_runs_from_result( result: TestResult, task: TestTask ) -> List[WorkloadRunData]: """Extract workload runs from a test result.""" runs = [] if not result.result_data or "nodes" not in result.result_data: return runs nodes = result.result_data["nodes"] workload_config = task.workload_config for node in nodes: # Extract parameters params = extract_workload_params(node, workload_config) # Parse performance metrics workload_type = extract_workload_type_from_template( task.workload_template, node.get("node_name", "") ) metrics = parse_performance_metrics( summary=node.get("summary", ""), workload_type=task.workload_template, node_name=node.get("node_name", ""), ) # Fetch client metrics (CPU, Memory, Network) client_metrics = None # Extract tiworkload execution ID and task ID from node metadata # metadata.tiworkload_execution_id is the correct ID (not TCMS workflow_execution_id) metadata = node.get("metadata", {}) tiworkload_exec_id = metadata.get("tiworkload_execution_id") workload_tasks = metadata.get("workload_tasks", []) tiworkload_task_id = workload_tasks[0].get("id") if workload_tasks else None if tiworkload_exec_id and tiworkload_task_id: try: client_metrics = get_task_metrics( exec_id=str(tiworkload_exec_id), task_id=tiworkload_task_id ) except Exception as e: # Log error but continue processing print( f"Warning: Failed to fetch client metrics for task {tiworkload_task_id}: {e}" ) # Create run data run = WorkloadRunData( node_name=node.get("node_name", ""), workload_type=workload_type, params=params, metrics=metrics, status=node.get("status", "UNKNOWN"), task_id=task.id, workflow_execution_id=task.workflow_execution_id, started_at=None, # Not available in current schema completed_at=None, # Not available in current schema client_metrics=client_metrics, ) runs.append(run) return runs def compare_spu_data( spu_id_1: str, spu_id_2: str, workflow_types: Optional[List[str]] = None ) -> Optional[SPUComparisonReport]: """Compare two SPU configurations. Args: spu_id_1: First SPU identifier spu_id_2: Second SPU identifier workflow_types: List of workflow types to compare (None = all) Returns: SPUComparisonReport or None if data not found """ # Get data for both SPUs spu_1_data = aggregate_spu_data(spu_id_1) spu_2_data = aggregate_spu_data(spu_id_2) if not spu_1_data or not spu_2_data: return None # Determine which workflows to compare if workflow_types is None: # Compare all common workflows wf_types_1 = {wf.workflow_type for wf in spu_1_data.workflows} wf_types_2 = {wf.workflow_type for wf in spu_2_data.workflows} workflow_types = list(wf_types_1.intersection(wf_types_2)) # Generate comparisons comparisons = [] for wf_type in workflow_types: comparison = compare_workflow_data(spu_1_data, spu_2_data, wf_type) if comparison: comparisons.append(comparison) return SPUComparisonReport( spu_1=spu_1_data, spu_2=spu_2_data, comparisons=comparisons ) def calculate_metric_stats(values: List[float]) -> MetricStats: """Calculate statistical summary for a list of values. Args: values: List of metric values Returns: MetricStats with avg, median, std, min, max """ if not values: return MetricStats() return MetricStats( avg=statistics.mean(values), median=statistics.median(values), std=statistics.stdev(values) if len(values) > 1 else 0.0, min=min(values), max=max(values), ) def extract_metric_value(run: WorkloadRunData, metric_path: str) -> Optional[float]: """Extract a metric value from a run using a path like 'tps' or 'latency.p95'. Args: run: WorkloadRunData containing metrics metric_path: Path to the metric (e.g., 'tps', 'latency.p95', 'tpmC') Returns: Metric value or None if not found """ parts = metric_path.split(".") value = run.metrics for part in parts: if isinstance(value, dict): value = value.get(part) else: return None if isinstance(value, (int, float)): return float(value) return None def get_all_metric_names(runs: List[WorkloadRunData]) -> List[str]: """Get all unique metric names from a list of runs. Args: runs: List of workload runs Returns: List of metric paths (e.g., ['tps', 'tpmC', 'qps', 'latency.avg', 'latency.p95', 'latency.p99']) """ metric_names = set() for run in runs: # Add top-level metrics for key in run.metrics: if isinstance(run.metrics[key], (int, float)): metric_names.add(key) elif key == "latency" and isinstance(run.metrics[key], dict): # Add latency sub-metrics for latency_key in run.metrics[key]: if isinstance(run.metrics[key][latency_key], (int, float)): metric_names.add(f"latency.{latency_key}") return sorted(metric_names) def compare_workflow_data( spu_1_data: SPUReportData, spu_2_data: SPUReportData, workflow_type: str ) -> Optional[ComparisonData]: """Compare workflow results between two SPUs. This function creates a comprehensive comparison including: - All runs from both configurations - Summary statistics for all metrics - Improvement percentages Args: spu_1_data: First SPU data spu_2_data: Second SPU data workflow_type: Workflow type to compare Returns: ComparisonData or None if not comparable """ # Get runs for this workflow type from both SPUs runs_1 = spu_1_data.get_all_runs(workflow_type) runs_2 = spu_2_data.get_all_runs(workflow_type) if not runs_1 or not runs_2: return None # Group runs by thread count (keep all runs, don't overwrite!) runs_1_by_thread: Dict[int, List[WorkloadRunData]] = {} for run in runs_1: thread_count = run.params.get("thread") if thread_count: if thread_count not in runs_1_by_thread: runs_1_by_thread[thread_count] = [] runs_1_by_thread[thread_count].append(run) runs_2_by_thread: Dict[int, List[WorkloadRunData]] = {} for run in runs_2: thread_count = run.params.get("thread") if thread_count: if thread_count not in runs_2_by_thread: runs_2_by_thread[thread_count] = [] runs_2_by_thread[thread_count].append(run) # Get all thread counts from both configs (not just common ones) all_threads = set(runs_1_by_thread.keys()).union(set(runs_2_by_thread.keys())) if not all_threads: return None # Create comparison comparison = ComparisonData( workflow_type=workflow_type, spu_1_identifier=spu_1_data.spu_identifier, spu_2_identifier=spu_2_data.spu_identifier, ) # Get all metrics from both sets of runs all_runs = runs_1 + runs_2 all_metrics = get_all_metric_names(all_runs) # Process each thread count # Track improvements by metric for calculating per-metric averages metric_improvements_by_thread: Dict[str, List[float]] = {} for thread_count in sorted(all_threads): thread_comp = ThreadComparisonData(thread_count=thread_count) # Get runs for this thread count config1_runs = runs_1_by_thread.get(thread_count, []) config2_runs = runs_2_by_thread.get(thread_count, []) # Store runs with task IDs for run in config1_runs: thread_comp.config1_runs.append( TaskComparisonRun(task_id=run.task_id or 0, run=run) ) for run in config2_runs: thread_comp.config2_runs.append( TaskComparisonRun(task_id=run.task_id or 0, run=run) ) # Calculate summary statistics for each metric for metric_name in all_metrics: # Extract values for config 1 config1_values = [] for run in config1_runs: value = extract_metric_value(run, metric_name) if value is not None: config1_values.append(value) # Extract values for config 2 config2_values = [] for run in config2_runs: value = extract_metric_value(run, metric_name) if value is not None: config2_values.append(value) # Calculate stats if we have values if config1_values or config2_values: stats1 = calculate_metric_stats(config1_values) stats2 = calculate_metric_stats(config2_values) # Calculate improvement (higher is better for throughput, lower is better for latency) improvement = 0.0 if stats1.avg > 0 and stats2.avg > 0: if "latency" in metric_name.lower(): # For latency, lower is better improvement = ( -((stats2.avg - stats1.avg) / (stats1.avg + 1e-6)) * 100 ) else: # For throughput, higher is better improvement = ( (stats2.avg - stats1.avg) / (stats1.avg + 1e-6) ) * 100 # Collect improvement for this metric if metric_name not in metric_improvements_by_thread: metric_improvements_by_thread[metric_name] = [] metric_improvements_by_thread[metric_name].append(improvement) thread_comp.summary_stats[metric_name] = (stats1, stats2, improvement) # Store thread comparison data comparison.thread_data[thread_count] = thread_comp # Also populate legacy fields for backward compatibility (TPS and P95 only) if config1_runs and config2_runs: # TPS comparison (use tps or tpmC) tps_metric = "tps" if "tps" in all_metrics else "tpmC" if tps_metric in thread_comp.summary_stats: stats1, stats2, improvement = thread_comp.summary_stats[tps_metric] comparison.tps_comparison[thread_count] = ( stats1.avg, stats2.avg, improvement, ) # P95 comparison if "latency.p95" in thread_comp.summary_stats: stats1, stats2, improvement = thread_comp.summary_stats["latency.p95"] comparison.p95_comparison[thread_count] = ( stats1.avg, stats2.avg, improvement, ) # Calculate average improvement per metric (across all thread counts) for metric_name, improvements in metric_improvements_by_thread.items(): if improvements: comparison.metric_improvements[metric_name] = statistics.mean(improvements) # Calculate overall average improvement (for backward compatibility) all_improvements = [] for improvements in metric_improvements_by_thread.values(): all_improvements.extend(improvements) if all_improvements: comparison.average_improvement = statistics.mean(all_improvements) return comparison def aggregate_tasks_by_ids(task_ids: List[int]) -> Optional[SPUReportData]: """Aggregate test data from specified task IDs into a single report. Validation rules: - TiDB SPU must match across all tasks (raises ValueError if different) - TiDB Count must match across all tasks (raises ValueError if diff) - Cluster Version must match across all tasks (raises ValueError if diff) - TiProxy status must match across all tasks (raises ValueError if diff) - Different workload types are allowed and will be aggregated together Args: task_ids: List of task IDs to aggregate Returns: SPUReportData or None if no tasks found Raises: ValueError: If configuration validation fails """ with get_db() as db: # Query all tasks by IDs tasks = ( db.query(TestTask) .filter(TestTask.id.in_(task_ids), TestTask.deleted.is_not(True)) .all() ) if not tasks: return None # Validate configuration consistency base_task = tasks[0] base_tiproxy_enabled = ( base_task.tiproxy_count is not None and base_task.tiproxy_count > 0 ) for task in tasks[1:]: # Check TiDB SPU if task.tidb_spu != base_task.tidb_spu: raise ValueError( f"TiDB SPU mismatch: task {base_task.id} has " f"'{base_task.tidb_spu}', but task {task.id} has " f"'{task.tidb_spu}'" ) # Check TiDB Count if task.tidb_count != base_task.tidb_count: raise ValueError( f"TiDB Count mismatch: task {base_task.id} has " f"{base_task.tidb_count}, but task {task.id} has " f"{task.tidb_count}" ) # Check Cluster Version if task.cluster_version != base_task.cluster_version: raise ValueError( f"Cluster Version mismatch: task {base_task.id} has " f"'{base_task.cluster_version}', but task {task.id} " f"has '{task.cluster_version}'" ) # Check TiProxy status task_tiproxy_enabled = ( task.tiproxy_count is not None and task.tiproxy_count > 0 ) if task_tiproxy_enabled != base_tiproxy_enabled: base_status = "enabled" if base_tiproxy_enabled else "disabled" task_status = "enabled" if task_tiproxy_enabled else "disabled" raise ValueError( f"TiProxy status mismatch: task {base_task.id} has " f"TiProxy {base_status}, but task {task.id} has " f"TiProxy {task_status}" ) # Extract cluster config from first task cluster_config = extract_cluster_config_from_task(base_task) # Compute SPU identifier spu_identifier = compute_spu_identifier( tidb_spu=base_task.tidb_spu, tidb_count=base_task.tidb_count, tikv_spu=base_task.tikv_spu, tikv_count=base_task.tikv_count, cloud_provider=base_task.cloud_provider, region=base_task.region, cluster_version=base_task.cluster_version, workload_config=base_task.workload_config, ) # Group tasks by workflow type and aggregate workflow_groups = {} for task in tasks: wf_type = task.workload_template if wf_type not in workflow_groups: workflow_groups[wf_type] = [] workflow_groups[wf_type].append(task) # Process each workflow type workflows = [] for wf_type, task_list in workflow_groups.items(): for task in task_list: # Get result for this task result = ( db.query(TestResult) .filter_by(task_id=task.id) .filter(TestResult.deleted.is_not(True)) .order_by(TestResult.id.desc()) .first() ) if not result or not result.result_data: continue # Extract runs from result runs = extract_runs_from_result(result, task) # Create workflow data workflow = WorkflowData( workflow_type=wf_type, workflow_execution_id=task.workflow_execution_id or "", task_id=task.id, status=task.status, runs=runs, ) workflows.append(workflow) return SPUReportData( spu_identifier=spu_identifier, cluster_config=cluster_config, workflows=workflows, ) def list_available_spus() -> List[Dict]: """List all available SPU configurations in the database. Returns: List of dictionaries with SPU information """ with get_db() as db: # Get all unique SPU configurations all_tasks = db.query(TestTask).filter(TestTask.deleted == False).all() spu_map = {} spu_test_count = {} for task in all_tasks: spu_id = compute_spu_identifier( tidb_spu=task.tidb_spu, tidb_count=task.tidb_count, tikv_spu=task.tikv_spu, tikv_count=task.tikv_count, cloud_provider=task.cloud_provider, region=task.region, cluster_version=task.cluster_version, workload_config=task.workload_config, ) if spu_id not in spu_map: spu_map[spu_id] = { "spu_identifier": spu_id, "tidb_spu": task.tidb_spu, "tidb_count": task.tidb_count, "tikv_spu": task.tikv_spu, "tikv_count": task.tikv_count, "cloud_provider": task.cloud_provider, "region": task.region, "cluster_version": task.cluster_version, } spu_test_count[spu_id] = 0 spu_test_count[spu_id] += 1 # Add test count to each SPU result = [] for spu_id, spu_info in spu_map.items(): spu_info["test_count"] = spu_test_count[spu_id] result.append(spu_info) return result