trinity.common.models.sglang_model module

Contents

trinity.common.models.sglang_model module#

class trinity.common.models.sglang_model.SGLangClient(server_url: str, api_key: str | None, logger: Logger)[source]#

Bases: object

A simple http client to interact with the SGLang API server.

__init__(server_url: str, api_key: str | None, logger: Logger)[source]#

async health_check() → bool[source]#

async init_weights_update_group(master_address: str, master_port: int, rank_offset: int, world_size: int, group_name: str, backend: str = 'nccl', timeout: int = 1200) → bool[source]#

async destroy_weights_update_group(group_name: str) → bool[source]#

async flush_cache(timeout: float = 30) → bool[source]#: Flush KV and Mamba cache to free GPU memory before weight sync.

async update_weights_from_distributed(state_dict_meta_list: List[Tuple[str, str, Tuple]], group_name: str, flush_cache: bool = True, abort_all_requests: bool = False, weight_version: str | None = None, timeout: float = 300) → bool[source]#

async update_weights_from_disk(model_path: str, abort_all_requests: bool = True, weight_version: str | None = None, is_async: bool = False, timeout: float = 300) → bool[source]#

async generate(input_ids: List[int], **kwargs) → Sequence[dict[str, Any]][source]#

class trinity.common.models.sglang_model.SGLangRolloutModel(config: InferenceModelConfig)[source]#

Bases: BaseInferenceModel

Wrapper around the SGLang engine to handle async requests.

Parameters:: config (Config) – The config.

__init__(config: InferenceModelConfig) → None[source]#

async init_process_group(master_address: str, master_port: int, rank_offset: int, world_size: int, group_name: str, backend: str = 'nccl', timeout: int = 1200)[source]#: Initialize the process group for model weight synchronization.

async set_state_dict_meta(state_dict_meta: List[Tuple[str, str, Tuple]])[source]#: Set the state_dict meta for NCCL weight sync.

async teardown_process_group()[source]#

Destroy the weight update group via the SGLang API.

Only the main node (node_rank=0) issues the API call; other nodes just clear local state.

async prepare() → None[source]#: Prepare the model before inference.

async generate(prompt: str, lora_request=None, **kwargs) → Sequence[Experience][source]#: Generate a responses from a prompt in async.

async chat(messages: List[dict], lora_request=None, **kwargs) → Sequence[Experience][source]#: Generate experiences from a list of history chat messages in async.

async logprobs(token_ids: List[int], **kwargs) → Tensor[source]#: Generate logprobs for a list of tokens in async.

async convert_messages_to_experience(messages: List[dict], tools=None, temperature: float | None = None) → Experience[source]#

Convert a list of messages into an experience in async.

Parameters:

messages – List of message dictionaries
tools – Optional list of tools
temperature – Optional temperature for logprobs calculation

async run_api_server() → bool[source]#

async shutdown() → None[source]#: Shutdown the model and release resources.

async sync_model_weights(model_version: int, method: SyncMethod, timeout: float = 1200) → int[source]#: Sync the model with the latest model_version.

get_model_version() → int[source]#: Get the checkpoint version.

get_api_server_url() → str | None[source]#: Get the API server URL if available.