Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/session/session.py: 85%
169 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# -----------------------------------------------------------------------------------------------------------
5# Copyright (c) 2025 Huawei Technologies Co., Ltd.
6# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
7# CANN Open Software License Agreement Version 2.0 (the "License").
8# Please refer to the License for details. You may not use this file except in compliance with the License.
9# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
10# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
11# See LICENSE in the root of the software repository for the full text of the License.
12# -----------------------------------------------------------------------------------------------------------
14"""Session module for GraphEngine graph operations."""
16import ctypes
17import threading
18import weakref
19from typing import List, Optional
21from ge._capi._allocator_callback_adapter import (
22 create_allocator_c_callbacks,
23 rollback_allocator_c_callbacks,
24)
25from ge._capi.pysession_wrapper import session_lib
26from ge.allocator import Allocator
27from ge.error import raise_ge_error
28from ge.graph.graph import Graph
29from ge.graph.tensor import Tensor
31_default_allocator_lock = threading.RLock()
32# stream -> weakref(session) that currently holds the default allocator for that stream.
33_stream_to_default_allocator_owner: dict = {}
36def _is_default_allocator_owner(session, stream: int) -> bool:
37 owner_ref = _stream_to_default_allocator_owner.get(stream)
38 if owner_ref is None:
39 return False
41 owner = owner_ref()
42 if owner is None:
43 _stream_to_default_allocator_owner.pop(stream, None)
44 return False
46 return owner is session
49def _str_list_to_c_array(python_list: list):
50 """
51 Convert python list to c array
53 Parameters:
54 python_list: python list.
56 Returns:
57 c_array: c array_ptr.
58 """
59 size = len(python_list)
60 c_array = (ctypes.c_char_p * size)()
61 for i, item in enumerate(python_list):
62 c_array[i] = item.encode("utf-8")
63 return c_array
66def _tensor_list_to_c_array(tensors: List[Tensor]):
67 handles = [ctypes.cast(tensor._handle, ctypes.c_void_p) for tensor in tensors]
68 arr_type = ctypes.c_void_p * len(handles)
69 return arr_type(*handles)
72class Session:
73 """Session class for session operations.
75 This class provides a Pythonic interface for session operations
76 using the session C API.
77 """
79 def __init__(self, options: Optional[dict] = None) -> None:
80 """Session Initialize a Session"""
81 self._handle = None
82 self._owns_handle = False
83 self._default_allocator_streams = set()
84 if options is None:
85 self._handle = session_lib.GeApiWrapper_Session_CreateSession()
86 elif isinstance(options, dict):
87 keys = [k for k in options.keys()]
88 values = [v for v in options.values()]
89 c_array_key = _str_list_to_c_array(keys)
90 c_array_value = _str_list_to_c_array(values)
91 c_array_key_ptr = ctypes.cast(c_array_key, ctypes.POINTER(ctypes.c_char_p))
92 c_value_key_ptr = ctypes.cast(c_array_value, ctypes.POINTER(ctypes.c_char_p))
93 self._handle = session_lib.GeApiWrapper_Session_CreateSessionWithOptions(
94 c_array_key_ptr, c_value_key_ptr, len(keys)
95 )
96 else:
97 raise TypeError("option must be a dictionary")
98 if not self._handle:
99 raise_ge_error("CreateSession")
100 self._owns_handle = True
102 def __del__(self) -> None:
103 """Clean up resources."""
104 self._unregister_default_allocators()
105 self._default_allocator_streams.clear()
106 if self._owns_handle:
107 session_lib.GeApiWrapper_Session_DestroySession(self._handle)
108 self._handle = None
110 def __copy__(self) -> None:
111 """Copy is not supported."""
112 raise RuntimeError("Session does not support copy")
114 def __deepcopy__(self, session) -> None:
115 """Deep copy is not supported."""
116 raise RuntimeError("Session does not support deepcopy")
118 def register_external_allocator(self, stream: int, allocator: Allocator) -> None:
119 """Register an external allocator for the given stream.
121 Args:
122 stream: Stream address.
123 allocator: An Allocator subclass instance.
124 """
125 if not isinstance(stream, int):
126 raise TypeError("stream must be an integer")
127 if not isinstance(allocator, Allocator):
128 raise TypeError("allocator must be an Allocator instance")
129 cb, prevent_gc_key, c_on_allocator_destroy = create_allocator_c_callbacks(allocator)
130 with _default_allocator_lock:
131 ret = session_lib.GeApiWrapper_Session_RegisterExternalAllocator(
132 self._handle,
133 ctypes.c_void_p(stream),
134 cb.c_malloc,
135 cb.c_free,
136 cb.c_get_addr,
137 c_on_allocator_destroy,
138 ctypes.c_void_p(prevent_gc_key),
139 )
140 if ret == 0:
141 self._default_allocator_streams.discard(stream)
142 _stream_to_default_allocator_owner.pop(stream, None)
143 if ret != 0:
144 rollback_allocator_c_callbacks(prevent_gc_key)
145 raise_ge_error("RegisterExternalAllocator", ret, stream=f"0x{stream:x}")
147 def unregister_external_allocator(self, stream: int) -> None:
148 """Unregister the external allocator for the given stream.
150 Args:
151 stream: Stream address.
152 """
153 if not isinstance(stream, int):
154 raise TypeError("stream must be an integer")
155 with _default_allocator_lock:
156 ret = session_lib.GeApiWrapper_Session_UnregisterExternalAllocator(self._handle, ctypes.c_void_p(stream))
157 if ret == 0:
158 self._default_allocator_streams.discard(stream)
159 _stream_to_default_allocator_owner.pop(stream, None)
160 if ret != 0:
161 raise_ge_error("UnregisterExternalAllocator", ret, stream=f"0x{stream:x}")
163 def add_graph(self, graph_id: int, add_graph: Graph, options: dict = None) -> None:
164 if not isinstance(graph_id, int):
165 raise TypeError("Graph_id must be an integer")
166 if not isinstance(add_graph, Graph):
167 raise TypeError("Add_graph must be a Graph")
168 if options is None:
169 ret = session_lib.GeApiWrapper_Session_AddGraph(self._handle, ctypes.c_uint32(graph_id), add_graph._handle)
170 elif not isinstance(options, dict):
171 raise TypeError("options must be a dictionary")
172 else:
173 keys = [k for k in options.keys()]
174 values = [v for v in options.values()]
175 c_array_key = _str_list_to_c_array(keys)
176 c_array_value = _str_list_to_c_array(values)
177 c_array_key_ptr = ctypes.cast(c_array_key, ctypes.POINTER(ctypes.c_char_p))
178 c_value_key_ptr = ctypes.cast(c_array_value, ctypes.POINTER(ctypes.c_char_p))
179 ret = session_lib.GeApiWrapper_Session_AddGraphWithOptions(
180 self._handle,
181 ctypes.c_uint32(graph_id),
182 add_graph._handle,
183 c_array_key_ptr,
184 c_value_key_ptr,
185 len(keys),
186 )
187 if ret != 0:
188 raise_ge_error("AddGraph", ret, graph_id=graph_id)
189 return ret
191 def remove_graph(self, graph_id: int) -> None:
192 if not isinstance(graph_id, int):
193 raise TypeError("Graph_id must be an integer")
194 ret = session_lib.GeApiWrapper_Session_RemoveGraph(self._handle, ctypes.c_uint32(graph_id))
195 if ret != 0:
196 raise_ge_error("RemoveGraph", ret, graph_id=graph_id)
198 def run_graph(self, graph_id: int, inputs: List[Tensor]) -> List[Tensor]:
199 if not isinstance(graph_id, int):
200 raise TypeError("Graph_id must be an integer")
201 if not isinstance(inputs, list):
202 raise TypeError("inputs must be a list of Tensor")
203 if not all(isinstance(input_tensor, Tensor) for input_tensor in inputs):
204 raise TypeError("All elements in inputs must be the type of Tensor")
205 arr = _tensor_list_to_c_array(inputs)
206 tensor_num = ctypes.c_size_t()
207 output_tensors = ctypes.POINTER(ctypes.c_void_p)()
208 ret = session_lib.GeApiWrapper_Session_RunGraph(
209 self._handle,
210 ctypes.c_uint32(graph_id),
211 arr,
212 len(inputs),
213 ctypes.byref(output_tensors),
214 ctypes.byref(tensor_num),
215 )
216 try:
217 if ret != 0:
218 raise_ge_error("RunGraph", ret, graph_id=graph_id)
219 return [Tensor._create_from(output_tensors[i]) for i in range(tensor_num.value)]
220 finally:
221 if output_tensors:
222 session_lib.GeApiWrapper_Session_FreeTensorArray(output_tensors)
224 def run_graph_with_stream_async(self, graph_id: int, stream: int, inputs: List[Tensor]) -> List[Tensor]:
225 """Run the graph asynchronously on the given stream and return output tensors.
227 Output tensor memory is allocated according to the following priority:
228 1. The external allocator registered via register_external_allocator(stream, allocator).
229 2. If no external allocator is registered, GE uses a built-in allocator automatically.
231 Args:
232 graph_id: Graph ID.
233 stream: Stream address.
234 inputs: List of input tensors.
236 Returns:
237 List of output tensors.
238 """
239 if not isinstance(graph_id, int):
240 raise TypeError("Graph_id must be an integer")
241 if not isinstance(stream, int):
242 raise TypeError("Stream must be an integer")
243 if not isinstance(inputs, list):
244 raise TypeError("inputs must be a list of Tensor")
245 if not all(isinstance(input_tensor, Tensor) for input_tensor in inputs):
246 raise TypeError("All elements in inputs must be the type of Tensor")
248 self._ensure_default_allocator(stream)
249 arr = _tensor_list_to_c_array(inputs)
250 tensor_num = ctypes.c_size_t()
251 output_tensors = ctypes.POINTER(ctypes.c_void_p)()
252 ret = session_lib.GeApiWrapper_Session_RunGraphWithStreamAsync(
253 self._handle,
254 ctypes.c_uint32(graph_id),
255 ctypes.c_void_p(stream),
256 arr,
257 len(inputs),
258 ctypes.byref(output_tensors),
259 ctypes.byref(tensor_num),
260 )
261 try:
262 if ret != 0:
263 raise_ge_error(
264 "RunGraphWithStreamAsync",
265 ret,
266 graph_id=graph_id,
267 stream=f"0x{stream:x}",
268 )
269 return [Tensor._create_from(output_tensors[i]) for i in range(tensor_num.value)]
270 finally:
271 if output_tensors:
272 session_lib.GeApiWrapper_Session_FreeTensorArray(output_tensors)
274 def _ensure_default_allocator(self, stream: int) -> None:
275 """Register a default allocator if none exists."""
276 with _default_allocator_lock:
277 has_external = session_lib.GeApiWrapper_HasExternalAllocator(ctypes.c_void_p(stream))
278 if stream in self._default_allocator_streams or has_external:
279 return
280 ret = session_lib.GeApiWrapper_Session_RegisterDefaultAllocator(self._handle, ctypes.c_void_p(stream))
281 if ret == 0:
282 self._default_allocator_streams.add(stream)
283 _stream_to_default_allocator_owner[stream] = weakref.ref(self)
284 if ret != 0:
285 raise_ge_error("RegisterDefaultAllocator", ret, stream=f"0x{stream:x}")
287 def _unregister_default_allocators(self) -> None:
288 if not self._default_allocator_streams or not session_lib.GeApiWrapper_IsGEInitialized():
289 return
291 with _default_allocator_lock:
292 for stream in list(self._default_allocator_streams):
293 if not _is_default_allocator_owner(self, stream):
294 continue
295 if not session_lib.GeApiWrapper_HasDefaultAllocator(ctypes.c_void_p(stream)):
296 _stream_to_default_allocator_owner.pop(stream, None)
297 continue
299 session_lib.GeApiWrapper_Session_UnregisterExternalAllocator(self._handle, ctypes.c_void_p(stream))
300 _stream_to_default_allocator_owner.pop(stream, None)