diff --git a/heatmap_visualization/src/heatmap_visualization/cli.py b/heatmap_visualization/src/heatmap_visualization/cli.py new file mode 100644 index 0000000..e2df802 --- /dev/null +++ b/heatmap_visualization/src/heatmap_visualization/cli.py @@ -0,0 +1,70 @@ +import argparse +import json +import logging +from pathlib import Path +from typing import Any + +import pandas as pd +from heatmap_visualization import core + +# Configure logging for CI and debugging +logging.basicConfig(level=logging.INFO, format='[%(levelname)s] %(message)s') +logger = logging.getLogger(__name__) + + +def _validate_input_data(data: Any) -> pd.DataFrame: + """Validate that the input data conforms to the expected LogData structure.""" + if isinstance(data, list): + try: + df = pd.DataFrame(data) + except Exception as e: + raise ValueError(f"Could not convert input list to DataFrame: {e}") + elif isinstance(data, pd.DataFrame): + df = data.copy() + else: + raise TypeError("Input data must be a list of dicts or a pandas.DataFrame.") + + required_fields = {"worker_start_offset", "expires_at_dist_hours", "retry_total_overhead_ms"} + missing = required_fields - set(df.columns) + if missing: + raise ValueError(f"Input data missing required fields: {', '.join(missing)}") + + # Validate types + for field in required_fields: + if not pd.api.types.is_numeric_dtype(df[field]): + raise ValueError(f"Field '{field}' must be numeric.") + return df + + +def main() -> None: + """Command-line entry point for generating the Run 26 heatmap visualization.""" + parser = argparse.ArgumentParser(description="Generate a latency heatmap from run log data.") + parser.add_argument("--input", required=True, help="Path to input log JSON file.") + parser.add_argument("--output", required=True, help="Path to output PNG file.") + args = parser.parse_args() + + input_path = Path(args.input) + output_path = Path(args.output) + + if not input_path.exists(): + logger.error(f"Input file not found: {input_path}") + raise FileNotFoundError(f"Input file not found: {input_path}") + + logger.info(f"Loading JSON data from {input_path}") + with input_path.open("r", encoding="utf-8") as f: + data = json.load(f) + + df = _validate_input_data(data) + + logger.info("Generating heatmap figure...") + fig = core.generate_heatmap(df) + + logger.info(f"Saving heatmap to {output_path}") + output_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(output_path, dpi=300, bbox_inches="tight") + + logger.info("Heatmap generation completed successfully.") + + +if __name__ == "__main__": + main() \ No newline at end of file