dockyard_rl.models.generation.sglang.sglang_generation

Module Contents

Classes

SGLangGeneration

Data

TOP_K_THRESHOLD
TOP_P_THRESHOLD
logger

API

dockyard_rl.models.generation.sglang.sglang_generation.TOP_K_THRESHOLD

8000

dockyard_rl.models.generation.sglang.sglang_generation.TOP_P_THRESHOLD

0.99

dockyard_rl.models.generation.sglang.sglang_generation.logger

‘getLogger(…)’

class dockyard_rl.models.generation.sglang.sglang_generation.SGLangGeneration(cluster: dockyard_rl.distributed.virtual_cluster.RayVirtualCluster, config: dockyard_rl.models.generation.sglang.config.SGLangConfig, name_prefix: str = 'sglang_policy', workers_per_node: Optional[Union[int, list[int]]] = None)

Bases: dockyard_rl.models.generation.interfaces.GenerationInterface

init_collective(ip: str, port: int, world_size: int, *, train_world_size: int) → list[ray.ObjectRef]

No-op for the SGLang backend.

SGLang receives refreshed weights over its HTTP /update_weights_from_tensor endpoint using CUDA IPC handles gathered on the trainer side (see stream_weights_via_http_impl in models/policy/utils.py). That path is self-contained and needs no pre-established NCCL weight-sync communicator, so there is nothing to initialise here. The method is required by GenerationInterface; it intentionally starts no collective and returns an empty list of object refs (no remote work is scheduled).

generate(data: dockyard_rl.distributed.batched_data_dict.BatchedDataDict[dockyard_rl.models.generation.interfaces.GenerationDatumSpec], greedy: bool = False) → dockyard_rl.distributed.batched_data_dict.BatchedDataDict[dockyard_rl.models.generation.interfaces.GenerationOutputSpec]

Generate a batch of data using SGLang.

async generate_async(data: dockyard_rl.distributed.batched_data_dict.BatchedDataDict[dockyard_rl.models.generation.interfaces.GenerationDatumSpec], greedy: bool = False) → AsyncGenerator[tuple[int, dockyard_rl.distributed.batched_data_dict.BatchedDataDict[dockyard_rl.models.generation.interfaces.GenerationOutputSpec]], None]

Per-sample streaming generation; yields (original_idx, result) as ready.

prepare_refit_info(state_dict_info: dict[str, Any]) → None

update_weights_via_ipc_zmq() → list[ray.ObjectRef]

update_weights_from_collective() → list[ray.ObjectRef]

get_sglang_server_urls() → list[str]

Get base URLs of all SGLang servers.

Returns: List of base URLs (e.g., [“http://localhost:30000”, “http://localhost:30001”])

get_sglang_url_to_gpu_uuids() → dict[str, list[str]]

Get mapping from SGLang server URL to list of GPU UUIDs it uses.

Returns: Dict mapping server URL to list of GPU UUIDs e.g., {“http://localhost:30000”: [“GPU-aaa”, “GPU-bbb”], …}

prepare_for_generation(*args: Any, **kwargs: Any) → bool

Reacquire the colocated server’s GPU memory before a generation round.

SGLang is always colocated in dockyard (it shares GPUs with the trainer), so with memory-saver enabled the server releases its GPU memory during the training phase (see finish_generation) and must resume it here. An optional tags kwarg (e.g. [“weights”], [“kv_cache”]) selects a partial resume: the HTTP weight-sync lifecycle calls this twice — once to stage refreshed weights, once to rebuild the KV cache. With memory-saver off the server is always-on and this is a no-op.

finish_generation(*args: Any, **kwargs: Any) → bool

Release the colocated server’s GPU memory at the end of a round.

With memory-saver enabled the server releases all GPU memory occupation so the co-resident trainer can use it during the training phase; prepare_for_generation resumes it next round. With memory-saver off this is a no-op: the prefix/KV cache is deliberately preserved between rollout batches so SGLang’s RadixAttention prefix caching stays warm (it is flushed only at weight-sync time via invalidate_kv_cache()).

shutdown() → bool

Shut down all SGLang workers and clean up resources.

invalidate_kv_cache() → bool

Invalidate KV cache before weight updates.

This flushes the cache before weight updates to clear stale cache. Only primary workers (TP rank 0, model owners) will flush their cache.

Returns: bool: True if all caches were flushed successfully, False otherwise