Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/build_st/python_tests/v1/st/test_v1_enable_local_comm_res.py: 90%

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

12 

13import time 

14import unittest 

15 

16from llm_datadist_v1 import * 

17 

18 

19class LlmCacheManagerSt(unittest.TestCase): 

20 def setUp(self) -> None: 

21 print("Begin ", self._testMethodName) 

22 config = LlmConfig() 

23 config.device_id = 0 

24 config.enable_switch_role = True 

25 config.rdma_service_level = 100 

26 config.rdma_traffic_class = 100 

27 config.listen_ip_info = "127.0.0.1:26008" 

28 config.local_comm_res = """ 

29 { 

30 "server_count": "1", 

31 "server_list": [{ 

32 "device": [{ 

33 "device_id": "0", 

34 "device_ip": "1.1.1.1" 

35 }], 

36 "server_id": "127.0.0.1" 

37 }], 

38 "status": "completed", 

39 "version": "1.0" 

40 } 

41 """ 

42 engine_options = config.generate_options() 

43 self.llm_datadist = LLMDataDist(LLMRole.PROMPT, 1) 

44 self.llm_datadist.init(engine_options) 

45 time.sleep(1) # wait listen 

46 self.has_exception = False 

47 

48 def tearDown(self) -> None: 

49 print("End ", self._testMethodName) 

50 self.llm_datadist.finalize() 

51 

52 def create_link_cluster(self): 

53 cluster = LLMClusterInfo() 

54 cluster.remote_cluster_id = 1 

55 cluster.append_local_ip_info("127.0.0.1", 26008) 

56 cluster.append_remote_ip_info("127.0.0.1", 26008) 

57 ret, rets = self.llm_datadist.link_clusters([cluster], 5000) 

58 self.assertEqual(ret, LLMStatusCode.LLM_SUCCESS) 

59 

60 def test_unlink_cluster(self): 

61 self.llm_datadist.switch_role(LLMRole.DECODER) 

62 self.llm_datadist._enable_cache_mgr = False 

63 self.create_link_cluster() 

64 cluster = LLMClusterInfo() 

65 cluster.remote_cluster_id = 1 

66 cluster.append_local_ip_info("127.0.0.1", 26008) 

67 cluster.append_remote_ip_info("127.0.0.1", 26008) 

68 ret, rets = self.llm_datadist.unlink_clusters([cluster], 5000) 

69 self.assertEqual(ret, LLMStatusCode.LLM_SUCCESS) 

70 

71 def test_local_comm_res_switch_role(self): 

72 try: 

73 self.llm_datadist.switch_role(LLMRole.DECODER) 

74 options = {"llm.listenIpInfo": "127.0.0.1:26008"} 

75 self.llm_datadist.switch_role(LLMRole.PROMPT, options) 

76 except Exception as e: 

77 print(f"{type(e).__name__} - {str(e)}") 

78 import traceback 

79 

80 print(traceback.format_exc()) 

81 self.has_exception = True 

82 self.assertEqual(self.has_exception, False)