Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/graph/graph.py: 94%
195 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-27 10:02 +0800
1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3# -----------------------------------------------------------------------------------------------------------
4# Copyright (c) 2025 Huawei Technologies Co., Ltd.
5# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
6# CANN Open Software License Agreement Version 2.0 (the "License").
7# Please refer to the License for details. You may not use this file except in compliance with the License.
8# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
9# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
10# See LICENSE in the root of the software repository for the full text of the License.
11# -----------------------------------------------------------------------------------------------------------
13"""Graph module for GraphEngine graph operations."""
15import ctypes
16from enum import Enum
17from typing import TYPE_CHECKING, Any, List, Optional
19from ge._capi.pygraph_wrapper import graph_lib
21from ._attr import _AttrValue
22from .node import Node
24if TYPE_CHECKING:
25 from ge.es.graph_builder import GraphBuilder
28class DumpFormat(Enum):
29 kOnnx = 0
30 kTxt = 1
31 kReadable = 2
34class Graph:
35 """Graph class for GraphEngine graph operations.
37 This class provides a Pythonic interface for graph operations
38 using the GraphEngine C API.
40 OWNERSHIP MANAGEMENT:
41 Graph objects can be in two ownership states:
42 1. Python-owned (default): Python is responsible for releasing C++ resources
43 2. C++-owned: C++ side manages the resource, Python is only a snapshot of the graph
45 When a Graph is passed as a subgraph parameter to operators (e.g., If, While, Case),
46 ownership is automatically transferred to the C++ side to avoid double-free issues.
48 Example:
49 >>> graph = Graph("my_graph")
50 >>> nodes = graph.get_all_nodes()
51 >>> attr = graph.get_attr("some_attr")
52 >>> graph.set_attr("new_attr", attr_value)
53 """
55 def __init__(self, name: Optional[str] = "graph") -> None:
56 """Initialize a Graph.
58 Args:
59 name: Graph name (optional).
61 Raises:
62 TypeError: If name is not a string.
63 RuntimeError: If graph creation fails.
64 """
65 # Init
66 self._owns_handle = False
67 self._owner = None
68 self._handle = None
69 if not isinstance(name, str):
70 raise TypeError("Graph name must be a string")
72 name_bytes = name.encode("utf-8")
73 self._handle = graph_lib.GeApiWrapper_Graph_CreateGraph(name_bytes)
74 if not self._handle:
75 raise RuntimeError("Failed to create Graph")
76 self._owns_handle = True # Python owns the resource by default
77 self._owner = None # Track the owner if ownership is transferred
79 def __del__(self) -> None:
80 """Clean up resources only if we own them."""
81 # Only destroy the graph if Python owns the handle
82 if self._owns_handle:
83 graph_lib.GeApiWrapper_Graph_DestroyGraph(self._handle)
84 self._handle = None
86 def __copy__(self) -> None:
87 """Copy is not supported."""
88 raise RuntimeError("Graph does not support copy")
90 def __deepcopy__(self, memodict) -> None:
91 """Deep copy is not supported."""
92 raise RuntimeError("Graph does not support deepcopy")
94 def __str__(self) -> str:
95 """Dump graph to stream in readable format."""
96 return self.dump_to_stream(DumpFormat.kReadable)
98 @classmethod
99 def _create_from(
100 cls,
101 handle: ctypes.c_void_p,
102 owns_handle: bool = True,
103 owner: Optional[Any] = None,
104 ) -> "Graph":
105 """Create Graph object from C++ pointer.
107 Internal helper used by GraphBuilder and Python pass bridge borrowed views.
108 Do not call this method directly unless you also understand the underlying
109 handle ownership model.
111 Args:
112 handle: C++ Graph object pointer.
113 owns_handle: Whether Python owns the handle and should destroy it.
114 owner: Optional Python object that keeps the real owner alive when
115 this Graph is only a borrowed view.
117 Returns:
118 Graph object.
120 Raises:
121 ValueError: If pointer is None.
122 """
123 if not handle:
124 raise ValueError("Graph pointer cannot be None")
125 instance = cls.__new__(cls)
126 instance._handle = handle
127 instance._owns_handle = owns_handle
128 instance._owner = owner
129 return instance
131 def _transfer_ownership_when_pass_as_subgraph(self, new_owner: "GraphBuilder") -> None:
132 """Transfer ownership of the C++ resource to new_owner.
134 After calling this method, Python will no longer destroy the underlying
135 C++ Graph object. This is called automatically when a Graph is passed
136 as a subgraph parameter to operators(i.e., If, While, Case inner_graph).
138 Args:
139 new_owner: The object that will manage the C++ resource (typically a GraphBuilder).
140 this Graph will hold a reference to keep the new_owner alive.
142 Example:
143 >>> sub_graph = create_subgraph()
144 >>> result = If(..., then_graph=sub_graph, ...)
145 >>> # sub_graph._transfer_ownership_when_pass_as_subgraph(main_builder) is called automatically
146 >>> # sub_graph._owner now holds a reference to main_builder
147 >>> # As long as sub_graph exists, main_builder won't be GC'd
148 """
149 if self._owner is not None:
150 raise RuntimeError(
151 "Graph :{} already has an new owner builder :{}, cannot transfer ownership again".format(
152 self.name, self._owner.name
153 )
154 )
155 self._owns_handle = False
156 self._owner = new_owner # Keep reference to prevent premature GC
158 @property
159 def name(self) -> str:
160 """Get graph name.
162 Returns:
163 Graph name.
165 Raises:
166 RuntimeError: If name retrieval fails.
167 """
168 c_str = graph_lib.GeApiWrapper_Graph_GetName(self._handle)
169 if not c_str:
170 raise RuntimeError("Failed to get Graph name")
172 try:
173 return ctypes.string_at(c_str).decode("utf-8")
174 finally:
175 graph_lib.GeApiWrapper_FreeString(c_str)
177 def get_all_nodes(self) -> List[Node]:
178 """Get all nodes in the graph.
180 Returns:
181 List of Node objects.
182 """
183 node_num = ctypes.c_size_t()
184 nodes = graph_lib.GeApiWrapper_Graph_GetAllNodes(self._handle, ctypes.byref(node_num))
185 if not nodes:
186 return []
188 try:
189 return [Node._create_from(nodes[i]) for i in range(node_num.value)]
190 finally:
191 graph_lib.GeApiWrapper_GNode_FreeGNodeArray(nodes)
193 def get_direct_nodes(self) -> List[Node]:
194 """Get direct nodes in the graph.
196 Returns:
197 List of Node objects.
198 """
199 node_num = ctypes.c_size_t()
200 nodes = graph_lib.GeApiWrapper_Graph_GetDirectNode(self._handle, ctypes.byref(node_num))
201 if not nodes:
202 return []
204 try:
205 return [Node._create_from(nodes[i]) for i in range(node_num.value)]
206 finally:
207 graph_lib.GeApiWrapper_GNode_FreeGNodeArray(nodes)
209 def get_attr(self, key: str) -> Any:
210 """Get graph attribute.
212 Args:
213 key: Attribute name.
215 Returns:
216 AttrValue object.
218 Raises:
219 TypeError: If key is not a string.
220 RuntimeError: If attribute retrieval fails.
221 """
222 if not isinstance(key, str):
223 raise TypeError("Attribute key must be a string")
225 attr_value = _AttrValue()
226 key_bytes = key.encode("utf-8")
228 ret = graph_lib.GeApiWrapper_Graph_GetAttr(self._handle, key_bytes, attr_value._av_ptr)
229 if ret != 0: # GRAPH_SUCCESS
230 raise RuntimeError(f"Failed to get attribute '{key}' from Graph {self.name}")
231 return attr_value.get_value()
233 def set_attr(self, key: str, value: Any) -> None:
234 """Set graph attribute.
236 Args:
237 key: Attribute name.
238 value: Attribute value.
240 Raises:
241 TypeError: If arguments have wrong types.
242 RuntimeError: If attribute setting fails.
243 """
244 if not isinstance(key, str):
245 raise TypeError("Attribute key must be a string")
247 key_bytes = key.encode("utf-8")
248 attr_value = _AttrValue()
249 attr_value.set_value(value)
250 ret = graph_lib.GeApiWrapper_Graph_SetAttr(self._handle, key_bytes, attr_value._av_ptr)
251 if ret != 0: # GRAPH_SUCCESS
252 raise RuntimeError(f"Failed to set attribute '{key}' on Graph {self.name}")
254 def dump_to_file(self, format: DumpFormat = DumpFormat.kReadable, suffix: str = "") -> None:
255 """Dump graph to file.
257 Args:
258 format: Format of the file. Defaults to "kReadable". Can be "kOnnx" or "kTxt" or "kReadable".
259 suffix: Suffix to append to the filename. Defaults to empty string. eg: xxxx
260 If path and suffix is not empty, the file name will be like: path/ge_<onnx/txt/readable>_00000_graph_0_xxxx.<txt/pbtxt>
261 Note:
262 pbtxt has only graph structure, no weights data or other attributes.
263 Raises:
264 TypeError: If format is not in [DumpFormat.kOnnx, DumpFormat.kTxt, DumpFormat.kReadable].
265 TypeError: If suffix is not a string.
266 RuntimeError: If dump operation fails.
267 """
269 if not isinstance(suffix, str):
270 raise TypeError("Suffix must be a string")
272 if format not in [DumpFormat.kOnnx, DumpFormat.kTxt, DumpFormat.kReadable]:
273 raise TypeError("Format must be in [DumpFormat.kOnnx, DumpFormat.kTxt, DumpFormat.kReadable]")
275 suffix_bytes = suffix.encode("utf-8")
276 ret = graph_lib.GeApiWrapper_Graph_Dump_To_File(self._handle, format.value, suffix_bytes)
277 if ret != 0:
278 raise RuntimeError(f"Failed to dump graph: {self.name} to file format {format} with suffix {suffix}")
280 def dump_to_stream(self, format: DumpFormat = DumpFormat.kReadable) -> str:
281 """Dump graph to stream.
283 Args:
284 format: Format of the stream. Defaults to "kReadable". Can be "kOnnx" or "kTxt" or "kReadable".
286 Returns:
287 Stream of the graph.
288 """
289 if format not in [DumpFormat.kOnnx, DumpFormat.kTxt, DumpFormat.kReadable]:
290 raise TypeError("Format must be in [DumpFormat.kOnnx, DumpFormat.kTxt, DumpFormat.kReadable]")
292 c_str = graph_lib.GeApiWrapper_Graph_Dump_To_Stream(self._handle, format.value)
293 if not c_str:
294 raise RuntimeError(f"Failed to dump graph: {self.name} to stream")
296 try:
297 return ctypes.string_at(c_str).decode("utf-8")
298 finally:
299 graph_lib.GeApiWrapper_FreeString(c_str)
301 def save_to_air(self, file_path: str) -> None:
302 """Save graph to AIR format.
304 Args:
305 path: file path to save the AIR file.
307 Raises:
308 TypeError: If file_path is not a string.
309 RuntimeError: If save operation fails.
310 """
311 if not isinstance(file_path, str):
312 raise TypeError("file_path must be a string")
313 ret = graph_lib.GeApiWrapper_Graph_SaveToAir(self._handle, file_path.encode("utf-8"))
314 if ret != 0:
315 raise RuntimeError("Failed to save graph to AIR format")
317 def load_from_air(self, file_path: str) -> None:
318 """Load graph from AIR format.
320 Args:
321 path: file path to load the AIR file.
323 Raises:
324 TypeError: If file_path is not a string.
325 RuntimeError: If load operation fails.
326 """
327 if not isinstance(file_path, str):
328 raise TypeError("file_path must be a string")
329 ret = graph_lib.GeApiWrapper_Graph_LoadFromAir(self._handle, file_path.encode("utf-8"))
330 if ret != 0:
331 raise RuntimeError("Failed to load graph from AIR format")
333 def remove_node(self, node: Node) -> None:
334 """Remove node.
336 Args:
337 node: node to be removed.
339 Raises:
340 TypeError: If node is not a Node.
341 RuntimeError: If removing node operation fails.
342 """
343 if not isinstance(node, Node):
344 raise TypeError("node must be a Node")
345 ret = graph_lib.GeApiWrapper_Graph_RemoveNode(self._handle, node._handle)
346 if ret != 0: # GRAPH_SUCCESS
347 raise RuntimeError(f"Failed to remove Node {node.name} from Graph {self.name}")
349 def remove_edge(self, src_node: Node, src_port_index: int, dst_node: Node, dst_port_index: int) -> None:
350 """Remove edge.
352 Args:
353 src_node: source node of edge
354 src_port_index: source port index of edge. If removing control edge, should be set to -1
355 dst_node: destination node of edge
356 dst_port_index: destination port index of edge. If removing control edge, should be set to -1
359 Raises:
360 TypeError: If src_node, dst_node are not nodes or src_port_index,dst_port_index are not integers.
361 RuntimeError: If removing edge operation fails.
363 Example:
364 >>> remove_edge(src_node, -1, dst_node, -1)
365 >>> remove_edge()
366 """
367 if not isinstance(src_node, Node):
368 raise TypeError("src_node must be a Node")
369 if not isinstance(src_port_index, int):
370 raise TypeError("src_port_index must be an integer")
371 if not isinstance(dst_node, Node):
372 raise TypeError("dst_node must be a Node")
373 if not isinstance(dst_port_index, int):
374 raise TypeError("dst_port_index must be an integer")
376 ret = graph_lib.GeApiWrapper_Graph_RemoveEdge(
377 self._handle,
378 src_node._handle,
379 src_port_index,
380 dst_node._handle,
381 dst_port_index,
382 )
383 if ret != 0: # GRAPH_SUCCESS
384 raise RuntimeError(
385 f"Failed to remove Edge from Node {src_node.name}, Port Index {src_port_index} to Node {dst_node.name}, Port Index {dst_port_index}"
386 )
388 def add_data_edge(self, src_node: Node, src_port_index: int, dst_node: Node, dst_port_index: int) -> None:
389 """Add data edge.
391 Args:
392 src_node: source node of edge
393 src_port_index: source port index of edge.
394 dst_node: destination node of edge
395 dst_port_index: destination port index of edge.
398 Raises:
399 TypeError: If src_node, dst_node are not nodes or src_port_index,dst_port_index are not integers.
400 RuntimeError: If adding data edge operation fails.
401 """
402 if not isinstance(src_node, Node):
403 raise TypeError("src_node must be a Node")
404 if not isinstance(src_port_index, int):
405 raise TypeError("src_port_index must be an integer")
406 if not isinstance(dst_node, Node):
407 raise TypeError("dst_node must be a Node")
408 if not isinstance(dst_port_index, int):
409 raise TypeError("dst_port_index must be an integer")
411 ret = graph_lib.GeApiWrapper_Graph_AddDataEdge(
412 self._handle,
413 src_node._handle,
414 src_port_index,
415 dst_node._handle,
416 dst_port_index,
417 )
418 if ret != 0: # GRAPH_SUCCESS
419 raise RuntimeError(
420 f"Failed to add DataEdge from Node {src_node.name}, Port Index {src_port_index} to Node {dst_node.name}, Port Index {dst_port_index}"
421 )
423 def add_control_edge(self, src_node: Node, dst_node: Node) -> None:
424 """Add control edge.
426 Args:
427 src_node: source node of edge
428 dst_node: destination node of edge
431 Raises:
432 TypeError: If src_node, dst_node are not nodes .
433 RuntimeError: If adding control edge operation fails.
434 """
435 if not isinstance(src_node, Node):
436 raise TypeError("src_node must be a Node")
437 if not isinstance(dst_node, Node):
438 raise TypeError("dst_node must be a Node")
440 ret = graph_lib.GeApiWrapper_Graph_AddControlEdge(self._handle, src_node._handle, dst_node._handle)
441 if ret != 0: # GRAPH_SUCCESS
442 raise RuntimeError(f"Failed to add Control from Node {src_node.name} to Node {dst_node.name}")
444 def find_node_by_name(self, name: str) -> Node:
445 """Find node by name
447 Args:
448 name: node name.
450 Returns:
451 Node found.
453 Raises:
454 TypeError: If name is not string.
455 RuntimeError: If finding node by name operation fails.
456 """
457 if not isinstance(name, str):
458 raise TypeError("name must be a string")
460 name_bytes = name.encode("utf-8")
461 node = ctypes.c_void_p()
462 ret = graph_lib.GeApiWrapper_Graph_FindNodeByName(self._handle, name_bytes, ctypes.byref(node))
463 if ret != 0: # GRAPH_SUCCESS
464 raise RuntimeError(f"Failed to find Node name {name}")
465 return Node._create_from(node)
467 def get_all_subgraphs(self) -> List["Graph"]:
468 """Get all subgraphs in the graph.
470 Returns:
471 List of Graph objects representing subgraphs.
472 """
473 subgraph_num = ctypes.c_size_t()
474 subgraphs = graph_lib.GeApiWrapper_Graph_GetAllSubgraphs(self._handle, ctypes.byref(subgraph_num))
475 if not subgraphs:
476 return []
478 try:
479 return [Graph._create_from(subgraphs[i]) for i in range(subgraph_num.value)]
480 finally:
481 graph_lib.GeApiWrapper_Graph_FreeGraphArray(subgraphs)
483 def get_subgraph(self, name: str) -> Optional["Graph"]:
484 """Get subgraph by name.
486 Args:
487 name: Subgraph name.
489 Returns:
490 Graph object representing the subgraph, or None if not found.
492 Raises:
493 TypeError: If name is not a string.
494 RuntimeError: If getting subgraph fails.
495 """
496 if not isinstance(name, str):
497 raise TypeError("name must be a string")
499 name_bytes = name.encode("utf-8")
500 subgraph_handle = graph_lib.GeApiWrapper_Graph_GetSubGraph(self._handle, name_bytes)
501 if not subgraph_handle:
502 return None
504 return Graph._create_from(subgraph_handle)
506 def add_subgraph(self, subgraph: "Graph") -> None:
507 """Add a subgraph to the graph.
509 The subgraph is indexed by its name. Subgraph names must be unique within the parent graph;
510 attempting to add a subgraph whose name already exists will fail.
512 Args:
513 subgraph: Graph object to be added as a subgraph.
515 Raises:
516 TypeError: If subgraph is not a Graph.
517 RuntimeError: If adding subgraph fails.
518 """
519 if not isinstance(subgraph, Graph):
520 raise TypeError("subgraph must be a Graph")
522 ret = graph_lib.GeApiWrapper_Graph_AddSubGraph(self._handle, subgraph._handle)
523 if ret != 0: # GRAPH_SUCCESS
524 raise RuntimeError(f"Failed to add subgraph '{subgraph.name}' to graph '{self.name}'")
526 def remove_subgraph(self, name: str) -> None:
527 """Remove subgraph by name.
529 Args:
530 name: Subgraph name.
532 Raises:
533 TypeError: If name is not a string.
534 RuntimeError: If removing subgraph fails.
535 """
536 if not isinstance(name, str):
537 raise TypeError("name must be a string")
539 name_bytes = name.encode("utf-8")
540 ret = graph_lib.GeApiWrapper_Graph_RemoveSubgraph(self._handle, name_bytes)
541 if ret != 0: # GRAPH_SUCCESS
542 raise RuntimeError(f"Failed to remove subgraph '{name}' from graph '{self.name}'")