flyteplugins-sglang 2.0.0b45__py3-none-any.whl

flyteplugins/sglang/__init__.py

@@ -0,0 +1,3 @@
+__all__ = ["SGLangAppEnvironment"]
+
+from flyteplugins.sglang._app_environment import SGLangAppEnvironment

flyteplugins/sglang/_app_environment.py

@@ -0,0 +1,212 @@
+from __future__ import annotations
+
+import shlex
+from dataclasses import dataclass, field, replace
+from typing import Any, Literal, Optional, Union
+
+import flyte.app
+import rich.repr
+from flyte import Environment, Image, Resources, SecretRequest
+from flyte.app import Parameter, RunOutput
+from flyte.app._types import Port
+from flyte.models import SerializationContext
+
+from flyteplugins.sglang._constants import SGLANG_MIN_VERSION_STR
+
+DEFAULT_SGLANG_IMAGE = (
+    flyte.Image.from_debian_base(name="sglang-app-image")
+    # install system dependencies, including CUDA toolkit, which is needed by sglang for compiling the model
+    .with_apt_packages("libnuma-dev", "wget")
+    .with_commands(
+        [
+            "wget https://developer.download.nvidia.com/compute/cuda/repos/debian12/x86_64/cuda-keyring_1.1-1_all.deb",
+            "dpkg -i cuda-keyring_1.1-1_all.deb",
+            "apt-get update",
+            "apt-get install -y cuda-toolkit-12-8",
+        ]
+    )
+    # install flash-infer
+    .with_pip_packages("flashinfer-python", "flashinfer-cubin")
+    .with_pip_packages("flashinfer-jit-cache", index_url="https://flashinfer.ai/whl/cu128")
+    .with_pip_packages("flyteplugins-sglang", pre=True)
+    .with_pip_packages(f"sglang>={SGLANG_MIN_VERSION_STR}")
+    .with_env_vars({"CUDA_HOME": "/usr/local/cuda-12.8"})
+)
+
+
+@rich.repr.auto
+@dataclass(kw_only=True, repr=True)
+class SGLangAppEnvironment(flyte.app.AppEnvironment):
+    """
+    App environment backed by SGLang for serving large language models.
+
+    This environment sets up an SGLang server with the specified model and configuration.
+
+    :param name: The name of the application.
+    :param container_image: The container image to use for the application.
+    :param port: Port application listens to. Defaults to 8000 for SGLang.
+    :param requests: Compute resource requests for application.
+    :param secrets: Secrets that are requested for application.
+    :param limits: Compute resource limits for application.
+    :param env_vars: Environment variables to set for the application.
+    :param scaling: Scaling configuration for the app environment.
+    :param domain: Domain to use for the app.
+    :param cluster_pool: The target cluster_pool where the app should be deployed.
+    :param requires_auth: Whether the public URL requires authentication.
+    :param type: Type of app.
+    :param extra_args: Extra args to pass to `python -m sglang.launch_server`. See
+        https://docs.sglang.io/advanced_features/server_arguments.html for details.
+    :param model_path: Remote path to model (e.g., s3://bucket/path/to/model).
+    :param model_hf_path: Hugging Face path to model (e.g., Qwen/Qwen3-0.6B).
+    :param model_id: Model id that is exposed by SGLang.
+    :param stream_model: Set to True to stream model from blob store to the GPU directly.
+        If False, the model will be downloaded to the local file system first and then loaded
+        into the GPU.
+    """
+
+    port: int | Port = 8080
+    type: str = "SGLang"
+    extra_args: str | list[str] = ""
+    model_path: str | RunOutput = ""
+    model_hf_path: str = ""
+    model_id: str = ""
+    stream_model: bool = True
+    image: str | Image | Literal["auto"] = DEFAULT_SGLANG_IMAGE
+    _model_mount_path: str = field(default="/root/flyte", init=False)
+
+    def __post_init__(self):
+        if self.env_vars is None:
+            self.env_vars = {}
+
+        if self._server is not None:
+            raise ValueError("server function cannot be set for SGLangAppEnvironment")
+
+        if self._on_startup is not None:
+            raise ValueError("on_startup function cannot be set for SGLangAppEnvironment")
+
+        if self._on_shutdown is not None:
+            raise ValueError("on_shutdown function cannot be set for SGLangAppEnvironment")
+
+        if self.model_id == "":
+            raise ValueError("model_id must be defined")
+
+        if self.model_path == "" and self.model_hf_path == "":
+            raise ValueError("model_path or model_hf_path must be defined")
+        if self.model_path != "" and self.model_hf_path != "":
+            raise ValueError("model_path and model_hf_path cannot be set at the same time")
+
+        if self.model_hf_path:
+            self._model_mount_path = self.model_hf_path
+
+        if self.args:
+            raise ValueError("args cannot be set for SGLangAppEnvironment. Use `extra_args` to add extra arguments.")
+
+        if isinstance(self.extra_args, str):
+            extra_args = shlex.split(self.extra_args)
+        else:
+            extra_args = self.extra_args
+
+        self.args = [
+            "sglang-fserve",
+            "--model-path",
+            self._model_mount_path,
+            "--served-model-name",
+            self.model_id,
+            "--port",
+            str(self.get_port().port),
+            *extra_args,
+        ]
+
+        if self.parameters:
+            raise ValueError("parameters cannot be set for SGLangAppEnvironment")
+
+        input_kwargs = {}
+        if self.stream_model:
+            self.env_vars["FLYTE_MODEL_LOADER_STREAM_SAFETENSORS"] = "true"
+            input_kwargs["env_var"] = "FLYTE_MODEL_LOADER_REMOTE_MODEL_PATH"
+            input_kwargs["download"] = False
+        else:
+            self.env_vars["FLYTE_MODEL_LOADER_STREAM_SAFETENSORS"] = "false"
+            input_kwargs["download"] = True
+            input_kwargs["mount"] = self._model_mount_path
+
+        if self.model_path:
+            self.parameters = [Parameter(name="model_path", value=self.model_path, **input_kwargs)]
+
+        self.env_vars["FLYTE_MODEL_LOADER_LOCAL_MODEL_PATH"] = self._model_mount_path
+        self.links = [flyte.app.Link(path="/docs", title="SGLang OpenAPI Docs", is_relative=True)]
+
+        if self.image is None or self.image == "auto":
+            self.image = DEFAULT_SGLANG_IMAGE
+
+        super().__post_init__()
+
+    def container_args(self, serialization_context: SerializationContext) -> list[str]:
+        """Return the container arguments for SGLang."""
+        if isinstance(self.args, str):
+            return shlex.split(self.args)
+        return self.args or []
+
+    def clone_with(
+        self,
+        name: str,
+        image: Optional[Union[str, Image, Literal["auto"]]] = None,
+        resources: Optional[Resources] = None,
+        env_vars: Optional[dict[str, str]] = None,
+        secrets: Optional[SecretRequest] = None,
+        depends_on: Optional[list[Environment]] = None,
+        description: Optional[str] = None,
+        interruptible: Optional[bool] = None,
+        **kwargs: Any,
+    ) -> SGLangAppEnvironment:
+        port = kwargs.pop("port", None)
+        extra_args = kwargs.pop("extra_args", None)
+        if "model_path" in kwargs:
+            set_model_path = True
+            model_path = kwargs.pop("model_path", "") or ""
+        else:
+            set_model_path = False
+            model_path = self.model_path
+        if "model_hf_path" in kwargs:
+            set_model_hf_path = True
+            model_hf_path = kwargs.pop("model_hf_path", "") or ""
+        else:
+            set_model_hf_path = False
+            model_hf_path = self.model_hf_path
+        model_id = kwargs.pop("model_id", None)
+        stream_model = kwargs.pop("stream_model", None)
+
+        if kwargs:
+            raise TypeError(f"Unexpected keyword arguments: {list(kwargs.keys())}")
+
+        kwargs = self._get_kwargs()
+        kwargs["name"] = name
+        kwargs["args"] = None
+        kwargs["parameters"] = None
+        if image is not None:
+            kwargs["image"] = image
+        if resources is not None:
+            kwargs["resources"] = resources
+        if env_vars is not None:
+            kwargs["env_vars"] = env_vars
+        if secrets is not None:
+            kwargs["secrets"] = secrets
+        if depends_on is not None:
+            kwargs["depends_on"] = depends_on
+        if description is not None:
+            kwargs["description"] = description
+        if interruptible is not None:
+            kwargs["interruptible"] = interruptible
+        if port is not None:
+            kwargs["port"] = port
+        if extra_args is not None:
+            kwargs["extra_args"] = extra_args
+        if set_model_path:
+            kwargs["model_path"] = model_path
+        if set_model_hf_path:
+            kwargs["model_hf_path"] = model_hf_path
+        if model_id is not None:
+            kwargs["model_id"] = model_id
+        if stream_model is not None:
+            kwargs["stream_model"] = stream_model
+        return replace(self, **kwargs)

