trinity.common.models.allocator module

trinity.common.models.allocator module#

Allocator module for managing inference engines.

class trinity.common.models.allocator.BundleResult(bundles: List[Dict[str, float]], actor_bundle_map: Dict[str, int], bundle_actor_map: Dict[int, str])[source]#

Bases: object

Data class for storing the result of bundle allocation.

bundles: List[Dict[str, float]]#
actor_bundle_map: Dict[str, int]#
bundle_actor_map: Dict[int, str]#
__init__(bundles: List[Dict[str, float]], actor_bundle_map: Dict[str, int], bundle_actor_map: Dict[int, str]) → None#
class trinity.common.models.allocator.Allocator(config: ExplorerConfig)[source]#

Bases: object

Allocator class for managing inference engines.

__init__(config: ExplorerConfig)[source]#

Initialize the Allocator.

get_actor_name(role: str, engine_id: int, node_id: int) → str[source]#

Generate a unique actor name based on the model config, engine ID, and node ID.

property placement_group_name: str#

Generate a descriptive name for the inference model placement group.

allocate_bundles() → BundleResult[source]#

Allocate bundles for the rollout model and auxiliary models based on the configuration.

analyze_placement_group(pg: PlacementGroup, bundle_result: BundleResult)[source]#
async create_engine(config: InferenceModelConfig, role: str, engine_id: int) → ModelWrapper[source]#
async create_all_models() → Tuple[List[ModelWrapper], List[List[ModelWrapper]]][source]#

Create all model actors for the rollout model and auxiliary models based on the configuration.

get_model(config: InferenceModelConfig, role: str, engine_id: int) → ModelWrapper[source]#

Get the model actor for the given role and engine ID.

async trinity.common.models.allocator.get_model_wrapper(actor_cls, config: InferenceModelConfig, pg: PlacementGroup, actor_bundle_list: List[Tuple[str, int]]) → ModelWrapper[source]#

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.

Parameters:

  • 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:

A wrapper for the model actors with distributed communication setup.

Return type:

ModelWrapper

async trinity.common.models.allocator.get_external_model_wrapper(config: InferenceModelConfig) → ModelWrapper[source]#

Get the ModelWrapper for the external model.