21798755656

Committed 08 Feb 2026 01:13PM UTC coverage: 66.496% (+0.06%) from 66.434%

Build # 21798755656

Build Type

Pull #511

github

Committed by

web-flow

Commit Message

Merge cbe5c4a65 into fdf653f6d

Pull Request Pull Request #511: adds ONNX dispatcher for tensor nodes

Coverage Stats

85 of 94 new or added lines in 2 files covered. (90.43%)

23253 of 34969 relevant lines covered (66.5%)

373.95 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

89.66

/python/docc/python/onnx_model_builder.py

"""
ONNX Model Builder

Converts JSON graph descriptions emitted by the ONNX tensor dispatchers
into valid ONNX model files that can be loaded by ONNX Runtime.
"""

import json
import os
from pathlib import Path
from typing import List, Dict, Any, Set, Tuple


def build_onnx_model(json_path: str, output_path: str) -> str:
    """
    Build an ONNX model from a JSON graph description.

    Args:
        json_path: Path to the JSON file containing node definitions
        output_path: Path where the .onnx model should be written

    Returns:
        Path to the generated ONNX model file
    """
    try:
        import onnx
        from onnx import helper, TensorProto
    except ImportError:
        raise ImportError(
            "The 'onnx' package is required for ONNX model generation. "
            "Install it with: pip install onnx"
        )

    # Read the JSON file
    with open(json_path, "r") as f:
        content = f.read().strip()

    # The JSON file contains comma-separated node definitions
    # Wrap in array brackets to make it valid JSON
    if content.endswith(","):
        content = content[:-1]  # Remove trailing comma
    json_content = f"[{content}]"

    try:
        nodes_data = json.loads(json_content)
    except json.JSONDecodeError as e:
        raise ValueError(
            f"Failed to parse ONNX graph JSON: {e}\nContent: {json_content[:500]}"
        )

    # Collect all inputs and outputs
    all_inputs: Set[str] = set()
    all_outputs: Set[str] = set()
    intermediate: Set[str] = set()
    elem_type = TensorProto.FLOAT  # Default

    # Build ONNX nodes
    onnx_nodes = []
    for node_data in nodes_data:
        op_type = node_data.get("op_type", "Unknown")
        name = node_data.get("name", "unnamed")
        inputs = node_data.get("inputs", [])
        outputs = node_data.get("outputs", [])
        attributes = node_data.get("attributes", {})

        if "elem_type" in node_data:
            elem_type = node_data["elem_type"]

        # Track inputs/outputs
        for inp in inputs:
            if inp not in intermediate:
                all_inputs.add(inp)
        for out in outputs:
            intermediate.add(out)
            all_outputs.add(out)

        # Build attribute list
        onnx_attrs = []
        for attr_name, attr_value in attributes.items():
            if isinstance(attr_value, list):
                if all(isinstance(x, int) for x in attr_value):
                    onnx_attrs.append(helper.make_attribute(attr_name, attr_value))
                elif all(isinstance(x, float) for x in attr_value):
                    onnx_attrs.append(helper.make_attribute(attr_name, attr_value))
            elif isinstance(attr_value, int):
                onnx_attrs.append(helper.make_attribute(attr_name, attr_value))
            elif isinstance(attr_value, float):
                onnx_attrs.append(helper.make_attribute(attr_name, attr_value))
            elif isinstance(attr_value, str):
                onnx_attrs.append(helper.make_attribute(attr_name, attr_value))

        # Create the ONNX node
        node = helper.make_node(
            op_type,
            inputs=inputs,
            outputs=outputs,
            name=name,
        )
        # Add attributes
        node.attribute.extend(onnx_attrs)
        onnx_nodes.append(node)

    # Remove intermediate values from outputs (only keep final outputs)
    # Final outputs are those not consumed by any other node
    consumed = set()
    for node_data in nodes_data:
        for inp in node_data.get("inputs", []):
            consumed.add(inp)

    final_outputs = all_outputs - consumed
    graph_inputs = all_inputs - intermediate

    # If no final outputs, use all outputs
    if not final_outputs:
        final_outputs = all_outputs

    # Create graph inputs (with dynamic shapes using dim_param)
    input_tensors = []
    for inp_name in sorted(graph_inputs):
        # Use dynamic shape with symbolic dimension
        input_tensor = helper.make_tensor_value_info(
            inp_name, elem_type, None  # Dynamic shape
        )
        input_tensors.append(input_tensor)

    # Create graph outputs
    output_tensors = []
    for out_name in sorted(final_outputs):
        output_tensor = helper.make_tensor_value_info(
            out_name, elem_type, None  # Dynamic shape
        )
        output_tensors.append(output_tensor)

    # Create the graph
    graph = helper.make_graph(
        onnx_nodes,
        "docc_onnx_graph",
        input_tensors,
        output_tensors,
    )

    # Create the model
    model = helper.make_model(
        graph, opset_imports=[helper.make_opsetid("", 13)]  # ONNX opset 13
    )

    # Set IR version
    model.ir_version = 7

    # Validate the model
    try:
        onnx.checker.check_model(model)
    except onnx.checker.ValidationError as e:
        # Log warning but continue - dynamic shapes may cause validation issues
        print(f"ONNX validation warning: {e}")

    # Save the model
    onnx.save(model, output_path)

    return output_path


