Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/build_st/python_tests/v1/st/test_v1_llm_datadist.py: 99%
412 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-27 10:03 +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# -----------------------------------------------------------------------------------------------------------
14import json
15import os
16import time
17import unittest
18from typing import Optional
20import numpy as np
21from llm_datadist_v1 import *
22from llm_datadist_v1.config import EngineConfig
23from llm_datadist_v1.llm_datadist import _shutdown_handler
24from llm_datadist_v1.llm_types import BlocksCacheKey, KvCache, Placement
25from llm_datadist_v1.llm_utils import TransferCacheJob, TransferCacheParameters
27_INVALID_ID = 2**64 - 1
29_TEST_BASE_DIR = "../tests/dflow/llm_datadist/st/testcase/llm_datadist"
32class LayerSynchronizerImpl(LayerSynchronizer):
33 def __init__(self, ret=True):
34 self._ret = ret
36 def synchronize_layer(self, layer_index: int, timeout_in_millis: Optional[int]) -> bool:
37 if layer_index == 0:
38 time.sleep(0.3)
39 return self._ret
42class MockTransferCacheJob(TransferCacheJob):
43 def __init__(self, params: TransferCacheParameters) -> None:
44 super().__init__(params, LayerSynchronizerImpl(True), None)
46 def transfer_layer(self, src_layer_index: int, dst_layer_idx, transfer_config: TransferConfig) -> LLMStatusCode:
47 return LLMStatusCode.LLM_WAIT_PROCESS_TIMEOUT
50class LlmEngineV2St(unittest.TestCase):
51 def setUp(self) -> None:
52 os.environ["ASCEND_GLOBAL_LOG_LEVEL"] = "1"
53 print("Begin ", self._testMethodName)
55 def tearDown(self) -> None:
56 os.environ.pop("RESOURCE_CONFIG_PATH", None)
57 _shutdown_handler()
58 print("End ", self._testMethodName)
60 @staticmethod
61 def _engine_options(is_prompt: bool, cluster_id: int = 0, rank_id: int = -1, resource_path: str = ""):
62 cluster_info = {
63 "cluster_id": cluster_id,
64 "logic_device_id": ["0:0:0:0", "0:0:1:0", "0:0:2:0", "0:0:3:0"],
65 }
66 if is_prompt:
67 cluster_info["listen_ip_info"] = [
68 {"ip": 0, "port": 26000},
69 {"ip": 1, "port": 26000},
70 {"ip": 2, "port": 26000},
71 {"ip": 3, "port": 26000},
72 ]
73 engine_options = {"llm.ClusterInfo": json.dumps(cluster_info)}
74 if rank_id != -1:
75 engine_options["ge.exec.rankId"] = str(rank_id)
76 if resource_path != "":
77 engine_options["ge.resourceConfigPath"] = resource_path
78 return engine_options
80 def test_prompt_cache_ops(self):
81 cluster_id = 0
82 engine = LLMDataDist(LLMRole.PROMPT, cluster_id)
83 engine.init(LlmEngineV2St._engine_options(True, cluster_id))
84 kv_cache_manager = engine.kv_cache_manager
85 cache_desc = CacheDesc(80, [2, 8], DataType.DT_FLOAT16)
86 cache_key = CacheKey(0, 0, 1)
87 kv_cache = kv_cache_manager.allocate_cache(cache_desc, [cache_key])
88 with self.assertRaises(LLMException):
89 tensors = kv_cache_manager.get_cache_tensors(kv_cache, -1)
90 tensors = kv_cache_manager.get_cache_tensors(kv_cache, 0)
91 print(f"kv_cache: {kv_cache}")
92 print(f"tensors: {Tensor(tensors[0]).numpy()}")
93 kv_cache_manager.deallocate_cache(kv_cache)
94 kv_cache_manager.remove_cache_key(cache_key)
95 engine.finalize()
97 def test_decoder_cache_ops(self):
98 cluster_id = 0
99 llm_engine = LLMDataDist(LLMRole.DECODER, cluster_id)
100 llm_engine.init(LlmEngineV2St._engine_options(False, cluster_id))
101 kv_cache_manager = llm_engine.kv_cache_manager
102 cache_desc = CacheDesc(80, [2, 8], DataType.DT_FLOAT16)
103 cache_key = CacheKey(0, 1, 1)
104 # mock prompt allocate kv
105 kv_cache_manager._role = LLMRole.PROMPT
106 kv_cache = kv_cache_manager.allocate_cache(cache_desc, [cache_key])
107 kv_cache_manager._role = LLMRole.DECODER
109 dst_kv_cache = kv_cache_manager.allocate_cache(cache_desc)
110 print(f"kv_cache: {kv_cache}")
111 with self.assertRaises(LLMException):
112 kv_cache_manager.pull_cache(cache_key, dst_kv_cache, 0, -2)
113 kv_cache_manager.pull_cache(cache_key, dst_kv_cache, src_cache_offset=0, dst_cache_offset=0)
114 cache_key_id = CacheKeyByIdAndIndex(cluster_id, kv_cache.cache_id, 0)
115 kv_cache_manager.pull_cache(cache_key_id, dst_kv_cache, 0)
116 kv_cache_manager.copy_cache(dst_kv_cache, kv_cache)
117 with self.assertRaises(LLMException):
118 kv_cache_manager.copy_cache(dst_kv_cache, kv_cache, 0, 0, 0, -2)
119 with self.assertRaises(LLMException):
120 kv_cache_manager.copy_cache(dst_kv_cache, kv_cache, 0, 0, 0, 0)
121 kv_cache_manager.deallocate_cache(kv_cache)
122 # test use after deallocated
123 with self.assertRaises(LLMException) as ex:
124 kv_cache_manager.pull_cache(cache_key, kv_cache)
125 self.assertEqual(ex.exception.status_code, LLMStatusCode.LLM_KV_CACHE_NOT_EXIST)
126 kv_cache_manager.deallocate_cache(dst_kv_cache)
127 llm_engine.finalize()
129 def test_cluster_ops(self):
130 cluster_id = 0
131 llm_engine = LLMDataDist(LLMRole.DECODER, cluster_id)
132 cluster_info = {"cluster_id": 1, "logic_device_id": ["0:0:0:0"]}
133 engine_options = {"llm.ClusterInfo": json.dumps(cluster_info)}
134 llm_engine.init(engine_options)
136 cluster = LLMClusterInfo()
137 cluster.remote_cluster_id = 1
138 cluster.append_local_ip_info(1, 26000)
139 cluster.append_remote_ip_info(1, 26000)
140 ret, rets = llm_engine.link_clusters([cluster], 5000)
141 self.assertEqual(ret, LLMStatusCode.LLM_SUCCESS)
142 ret, rets = llm_engine.unlink_clusters([cluster], 5000)
143 self.assertEqual(ret, LLMStatusCode.LLM_SUCCESS)
144 llm_engine.check_link_status(1)
146 def test_init_and_finalize(self):
147 cluster_id = 0
148 decoder_engine = LLMDataDist(LLMRole.DECODER, cluster_id)
149 decoder_engine.finalize()
150 decoder_engine.init(self._engine_options(False, cluster_id, 1))
151 decoder_engine.init(self._engine_options(False, cluster_id, 1))
152 decoder_engine.finalize()
153 decoder_engine.finalize()
154 decoder_engine_2 = LLMDataDist(LLMRole.DECODER, cluster_id)
155 decoder_engine_2.init(self._engine_options(False, cluster_id, 1))
157 def test_init_witch_device_id_and_rank_id(self):
158 llm_config = LLMConfig()
159 llm_config.device_id = 1
160 llm_config.listen_ip_info = "127.0.0.1:26000"
161 llm_config.ge_options = {"ge.exec.rankId": "1"}
162 engine_options = llm_config.generate_options()
163 print("engine_options:", engine_options)
164 prompt_engine = LLMDataDist(LLMRole.PROMPT, 0)
165 prompt_engine.init(engine_options)
166 EngineConfig.from_engine_options(True, engine_options)
168 def test_simple_option(self):
169 cluster_id = 0
170 prompt_engine = LLMDataDist(LLMRole.PROMPT, cluster_id)
171 llm_config = LLMConfig()
172 llm_config.device_id = 1
173 llm_config.listen_ip_info = "127.0.0.1:26000"
174 llm_config.deploy_res_path = "./"
175 llm_config.ge_options = {"ge.flowGraphMemMaxSize": "10000000"}
176 engine_options = llm_config.generate_options()
177 print("engine_options:", engine_options)
178 prompt_engine.init(engine_options)
180 def test_shutdown_failed(self):
181 cluster_id = 0
182 decoder_engine = LLMDataDist(LLMRole.DECODER, cluster_id)
183 decoder_engine.init(self._engine_options(False, cluster_id, 1))
185 # 模拟被其它框架Finalize
186 decoder_engine._llm_datadist.finalize()
187 _shutdown_handler()
188 decoder_engine._is_initialized = False
189 LLMDataDist.llm_engine_instance = None
191 def test_switch_role_failed_option_not_set(self):
192 engine = LLMDataDist(LLMRole.MIX, 0)
193 options = {"ge.exec.deviceId": "0"}
194 engine.init(options)
195 try:
196 engine.switch_role(LLMRole.DECODER)
197 except LLMException as e:
198 self.assertEqual(e.status_code, LLMStatusCode.LLM_FEATURE_NOT_ENABLED)
200 def test_switch_role_failed_identical_role(self):
201 engine = LLMDataDist(LLMRole.MIX, 0)
202 options = {"llm.EnableSwitchRole": "1", "ge.exec.deviceId": "0"}
203 engine.init(options)
204 with self.assertRaises(LLMException):
205 engine.switch_role(LLMRole.MIX)
207 def test_switch_role(self):
208 engine = LLMDataDist(LLMRole.MIX, 0)
209 llm_config = LLMConfig()
210 llm_config.device_id = 0
211 llm_config.enable_switch_role = True
212 options = llm_config.generate_options()
213 os.environ["RESOURCE_CONFIG_PATH"] = _TEST_BASE_DIR + "/json_file/numa_config.json"
214 engine.init(options)
215 engine.switch_role(LLMRole.DECODER)
216 switch_options = {
217 "llm.listenIpInfo": "127.0.0.1:1111",
218 }
219 engine.switch_role(LLMRole.PROMPT, switch_options)
221 @staticmethod
222 def _allocate_npu_cache(kv_cache_manager, block_size, num_block, num_tensors):
223 npu_cache_desc = CacheDesc(
224 num_tensors=num_tensors,
225 shape=[num_block, block_size],
226 data_type=DataType.DT_FLOAT16,
227 )
228 npu_cache_key = BlocksCacheKey(0, 0)
229 cache = kv_cache_manager.allocate_blocks_cache(npu_cache_desc, npu_cache_key)
230 return cache, npu_cache_key
232 @staticmethod
233 def _allocate_cpu_cache(kv_cache_manager, block_size, num_block, num_tensors):
234 # DT没有友好的方式创建cpu tensor，用npu接口模拟
235 cpu_cache_desc = CacheDesc(
236 num_tensors=num_tensors,
237 shape=[num_block, block_size],
238 data_type=DataType.DT_FLOAT16,
239 )
240 cpu_cache_key = BlocksCacheKey(1, 1)
241 cache = kv_cache_manager.allocate_blocks_cache(cpu_cache_desc, cpu_cache_key)
242 cpu_cache_desc = CacheDesc(
243 num_tensors=num_tensors,
244 shape=[num_block, block_size],
245 data_type=DataType.DT_FLOAT16,
246 placement=Placement.HOST,
247 )
248 return KvCache.create_cpu_cache(cpu_cache_desc, cache.per_device_tensor_addrs[0]), cache
250 def test_swap_blocks(self):
251 cluster_id = 0
252 llm_engine = LLMDataDist(LLMRole.PROMPT, cluster_id)
253 llm_config = LLMConfig()
254 llm_config.device_id = 0
255 llm_config.listen_ip_info = "0.0.0.0:26000"
256 llm_config.ge_options = {"ge.flowGraphMemMaxSize": "10240"}
257 init_options = llm_config.generate_options()
258 llm_engine.init(init_options)
260 kv_cache_manager = llm_engine.kv_cache_manager
261 # allocate npu cache
262 npu_cache, npu_cache_key = self._allocate_npu_cache(kv_cache_manager, 64 * 1024, 10, 10)
263 cpu_cache, tmp_cache = self._allocate_cpu_cache(kv_cache_manager, 64 * 1024, 20, 10)
264 src_to_dst = {3: 4, 0: 0, 1: 1, 2: 2, 5: 6, 6: 7, 7: 8, 9: 9}
265 kv_cache_manager.swap_blocks(npu_cache, cpu_cache, src_to_dst)
266 kv_cache_manager.swap_blocks(cpu_cache, npu_cache, src_to_dst)
268 def test_create_cpu_cache_failed(self):
269 cluster_id = 0
270 llm_engine = LLMDataDist(LLMRole.PROMPT, cluster_id)
271 llm_config = LLMConfig()
272 llm_config.device_id = 0
273 llm_config.listen_ip_info = "0.0.0.0:26000"
274 llm_config.ge_options = {"ge.flowGraphMemMaxSize": "10240"}
275 init_options = llm_config.generate_options()
276 llm_engine.init(init_options)
278 has_err = False
279 try:
280 cpu_cache_desc = CacheDesc(
281 num_tensors=1,
282 shape=[1, 1],
283 data_type=DataType.DT_FLOAT16,
284 placement=Placement.HOST,
285 )
286 KvCache.create_cpu_cache(cpu_cache_desc, [[1], 1])
287 except LLMException as ex:
288 self.assertEqual("should be consistent" in ex.__str__(), True)
289 has_err = True
290 self.assertEqual(has_err, True)
292 def test_copy_blocks_validate(self):
293 cluster_id = 0
294 llm_datadist = LLMDataDist(LLMRole.DECODER, cluster_id)
295 llm_datadist.init(self._engine_options(False, cluster_id))
297 kv_cache_manager = llm_datadist.kv_cache_manager
298 cache_desc = CacheDesc(1, [2, 8], DataType.DT_FLOAT16)
299 kv_cache = kv_cache_manager.allocate_blocks_cache(cache_desc)
300 try:
301 kv_cache_manager.copy_blocks(kv_cache, {0: "1"})
302 except Exception as ex:
303 self.assertEqual("only support" in ex.__str__(), True)
304 print("copy_blocks validate2")
305 try:
306 kv_cache_manager.copy_blocks(kv_cache, {0: ["1"]})
307 except Exception as ex:
308 self.assertEqual("inner type only support" in ex.__str__(), True)
310 llm_datadist.finalize()
312 def test_transfer_cache(self):
313 cluster_id = 0
314 engine = LLMDataDist(LLMRole.PROMPT, cluster_id)
315 engine.init(LlmEngineV2St._engine_options(True, cluster_id))
316 kv_cache_manager = engine.kv_cache_manager
317 cache_desc = CacheDesc(10, [2, 8], DataType.DT_FLOAT16)
318 cache_key = CacheKey(0, 0, 1)
319 kv_cache = kv_cache_manager.allocate_cache(cache_desc, [cache_key])
320 dst_addrs_1 = [10000000, 20000000, 30000000, 40000000, 50000000, 60000000]
321 dst_addrs_2 = [10000000, 20000000, 30000000, 40000000]
322 transfer_config_1 = TransferConfig(1, dst_addrs_1, range(0, 3))
323 print(transfer_config_1)
324 transfer_config_2 = TransferConfig(2, dst_addrs_2, range(2, 4))
325 transfer_configs = (transfer_config_1, transfer_config_2)
326 cache_task = kv_cache_manager.transfer_cache_async(kv_cache, LayerSynchronizerImpl(True), transfer_configs)
327 ret = cache_task.synchronize(0)
328 self.assertEqual(ret, LLMStatusCode.LLM_WAIT_PROCESS_TIMEOUT)
329 rets = cache_task.get_results(0)
330 self.assertEqual(rets, [LLMStatusCode.LLM_WAIT_PROCESS_TIMEOUT] * 2)
331 ret = cache_task.synchronize(1000)
332 rets = cache_task.get_results(1000)
333 self.assertEqual(ret, LLMStatusCode.LLM_SUCCESS)
334 self.assertEqual(rets[0], LLMStatusCode.LLM_SUCCESS)
335 self.assertEqual(rets[1], LLMStatusCode.LLM_SUCCESS)
337 transfer_config_3 = TransferConfig(2, dst_addrs_1 + dst_addrs_2)
338 cache_task = kv_cache_manager.transfer_cache_async(kv_cache, LayerSynchronizerImpl(True), [transfer_config_3])
339 ret = cache_task.synchronize()
340 self.assertEqual(ret, LLMStatusCode.LLM_SUCCESS)
342 # test blocks suc
343 kv_cache = kv_cache_manager.allocate_blocks_cache(cache_desc)
344 dst_addrs_1 = [10000000, 20000000, 30000000, 40000000, 50000000, 60000000]
345 transfer_config_1 = TransferConfig(1, dst_addrs_1, range(0, 3))
346 transfer_configs = [transfer_config_1]
347 block_indices = [0, 1]
348 cache_task = kv_cache_manager.transfer_cache_async(
349 kv_cache,
350 LayerSynchronizerImpl(True),
351 transfer_configs,
352 block_indices,
353 block_indices,
354 dst_block_memory_size=32,
355 )
356 ret = cache_task.synchronize()
357 rets = cache_task.get_results()
358 self.assertEqual(ret, LLMStatusCode.LLM_SUCCESS)
359 self.assertEqual(rets[0], LLMStatusCode.LLM_SUCCESS)
361 # test sync layer failed
362 transfer_config_1 = TransferConfig(1, dst_addrs_1, range(0, 3))
363 transfer_config_2 = TransferConfig(2, dst_addrs_2, range(2, 4))
364 transfer_configs = (transfer_config_1, transfer_config_2)
365 cache_task = kv_cache_manager.transfer_cache_async(kv_cache, LayerSynchronizerImpl(False), transfer_configs)
366 ret = cache_task.synchronize()
367 rets = cache_task.get_results()
368 self.assertNotEqual(ret, LLMStatusCode.LLM_SUCCESS)
369 self.assertNotEqual(rets[0], LLMStatusCode.LLM_SUCCESS)
370 self.assertEqual(rets[1], None)
372 # test transfer failed
373 params = TransferCacheParameters(kv_cache, transfer_configs, None, None, None)
374 job = MockTransferCacheJob(params)
375 job.init()
376 job.transfer_layers()
377 self.assertEqual(job.get_results()[0], LLMStatusCode.LLM_WAIT_PROCESS_TIMEOUT)
378 self.assertEqual(job.get_results()[1], None)
380 src_block_indices = [1, 2]
381 with self.assertRaisesRegex(LLMException, "transfer from blocks to cache is not supported"):
382 _ = kv_cache_manager.transfer_cache_async(
383 kv_cache,
384 LayerSynchronizerImpl(False),
385 transfer_configs,
386 src_block_indices,
387 )
388 with self.assertRaises(TypeError):
389 _ = kv_cache_manager.transfer_cache_async(
390 kv_cache,
391 LayerSynchronizerImpl(False),
392 transfer_configs,
393 [0, 1],
394 [0, "1"],
395 )
396 with self.assertRaises(ValueError):
397 _ = kv_cache_manager.transfer_cache_async(
398 kv_cache,
399 LayerSynchronizerImpl(False),
400 transfer_configs,
401 [-1, 1],
402 [0, 1],
403 )
404 with self.assertRaises(ValueError):
405 _ = kv_cache_manager.transfer_cache_async(
406 kv_cache,
407 LayerSynchronizerImpl(False),
408 transfer_configs,
409 [0, 1],
410 [-1, 1],
411 )
412 with self.assertRaises(LLMException):
413 _ = kv_cache_manager.transfer_cache_async(
414 kv_cache, LayerSynchronizerImpl(False), transfer_configs, [0, 1], [0]
415 )
416 with self.assertRaises(LLMException):
417 _ = kv_cache_manager.transfer_cache_async(
418 kv_cache, LayerSynchronizerImpl(False), transfer_configs, [], [], 1
419 )
420 with self.assertRaises(ValueError):
421 _ = kv_cache_manager.transfer_cache_async(
422 kv_cache,
423 LayerSynchronizerImpl(False),
424 transfer_configs,
425 [0, 1],
426 [0, 1],
427 -1,
428 )
429 with self.assertRaises(TypeError):
430 _ = kv_cache_manager.transfer_cache_async("cache", LayerSynchronizerImpl(False), transfer_configs)
431 with self.assertRaises(TypeError):
432 _ = kv_cache_manager.transfer_cache_async(None, LayerSynchronizerImpl(False), transfer_configs)
433 kv_cache_manager.deallocate_cache(kv_cache)
434 kv_cache_manager.remove_cache_key(cache_key)
435 engine.finalize()
437 def test_transfer_config(self):
438 config = TransferConfig(2, [1000, 2000])
439 self.assertEqual(config.dst_cluster_id, 2)
440 self.assertEqual(config.dst_addrs, [1000, 2000])
441 self.assertIsNone(config.src_layer_range)
442 self.assertEqual(config.src_batch_index, 0)
444 with self.assertRaises(TypeError):
445 config.dst_cluster_id = None
446 with self.assertRaises(TypeError):
447 config.dst_cluster_id = "123"
448 with self.assertRaises(TypeError):
449 config.dst_addrs = None
450 with self.assertRaises(TypeError):
451 config.dst_addrs = "123"
452 with self.assertRaises(TypeError):
453 config.dst_addrs = ["123"]
454 with self.assertRaises(ValueError):
455 config.dst_addrs = [-1]
456 with self.assertRaises(TypeError):
457 config.src_layer_range = [0, 1]
458 with self.assertRaises(LLMException):
459 config.src_layer_range = range(0, 5, 2)
460 with self.assertRaises(LLMException):
461 config.src_layer_range = range(0, 1, -1)
462 with self.assertRaises(LLMException):
463 config.src_layer_range = range(4, 1)
464 with self.assertRaises(LLMException):
465 config.src_layer_range = range(0, 0)
466 with self.assertRaises(TypeError):
467 config.src_batch_index = None
468 with self.assertRaises(TypeError):
469 config.src_batch_index = "1"
470 with self.assertRaises(ValueError):
471 config.src_batch_index = -1
472 with self.assertRaises(ValueError):
473 config.src_batch_index = 2**32
475 with self.assertRaises(TypeError):
476 _ = TransferConfig(2, [1000, 2000], None, "1")
477 with self.assertRaises(TypeError):
478 _ = TransferConfig(2, [1000, 2000], [0, 1])
479 with self.assertRaises(TypeError):
480 _ = TransferConfig(2, [1000, "2000"])
481 with self.assertRaises(TypeError):
482 _ = TransferConfig(2, None)
483 with self.assertRaises(TypeError):
484 _ = TransferConfig("1", [1000, 2000])
486 def test_multiple_devices(self):
487 llm_config = LLMConfig()
488 llm_config.device_id = [1, 2]
489 llm_config.listen_ip_info = "127.0.0.1:26000;127.0.0.1:26000"
490 engine_options = llm_config.generate_options()
491 print("engine_options:", engine_options)
492 prompt_engine = LLMDataDist(LLMRole.PROMPT, 0)
493 prompt_engine.init(engine_options)
494 print(engine_options["llm.ClusterInfo"])
496 def test_tensor_bf16(self):
497 arr1 = np.array([1.875], np.float16)
498 tensor_desc = TensorDesc(DataType.DT_BF16, [1])
499 tensor = Tensor(arr1, tensor_desc)
500 print("generated numpy:", tensor.numpy())
501 self.assertEqual(tensor.numpy().dtype, np.float32)
502 self.assertEqual(int(tensor.numpy()[0]), 1)
504 def test_tensor_foat32_copy_true(self):
505 arr1 = np.array([[1.0, 2.0], [3.0, 4.0]], np.float32)
506 tensor_desc = TensorDesc(DataType.DT_FLOAT, (2, 2))
507 tensor = Tensor(arr1, tensor_desc)
508 res = tensor.numpy(copy=True)
509 print("generated numpy:", res)
510 np.testing.assert_array_equal(res, arr1)
511 self.assertEqual(res.dtype, np.float32)
512 self.assertTrue(res.flags.c_contiguous)
513 self.assertTrue(res.flags.writeable)
514 arr1[0, 0] = 999.0
515 self.assertNotEqual(res[0, 0], 999.0)
517 def test_str_tensor(self):
518 arr1 = np.array([["aaaa", "bbbb"], ["ccc", "ddd"]])
519 src_type = arr1.dtype
520 tensor = Tensor(arr1)
521 self.assertEqual(tensor.numpy(True).dtype, src_type)
522 print(tensor.numpy(True))
524 def test_check_flow_graph_mem_max_size(self):
525 cluster_id = 0
526 llm_engine = LLMDataDist(LLMRole.PROMPT, cluster_id)
527 llm_config = LLMConfig()
528 llm_config.device_id = 0
529 llm_config.listen_ip_info = "0.0.0.0:26000"
530 llm_config.ge_options = {"ge.flowGraphMemMaxSize": "-1"}
531 init_options = llm_config.generate_options()
533 has_err = False
534 try:
535 llm_engine.init(init_options)
536 except LLMException:
537 has_err = True
538 self.assertEqual(has_err, True)
540 def test_check_flow_graph_mem_max_size2(self):
541 cluster_id = 0
542 llm_engine = LLMDataDist(LLMRole.PROMPT, cluster_id)
543 llm_config = LLMConfig()
544 llm_config.device_id = 0
545 llm_config.listen_ip_info = "0.0.0.0:26000"
546 llm_config.ge_options = {"llm.EnableCacheManager": "0"}
547 init_options = llm_config.generate_options()
548 has_err = False
549 try:
550 llm_engine.init(init_options)
551 except LLMException:
552 has_err = True
553 self.assertEqual(has_err, False)