Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/build_ut/python_tests/v1/ut/test_v1_llm_datadist.py: 83%

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# ----------------------------------------------------------------------------------------------------------- 

13 

14import json 

15import os 

16import unittest 

17 

18from llm_datadist_v1 import * 

19from llm_datadist_v1.llm_datadist import _shutdown_handler 

20 

21_INVALID_ID = 2**64 - 1 

22 

23_TEST_BASE_DIR = "../tests/dflow/llm_datadist/st/testcase/llm_datadist" 

24 

25 

26class LlmEngineV2Ut(unittest.TestCase): 

27 def setUp(self) -> None: 

28 os.environ["ASCEND_GLOBAL_LOG_LEVEL"] = "1" 

29 print("Begin ", self._testMethodName) 

30 

31 def tearDown(self) -> None: 

32 os.environ.pop("RESOURCE_CONFIG_PATH", None) 

33 _shutdown_handler() 

34 print("End ", self._testMethodName) 

35 

36 @staticmethod 

37 def _engine_options(is_prompt: bool, cluster_id: int = 0, rank_id: int = -1, resource_path: str = ""): 

38 cluster_info = { 

39 "cluster_id": cluster_id, 

40 "logic_device_id": ["0:0:0:0", "0:0:1:0", "0:0:2:0", "0:0:3:0"], 

41 } 

42 if is_prompt: 

43 cluster_info["listen_ip_info"] = [ 

44 {"ip": 0, "port": 26000}, 

45 {"ip": 1, "port": 26000}, 

46 {"ip": 2, "port": 26000}, 

47 {"ip": 3, "port": 26000}, 

48 ] 

49 engine_options = {"llm.ClusterInfo": json.dumps(cluster_info)} 

50 if rank_id != -1: 

51 engine_options["ge.exec.rankId"] = str(rank_id) 

52 if resource_path != "": 

53 engine_options["ge.resourceConfigPath"] = resource_path 

54 return engine_options 

55 

56 def test_simple_option(self): 

57 cluster_id = 0 

58 prompt_engine = LLMDataDist(LLMRole.PROMPT, cluster_id) 

59 llm_config = LLMConfig() 

60 llm_config.device_id = 1 

61 llm_config.listen_ip_info = "127.0.0.1:26000" 

62 llm_config.deploy_res_path = "./" 

63 llm_config.ge_options = {"ge.flowGraphMemMaxSize": "10000000"} 

64 engine_options = llm_config.generate_options() 

65 print("engine_options:", engine_options) 

66 prompt_engine.init(engine_options) 

67 

68 def test_check_flow_graph_mem_max_size(self): 

69 cluster_id = 0 

70 llm_engine = LLMDataDist(LLMRole.PROMPT, cluster_id) 

71 llm_config = LLMConfig() 

72 llm_config.device_id = 0 

73 llm_config.listen_ip_info = "0.0.0.0:26000" 

74 llm_config.ge_options = {"ge.flowGraphMemMaxSize": "-1"} 

75 init_options = llm_config.generate_options() 

76 

77 has_err = False 

78 try: 

79 llm_engine.init(init_options) 

80 except LLMException: 

81 has_err = True 

82 self.assertEqual(has_err, True) 

83 

84 def test_check_flow_graph_mem_max_size2(self): 

85 cluster_id = 0 

86 llm_engine = LLMDataDist(LLMRole.PROMPT, cluster_id) 

87 llm_config = LLMConfig() 

88 llm_config.device_id = 0 

89 llm_config.listen_ip_info = "0.0.0.0:26000" 

90 llm_config.ge_options = {"llm.EnableCacheManager": "0"} 

91 init_options = llm_config.generate_options() 

92 has_err = False 

93 try: 

94 llm_engine.init(init_options) 

95 except LLMException: 

96 has_err = True 

97 self.assertEqual(has_err, False)