sglang 0.5.1.post2__py3-none-any.whl → 0.5.2__py3-none-any.whl

sglang/srt/disaggregation/launch_lb.py

@@ -1,131 +0,0 @@
-import argparse
-import dataclasses
-
-from sglang.srt.disaggregation.mini_lb import PrefillConfig, run
-
-
-@dataclasses.dataclass
-class LBArgs:
-    rust_lb: bool = False
-    host: str = "0.0.0.0"
-    port: int = 8000
-    policy: str = "random"
-    prefill_infos: list = dataclasses.field(default_factory=list)
-    decode_infos: list = dataclasses.field(default_factory=list)
-    log_interval: int = 5
-    timeout: int = 600
-
-    @staticmethod
-    def add_cli_args(parser: argparse.ArgumentParser):
-        parser.add_argument(
-            "--rust-lb",
-            action="store_true",
-            help="Deprecated, please use SGLang Router instead, this argument will have no effect.",
-        )
-        parser.add_argument(
-            "--host",
-            type=str,
-            default=LBArgs.host,
-            help=f"Host to bind the server (default: {LBArgs.host})",
-        )
-        parser.add_argument(
-            "--port",
-            type=int,
-            default=LBArgs.port,
-            help=f"Port to bind the server (default: {LBArgs.port})",
-        )
-        parser.add_argument(
-            "--policy",
-            type=str,
-            default=LBArgs.policy,
-            choices=["random", "po2"],
-            help=f"Policy to use for load balancing (default: {LBArgs.policy})",
-        )
-        parser.add_argument(
-            "--prefill",
-            type=str,
-            default=[],
-            nargs="+",
-            help="URLs for prefill servers",
-        )
-        parser.add_argument(
-            "--decode",
-            type=str,
-            default=[],
-            nargs="+",
-            help="URLs for decode servers",
-        )
-        parser.add_argument(
-            "--prefill-bootstrap-ports",
-            type=int,
-            nargs="+",
-            help="Bootstrap ports for prefill servers",
-        )
-        parser.add_argument(
-            "--log-interval",
-            type=int,
-            default=LBArgs.log_interval,
-            help=f"Log interval in seconds (default: {LBArgs.log_interval})",
-        )
-        parser.add_argument(
-            "--timeout",
-            type=int,
-            default=LBArgs.timeout,
-            help=f"Timeout in seconds (default: {LBArgs.timeout})",
-        )
-
-    @classmethod
-    def from_cli_args(cls, args: argparse.Namespace) -> "LBArgs":
-        bootstrap_ports = args.prefill_bootstrap_ports
-        if bootstrap_ports is None:
-            bootstrap_ports = [None] * len(args.prefill)
-        elif len(bootstrap_ports) == 1:
-            bootstrap_ports = bootstrap_ports * len(args.prefill)
-        else:
-            if len(bootstrap_ports) != len(args.prefill):
-                raise ValueError(
-                    "Number of prefill URLs must match number of bootstrap ports"
-                )
-
-        prefill_infos = [
-            (url, port) for url, port in zip(args.prefill, bootstrap_ports)
-        ]
-
-        return cls(
-            rust_lb=args.rust_lb,
-            host=args.host,
-            port=args.port,
-            policy=args.policy,
-            prefill_infos=prefill_infos,
-            decode_infos=args.decode,
-            log_interval=args.log_interval,
-            timeout=args.timeout,
-        )
-
-    def __post_init__(self):
-        if not self.rust_lb:
-            assert (
-                self.policy == "random"
-            ), "Only random policy is supported for Python load balancer"
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="PD Disaggregation Load Balancer Server"
-    )
-    LBArgs.add_cli_args(parser)
-    args = parser.parse_args()
-    lb_args = LBArgs.from_cli_args(args)
-
-    prefill_configs = [PrefillConfig(url, port) for url, port in lb_args.prefill_infos]
-    run(
-        prefill_configs,
-        lb_args.decode_infos,
-        lb_args.host,
-        lb_args.port,
-        lb_args.timeout,
-    )
-
-
-if __name__ == "__main__":
-    main()

sglang/srt/mem_cache/storage/mooncake_store/unit_test.py

@@ -1,40 +0,0 @@
-import torch
-from mooncake_store import MooncakeStore
-
-
-def test_init_and_warmup():
-    store = MooncakeStore()
-    assert store.store is not None
-
-
-def test_register_buffer():
-    store = MooncakeStore()
-    tensor = torch.zeros(1024, dtype=torch.float32)
-    store.register_buffer(tensor)
-
-
-def test_set_and_get():
-    store = MooncakeStore()
-
-    key = ["test_key_" + str(i) for i in range(2)]
-    tensor = torch.arange(256, dtype=torch.float32).cuda()
-    ptrs = [tensor.data_ptr(), tensor.data_ptr()]
-    sizes = [tensor.numel() * tensor.element_size()] * 2
-
-    store.set(key, target_location=ptrs, target_sizes=sizes)
-    store.get(key, target_location=ptrs, target_sizes=sizes)
-
-
-def test_exists():
-    store = MooncakeStore()
-    keys = ["test_key_0", "non_existent_key"]
-    result = store.exists(keys)
-    assert isinstance(result, dict)
-    assert "test_key_0" in result
-
-
-if __name__ == "__main__":
-    test_init_and_warmup()
-    test_register_buffer()
-    test_set_and_get()
-    test_exists()