def convert_json_to_onnx(build_path: str) -> str:
    """
    Find and convert ONNX JSON files in the build directory.

    Args:
        build_path: Path to the build directory

    Returns:
        Path to the generated ONNX model, or None if no JSON found
    """
    json_paths = Path(build_path).glob("*.onnx.json")
    models = []
    for json_path in json_paths:
        onnx_path = json_path.parent / (json_path.name.replace(".onnx.json", ".onnx"))
        build_onnx_model(str(json_path), str(onnx_path))
        models.append(str(onnx_path))

    return models

1	"""
2	ONNX Model Builder
3
4	Converts JSON graph descriptions emitted by the ONNX tensor dispatchers
5	into valid ONNX model files that can be loaded by ONNX Runtime.
6	"""
7
8	import json	4✔
9	import os	4✔
10	from pathlib import Path	4✔
11	from typing import List, Dict, Any, Set, Tuple	4✔
12
13
14	def build_onnx_model(json_path: str, output_path: str) -> str:	4✔
15	"""
16	Build an ONNX model from a JSON graph description.
17
18	Args:
19	json_path: Path to the JSON file containing node definitions
20	output_path: Path where the .onnx model should be written
21
22	Returns:
23	Path to the generated ONNX model file
24	"""
25	try:	4✔
26	import onnx	4✔
27	from onnx import helper, TensorProto	4✔
NEW 28	except ImportError:	×
NEW 29	raise ImportError(	×
30	"The 'onnx' package is required for ONNX model generation. "
31	"Install it with: pip install onnx"
32	)
33
34	# Read the JSON file
35	with open(json_path, "r") as f:	4✔
36	content = f.read().strip()	4✔
37
38	# The JSON file contains comma-separated node definitions
39	# Wrap in array brackets to make it valid JSON
40	if content.endswith(","):	4✔
41	content = content[:-1] # Remove trailing comma	4✔
42	json_content = f"[{content}]"	4✔
43
44	try:	4✔
45	nodes_data = json.loads(json_content)	4✔
NEW 46	except json.JSONDecodeError as e:	×
NEW 47	raise ValueError(	×
48	f"Failed to parse ONNX graph JSON: {e}\nContent: {json_content[:500]}"
49	)
50
51	# Collect all inputs and outputs
52	all_inputs: Set[str] = set()	4✔
53	all_outputs: Set[str] = set()	4✔
54	intermediate: Set[str] = set()	4✔
55	elem_type = TensorProto.FLOAT # Default	4✔
56
57	# Build ONNX nodes
58	onnx_nodes = []	4✔
59	for node_data in nodes_data:	4✔
60	op_type = node_data.get("op_type", "Unknown")	4✔
61	name = node_data.get("name", "unnamed")	4✔
62	inputs = node_data.get("inputs", [])	4✔
63	outputs = node_data.get("outputs", [])	4✔
64	attributes = node_data.get("attributes", {})	4✔
65
66	if "elem_type" in node_data:	4✔
67	elem_type = node_data["elem_type"]	4✔
68
69	# Track inputs/outputs
70	for inp in inputs:	4✔
71	if inp not in intermediate:	4✔
72	all_inputs.add(inp)	4✔
73	for out in outputs:	4✔
74	intermediate.add(out)	4✔
75	all_outputs.add(out)	4✔
76
77	# Build attribute list
78	onnx_attrs = []	4✔
79	for attr_name, attr_value in attributes.items():	4✔
80	if isinstance(attr_value, list):	4✔
81	if all(isinstance(x, int) for x in attr_value):	4✔
82	onnx_attrs.append(helper.make_attribute(attr_name, attr_value))	4✔
NEW 83	elif all(isinstance(x, float) for x in attr_value):	×
NEW 84	onnx_attrs.append(helper.make_attribute(attr_name, attr_value))	×
85	elif isinstance(attr_value, int):	4✔
86	onnx_attrs.append(helper.make_attribute(attr_name, attr_value))	4✔
87	elif isinstance(attr_value, float):	4✔
88	onnx_attrs.append(helper.make_attribute(attr_name, attr_value))	4✔
NEW 89	elif isinstance(attr_value, str):	×
NEW 90	onnx_attrs.append(helper.make_attribute(attr_name, attr_value))	×
91
92	# Create the ONNX node
93	node = helper.make_node(	4✔
94	op_type,
95	inputs=inputs,
96	outputs=outputs,
97	name=name,
98	)
99	# Add attributes
100	node.attribute.extend(onnx_attrs)	4✔
101	onnx_nodes.append(node)	4✔
102
103	# Remove intermediate values from outputs (only keep final outputs)
104	# Final outputs are those not consumed by any other node
105	consumed = set()	4✔
106	for node_data in nodes_data:	4✔
107	for inp in node_data.get("inputs", []):	4✔
108	consumed.add(inp)	4✔
109
110	final_outputs = all_outputs - consumed	4✔
111	graph_inputs = all_inputs - intermediate	4✔
112
113	# If no final outputs, use all outputs
114	if not final_outputs:	4✔
NEW 115	final_outputs = all_outputs	×
116
117	# Create graph inputs (with dynamic shapes using dim_param)
118	input_tensors = []	4✔
119	for inp_name in sorted(graph_inputs):	4✔
120	# Use dynamic shape with symbolic dimension
121	input_tensor = helper.make_tensor_value_info(	4✔
122	inp_name, elem_type, None # Dynamic shape
123	)
124	input_tensors.append(input_tensor)	4✔
125
126	# Create graph outputs
127	output_tensors = []	4✔
128	for out_name in sorted(final_outputs):	4✔
129	output_tensor = helper.make_tensor_value_info(	4✔
130	out_name, elem_type, None # Dynamic shape
131	)
132	output_tensors.append(output_tensor)	4✔
133
134	# Create the graph
135	graph = helper.make_graph(	4✔
136	onnx_nodes,
137	"docc_onnx_graph",
138	input_tensors,
139	output_tensors,
140	)
141
142	# Create the model
143	model = helper.make_model(	4✔
144	graph, opset_imports=[helper.make_opsetid("", 13)] # ONNX opset 13
145	)
146
147	# Set IR version
148	model.ir_version = 7	4✔
149
150	# Validate the model
151	try:	4✔
152	onnx.checker.check_model(model)	4✔
153	except onnx.checker.ValidationError as e:	4✔
154	# Log warning but continue - dynamic shapes may cause validation issues
155	print(f"ONNX validation warning: {e}")	4✔
156
157	# Save the model
158	onnx.save(model, output_path)	4✔
159
160	return output_path	4✔
161
162
163	def convert_json_to_onnx(build_path: str) -> str:	4✔
164	"""
165	Find and convert ONNX JSON files in the build directory.
166
167	Args:
168	build_path: Path to the build directory
169
170	Returns:
171	Path to the generated ONNX model, or None if no JSON found
172	"""
173	json_paths = Path(build_path).glob("*.onnx.json")	4✔
174	models = []	4✔
175	for json_path in json_paths:	4✔
176	onnx_path = json_path.parent / (json_path.name.replace(".onnx.json", ".onnx"))	4✔
177	build_onnx_model(str(json_path), str(onnx_path))	4✔
178	models.append(str(onnx_path))	4✔
179
180	return models	4✔

daisytuner / docc / 21798755656

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous