平衡本地机器创建的数据和任务，并与全局同步。
| 245 | logger.error(f"Error sending reply for unknown message type: {e}") |
| 246 | |
| 247 | class LocalBalanceManager: |
| 248 | """平衡本地机器创建的数据和任务，并与全局同步。""" |
| 249 | |
| 250 | def __init__( |
| 251 | self, |
| 252 | local_collect_address: str, |
| 253 | local_provider_address: str, |
| 254 | local_steal_port: int, |
| 255 | global_sync_address: str, |
| 256 | global_result_collect_address: str, |
| 257 | global_data_dispatch_address: str, |
| 258 | chunk_size: int, |
| 259 | mt_max_beam_width: int, |
| 260 | max_cache_size: int, |
| 261 | processing_class_name_or_path: str, |
| 262 | max_prompt_length: int, |
| 263 | steal_threshold: int = 16, |
| 264 | tp_size: int = 1, |
| 265 | timeout: int = 3600 |
| 266 | ): |
| 267 | """初始化本地平衡管理器。 |
| 268 | |
| 269 | Args: |
| 270 | local_collect_address: 本地收集地址 |
| 271 | local_provider_address: 本地提供者地址 |
| 272 | local_steal_port: 本地窃取端口 |
| 273 | global_sync_address: 全局同步地址 |
| 274 | global_result_collect_address: 全局结果收集地址 |
| 275 | global_data_dispatch_address: 全局数据分发地址 |
| 276 | chunk_size: 块大小 |
| 277 | mt_max_beam_width: 最大束宽度 |
| 278 | max_cache_size: 最大缓存大小 |
| 279 | processing_class_name_or_path: 处理器类名或路径 |
| 280 | max_prompt_length: 最大提示长度 |
| 281 | steal_threshold: 窃取阈值 |
| 282 | tp_size: 张量并行大小 |
| 283 | timeout: 超时时间 |
| 284 | """ |
| 285 | self.local_collect_address = local_collect_address |
| 286 | self.local_provider_address = local_provider_address |
| 287 | self.local_steal_port = local_steal_port |
| 288 | self.global_sync_address = global_sync_address |
| 289 | self.global_result_collect_address = global_result_collect_address |
| 290 | self.global_data_dispatch_address = global_data_dispatch_address |
| 291 | self.chunk_size = chunk_size |
| 292 | self.mt_max_beam_width = mt_max_beam_width |
| 293 | self.max_cache_size = max_cache_size |
| 294 | self.processing_class_name_or_path = processing_class_name_or_path |
| 295 | self.max_prompt_length = max_prompt_length |
| 296 | self.steal_threshold = steal_threshold |
| 297 | self.tp_size = tp_size |
| 298 | self.timeout = timeout |
| 299 | |
| 300 | # 初始化ZMQ上下文 |
| 301 | self.zmqctx = zmq.Context(16) |
| 302 | |
| 303 | # 初始化所有ZMQ套接字 |
| 304 | self._init_sockets() |
no outgoing calls
no test coverage detected