flyteplugins/sglang/_constants.py

@@ -0,0 +1,2 @@
+SGLANG_MIN_VERSION = (0, 5, 2)
+SGLANG_MIN_VERSION_STR = ".".join(map(str, SGLANG_MIN_VERSION))

flyteplugins/sglang/_model_loader/shim.py

@@ -0,0 +1,171 @@
+import logging
+import os
+import sys
+from typing import Generator
+
+import torch
+from flyte.app.extras._model_loader.config import (
+    LOCAL_MODEL_PATH,
+    REMOTE_MODEL_PATH,
+    STREAM_SAFETENSORS,
+)
+from flyte.app.extras._model_loader.loader import SafeTensorsStreamer, prefetch
+
+from flyteplugins.sglang._constants import SGLANG_MIN_VERSION, SGLANG_MIN_VERSION_STR
+
+try:
+    import sglang
+except ImportError:
+    raise ImportError(
+        f"sglang is not installed. Please install 'sglang>={SGLANG_MIN_VERSION_STR}', to use the model loader."
+    )
+
+if tuple([int(part) for part in sglang.__version__.split(".") if part.isdigit()]) < SGLANG_MIN_VERSION:
+    raise ImportError(
+        f"sglang version >={SGLANG_MIN_VERSION_STR} required, but found {sglang.__version__}. Please upgrade sglang."
+    )
+
+import sglang.srt.model_loader.loader
+from sglang.srt.configs.device_config import DeviceConfig
+from sglang.srt.configs.model_config import ModelConfig
+from sglang.srt.server_args import prepare_server_args
+from sglang.srt.utils import kill_process_tree
+
+try:
+    from sglang.srt.server import launch_server
+except (ModuleNotFoundError, AttributeError, ImportError):
+    try:
+        from sglang.launch_server import launch_server
+    except (ModuleNotFoundError, AttributeError, ImportError):
+        from sglang.launch_server import run_server as launch_server
+
+
+logger = logging.getLogger(__name__)
+
+_OrigDefaultModelLoader = sglang.srt.model_loader.loader.DefaultModelLoader
+
+
+class FlyteModelLoader(_OrigDefaultModelLoader):
+    """Custom model loader for streaming model weights from object storage."""
+
+    def _get_weights_iterator(self, source) -> Generator[tuple[str, torch.Tensor], None, None]:
+        # Try to load weights using the Flyte SafeTensorsStreamer. Fallback to the default loader otherwise.
+        try:
+            streamer = SafeTensorsStreamer(REMOTE_MODEL_PATH, LOCAL_MODEL_PATH)
+        except ValueError:
+            return super()._get_weights_iterator(source)
+        else:
+            for name, tensor in streamer.get_tensors():
+                yield source.prefix + name, tensor
+
+    def download_model(self, model_config: ModelConfig) -> None:
+        # This model loader supports streaming only
+        pass
+
+    def _load_sharded_model(self, *, model_config: ModelConfig, device_config: DeviceConfig) -> torch.nn.Module:
+        # Forked from: https://github.com/sgl-project/sglang/blob/1c4e0d2445311f2e635e9dab5a660d982731ad20/python/sglang/srt/model_loader/loader.py#L564
+        from sglang.srt.distributed import (
+            get_tensor_model_parallel_rank,
+            get_tensor_model_parallel_world_size,
+            model_parallel_is_initialized,
+        )
+        from sglang.srt.model_loader.loader import (
+            ShardedStateLoader,
+            _initialize_model,
+        )
+        from sglang.srt.model_loader.utils import set_default_torch_dtype
+
+        # Sanity checks
+        if model_parallel_is_initialized():
+            tensor_parallel_size = get_tensor_model_parallel_world_size()
+            rank = get_tensor_model_parallel_rank()
+        else:
+            tensor_parallel_size = 1
+            rank = 0
+        if rank >= tensor_parallel_size:
+            raise ValueError(f"Invalid rank {rank} for tensor parallel size {tensor_parallel_size}")
+        with set_default_torch_dtype(model_config.dtype):
+            with torch.device(device_config.device):
+                model = _initialize_model(model_config, self.load_config)
+                for _, module in model.named_modules():
+                    quant_method = getattr(module, "quant_method", None)
+                    if quant_method is not None:
+                        quant_method.process_weights_after_loading(module)
+            state_dict = ShardedStateLoader._filter_subtensors(model.state_dict())
+            streamer = SafeTensorsStreamer(
+                REMOTE_MODEL_PATH,
+                LOCAL_MODEL_PATH,
+                rank=rank,
+                tensor_parallel_size=tensor_parallel_size,
+            )
+            for name, tensor in streamer.get_tensors():
+                # If loading with LoRA enabled, additional padding may
+                # be added to certain parameters. We only load into a
+                # narrowed view of the parameter data.
+                param_data = state_dict[name].data
+                param_shape = state_dict[name].shape
+                for dim, size in enumerate(tensor.shape):
+                    if size < param_shape[dim]:
+                        param_data = param_data.narrow(dim, 0, size)
+                if tensor.shape != param_shape:
+                    logger.warning(
+                        "loading tensor of shape %s into parameter '%s' of shape %s",
+                        tensor.shape,
+                        name,
+                        param_shape,
+                    )
+                param_data.copy_(tensor)
+                state_dict.pop(name)
+            if state_dict:
+                raise ValueError(f"Missing keys {tuple(state_dict)} in loaded state!")
+        return model.eval()
+
+    def load_model(self, *, model_config: ModelConfig, device_config: DeviceConfig) -> torch.nn.Module:
+        from sglang.srt.distributed import get_tensor_model_parallel_world_size
+
+        logger.info("Loading model with FlyteModelLoader")
+        if get_tensor_model_parallel_world_size() > 1:
+            return self._load_sharded_model(model_config=model_config, device_config=device_config)
+        else:
+            return super().load_model(model_config=model_config, device_config=device_config)
+
+
+# Monkeypatch the default model loader when streaming is enabled
+if REMOTE_MODEL_PATH and STREAM_SAFETENSORS:
+    sglang.srt.model_loader.loader.DefaultModelLoader = FlyteModelLoader
+
+
+async def _get_model_files():
+    import flyte.storage as storage
+
+    if not await storage.exists(REMOTE_MODEL_PATH):
+        raise FileNotFoundError(f"Model path not found: {REMOTE_MODEL_PATH}")
+
+    await prefetch(
+        REMOTE_MODEL_PATH,
+        LOCAL_MODEL_PATH,
+        exclude_safetensors=STREAM_SAFETENSORS,
+    )
+
+
+def main():
+    import asyncio
+
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    )
+    logger.info(f"REMOTE_MODEL_PATH: {REMOTE_MODEL_PATH}")
+    logger.info(f"LOCAL_MODEL_PATH: {LOCAL_MODEL_PATH}")
+    logger.info(f"STREAM_SAFETENSORS: {STREAM_SAFETENSORS}")
+
+    # Prefetch the model
+    if REMOTE_MODEL_PATH:
+        logger.info("Prefetching model files...")
+        asyncio.run(_get_model_files())
+
+    server_args = prepare_server_args(sys.argv[1:])
+    try:
+        launch_server(server_args)
+    finally:
+        kill_process_tree(os.getpid(), include_parent=False)

