Function get_model_wrapper

trinity/common/models/allocator.py:165–213

Get the Ray actor wrapper for the inference model. Creates Ray actors with appropriate GPU allocation and distributed communication setup based on the inferred launch mode. Args: actor_cls: The actor class to instantiate (e.g., vLLMRolloutModel, SGLangRolloutModel). con

(
    actor_cls,
    config: "InferenceModelConfig",
    pg: PlacementGroup,
    actor_bundle_list: List[Tuple[str, int]],
)

async def get_model_wrapper(
    actor_cls,
    config: "InferenceModelConfig",
    pg: PlacementGroup,
    actor_bundle_list: List[Tuple[str, int]],
) -> ModelWrapper:
    """Get the Ray actor wrapper for the inference model.

    Creates Ray actors with appropriate GPU allocation and distributed communication
    setup based on the inferred launch mode.

    Args:
        actor_cls: The actor class to instantiate (e.g., vLLMRolloutModel, SGLangRolloutModel).
        config (InferenceModelConfig): The model config.
        pg (PlacementGroup): The placement group for the actors.
        actor_bundle_list (List[Tuple[str, int]]): The list of (actor_name, bundle_id) tuples
            for distributed setup.
    Returns:
        ModelWrapper: A wrapper for the model actors with distributed communication setup.
    """
    handlers = []
    for i, (actor_name, bundle_id) in enumerate(actor_bundle_list):
        engine_config = deepcopy(config)
        engine_config.ray_actor_name = actor_name
        engine_config.node_rank = i
        handlers.append(
            ray.remote(actor_cls)
            .options(
                name=actor_name,
                num_gpus=engine_config.gpu_per_engine / engine_config.nnodes,
                namespace=engine_config.ray_namespace,
                scheduling_strategy=PlacementGroupSchedulingStrategy(
                    placement_group=pg,
                    placement_group_capture_child_tasks=True,
                    placement_group_bundle_index=bundle_id,
                ),
            )
            .remote(config=engine_config)
        )
    if len(actor_bundle_list) > 1:
        # get master address and port from the first handler and set it to all handlers for distributed communication
        master_addr, master_port = await handlers[0].get_available_address.remote(random_port=True)
        for handler in handlers:
            await handler.set_master_addr_port.remote(master_addr, master_port)
    await asyncio.gather(*[handler.prepare.remote() for handler in handlers])
    server_address = await handlers[0].get_api_server_url.remote()
    wrapper = ModelWrapper(models=handlers, config=config, api_address=server_address)
    await wrapper.prepare()
    return wrapper


async def get_external_model_wrapper(config: InferenceModelConfig) -> ModelWrapper: