haoline 0.3.0__py3-none-any.whl

haoline/tests/test_hardware_recommender.py

@@ -0,0 +1,90 @@
+# Copyright (c) 2025 HaoLine Contributors
+# SPDX-License-Identifier: MIT
+
+import unittest
+
+from ..hardware import (
+    NVIDIA_JETSON_NANO,
+    NVIDIA_RTX_3060_12GB,
+    BatchSizeSweeper,
+    HardwareEstimator,
+    SystemRequirementsRecommender,
+)
+
+
+class TestSystemRequirements(unittest.TestCase):
+    def setUp(self):
+        self.estimator = HardwareEstimator()
+        self.recommender = SystemRequirementsRecommender(self.estimator)
+
+    def test_recommend_tiny_model(self):
+        # Tiny model (10MB weights) should fit on Nano
+        reqs = self.recommender.recommend(
+            model_params=2_500_000,  # ~10MB FP32
+            model_flops=1_000_000_000,  # 1 GFLOP
+            peak_activation_bytes=10_000_000,  # 10MB
+            target_batch_size=1,
+        )
+
+        self.assertEqual(reqs.minimum_gpu.name, NVIDIA_JETSON_NANO.name)
+        self.assertLess(reqs.minimum_vram_gb, 4.0)
+
+    def test_recommend_large_model(self):
+        # Large model (10GB weights) needs bigger GPU
+        reqs = self.recommender.recommend(
+            model_params=5_000_000_000,  # ~10GB FP16
+            model_flops=10_000_000_000_000,  # 10 TFLOPs
+            peak_activation_bytes=1_000_000_000,  # 1GB
+            target_batch_size=1,
+            precision="fp16",
+        )
+
+        # Should definitely NOT be Nano
+        self.assertNotEqual(reqs.minimum_gpu.name, NVIDIA_JETSON_NANO.name)
+        # Should be at least a 12GB card
+        self.assertGreaterEqual(reqs.minimum_gpu.vram_bytes, 10 * 1024**3)
+
+
+class TestBatchSizeSweep(unittest.TestCase):
+    def setUp(self):
+        self.estimator = HardwareEstimator()
+        self.sweeper = BatchSizeSweeper(self.estimator)
+        self.gpu = NVIDIA_RTX_3060_12GB
+
+    def test_sweep_basic(self):
+        # Medium model
+        sweep = self.sweeper.sweep(
+            model_params=100_000_000,  # 400MB
+            model_flops=10_000_000_000,  # 10 GFLOPs
+            peak_activation_bytes=50_000_000,  # 50MB
+            hardware=self.gpu,
+            max_batch_size=8,
+        )
+
+        self.assertEqual(len(sweep.batch_sizes), 4)  # 1, 2, 4, 8
+        self.assertEqual(sweep.batch_sizes[-1], 8)
+        self.assertTrue(all(lat > 0 for lat in sweep.latencies))
+
+        # Throughput should generally increase (or saturate)
+        self.assertGreater(sweep.throughputs[-1], sweep.throughputs[0])
+
+    def test_sweep_oom(self):
+        # Huge activation model that OOMs quickly
+        sweep = self.sweeper.sweep(
+            model_params=100_000_000,
+            model_flops=10_000_000_000,
+            peak_activation_bytes=4 * 1024**3,  # 4GB activations per sample
+            hardware=self.gpu,  # 12GB VRAM
+            max_batch_size=8,
+        )
+
+        # Should fit batch 1 (4GB < 12GB)
+        # batch 2 (8GB < 12GB)
+        # batch 4 (16GB > 12GB) -> OOM
+        # So we expect 2 results
+        self.assertLess(len(sweep.batch_sizes), 4)
+        self.assertTrue(len(sweep.batch_sizes) >= 1)
+
+
+if __name__ == "__main__":
+    unittest.main()

haoline/tests/test_hierarchical_graph.py

@@ -0,0 +1,326 @@
+# Copyright (c) 2025 HaoLine Contributors
+# SPDX-License-Identifier: MIT
+
+"""Tests for hierarchical graph view."""
+
+from __future__ import annotations
+
+import json
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import onnx
+import pytest
+from onnx import TensorProto, helper
+
+from ..analyzer import ONNXGraphLoader
+from ..hierarchical_graph import (
+    HierarchicalGraph,
+    HierarchicalGraphBuilder,
+    HierarchicalNode,
+    generate_summary,
+)
+from ..patterns import PatternAnalyzer
+
+
+def create_test_model() -> onnx.ModelProto:
+    """Create a simple model for hierarchical testing."""
+    X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 64, 56, 56])
+    Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 64, 56, 56])
+
+    weight1 = helper.make_tensor(
+        "w1",
+        TensorProto.FLOAT,
+        [64, 64, 3, 3],
+        np.random.randn(64, 64, 3, 3).astype(np.float32).flatten().tolist(),
+    )
+    weight2 = helper.make_tensor(
+        "w2",
+        TensorProto.FLOAT,
+        [64, 64, 3, 3],
+        np.random.randn(64, 64, 3, 3).astype(np.float32).flatten().tolist(),
+    )
+
+    nodes = [
+        helper.make_node("Conv", ["X", "w1"], ["c1"], kernel_shape=[3, 3], pads=[1, 1, 1, 1]),
+        helper.make_node("Relu", ["c1"], ["r1"]),
+        helper.make_node("Conv", ["r1", "w2"], ["c2"], kernel_shape=[3, 3], pads=[1, 1, 1, 1]),
+        helper.make_node("Add", ["X", "c2"], ["add_out"]),
+        helper.make_node("Relu", ["add_out"], ["Y"]),
+    ]
+
+    graph = helper.make_graph(nodes, "test_model", [X], [Y], [weight1, weight2])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
+    return model
+
+
+class TestHierarchicalNode:
+    """Tests for HierarchicalNode."""
+
+    def test_basic_creation(self):
+        """Test basic node creation."""
+        node = HierarchicalNode(
+            id="test_1",
+            name="TestNode",
+            node_type="op",
+            op_type="Conv",
+        )
+        assert node.id == "test_1"
+        assert node.name == "TestNode"
+        assert node.is_leaf()
+        assert node.is_collapsed
+
+    def test_display_name_with_repeat(self):
+        """Test xN notation for repeated blocks."""
+        node = HierarchicalNode(
+            id="block_1",
+            name="TransformerLayer",
+            node_type="block",
+            repeat_count=12,
+            is_repeated=True,
+        )
+        assert node.get_display_name() == "TransformerLayer x12"
+
+    def test_collapse_expand(self):
+        """Test collapse/expand state management."""
+        parent = HierarchicalNode(
+            id="parent",
+            name="Parent",
+            node_type="block",
+            children=[
+                HierarchicalNode(id="child1", name="Child1", node_type="op"),
+                HierarchicalNode(id="child2", name="Child2", node_type="op"),
+            ],
+        )
+
+        assert parent.is_collapsed
+        parent.expand()
+        assert not parent.is_collapsed
+        parent.toggle()
+        assert parent.is_collapsed
+
+    def test_aggregate_stats(self):
+        """Test stats aggregation from children."""
+        child1 = HierarchicalNode(
+            id="c1",
+            name="C1",
+            node_type="op",
+            total_flops=1000,
+            total_params=100,
+            total_memory_bytes=400,
+        )
+        child2 = HierarchicalNode(
+            id="c2",
+            name="C2",
+            node_type="op",
+            total_flops=2000,
+            total_params=200,
+            total_memory_bytes=800,
+        )
+        parent = HierarchicalNode(
+            id="parent",
+            name="Parent",
+            node_type="block",
+            children=[child1, child2],
+        )
+
+        parent.aggregate_stats()
+
+        assert parent.total_flops == 3000
+        assert parent.total_params == 300
+        assert parent.total_memory_bytes == 1200
+        assert parent.node_count == 2
+
+    def test_aggregate_with_repeat(self):
+        """Test stats aggregation with repeat multiplier."""
+        child = HierarchicalNode(
+            id="c1",
+            name="C1",
+            node_type="op",
+            total_flops=1000,
+            total_params=100,
+            total_memory_bytes=400,
+        )
+        parent = HierarchicalNode(
+            id="parent",
+            name="Parent",
+            node_type="block",
+            children=[child],
+            repeat_count=12,
+        )
+
+        parent.aggregate_stats()
+
+        assert parent.total_flops == 12000  # Multiplied
+        assert parent.total_params == 100  # Shared, not multiplied
+        assert parent.total_memory_bytes == 4800  # Multiplied
+
+    def test_to_dict(self):
+        """Test JSON export."""
+        node = HierarchicalNode(
+            id="test",
+            name="TestNode",
+            node_type="op",
+            op_type="Conv",
+        )
+        d = node.to_dict()
+
+        assert d["id"] == "test"
+        assert d["name"] == "TestNode"
+        assert d["op_type"] == "Conv"
+        assert "children" not in d  # No children for leaf
+
+
+class TestHierarchicalGraph:
+    """Tests for HierarchicalGraph."""
+
+    def test_graph_creation(self):
+        """Test creating a hierarchical graph."""
+        root = HierarchicalNode(id="root", name="Model", node_type="model")
+        graph = HierarchicalGraph(root=root, nodes_by_id={"root": root})
+
+        assert graph.root.name == "Model"
+        assert graph.get_node("root") == root
+
+    def test_visible_nodes_collapsed(self):
+        """Test visible nodes when parent is collapsed."""
+        child = HierarchicalNode(id="child", name="Child", node_type="op")
+        root = HierarchicalNode(
+            id="root",
+            name="Root",
+            node_type="model",
+            children=[child],
+            is_collapsed=True,
+        )
+        graph = HierarchicalGraph(
+            root=root,
+            nodes_by_id={"root": root, "child": child},
+        )
+
+        visible = graph.get_visible_nodes()
+        assert len(visible) == 1
+        assert visible[0].id == "root"
+
+    def test_visible_nodes_expanded(self):
+        """Test visible nodes when parent is expanded."""
+        child = HierarchicalNode(id="child", name="Child", node_type="op")
+        root = HierarchicalNode(
+            id="root",
+            name="Root",
+            node_type="model",
+            children=[child],
+            is_collapsed=False,
+        )
+        graph = HierarchicalGraph(
+            root=root,
+            nodes_by_id={"root": root, "child": child},
+        )
+
+        visible = graph.get_visible_nodes()
+        assert len(visible) == 2
+
+    def test_expand_to_depth(self):
+        """Test expanding to a specific depth."""
+        grandchild = HierarchicalNode(id="gc", name="GC", node_type="op", depth=2)
+        child = HierarchicalNode(
+            id="child",
+            name="Child",
+            node_type="block",
+            children=[grandchild],
+            depth=1,
+        )
+        root = HierarchicalNode(
+            id="root",
+            name="Root",
+            node_type="model",
+            children=[child],
+            depth=0,
+        )
+        graph = HierarchicalGraph(
+            root=root,
+            nodes_by_id={"root": root, "child": child, "gc": grandchild},
+        )
+
+        graph.expand_to_depth(1)
+
+        assert not root.is_collapsed
+        assert not child.is_collapsed
+        assert grandchild.is_collapsed
+
+    def test_to_json(self):
+        """Test JSON export."""
+        root = HierarchicalNode(id="root", name="Model", node_type="model")
+        graph = HierarchicalGraph(root=root, total_nodes=1, depth=0)
+
+        json_str = graph.to_json()
+        data = json.loads(json_str)
+
+        assert data["root"]["name"] == "Model"
+        assert data["total_nodes"] == 1
+
+
+class TestHierarchicalGraphBuilder:
+    """Tests for building hierarchical graphs from ONNX."""
+
+    def test_build_from_onnx(self):
+        """Test building hierarchy from ONNX model."""
+        model = create_test_model()
+
+        with tempfile.NamedTemporaryFile(suffix=".onnx", delete=False) as f:
+            onnx.save(model, f.name)
+            model_path = Path(f.name)
+
+        try:
+            loader = ONNXGraphLoader()
+            _, graph_info = loader.load(model_path)
+
+            pattern_analyzer = PatternAnalyzer()
+            blocks = pattern_analyzer.group_into_blocks(graph_info)
+
+            builder = HierarchicalGraphBuilder()
+            graph = builder.build(graph_info, blocks, "TestModel")
+
+            assert graph.root.name == "TestModel"
+            assert graph.total_nodes > 0
+            assert len(graph.root.children) > 0
+        finally:
+            model_path.unlink()
+
+
+class TestSummaryGeneration:
+    """Tests for summary generation."""
+
+    def test_generate_summary(self):
+        """Test multi-level summary generation."""
+        child1 = HierarchicalNode(
+            id="c1",
+            name="ConvBlock",
+            node_type="block",
+            attributes={"block_type": "ConvRelu"},
+        )
+        child2 = HierarchicalNode(
+            id="c2",
+            name="Layer",
+            node_type="layer",
+            is_repeated=True,
+            repeat_count=12,
+        )
+        root = HierarchicalNode(
+            id="root",
+            name="Model",
+            node_type="model",
+            children=[child1, child2],
+            is_collapsed=False,
+        )
+        graph = HierarchicalGraph(root=root, total_nodes=3, depth=1)
+
+        summary = generate_summary(graph)
+
+        assert "Model" in summary
+        assert "ConvRelu" in summary
+        assert "x12" in summary
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

