A wrapper for the InferenceModel Ray Actor
| 376 | |
| 377 | |
| 378 | class ModelWrapper: |
| 379 | """A wrapper for the InferenceModel Ray Actor""" |
| 380 | |
| 381 | def __init__( |
| 382 | self, |
| 383 | model: Optional[ActorHandle[InferenceModel]] = None, |
| 384 | models: Optional[List[ActorHandle[InferenceModel]]] = None, |
| 385 | config: Optional[InferenceModelConfig] = None, |
| 386 | api_address: Optional[str] = None, |
| 387 | ): |
| 388 | """Initialize the ModelWrapper. |
| 389 | |
| 390 | Args: |
| 391 | model (InferenceModel): The inference model Ray actor. |
| 392 | models (List[InferenceModel]): A list of inference model Ray actors for ensemble. The first model will be used as the main model for generation and other models will be used for auxiliary purposes such as logprobs calculation. If `model` is provided, `models` will be ignored. |
| 393 | config (InferenceModelConfig): The configuration for the inference model. |
| 394 | api_address (str, optional): The API address for the model. Required if `enable_openai_api` is True in the config. |
| 395 | """ |
| 396 | if config is None: |
| 397 | raise ValueError("Model config must be provided.") |
| 398 | if model is None and models is None and config.engine_type != "external": |
| 399 | raise ValueError("Either model or models must be provided.") |
| 400 | if model is not None: |
| 401 | self.model = model |
| 402 | self.models = [model] |
| 403 | elif models is not None and len(models) > 0: |
| 404 | self.model = models[0] |
| 405 | self.models = models |
| 406 | else: |
| 407 | self.model = None |
| 408 | self.models = [] |
| 409 | self.config: InferenceModelConfig = config |
| 410 | if self.config.model_path is None: |
| 411 | raise ValueError("model_path must be provided in the config.") |
| 412 | self._model_path = self.config.model_path |
| 413 | self._engine_type = config.engine_type |
| 414 | self._generate_kwargs = { |
| 415 | "temperature": self.config.temperature, |
| 416 | "top_p": self.config.top_p, |
| 417 | "max_tokens": self.config.max_response_tokens, |
| 418 | } |
| 419 | if self.config.enable_thinking is not None: |
| 420 | self._generate_kwargs["extra_body"] = { |
| 421 | "chat_template_kwargs": {"enable_thinking": self.config.enable_thinking} |
| 422 | } |
| 423 | self.api_address: Optional[str] = api_address |
| 424 | self._api_key: str = self.config.api_key |
| 425 | self.openai_client: openai.OpenAI = None |
| 426 | self.openai_async_client: openai.AsyncOpenAI = None |
| 427 | self.logger = get_logger(__name__) |
| 428 | self.enable_lora = config.enable_lora |
| 429 | self.enable_history = config.enable_history |
| 430 | self.history = [] |
| 431 | self.status = RunningStatus.RUNNING |
| 432 | self.workflow_state: Dict = {} |
| 433 | self.request_count = 0 |
| 434 | self.state_lock = asyncio.Lock() |
| 435 | self._routed_experts_layout: Optional[Tuple[int, int]] = None |
no outgoing calls