flyteplugins_sglang-2.0.0b45.dist-info/METADATA

@@ -0,0 +1,69 @@
+Metadata-Version: 2.4
+Name: flyteplugins-sglang
+Version: 2.0.0b45
+Summary: SGLang plugin for flyte
+Author-email: Niels Bantilan <cosmicbboy@users.noreply.github.com>
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: flyte>=2.0.0b43
+
+# Flyte SGLang Plugin
+
+Serve large language models using SGLang with Flyte Apps.
+
+This plugin provides the `SGLangAppEnvironment` class for deploying and serving LLMs using [SGLang](https://docs.sglang.ai/).
+
+## Installation
+
+```bash
+pip install --pre flyteplugins-sglang
+```
+
+## Usage
+
+```python
+import flyte
+import flyte.app
+from flyteplugins.sglang import SGLangAppEnvironment
+
+# Define the SGLang app environment
+sglang_app = SGLangAppEnvironment(
+    name="my-llm-app",
+    model="s3://your-bucket/models/your-model",
+    model_id="your-model-id",
+    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="L40s:1"),
+    stream_model=True,  # Stream model directly from blob store to GPU
+    scaling=flyte.app.Scaling(
+        replicas=(0, 1),
+        scaledown_after=300,
+    ),
+)
+
+if __name__ == "__main__":
+    flyte.init_from_config()
+    app = flyte.serve(sglang_app)
+    print(f"Deployed SGLang app: {app.url}")
+```
+
+## Features
+
+- **Streaming Model Loading**: Stream model weights directly from object storage to GPU memory, reducing startup time and disk requirements.
+- **OpenAI-Compatible API**: The deployed app exposes an OpenAI-compatible API for chat completions.
+- **Auto-scaling**: Configure scaling policies to scale up/down based on traffic.
+- **Tensor Parallelism**: Support for distributed inference across multiple GPUs.
+
+## Extra Arguments
+
+You can pass additional arguments to the SGLang server using the `extra_args` parameter:
+
+```python
+sglang_app = SGLangAppEnvironment(
+    name="my-llm-app",
+    model="s3://your-bucket/models/your-model",
+    model_id="your-model-id",
+    extra_args="--max-model-len 8192 --enforce-eager",
+)
+```
+
+See the [SGLang server arguments documentation](https://docs.sglang.ai/backend/server_arguments.html) for available options.
+

flyteplugins_sglang-2.0.0b45.dist-info/RECORD

@@ -0,0 +1,11 @@
+flyteplugins/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+flyteplugins/sglang/__init__.py,sha256=xK4xPyWoYdYr8tuIhusQh-nAXfvwW5cOO0pggY0jvZ8,106
+flyteplugins/sglang/_app_environment.py,sha256=pc6nuNpXXIBYuEwPVaA_ed_MIeuaFFoeOqglL-_tzfY,8510
+flyteplugins/sglang/_constants.py,sha256=gCHEQipHhHY7J8L829C6e3XZ8sxX3ae0F7AJEcVzWBg,95
+flyteplugins/sglang/_model_loader/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+flyteplugins/sglang/_model_loader/shim.py,sha256=U6mHvQBQ3bRjWsVZQG9IxmikMeSTWlJCus6hrydw2zo,6775
+flyteplugins_sglang-2.0.0b45.dist-info/METADATA,sha256=5H4qZPEEufNajOceJcDIfHTRD9Av_yt77eg4_YlfARM,2054
+flyteplugins_sglang-2.0.0b45.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+flyteplugins_sglang-2.0.0b45.dist-info/entry_points.txt,sha256=6eEFwDOLxQIAgFc03TStKyreECu3gycDlBRCTQ0pn0Y,78
+flyteplugins_sglang-2.0.0b45.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
+flyteplugins_sglang-2.0.0b45.dist-info/RECORD,,

flyteplugins_sglang-2.0.0b45.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

flyteplugins_sglang-2.0.0b45.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+sglang-fserve = flyteplugins.sglang._model_loader.shim:main

flyteplugins_sglang-2.0.0b45.dist-info/top_level.txt

@@ -0,0 +1 @@
+flyteplugins