[lite] Support RemoteSGLangEngine with test cases.

arealite/README.md

@@ -83,7 +83,6 @@ def main_grpo():
             # or asynchronous rollout with filtering and off-policyness control
             # rollout_batch = rollout.prepare_batch(batch, 
             #                                       workflow=MyRolloutWorkflow(rollout_config.workflow),
-            #                                       offpolicyness=4,
             #                                       should_accept=lambda x: x['rewards'].mean() > 0)
 
             # In the single-controller mode

arealite/api/cli_args.py

@@ -1,5 +1,5 @@
 from dataclasses import asdict, dataclass, field
-from typing import List
+from typing import Dict, List, Optional
 
 
 @dataclass
@@ -70,8 +70,160 @@ def new(self, **kwargs):
         return GenerationHyperparameters(**args)
 
 
+@dataclass
+class SGLangConfig:
+    """Configuration for SGLang runtime. Refer to:
+    https://github.com/sgl-project/sglang for detailed documentation.
+    """
+
+    disable_cuda_graph: bool = False
+    disable_radix_cache: bool = False
+    disable_cuda_graph_padding: bool = False
+    enable_nccl_nvls: bool = False
+    disable_outlines_disk_cache: bool = False
+    disable_custom_all_reduce: bool = False
+    disable_overlap_schedule: bool = False
+    enable_mixed_chunk: bool = False
+    enable_dp_attention: bool = False
+    enable_ep_moe: bool = False
+    enable_torch_compile: bool = False
+    torch_compile_max_bs: int = 32
+    cuda_graph_max_bs: Optional[int] = None
+    cuda_graph_bs: Optional[List[int]] = None
+    torchao_config: str = ""
+    enable_nan_detection: bool = False
+    enable_p2p_check: bool = False
+    triton_attention_reduce_in_fp32: bool = False
+    triton_attention_num_kv_splits: int = 8
+    num_continuous_decode_steps: int = 1
+    enable_memory_saver: bool = False
+    allow_auto_truncate: bool = False
+    # NOTE: to avoid the illegal memory access error
+    attention_backend: Optional[str] = "flashinfer"
+    sampling_backend: Optional[str] = None
+    context_length: Optional[int] = 32768
+    mem_fraction_static: Optional[float] = 0.9
+    max_running_requests: Optional[int] = None
+    # NOTE: chunked_prefill_size is by default 8192 on GPUs with 80GB mem in SGLang,
+    # but we disable it to avoid precision issues
+    chunked_prefill_size: Optional[int] = -1
+    max_prefill_tokens: int = 32768
+    schedule_policy: str = "lpm"
+    schedule_conservativeness: float = 1.0
+    cpu_offload_gb: int = 0
+
+    dtype: str = "float16"
+    kv_cache_dtype: str = "auto"
+
+    # logging
+    log_level: str = "warning"
+    log_level_http: Optional[str] = "warning"
+    log_requests: bool = False
+    log_requests_level: int = 0
+    show_time_cost: bool = False
+    enable_metrics: bool = True  # Exports Prometheus-like metrics
+    # The interval (in decoding iterations) to log throughput
+    # and update prometheus metrics
+    decode_log_interval: int = 1
+
+    # Use staticmethod to make OmegaConf happy.
+    @staticmethod
+    def build_cmd(
+        sglang_config: "SGLangConfig",
+        model_path,
+        tp_size,
+        base_gpu_id,
+        dist_init_addr: Optional[str] = None,
+        served_model_name: Optional[str] = None,
+        skip_tokenizer_init: bool = True,
+    ):
+        from realhf.base import network, pkg_version, seeding
+        from realhf.experiments.common.utils import asdict as conf_as_dict
+
+        args: Dict = conf_as_dict(sglang_config)
+        args["random_seed"] = seeding.get_seed()
+
+        if served_model_name is None:
+            served_model_name = model_path
+        host_ip = network.gethostip()
+        host = "localhost" if not sglang_config.enable_metrics else host_ip
+        args = dict(
+            host=host,
+            model_path=model_path,
+            # Model and tokenizer
+            tokenizer_path=model_path,
+            tokenizer_mode="auto",
+            load_format="auto",
+            trust_remote_code=True,
+            device="cuda",
+            served_model_name=served_model_name,
+            is_embedding=False,
+            skip_tokenizer_init=skip_tokenizer_init,
+            # Other runtime options
+            tp_size=tp_size,
+            # Because we have set CUDA_VISIBLE_DEVICES to a single GPU in each process
+            base_gpu_id=base_gpu_id,
+            nnodes=1,
+            node_rank=0,
+            dist_init_addr=dist_init_addr,
+            **args,
+        )
+
+        if pkg_version.is_version_less("sglang", "0.4.4"):
+            args.pop("log_requests_level")
+        if pkg_version.is_version_less("sglang", "0.4.3"):
+            args.pop("enable_nccl_nvls")
+            args.pop("triton_attention_num_kv_splits")
+            args.pop("cuda_graph_bs")
+            args.pop("enable_memory_saver")
+            args.pop("allow_auto_truncate")
+            args.pop("file_storage_path")
+
+        flags = []
+        for k, v in args.items():
+            if v is None or v is False or v == "":
+                continue
+            if v is True:
+                flags.append(f"--{k.replace('_','-')} ")
+                continue
+            if isinstance(v, list):
+                values = " ".join(map(str, v))
+                flags.append(f"--{k.replace('_','-')} {values}")
+                continue
+            flags.append(f"--{k.replace('_','-')} {v}")
+        flags = " ".join(flags)
+        return f"python3 -m sglang.launch_server {flags}"
+
+
 @dataclass
 class InferenceEngineConfig:
