# Copyright (c) ONNX Project Contributors # # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations import os import tarfile import onnx.checker import onnx.helper import onnx.shape_inference from onnx import FunctionProto, ModelProto, NodeProto, TensorProto, ValueInfoProto class Extractor: def __init__(self, model: ModelProto) -> None: self.model = onnx.shape_inference.infer_shapes(model) self.graph = self.model.graph self.wmap = self._build_name2obj_dict(self.graph.initializer) self.vimap = self._build_name2obj_dict(self.graph.value_info) @staticmethod def _build_name2obj_dict(objs): # type: ignore return {obj.name: obj for obj in objs} def _collect_new_io_core(self, original_io, io_names_to_extract): # type: ignore original_io_map = self._build_name2obj_dict(original_io) original_io_names = set(original_io_map) s_io_names_to_extract = set(io_names_to_extract) io_names_to_keep = s_io_names_to_extract & original_io_names new_io_names_to_add = s_io_names_to_extract - original_io_names new_io_tensors = [original_io_map[name] for name in io_names_to_keep] # activation become input or output new_io_tensors.extend(self.vimap[name] for name in new_io_names_to_add) # adjust sequence new_io_tensors_map = self._build_name2obj_dict(new_io_tensors) return [new_io_tensors_map[name] for name in io_names_to_extract] def _collect_new_inputs(self, names: list[str]) -> list[ValueInfoProto]: return self._collect_new_io_core(self.graph.input, names) # type: ignore def _collect_new_outputs(self, names: list[str]) -> list[ValueInfoProto]: return self._collect_new_io_core(self.graph.output, names) # type: ignore def _dfs_search_reachable_nodes( self, node_output_name: str, graph_input_names: set[str], nodes: list[NodeProto], reachable: set[int], unreachable: set[int], ) -> None: """Helper function to find nodes which are connected to an output Arguments: node_output_name (str): The name of the output graph_input_names (set of string): The names of all inputs of the graph nodes (list of nodes): The list of all nodes of the graph reachable (set of int): The set of indexes to reachable nodes in `nodes` unreachable (set of int): The set of indexes to unreachable nodes in `nodes` """ # finish search at inputs if node_output_name in graph_input_names: return # find nodes connected to this output nodes_to_search = [ index for index in unreachable if node_output_name in nodes[index].output ] # add nodes connected to this output to sets for node_index in nodes_to_search: reachable.add(node_index) unreachable.remove(node_index) # recurse on inputs for node_index in nodes_to_search: for name in nodes[node_index].input: self._dfs_search_reachable_nodes( name, graph_input_names, nodes, reachable, unreachable ) def _collect_reachable_nodes( self, input_names: list[str], output_names: list[str], ) -> list[NodeProto]: _input_names = set(input_names) nodes = list(self.graph.node) reachable: set[int] = set() unreachable: set[int] = set(range(len(nodes))) for name in output_names: self._dfs_search_reachable_nodes( name, _input_names, nodes, reachable, unreachable ) # needs to be topologically sorted nodes = [nodes[node_index] for node_index in sorted(reachable)] return nodes def _collect_referred_local_functions( self, nodes, # type: list[NodeProto] ): # type: (...) -> list[FunctionProto] # a node in a model graph may refer a function. # a function contains nodes, some of which may in turn refer a function. # we need to find functions referred by graph nodes and # by nodes used to define functions. def find_referred_funcs(nodes, referred_local_functions): # type: ignore new_nodes = [] # type: list[NodeProto] for node in nodes: # check if the node is a function op match_function = next( ( f for f in self.model.functions if f.name == node.op_type and f.domain == node.domain ), None, ) if match_function and match_function not in referred_local_functions: referred_local_functions.append(match_function) new_nodes.extend(match_function.node) return new_nodes referred_local_functions = [] # type: list[FunctionProto] new_nodes = find_referred_funcs(nodes, referred_local_functions) while new_nodes: new_nodes = find_referred_funcs(new_nodes, referred_local_functions) return referred_local_functions def _collect_reachable_tensors( self, nodes: list[NodeProto], ) -> tuple[list[TensorProto], list[ValueInfoProto]]: all_tensors_names: set[str] = set() for node in nodes: all_tensors_names.update(node.input) all_tensors_names.update(node.output) initializer = [self.wmap[t] for t in self.wmap if t in all_tensors_names] value_info = [self.vimap[t] for t in self.vimap if t in all_tensors_names] len_sparse_initializer = len(self.graph.sparse_initializer) if len_sparse_initializer != 0: raise ValueError( f"len_sparse_initializer is {len_sparse_initializer}, it must be 0." ) len_quantization_annotation = len(self.graph.quantization_annotation) if len_quantization_annotation != 0: raise ValueError( f"len_quantization_annotation is {len_quantization_annotation}, it must be 0." ) return initializer, value_info def _make_model( self, nodes: list[NodeProto], inputs: list[ValueInfoProto], outputs: list[ValueInfoProto], initializer: list[TensorProto], value_info: list[ValueInfoProto], local_functions: list[FunctionProto], ) -> ModelProto: name = "Extracted from {" + self.graph.name + "}" graph = onnx.helper.make_graph( nodes, name, inputs, outputs, initializer=initializer, value_info=value_info ) meta = { "ir_version": self.model.ir_version, "opset_imports": self.model.opset_import, "producer_name": "onnx.utils.extract_model", "functions": local_functions, } return onnx.helper.make_model(graph, **meta) def extract_model( self, input_names: list[str], output_names: list[str], ) -> ModelProto: inputs = self._collect_new_inputs(input_names) outputs = self._collect_new_outputs(output_names) nodes = self._collect_reachable_nodes(input_names, output_names) initializer, value_info = self._collect_reachable_tensors(nodes) local_functions = self._collect_referred_local_functions(nodes) model = self._make_model( nodes, inputs, outputs, initializer, value_info, local_functions ) return model def extract_model( input_path: str | os.PathLike, output_path: str | os.PathLike, input_names: list[str], output_names: list[str], check_model: bool = True, ) -> None: """Extracts sub-model from an ONNX model. The sub-model is defined by the names of the input and output tensors *exactly*. Note: For control-flow operators, e.g. If and Loop, the _boundary of sub-model_, which is defined by the input and output tensors, should not _cut through_ the subgraph that is connected to the _main graph_ as attributes of these operators. Arguments: input_path (str | os.PathLike): The path to original ONNX model. output_path (str | os.PathLike): The path to save the extracted ONNX model. input_names (list of string): The names of the input tensors that to be extracted. output_names (list of string): The names of the output tensors that to be extracted. check_model (bool): Whether to run model checker on the extracted model. """ if not os.path.exists(input_path): raise ValueError(f"Invalid input model path: {input_path}") if not output_path: raise ValueError("Output model path shall not be empty!") if not output_names: raise ValueError("Output tensor names shall not be empty!") onnx.checker.check_model(input_path) model = onnx.load(input_path) e = Extractor(model) extracted = e.extract_model(input_names, output_names) onnx.save(extracted, output_path) if check_model: onnx.checker.check_model(output_path) def _tar_members_filter( tar: tarfile.TarFile, base: str | os.PathLike ) -> list[tarfile.TarInfo]: """Check that the content of ``tar`` will be extracted safely Args: tar: The tarball file base: The directory where the tarball will be extracted Returns: list of tarball members """ result = [] for member in tar: member_path = os.path.join(base, member.name) abs_base = os.path.abspath(base) abs_member = os.path.abspath(member_path) if not abs_member.startswith(abs_base): raise RuntimeError( f"The tarball member {member_path} in downloading model contains " f"directory traversal sequence which may contain harmful payload." ) elif member.issym() or member.islnk(): raise RuntimeError( f"The tarball member {member_path} in downloading model contains " f"symbolic links which may contain harmful payload." ) result.append(member) return result def _extract_model_safe( model_tar_path: str | os.PathLike, local_model_with_data_dir_path: str | os.PathLike ) -> None: """Safely extracts a tar file to a specified directory. This function ensures that the extraction process mitigates against directory traversal vulnerabilities by validating or sanitizing paths within the tar file. It also provides compatibility for different versions of the tarfile module by checking for the availability of certain attributes or methods before invoking them. Args: model_tar_path: The path to the tar file to be extracted. local_model_with_data_dir_path: The directory path where the tar file contents will be extracted to. """ with tarfile.open(model_tar_path) as model_with_data_zipped: # Mitigate tarball directory traversal risks if hasattr(tarfile, "data_filter"): model_with_data_zipped.extractall( path=local_model_with_data_dir_path, filter="data" ) else: model_with_data_zipped.extractall( path=local_model_with_data_dir_path, members=_tar_members_filter( model_with_data_zipped, local_model_with_data_dir_path ), )