haoline/tests/test_html_export.py

@@ -0,0 +1,180 @@
+# Copyright (c) 2025 HaoLine Contributors
+# SPDX-License-Identifier: MIT
+
+"""Tests for HTML export."""
+
+from __future__ import annotations
+
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from ..edge_analysis import EdgeAnalysisResult
+from ..hierarchical_graph import HierarchicalGraph, HierarchicalNode
+from ..html_export import HTMLExporter, generate_html
+
+
+def create_test_graph() -> HierarchicalGraph:
+    """Create a test hierarchical graph."""
+    child1 = HierarchicalNode(
+        id="op_1",
+        name="Conv1",
+        node_type="op",
+        op_type="Conv",
+        total_flops=1000000,
+        total_memory_bytes=1024 * 1024,
+    )
+    child2 = HierarchicalNode(
+        id="op_2",
+        name="Relu1",
+        node_type="op",
+        op_type="Relu",
+    )
+    block = HierarchicalNode(
+        id="block_1",
+        name="ConvBlock",
+        node_type="block",
+        children=[child1, child2],
+        is_collapsed=False,
+        attributes={"block_type": "ConvRelu"},
+    )
+    root = HierarchicalNode(
+        id="root",
+        name="TestModel",
+        node_type="model",
+        children=[block],
+        is_collapsed=False,
+    )
+
+    return HierarchicalGraph(
+        root=root,
+        nodes_by_id={"root": root, "block_1": block, "op_1": child1, "op_2": child2},
+        total_nodes=4,
+        depth=2,
+    )
+
+
+def create_test_edge_result() -> EdgeAnalysisResult:
+    """Create a test edge analysis result."""
+    return EdgeAnalysisResult(
+        edges=[],
+        total_activation_bytes=10 * 1024 * 1024,
+        peak_activation_bytes=5 * 1024 * 1024,
+        peak_activation_node="Conv1",
+        bottleneck_edges=["conv_out"],
+        attention_edges=[],
+        skip_connection_edges=[],
+        memory_profile=[("Conv1", 5 * 1024 * 1024)],
+    )
+
+
+class TestHTMLGeneration:
+    """Tests for HTML generation."""
+
+    def test_generate_html_basic(self):
+        """Test basic HTML generation."""
+        graph = create_test_graph()
+        html = generate_html(graph, title="Test Model")
+
+        # Should contain key elements
+        assert "<!DOCTYPE html>" in html
+        assert "<title>Test Model - Neural Architecture</title>" in html
+        assert "d3.v7.min.js" in html
+        assert "graphData" in html
+
+    def test_generate_html_with_edges(self):
+        """Test HTML generation with edge data."""
+        graph = create_test_graph()
+        edge_result = create_test_edge_result()
+        html = generate_html(graph, edge_result, title="Test Model")
+
+        assert "edgeData" in html
+        assert "peak_activation_bytes" in html
+
+    def test_generate_html_to_file(self):
+        """Test HTML generation to file."""
+        graph = create_test_graph()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            output_path = Path(tmpdir) / "test_model.html"
+            html = generate_html(graph, title="Test", output_path=output_path)
+
+            assert output_path.exists()
+            content = output_path.read_text(encoding="utf-8")
+            assert content == html
+
+    def test_html_contains_controls(self):
+        """Test that HTML contains interactive controls."""
+        graph = create_test_graph()
+        html = generate_html(graph)
+
+        assert "expandAll()" in html
+        assert "collapseAll()" in html
+        assert "resetZoom()" in html
+        assert "fitToScreen()" in html
+
+    def test_html_contains_legend(self):
+        """Test that HTML contains op type legend."""
+        graph = create_test_graph()
+        html = generate_html(graph)
+
+        # New design uses "Op Types" instead of "Legend"
+        assert "Op Types" in html
+        assert "Convolution" in html
+        assert "Attention" in html
+        assert "legend-item" in html  # CSS class for legend items
+
+    def test_html_contains_stats_panel(self):
+        """Test that HTML contains stats panel."""
+        graph = create_test_graph()
+        html = generate_html(graph)
+
+        assert "node-count" in html
+        assert "edge-count" in html
+        assert "peak-memory" in html
+
+
+class TestHTMLExporter:
+    """Tests for HTMLExporter class."""
+
+    def test_export_basic(self):
+        """Test basic export."""
+        graph = create_test_graph()
+        exporter = HTMLExporter()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            output_path = Path(tmpdir) / "model.html"
+            result = exporter.export(graph, output_path=output_path)
+
+            assert result == output_path
+            assert output_path.exists()
+
+    def test_export_with_edges(self):
+        """Test export with edge data."""
+        graph = create_test_graph()
+        edge_result = create_test_edge_result()
+        exporter = HTMLExporter()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            output_path = Path(tmpdir) / "model.html"
+            exporter.export(graph, edge_result, output_path=output_path)
+
+            content = output_path.read_text(encoding="utf-8")
+            assert "peak_activation_bytes" in content
+
+    def test_export_custom_title(self):
+        """Test export with custom title."""
+        graph = create_test_graph()
+        exporter = HTMLExporter()
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            output_path = Path(tmpdir) / "model.html"
+            exporter.export(graph, output_path=output_path, title="My Custom Model")
+
+            content = output_path.read_text(encoding="utf-8")
+            assert "My Custom Model" in content
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])