Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/llm_datadist_v1/llm_datadist.py: 94%
140 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# 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# -----------------------------------------------------------------------------------------------------------
12import atexit
13from typing import Dict, List, Optional, Tuple, Union
15from llm_datadist_v1 import llm_wrapper
17from .config import EngineConfig
18from .configs import LLMClusterInfo, LLMRole
19from .kv_cache_manager import KvCacheManager
20from .status import (
21 LLMException,
22 LLMStatusCode,
23 code_2_status,
24 handle_llm_status,
25 raise_if_false,
26)
27from .utils import log
28from .utils.utils import check_int32, check_isinstance, check_uint64
30__all__ = ["LLMDataDist", "KvCacheManager"]
33class LLMDataDist(object):
34 llm_engine_instance = None
35 """
36 LLMDataDist
38 Args:
39 role: role of LLMDataDist
40 cluster_id: cluster_id of LLMDataDist
41 """
43 def __init__(self, role: LLMRole, cluster_id: int):
44 check_isinstance("role", role, LLMRole)
45 self._kv_cache_manager = None
46 self._cache_manager = None
47 self._role = role
48 check_uint64("cluster_id", cluster_id)
49 self._cluster_id = cluster_id
50 self._llm_datadist = None
51 self._engine_config = None
52 self._is_initialized = False
53 self._engine_options: Dict[str, str] = {}
54 self._enable_cache_mgr = False
55 self._enable_free_comm = False
56 self._enable_local_comm_res = False
58 def _check_flow_graph_max_size(self, options: Dict[str, str]) -> None:
59 mem_utilization = float(options.get("llm.MemoryUtilization", "0.95"))
60 value = options.get("ge.flowGraphMemMaxSize", None)
61 if value is None:
62 return
63 check_isinstance("ge.flowGraphMemMaxSize", value, str)
64 raise_if_false(
65 len(value.split(",")) == 1,
66 "ge.flowGraphMemMaxSize only support one mem pool in llm datadist",
67 )
68 raise_if_false(value.isdigit(), "ge.flowGraphMemMaxSize must be digit")
70 def init(self, options: Dict[str, str]) -> None:
71 """
72 初始化LLM Engine
74 Args:
75 options: Engine相关options
76 """
77 if self._is_initialized:
78 return
79 raise_if_false(
80 LLMDataDist.llm_engine_instance is None,
81 "Cannot init multiple LLM engines",
82 status_code=LLMStatusCode.LLM_FAILED,
83 )
84 check_isinstance("options", options, dict)
85 self._check_flow_graph_max_size(options)
86 self._engine_options = options
87 self._engine_options["llm.Role"] = self._role_to_str(self._role)
88 self._enable_local_comm_res = "llm.LocalCommRes" in options
89 if self._enable_local_comm_res and "llm.EnableCacheManager" not in options:
90 self._engine_options["llm.EnableCacheManager"] = "1"
91 if self._enable_local_comm_res and "llm.EnableRemoteCacheAccessible" not in options:
92 self._engine_options["llm.EnableRemoteCacheAccessible"] = "1"
93 self._enable_cache_mgr = (
94 "llm.EnableCacheManager" in self._engine_options and self._engine_options["llm.EnableCacheManager"] == "1"
95 )
96 log.info("options = %s", self._engine_options)
97 self._llm_datadist = llm_wrapper
98 EngineConfig.gen_cluster_info_if_not_exist(self._cluster_id, self._role, self._engine_options)
99 ret = self._llm_datadist.initialize(self._cluster_id, self._engine_options)
100 handle_llm_status(
101 ret,
102 "[LLMDataDist.init]",
103 f"Failed to initialize llm datadist, options = {options}",
104 )
105 self._kv_cache_manager = KvCacheManager(self._llm_datadist, self._role)
106 LLMDataDist.llm_engine_instance = self
107 self._is_initialized = True
109 def _check_is_cache_mgr_mode(self, func_name):
110 raise_if_false(
111 self._enable_cache_mgr,
112 "{0} is not supported when llm.EnableCacheManager is not configured.",
113 func_name,
114 )
116 def _check_is_not_cache_mgr_mode(self, func_name):
117 raise_if_false(
118 not self._enable_cache_mgr,
119 "{0} is not supported when llm.EnableCacheManager is configured.",
120 func_name,
121 )
123 def finalize(self) -> None:
124 """
125 释放LLM Engine相关资源
126 """
127 if not self._is_initialized:
128 return
129 self._llm_datadist.finalize()
130 if self._kv_cache_manager is not None:
131 self._kv_cache_manager._initialized = False
132 self._is_initialized = False
133 LLMDataDist.llm_engine_instance = None
135 def _cluster_config(self):
136 if self._engine_config is None:
137 self._engine_config = EngineConfig.from_engine_options(self._role == LLMRole.PROMPT, self._engine_options)
138 return self._engine_config.cluster_config
140 def check_link_status(self, remote_cluster_id: int):
141 self._check_is_inited()
142 self._check_is_not_cache_mgr_mode("check_link_status")
143 check_uint64("remote_cluster_id", remote_cluster_id)
144 ret = self._llm_datadist.check_link_status(remote_cluster_id)
145 handle_llm_status(ret, "[check_link_status]", f"remote_cluster_id is {remote_cluster_id}")
146 log.info("[check_link_status] success")
148 def link_clusters(self, clusters: Union[List[LLMClusterInfo], Tuple[LLMClusterInfo]], timeout=3000):
149 self._check_is_inited()
150 check_int32("timeout", timeout)
151 raise_if_false(timeout > 0, "Param timeout should be greater than 0.")
152 check_isinstance("clusters", clusters, [list, tuple], LLMClusterInfo)
153 cluster_list = [
154 (
155 cluster.remote_cluster_id,
156 0,
157 cluster.local_ip_info_list,
158 cluster.remote_ip_info_list,
159 )
160 for cluster in clusters
161 ]
162 ret, rets = self._llm_datadist.link_clusters(cluster_list, timeout)
163 return code_2_status(ret), [code_2_status(cluster_ret) for cluster_ret in rets]
165 def unlink_clusters(
166 self,
167 clusters: Union[List[LLMClusterInfo], Tuple[LLMClusterInfo]],
168 timeout=3000,
169 force=False,
170 ):
171 self._check_is_inited()
172 self._check_is_not_cache_mgr_mode("unlink_clusters")
173 check_int32("timeout", timeout)
174 raise_if_false(timeout > 0, "Param timeout should be greater than 0.")
175 check_isinstance("clusters", clusters, [list, tuple], LLMClusterInfo)
176 check_isinstance("force", force, bool)
177 cluster_list = [
178 (
179 cluster.remote_cluster_id,
180 0,
181 cluster.local_ip_info_list,
182 cluster.remote_ip_info_list,
183 )
184 for cluster in clusters
185 ]
186 ret, rets = self._llm_datadist.unlink_clusters(cluster_list, timeout, force)
187 return code_2_status(ret), [code_2_status(cluster_ret) for cluster_ret in rets]
189 def switch_role(self, role: LLMRole, switch_options: Optional[Dict[str, str]] = None):
190 self._check_is_inited()
191 check_isinstance("role", role, LLMRole)
192 raise_if_false(self._role != role, f"role not changed, role = {role.name}")
193 role_str = self._role_to_str(role)
194 log.info(f"[switch_role] [{self._role.name}->{role.name}] start, switch_options = {switch_options}")
195 check_isinstance("switch_options", switch_options, dict)
196 options = switch_options.copy() if switch_options is not None else {}
197 if role == LLMRole.PROMPT:
198 raise_if_false(
199 "llm.listenIpInfo" in switch_options,
200 'Failed to switch to Prompt, option "llm.listenIpInfo" was specified',
201 )
202 listen_ip_info = switch_options["llm.listenIpInfo"]
203 ip, port = EngineConfig.parse_listen_ip_info(listen_ip_info)
204 options["llm.ListenIp"] = str(ip)
205 options["llm.ListenPort"] = str(port)
206 ret = self._llm_datadist.switch_role(role_str, options)
207 handle_llm_status(
208 ret,
209 "[LLMEngine.switch_role]",
210 f"Failed to switch role, role = {role}, options = {options}",
211 )
212 self._kv_cache_manager._switch_role(role)
213 log.info(f"[switch_role] [{self._role.name}->{role.name}] success")
214 self._role = role
216 @staticmethod
217 def _role_to_str(role: LLMRole) -> str:
218 role_mapping = {
219 LLMRole.PROMPT: "Prompt",
220 LLMRole.DECODER: "Decoder",
221 LLMRole.MIX: "Mix",
222 }
223 return role_mapping[role]
225 def _check_is_inited(self):
226 if not self._is_initialized:
227 raise RuntimeError("llm datadist is not initialized")
229 @property
230 def kv_cache_manager(self) -> KvCacheManager:
231 """
232 获取KvCacheManager
234 Returns:
235 KvCacheManager
236 """
237 self._check_is_inited()
238 self._check_is_not_cache_mgr_mode("kv_cache_manager")
239 return self._kv_cache_manager
241 @property
242 def cluster_id(self):
243 return self._cluster_id
246def _shutdown_handler():
247 if LLMDataDist.llm_engine_instance is not None:
248 log.info("[shutdown_handler] finalize llm datadist")
249 try:
250 LLMDataDist.llm_engine_instance.finalize()
251 except LLMException as e:
252 log.warn(
253 f"error occurred while finalize llm datadist: {e} may cause by already finalized by another framework"
254 )
257atexit.register(_shutdown_handler)