+    experiment_name: str
+    trial_name: str
+    max_concurrent_rollouts: None | int = field(
+        default=None,
+        metadata={
+            "help": "Maximum number of concurrent requests to the inference engine."
+        },
+    )
+    queue_size: None | int = field(
+        default=None,
+        metadata={"help": "Input/Output queue size for async rollout."},
+    )
+    consumer_batch_size: int = field(
+        default=1,
+        metadata={
+            "help": "Batch size for consuming requests from the queue. "
+            "If None, it will be set to max_concurrent_requests.",
+        },
+    )
+    max_head_offpolicyness: int = field(
+        default=0,
+        metadata={
+            "help": "Maximum off-policyness for the head. "
+            "If the current version is more than this many versions behind, "
+            "the request will not be accepted.",
+        },
+    )
     # Used by remote inference engines.
     server_addrs: List[str] = field(
         default_factory=list,

arealite/api/engine_api.py

@@ -109,6 +109,10 @@ def initialize(self, addr: str | None, ft_spec):
         """Initialize environments for distributed inference and load models."""
         raise NotImplementedError()
 
+    def destroy(self):
+        """Destroy the engine and release GPU memory."""
+        pass
+
     def update_weights(self, meta: WeightUpdateMeta) -> Future:
         """Update weights in the inference engine."""
         raise NotImplementedError()
@@ -121,7 +125,7 @@ def submit(self, data: Dict[str, Any], workflow: "RolloutWorkflow") -> None:
         """Asynchronously submit a request to the inference engine. Exits immediately."""
         raise NotImplementedError()
 
-    def wait(self, count: int, timeout: int) -> TensorDict:
+    def wait(self, count: int, timeout: float) -> TensorDict:
         """Wait for a specified number of requests to complete, with a timeout."""
         raise NotImplementedError()
 

arealite/api/workflow_api.py

@@ -9,7 +9,7 @@
 class RolloutWorkflow:
 
     async def arun_episode(
-        self, engine: InferenceEngine, data: Dict[str, Any]
+        self, engine: "InferenceEngine", data: Dict[str, Any]
     ) -> TensorDict:
         """Run a single episode of the workflow.