Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/llm_datadist_v1/llm_utils.py: 22%
236 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# -----------------------------------------------------------------------------------------------------------
12from dataclasses import dataclass
13from threading import Thread
14from typing import Dict, List, Optional, Tuple, Union
16from llm_datadist_v1 import llm_wrapper
18from .data_type import DataType, python_dtype_2_dwrapper_dtype
19from .llm_types import (
20 BlocksCacheKey,
21 CacheDesc,
22 CacheKey,
23 CacheKeyByIdAndIndex,
24 CacheTask,
25 KvCache,
26 LayerSynchronizer,
27 MemInfo,
28 Placement,
29 PushType,
30 TransferConfig,
31 TransferWithCacheKeyConfig,
32)
33from .status import (
34 LLMException,
35 LLMStatusCode,
36 code_2_status,
37 raise_if_false,
38 raise_if_true,
39)
40from .utils import log
41from .utils.utils import (
42 check_isinstance,
43 check_list_uint64,
44 check_type,
45 check_uint32,
46 check_uint64,
47)
49# UINT64_MAX
50_INVALID_ID = 2**64 - 1
51_MAX_DISPLAYED_REQ_ID_COUNT = 8
52_NUM_TENSORS_PER_LAYER = 2
55def clone_cache_desc(cache_desc: CacheDesc) -> CacheDesc:
56 return CacheDesc(
57 cache_desc.num_tensors,
58 cache_desc.shape[:],
59 cache_desc.data_type,
60 cache_desc.placement,
61 cache_desc.batch_dim,
62 cache_desc.seq_len_dim_index,
63 cache_desc.kv_tensor_format,
64 )
67def verify_cache_shape(shape: List[int]) -> None:
68 raise_if_false(len(shape) > 0, "scalar is not supported")
69 raise_if_false(0 not in shape, f"empty dimension is not supported, shape = {shape}")
70 for dim in shape:
71 if dim < -1:
72 raise LLMException(f"dim {dim} is not supported, shape = {shape}")
75def to_data_type(data_type_str: str) -> DataType:
76 for dtype in DataType:
77 if str(dtype.value) == data_type_str:
78 return dtype
79 raise LLMException(
80 f"unsupported dtype: {data_type_str}",
81 status_code=LLMStatusCode.LLM_PARAM_INVALID,
82 )
85def is_invalid_id(req_id_or_prefix_id: int) -> bool:
86 return (req_id_or_prefix_id < 0) or (req_id_or_prefix_id == _INVALID_ID)
89def is_valid_id(req_id_or_prefix_id: int) -> bool:
90 return not is_invalid_id(req_id_or_prefix_id)
93def calc_tensor_size(shape: List[int], data_type: DataType) -> int:
94 ge_data_type = python_dtype_2_dwrapper_dtype.get(data_type)
95 tensor_size = llm_wrapper.calc_tensor_size(shape, ge_data_type)
96 raise_if_false(
97 tensor_size >= 0,
98 "Failed to calc tensor size, shape = {0}, data_type = {1}",
99 shape,
100 data_type,
101 )
102 return tensor_size
105class CacheDescParser(object):
106 @staticmethod
107 def parse_by_options(options: Dict[str, str]) -> CacheDesc:
108 raise_if_false("llm.RefInputShapes" in options, "llm.RefInputShapes not set")
109 kv_shape_str = options["llm.RefInputShapes"]
110 kv_shapes = kv_shape_str.split(";")
111 raise_if_false(
112 len(set(kv_shapes)) == 1,
113 f"zero or multiple different shapes is not supported, llm.RefInputShapes = {kv_shapes}",
114 )
116 raise_if_false("llm.RefInputDtypes" in options, "llm.RefInputDtypes not set")
117 kv_data_type_str = options["llm.RefInputDtypes"]
118 kv_data_types = kv_data_type_str.split(";")
119 raise_if_false(
120 len(set(kv_data_types)) == 1,
121 f"zero or multiple different data types is not supported, llm.RefInputDtypes = {kv_data_types}",
122 )
123 raise_if_false(
124 len(kv_shapes) == len(kv_data_types),
125 f"kv_shapes num ({len(kv_shapes)}) mismatches that of kv_data_type ({len(kv_data_types)})",
126 )
127 kv_shape = list((int(dim) for dim in kv_shapes[0].split(",")))
128 verify_cache_shape(kv_shape)
129 if "llm.RefInputSeqLenDimIndex" in options:
130 seq_len_dim_index = int(options["llm.RefInputSeqLenDimIndex"])
131 log.info(f"get seq_len_dim_index from option, value = {seq_len_dim_index}")
132 raise_if_false(
133 0 <= seq_len_dim_index < len(kv_shape),
134 f"seq_len_dim_index ({seq_len_dim_index}) out of range, kv_shape = {kv_shape}",
135 )
136 elif -1 in kv_shape[1:]:
137 raise_if_false(
138 kv_shape[1:].count(-1) == 1,
139 f"can only have one dynamic dim apart from batch_dim(0), "
140 f"but got {kv_shape[1:].count(-1)}, shape = {kv_shape}",
141 )
142 seq_len_dim_index = kv_shape[1:].index(-1) + 1
143 log.info(f"seq_len_dim_index inferred by shape, value = {seq_len_dim_index}")
144 else:
145 seq_len_dim_index = -1
146 log.info(
147 f"llm.RefInputSeqLenDimIndex not set and can not infer by shape,seq_len_dim_index = {seq_len_dim_index}"
148 )
149 kv_tensor_format = options.get("llm.kvTensorFormat", None)
150 kv_data_type = to_data_type(kv_data_types[0])
151 cache_desc = CacheDesc(
152 len(kv_shapes),
153 kv_shape,
154 kv_data_type,
155 Placement.DEVICE,
156 0,
157 seq_len_dim_index,
158 kv_tensor_format,
159 )
160 log.info(f"parse cache_desc from option, value = {cache_desc}")
161 return cache_desc
164def build_npu_tensors(kv_cache: KvCache, device_index: int = 0) -> List[int]:
165 return llm_wrapper.build_npu_tensors(
166 kv_cache.cache_desc.shape,
167 kv_cache.cache_desc.data_type.value,
168 kv_cache.cache_desc.size,
169 kv_cache.per_device_tensor_addrs[device_index],
170 )
173def pack_cache_desc(cache_desc: CacheDesc) -> Tuple[int, int, int, List[int], int, int]:
174 return (
175 cache_desc.num_tensors,
176 cache_desc.data_type.value,
177 cache_desc.seq_len_dim_index,
178 cache_desc.shape,
179 cache_desc.placement.value,
180 cache_desc._is_blocks,
181 )
184def pack_cache_key(cache_key: CacheKey) -> Tuple[int, int, int, int, int, int, bool]:
185 return (
186 cache_key.prompt_cluster_id,
187 -1,
188 0,
189 cache_key.req_id,
190 cache_key.prefix_id,
191 cache_key.model_id,
192 False,
193 )
196def pack_cache_key_by_id(
197 cache_key: CacheKeyByIdAndIndex,
198) -> Tuple[int, int, int, int, int, int, bool]:
199 return (
200 cache_key.cluster_id,
201 cache_key.cache_id,
202 cache_key.batch_index,
203 _INVALID_ID,
204 _INVALID_ID,
205 0,
206 False,
207 )
210def pack_block_cache_key(
211 cache_key: BlocksCacheKey,
212) -> Tuple[int, int, int, int, int, int, bool]:
213 return (
214 cache_key.prompt_cluster_id,
215 -1,
216 0,
217 _INVALID_ID,
218 _INVALID_ID,
219 cache_key.model_id,
220 True,
221 )
224def pack_mem_info(mem_info: MemInfo) -> Tuple[int, int, int]:
225 return (mem_info.mem_type.value, mem_info.addr, mem_info.size)
228@dataclass
229class TransferCacheParameters:
230 src_cache: "Union[Cache, KvCache]"
231 transfer_configs: Union[
232 List[Union[TransferConfig, TransferWithCacheKeyConfig]],
233 Tuple[Union[TransferConfig, TransferWithCacheKeyConfig]],
234 ]
235 src_block_indices: Optional[List[int]] = None
236 dst_block_indices: Optional[List[int]] = None
237 dst_block_memory_size: Optional[int] = None
240class TransferCacheJob:
241 task_id = 0
243 def __init__(
244 self,
245 params: TransferCacheParameters,
246 layer_synchronizer: LayerSynchronizer,
247 transfer_cache_func,
248 ):
249 self._transfer_configs = params.transfer_configs
250 self._src_block_indices = params.src_block_indices
251 self._dst_block_indices = params.dst_block_indices
252 self._dst_block_memory_size = params.dst_block_memory_size
253 self._cache_id = params.src_cache.cache_id
254 self._cache_desc = params.src_cache.cache_desc
255 self._num_layers = 0
256 self._rets: Dict[int, LLMStatusCode] = {}
257 self._timeout_in_millis: Optional[int] = None
258 self._layer_synchronizer = layer_synchronizer
259 self._transfer_cache_func = transfer_cache_func
261 def init(self):
262 raise_if_false(
263 self._cache_desc.num_tensors % 2 == 0,
264 "cache_desc.num_tensors ({0}) is not even",
265 self._cache_desc.num_tensors,
266 )
267 self._num_layers = self._cache_desc.num_tensors // 2
268 for transfer_config in self._transfer_configs:
269 if transfer_config.src_layer_range is None:
270 transfer_config.src_layer_range = range(0, self._num_layers)
271 self.check_transfer_config(transfer_config)
273 def transfer_layers(self):
274 for layer_i, src_layer_index in enumerate(range(self._num_layers)):
275 to_transfer = [config for config in self._transfer_configs if src_layer_index in config.src_layer_range]
276 if not to_transfer:
277 log.info("src_layer %d sends to no destination", src_layer_index)
278 continue
279 if not self._layer_synchronizer.synchronize_layer(src_layer_index, self._timeout_in_millis):
280 log.error(f"Failed to synchronize layer {src_layer_index}")
281 for config in to_transfer:
282 # 错误码待添加
283 self._rets[config.dst_cluster_id] = LLMStatusCode.LLM_PARAM_INVALID
284 return
285 for config in to_transfer:
286 dst_layer_index = (
287 src_layer_index if config.dst_layer_range is None else list(config.dst_layer_range)[layer_i]
288 )
289 ret = self.transfer_layer(src_layer_index, dst_layer_index, config)
290 if ret != LLMStatusCode.LLM_SUCCESS:
291 log.error(f"Failed to transfer layer {src_layer_index} to dst_cluster_id={config.dst_cluster_id}")
292 self._rets[config.dst_cluster_id] = ret
293 return
294 log.info(f"transfer layer {src_layer_index} to dst_cluster_id={config.dst_cluster_id} success")
295 if src_layer_index == config.src_layer_range.stop - 1:
296 self._rets[config.dst_cluster_id] = LLMStatusCode.LLM_SUCCESS
297 log.info(f"transfer all layers to dst_cluster_id={config.dst_cluster_id} finished")
298 log.info("transfer all layers finished")
300 def transfer_layer(
301 self,
302 src_layer_index: int,
303 dst_layer_idx,
304 transfer_config: Union[TransferConfig, TransferWithCacheKeyConfig],
305 ) -> LLMStatusCode:
306 if check_type(transfer_config, TransferConfig):
307 dst_layer_index = src_layer_index - transfer_config.src_layer_range.start
308 dst_addrs = transfer_config.dst_addrs[
309 dst_layer_index * _NUM_TENSORS_PER_LAYER : dst_layer_index * _NUM_TENSORS_PER_LAYER
310 + _NUM_TENSORS_PER_LAYER
311 ]
312 transfer_config = (
313 self._cache_id,
314 transfer_config.src_batch_index,
315 src_layer_index,
316 dst_addrs,
317 transfer_config.dst_cluster_id,
318 0,
319 0,
320 PushType.NO_CACHE_KEY.value,
321 src_layer_index,
322 2,
323 )
324 else:
325 if check_type(transfer_config.cache_key, BlocksCacheKey):
326 transfer_config = (
327 self._cache_id,
328 0,
329 src_layer_index,
330 [],
331 transfer_config.cache_key.cluster_id,
332 transfer_config.cache_key.model_id,
333 0,
334 PushType.BLOCKS_CACHE_KEY.value,
335 dst_layer_idx,
336 2,
337 )
338 elif check_type(transfer_config.cache_key, CacheKeyByIdAndIndex):
339 transfer_config = (
340 self._cache_id,
341 transfer_config.src_batch_index,
342 src_layer_index,
343 [],
344 transfer_config.cache_key.cluster_id,
345 transfer_config.cache_key.cache_id,
346 transfer_config.cache_key.batch_index,
347 PushType.CACHE_KEY_BY_ID.value,
348 dst_layer_idx,
349 2,
350 )
351 block_config = (
352 self._dst_block_memory_size if self._dst_block_memory_size is not None else 0,
353 self._src_block_indices if self._src_block_indices is not None else [],
354 self._dst_block_indices if self._dst_block_indices is not None else [],
355 )
356 ret = self._transfer_cache_func(TransferCacheJob.task_id, transfer_config, block_config)
357 TransferCacheJob.task_id += 1
358 return code_2_status(ret)
360 def check_transfer_config(self, transfer_config: Union[TransferConfig, TransferWithCacheKeyConfig]):
361 if self._src_block_indices:
362 raise_if_false(
363 transfer_config.src_batch_index == 0,
364 "Invalid TransferConfig, src_batch_index ({0}) != 0 while src is blocks",
365 transfer_config.src_batch_index,
366 )
367 raise_if_false(
368 0 <= transfer_config.src_batch_index < self._cache_desc.batch_size,
369 "Invalid TransferConfig, src_batch_index ({0}) out of range: [0, {1})",
370 transfer_config.src_batch_index,
371 self._cache_desc.batch_size,
372 )
373 raise_if_false(
374 0 <= transfer_config.src_layer_range.start < transfer_config.src_layer_range.stop <= self._num_layers,
375 "src_layer_range: {0} out of range, src_layer_num = {1}",
376 transfer_config.src_layer_range,
377 self._num_layers,
378 )
379 num_tensors_to_transfer = (
380 transfer_config.src_layer_range.stop - transfer_config.src_layer_range.start
381 ) * _NUM_TENSORS_PER_LAYER
382 if check_type(transfer_config, TransferConfig):
383 raise_if_false(
384 len(transfer_config.dst_addrs) == num_tensors_to_transfer,
385 "expect {0} dst_addrs, but len(dst_addrs) = {1}, range = {2}",
386 num_tensors_to_transfer,
387 len(transfer_config.dst_addrs),
388 transfer_config.src_layer_range,
389 )
391 def get_results(self) -> List[LLMStatusCode]:
392 rets = []
393 for config in self._transfer_configs:
394 ret = self._rets.get(config.dst_cluster_id)
395 rets.append(ret)
396 return rets
398 def num_transfer_configs(self):
399 return len(self._transfer_configs)
402class TransferAsyncThread(Thread):
403 def __init__(self, transfer_job: TransferCacheJob, default_err_code=LLMStatusCode.LLM_TIMEOUT):
404 super().__init__()
405 self._transfer_job = transfer_job
406 self._rets = []
407 self._default_err_code = default_err_code
409 def run(self):
410 self._transfer_job.transfer_layers()
412 def get_results(self, timeout) -> List[LLMStatusCode]:
413 self.join(timeout)
414 if self.is_alive():
415 return [self._default_err_code] * self._transfer_job.num_transfer_configs()
416 return self._transfer_job.get_results()
418 def get(self, timeout) -> LLMStatusCode:
419 self.join(timeout)
420 if self.is_alive():
421 return self._default_err_code
422 rets = self._transfer_job.get_results()
423 for ret in rets:
424 if ret != LLMStatusCode.LLM_SUCCESS and ret is not None:
425 return ret
426 return LLMStatusCode.LLM_SUCCESS
429def _check_block_indices(arg_name, arg_value):
430 if arg_value is not None:
431 check_isinstance(arg_name, arg_value, [list, tuple], int)
432 check_list_uint64(arg_name, arg_value)
435def transfer_cache_async(
436 params: TransferCacheParameters,
437 layer_synchronizer: LayerSynchronizer,
438 transfer_cache_func,
439 default_error_code=LLMStatusCode.LLM_TIMEOUT,
440 enable_remote_cache=False,
441) -> CacheTask:
442 _check_block_indices("dst_block_indices", params.dst_block_indices)
443 _check_block_indices("src_block_indices", params.src_block_indices)
444 if params.dst_block_memory_size is not None:
445 check_uint64("dst_block_memory_size", params.dst_block_memory_size)
446 if params.src_block_indices: # src is blocks
447 raise_if_false(params.dst_block_indices, "transfer from blocks to cache is not supported")
448 raise_if_false(
449 len(params.src_block_indices) == len(params.dst_block_indices),
450 "num_block_indices mismatches, src_num = {0}, dst_num = {1}",
451 len(params.src_block_indices),
452 len(params.dst_block_indices),
453 )
454 else: # src is cache
455 raise_if_true(
456 (transfer_cache_func == llm_wrapper.transfer_cache) and (params.dst_block_indices is not None),
457 "transfer from cache to blocks is not supported",
458 )
459 if params.dst_block_indices:
460 raise_if_false(
461 params.dst_block_memory_size is not None,
462 "dst_block_memory_size must be set when transfer from cache to blocks",
463 )
464 check_isinstance("layer_synchronizer", layer_synchronizer, LayerSynchronizer, allow_none=False)
465 if not enable_remote_cache:
466 check_isinstance(
467 "transfer_configs",
468 params.transfer_configs,
469 [list, tuple],
470 TransferConfig,
471 allow_none=False,
472 )
473 else:
474 check_isinstance(
475 "transfer_configs",
476 params.transfer_configs,
477 [list, tuple],
478 TransferWithCacheKeyConfig,
479 "While enable_remote_cache_accessible is True, ",
480 allow_none=False,
481 )
482 raise_if_false(
483 params.dst_block_indices or params.dst_block_memory_size in (None, 0),
484 "dst_block_memory_size ({0}) is neither None nor 0 while dst is not blocks",
485 params.dst_block_memory_size,
486 )
487 transfer_job = TransferCacheJob(params, layer_synchronizer, transfer_cache_func)
488 transfer_job.init()
489 transfer_thread = TransferAsyncThread(transfer_job, default_error_code)
490 transfer_thread.start()
491 log.info("[transfer_cache_async] async task start")
492 cache_task = CacheTask(transfer_thread)
493 return cache_task
496def layer_range_to_tensor_indices(src_layer_range: range, dst_layer_range: range, tensor_num_per_layer: int = 2):
497 check_isinstance("src_layer_range", src_layer_range, range)
498 check_isinstance("dst_layer_range", dst_layer_range, range)
499 raise_if_true(
500 (src_layer_range is not None) and (src_layer_range.step != 1),
501 "param check failed, src_layer_range step must be 1.",
502 )
503 raise_if_true(
504 (dst_layer_range is not None) and (dst_layer_range.step != 1),
505 "param check failed, dst_layer_range step must be 1.",
506 )
507 src_layer_indices = [] if src_layer_range is None else list(src_layer_range)
508 dst_layer_indices = [] if dst_layer_range is None else list(dst_layer_range)
510 # 默认一层有两个tensor
511 one_layer_tensor_num = tensor_num_per_layer
512 src_tensor_indices = []
513 if len(src_layer_indices) != 0:
514 raise_if_true(
515 src_layer_indices[0] < 0,
516 "src_layer_range is invalid, the start value:{0} is < 0",
517 src_layer_indices[0],
518 )
519 check_uint32("src_layer_range", src_layer_indices[0])
520 check_uint32("src_layer_range", src_layer_indices[-1])
521 src_tensor_start_index = src_layer_indices[0] * one_layer_tensor_num
522 for i in range(len(src_layer_indices) * one_layer_tensor_num):
523 src_tensor_indices.append(src_tensor_start_index + i)
525 dst_tensor_indices = []
526 if len(dst_layer_indices) != 0:
527 raise_if_true(
528 dst_layer_indices[0] < 0,
529 "dst_layer_range is invalid, the start value:{0} is < 0",
530 dst_layer_indices[0],
531 )
532 check_uint32("dst_layer_range", dst_layer_indices[0])
533 check_uint32("dst_layer_range", dst_layer_indices[-1])
534 dst_tensor_start_index = dst_layer_indices[0] * one_layer_tensor_num
535 for i in range(len(dst_layer_indices) * one_layer_tensor_num):
536 dst_tensor_indices.append(dst_tensor_start_index + i)
538 return src_tensor_indices, dst_tensor_indices
541def parse_listen_ip_info(listen_ip_info: str) -> (str, int):
542 check_isinstance("listen_ip_info", listen_ip_info, [str])
543 ip_and_port = listen_ip_info.split(":")
544 raise_if_false(
545 len(ip_and_port) == 2,
546 f'llm.listenIpInfo "{ip_and_port}" is invalid, format should be "ip:port"',
547 )
548 ip = ip_and_port[0]
549 port = int(ip_and_port[1])
550 return ip, port