haoline 0.3.0__py3-none-any.whl

haoline/schema.py

@@ -0,0 +1,523 @@
+# Copyright (c) 2025 HaoLine Contributors
+# SPDX-License-Identifier: MIT
+
+"""
+JSON Schema definition and validation for HaoLine reports.
+
+Provides:
+- INSPECTION_REPORT_SCHEMA: JSON Schema for InspectionReport (auto-generated from Pydantic)
+- validate_report(): Validate a report dict against the schema using Pydantic
+- ValidationError: Exception raised on validation failure
+
+The Pydantic models in pydantic_models.py were auto-generated from the original
+JSON Schema using datamodel-code-generator. Validation is now done via Pydantic
+for better error messages and type safety.
+"""
+
+from __future__ import annotations
+
+import warnings
+from typing import Any
+
+# Try to import Pydantic models (preferred)
+try:
+    from pydantic import ValidationError as PydanticValidationError
+
+    from .pydantic_models import HaolineInspectionReport
+
+    PYDANTIC_AVAILABLE = True
+except ImportError:
+    PYDANTIC_AVAILABLE = False
+    PydanticValidationError = None  # type: ignore
+    HaolineInspectionReport = None  # type: ignore
+
+# Fallback to jsonschema if Pydantic not available
+try:
+    from jsonschema import Draft7Validator
+    from jsonschema import ValidationError as JsonSchemaError
+
+    JSONSCHEMA_AVAILABLE = True
+except ImportError:
+    JSONSCHEMA_AVAILABLE = False
+    JsonSchemaError = None  # type: ignore
+
+
+class ValidationError(Exception):
+    """Raised when JSON schema validation fails."""
+
+    def __init__(self, message: str, errors: list[str] | None = None):
+        super().__init__(message)
+        self.errors = errors or []
+
+
+# JSON Schema for InspectionReport (Draft 7)
+INSPECTION_REPORT_SCHEMA: dict[str, Any] = {
+    "$schema": "http://json-schema.org/draft-07/schema#",
+    "$id": "https://github.com/mdayku/HaoLine/schema/inspection-report.schema.json",
+    "title": "HaoLine Inspection Report",
+    "description": "Schema for ONNX model inspection reports generated by HaoLine",
+    "type": "object",
+    "required": ["metadata", "generated_at", "autodoc_version"],
+    "properties": {
+        "metadata": {
+            "type": "object",
+            "description": "Model metadata extracted from ONNX proto",
+            "required": ["path", "ir_version", "producer_name", "opsets"],
+            "properties": {
+                "path": {
+                    "type": "string",
+                    "description": "Path to the ONNX model file",
+                },
+                "ir_version": {
+                    "type": "integer",
+                    "minimum": 1,
+                    "description": "ONNX IR version",
+                },
+                "producer_name": {
+                    "type": "string",
+                    "description": "Name of the tool that produced the model",
+                },
+                "producer_version": {
+                    "type": "string",
+                    "description": "Version of the producer tool",
+                },
+                "domain": {"type": "string", "description": "Model domain"},
+                "model_version": {
+                    "type": "integer",
+                    "description": "Model version number",
+                },
+                "doc_string": {
+                    "type": "string",
+                    "description": "Model documentation string",
+                },
+                "opsets": {
+                    "type": "object",
+                    "description": "Opset versions by domain",
+                    "additionalProperties": {"type": "integer", "minimum": 1},
+                },
+            },
+        },
+        "generated_at": {
+            "type": "string",
+            "format": "date-time",
+            "description": "ISO 8601 timestamp when report was generated",
+        },
+        "autodoc_version": {
+            "type": "string",
+            "pattern": "^[0-9]+\\.[0-9]+\\.[0-9]+",
+            "description": "Version of HaoLine that generated the report",
+        },
+        "graph_summary": {
+            "type": ["object", "null"],
+            "description": "Summary statistics about the ONNX graph",
+            "properties": {
+                "num_nodes": {
+                    "type": "integer",
+                    "minimum": 0,
+                    "description": "Total number of nodes in graph",
+                },
+                "num_inputs": {
+                    "type": "integer",
+                    "minimum": 0,
+                    "description": "Number of graph inputs",
+                },
+                "num_outputs": {
+                    "type": "integer",
+                    "minimum": 0,
+                    "description": "Number of graph outputs",
+                },
+                "num_initializers": {
+                    "type": "integer",
+                    "minimum": 0,
+                    "description": "Number of initializers (weights)",
+                },
+                "input_shapes": {
+                    "type": "object",
+                    "description": "Input tensor shapes by name",
+                    "additionalProperties": {
+                        "type": "array",
+                        "items": {"type": ["integer", "string"]},
+                    },
+                },
+                "output_shapes": {
+                    "type": "object",
+                    "description": "Output tensor shapes by name",
+                    "additionalProperties": {
+                        "type": "array",
+                        "items": {"type": ["integer", "string"]},
+                    },
+                },
+                "op_type_counts": {
+                    "type": "object",
+                    "description": "Count of each operator type",
+                    "additionalProperties": {"type": "integer", "minimum": 0},
+                },
+            },
+        },
+        "param_counts": {
+            "type": ["object", "null"],
+            "description": "Parameter count statistics",
+            "properties": {
+                "total": {
+                    "type": "integer",
+                    "minimum": 0,
+                    "description": "Total parameter count",
+                },
+                "trainable": {
+                    "type": "integer",
+                    "minimum": 0,
+                    "description": "Trainable parameter count",
+                },
+                "non_trainable": {
+                    "type": "integer",
+                    "minimum": 0,
+                    "description": "Non-trainable parameter count",
+                },
+                "by_op_type": {
+                    "type": "object",
+                    "description": "Parameters by operator type (fractional for shared weights)",
+                    "additionalProperties": {"type": "number", "minimum": 0},
+                },
+                "shared_weights": {
+                    "type": "object",
+                    "description": "Information about shared weights",
+                    "properties": {
+                        "count": {
+                            "type": "integer",
+                            "minimum": 0,
+                            "description": "Number of weights shared across 2+ nodes",
+                        },
+                        "details": {
+                            "type": "object",
+                            "description": "Shared weight name to list of node names using it",
+                            "additionalProperties": {
+                                "type": "array",
+                                "items": {"type": "string"},
+                            },
+                        },
+                    },
+                },
+                "precision_breakdown": {
+                    "type": "object",
+                    "description": "Parameter count by data type (fp32, fp16, int8, etc.)",
+                    "additionalProperties": {"type": "integer", "minimum": 0},
+                },
+                "is_quantized": {
+                    "type": "boolean",
+                    "description": "Whether model uses quantized weights or ops",
+                },
+                "quantized_ops": {
+                    "type": "array",
+                    "description": "List of quantized operation types detected",
+                    "items": {"type": "string"},
+                },
+            },
+        },
+        "flop_counts": {
+            "type": ["object", "null"],
+            "description": "FLOP estimation statistics",
+            "properties": {
+                "total": {
+                    "type": "integer",
+                    "minimum": 0,
+                    "description": "Total estimated FLOPs",
+                },
+                "by_node_type": {
+                    "type": "object",
+                    "description": "FLOPs by operator type",
+                    "additionalProperties": {"type": "integer", "minimum": 0},
+                },
+                "hotspots": {
+                    "type": "array",
+                    "description": "Top compute-intensive nodes",
+                    "items": {
+                        "type": "object",
+                        "properties": {
+                            "name": {"type": "string"},
+                            "op_type": {"type": "string"},
+                            "flops": {"type": "integer", "minimum": 0},
+                        },
+                    },
+                },
+            },
+        },
+        "memory_estimates": {
+            "type": ["object", "null"],
+            "description": "Memory usage estimates",
+            "properties": {
+                "model_size_bytes": {
+                    "type": "integer",
+                    "minimum": 0,
+                    "description": "Model size in bytes",
+                },
+                "peak_activation_bytes": {
+                    "type": "integer",
+                    "minimum": 0,
+                    "description": "Peak activation memory in bytes",
+                },
+                "kv_cache_bytes_per_token": {
+                    "type": "integer",
+                    "minimum": 0,
+                    "description": "KV cache memory per token (transformers)",
+                },
+                "kv_cache_bytes_full_context": {
+                    "type": "integer",
+                    "minimum": 0,
+                    "description": "KV cache for full context length",
+                },
+                "kv_cache_config": {
+                    "type": "object",
+                    "description": "KV cache configuration",
+                    "properties": {
+                        "num_layers": {"type": "integer", "minimum": 0},
+                        "hidden_dim": {"type": "integer", "minimum": 0},
+                        "seq_len": {"type": "integer", "minimum": 0},
+                        "bytes_per_element": {"type": "integer", "minimum": 1},
+                    },
+                },
+                "breakdown": {
+                    "type": ["object", "null"],
+                    "description": "Memory breakdown by component",
+                    "properties": {
+                        "weights_by_op_type": {
+                            "type": "object",
+                            "additionalProperties": {"type": "integer", "minimum": 0},
+                        },
+                        "activations_by_op_type": {
+                            "type": "object",
+                            "additionalProperties": {"type": "integer", "minimum": 0},
+                        },
+                    },
+                },
+            },
+        },
+        "detected_blocks": {
+            "type": "array",
+            "description": "Detected architectural blocks",
+            "items": {
+                "type": "object",
+                "required": ["block_type", "name"],
+                "properties": {
+                    "block_type": {
+                        "type": "string",
+                        "description": "Type of block (e.g., ResidualAdd, Attention)",
+                    },
+                    "name": {"type": "string", "description": "Block identifier"},
+                    "nodes": {
+                        "type": "array",
+                        "items": {"type": "string"},
+                        "description": "Node names in this block",
+                    },
+                    "start_node": {"type": "string"},
+                    "end_node": {"type": "string"},
+                    "attributes": {
+                        "type": "object",
+                        "description": "Block-specific attributes",
+                    },
+                },
+            },
+        },
+        "architecture_type": {
+            "type": "string",
+            "enum": ["transformer", "cnn", "mlp", "hybrid", "unknown"],
+            "description": "Detected architecture type",
+        },
+        "risk_signals": {
+            "type": "array",
+            "description": "Detected risk signals",
+            "items": {
+                "type": "object",
+                "required": ["id", "severity", "description"],
+                "properties": {
+                    "id": {
+                        "type": "string",
+                        "description": "Risk signal identifier",
+                    },
+                    "severity": {
+                        "type": "string",
+                        "enum": ["info", "warning", "high"],
+                        "description": "Severity level",
+                    },
+                    "description": {
+                        "type": "string",
+                        "description": "Human-readable description",
+                    },
+                    "nodes": {
+                        "type": "array",
+                        "items": {"type": "string"},
+                        "description": "Affected node names",
+                    },
+                    "recommendation": {
+                        "type": "string",
+                        "description": "Recommended action",
+                    },
+                },
+            },
+        },
+        "hardware_profile": {
+            "type": ["object", "null"],
+            "description": "Target hardware profile",
+            "properties": {
+                "name": {"type": "string"},
+                "vram_bytes": {"type": "integer", "minimum": 0},
+                "peak_fp32_tflops": {"type": "number", "minimum": 0},
+                "peak_fp16_tflops": {"type": "number", "minimum": 0},
+                "memory_bandwidth_gbps": {"type": "number", "minimum": 0},
+                "tdp_watts": {"type": ["integer", "null"]},
+            },
+        },
+        "hardware_estimates": {
+            "type": ["object", "null"],
+            "description": "Hardware-specific estimates",
+            "properties": {
+                "device": {"type": "string"},
+                "precision": {"type": "string"},
+                "batch_size": {"type": "integer", "minimum": 1},
+                "vram_required_bytes": {"type": "integer", "minimum": 0},
+                "fits_in_vram": {"type": "boolean"},
+                "theoretical_latency_ms": {"type": "number", "minimum": 0},
+                "bottleneck": {"type": "string"},
+                "compute_utilization_estimate": {"type": "number", "minimum": 0},
+                "gpu_saturation": {"type": "number", "minimum": 0},
+            },
+        },
+        "llm_summary": {
+            "type": ["object", "null"],
+            "description": "LLM-generated summary",
+            "properties": {
+                "success": {"type": "boolean"},
+                "short_summary": {"type": "string"},
+                "detailed_summary": {"type": "string"},
+                "model": {"type": "string"},
+                "error": {"type": "string"},
+            },
+        },
+        "dataset_info": {
+            "type": ["object", "null"],
+            "description": "Dataset and class information",
+            "properties": {
+                "task": {"type": ["string", "null"]},
+                "num_classes": {"type": ["integer", "null"], "minimum": 0},
+                "class_names": {
+                    "type": "array",
+                    "items": {"type": "string"},
+                },
+                "source": {"type": ["string", "null"]},
+            },
+        },
+    },
+}
+
+
+def validate_report(report_dict: dict[str, Any]) -> tuple[bool, list[str]]:
+    """
+    Validate a report dictionary against the schema.
+
+    Uses Pydantic for validation (preferred) or falls back to jsonschema.
+
+    Args:
+        report_dict: The report as a dictionary (from to_dict()).
+
+    Returns:
+        Tuple of (is_valid, list of error messages).
+        If neither pydantic nor jsonschema is installed, returns (True, []) with a warning.
+    """
+    # Prefer Pydantic validation
+    if PYDANTIC_AVAILABLE:
+        try:
+            HaolineInspectionReport.model_validate(report_dict)
+            return True, []
+        except PydanticValidationError as e:
+            error_messages = []
+            for error in e.errors():
+                loc = " -> ".join(str(x) for x in error["loc"])
+                error_messages.append(f"{loc}: {error['msg']}")
+            return False, error_messages
+
+    # Fallback to jsonschema
+    if JSONSCHEMA_AVAILABLE:
+        validator = Draft7Validator(INSPECTION_REPORT_SCHEMA)
+        errors = list(validator.iter_errors(report_dict))
+
+        if not errors:
+            return True, []
+
+        error_messages = []
+        for error in errors:
+            path = " -> ".join(str(p) for p in error.absolute_path) or "root"
+            error_messages.append(f"{path}: {error.message}")
+
+        return False, error_messages
+
+    # No validation library available
+    warnings.warn(
+        "Neither pydantic nor jsonschema installed. "
+        "Install with 'pip install pydantic' for validation.",
+        UserWarning,
+        stacklevel=2,
+    )
+    return True, []
+
+
+def validate_report_strict(report_dict: dict[str, Any]) -> None:
+    """
+    Validate a report dictionary, raising ValidationError on failure.
+
+    Args:
+        report_dict: The report as a dictionary (from to_dict()).
+
+    Raises:
+        ValidationError: If validation fails.
+    """
+    is_valid, errors = validate_report(report_dict)
+    if not is_valid:
+        raise ValidationError(
+            f"Report validation failed with {len(errors)} error(s)", errors=errors
+        )
+
+
+def get_schema() -> dict[str, Any]:
+    """
+    Return the JSON schema for InspectionReport.
+
+    If Pydantic is available, returns the auto-generated schema from the model.
+    Otherwise, returns the manually-defined schema.
+    """
+    if PYDANTIC_AVAILABLE:
+        schema: dict[str, Any] = HaolineInspectionReport.model_json_schema()
+        return schema
+    return INSPECTION_REPORT_SCHEMA.copy()
+
+
+def validate_with_pydantic(report_dict: dict[str, Any]) -> HaolineInspectionReport | None:
+    """
+    Validate and parse a report dict into a Pydantic model.
+
+    Args:
+        report_dict: The report as a dictionary.
+
+    Returns:
+        HaolineInspectionReport instance if valid, None if Pydantic not available.
+
+    Raises:
+        ValidationError: If validation fails.
+    """
+    if not PYDANTIC_AVAILABLE:
+        warnings.warn(
+            "Pydantic not installed. Install with 'pip install pydantic'.",
+            UserWarning,
+            stacklevel=2,
+        )
+        return None
+
+    try:
+        result: HaolineInspectionReport = HaolineInspectionReport.model_validate(report_dict)
+        return result
+    except PydanticValidationError as e:
+        error_messages = [
+            f"{' -> '.join(str(x) for x in err['loc'])}: {err['msg']}" for err in e.errors()
+        ]
+        raise ValidationError(
+            f"Report validation failed with {len(error_messages)} error(s)",
+            errors=error_messages,
+        ) from e