Direct model loading from object storage with Runai Model Streamer (#17948)

Signed-off-by: Noa Neria <noa@run.ai>

docs/advanced_features/object_storage.md

@@ -0,0 +1,108 @@
+# Loading Models from Object Storage
+
+SGLang supports direct loading of models from object storage (S3 and Google Cloud Storage) without requiring a full local download. This feature uses the `runai_streamer` load format to stream model weights directly from cloud storage, significantly reducing startup time and local storage requirements.
+
+## Overview
+
+When loading models from object storage, SGLang uses a two-phase approach:
+
+1. **Metadata Download** (once, before process launch): Configuration files and tokenizer files are downloaded to a local cache
+2. **Weight Streaming** (lazy, during model loading): Model weights are streamed directly from object storage as needed
+
+## Supported Storage Backends
+
+1. **Amazon S3**: `s3://bucket-name/path/to/model/`
+2. **Google Cloud Storage**: `gs://bucket-name/path/to/model/`
+3. **Azure Blob**: `az://some-azure-container/path/`
+4. **S3 compatible**: `s3://bucket-name/path/to/model/`
+
+## Quick Start
+
+### Basic Usage
+
+Simply provide an object storage URI as the model path:
+
+```bash
+# S3
+python -m sglang.launch_server \
+  --model-path s3://my-bucket/models/llama-3-8b/ \
+  --load-format runai_streamer
+
+# Google Cloud Storage
+python -m sglang.launch_server \
+  --model-path gs://my-bucket/models/llama-3-8b/ \
+  --load-format runai_streamer
+```
+
+**Note**: The `--load-format runai_streamer` is automatically detected when using object storage URIs, so you can omit it:
+
+```bash
+python -m sglang.launch_server \
+  --model-path s3://my-bucket/models/llama-3-8b/
+```
+
+### With Tensor Parallelism
+
+```bash
+python -m sglang.launch_server \
+  --model-path gs://my-bucket/models/llama-70b/ \
+  --tp 4 \
+  --model-loader-extra-config '{"distributed": true}'
+```
+
+## Configuration
+
+### Load Format
+
+The `runai_streamer` load format is specifically designed for object storage, ssd and shared file systems
+
+```bash
+python -m sglang.launch_server \
+  --model-path s3://bucket/model/ \
+  --load-format runai_streamer
+```
+
+### Extended Configuration Parameters
+
+Use `--model-loader-extra-config` to pass additional configuration as a JSON string:
+
+```bash
+python -m sglang.launch_server \
+  --model-path s3://bucket/model/ \
+  --model-loader-extra-config '{
+    "distributed": true,
+    "concurrency": 8,
+    "memory_limit": 2147483648
+  }'
+```
+
+#### Available Parameters
+
+| Parameter | Type | Description | Default |
+|-----------|------|-------------|---------|
+| `distributed` | bool | Enable distributed streaming for multi-GPU setups. Automatically set to `true` for object storage paths and cuda alike devices. | Auto-detected |
+| `concurrency` | int | Number of concurrent download streams. Higher values can improve throughput for large models. | 4 |
+| `memory_limit` | int | Memory limit (in bytes) for the streaming buffer. | System-dependent |
+
+
+## Performance Considerations
+
+### Distributed Streaming
+
+For multi-GPU setups, enable distributed streaming to parallelize weight loading between the processes:
+
+```bash
+python -m sglang.launch_server \
+  --model-path s3://bucket/model/ \
+  --tp 8 \
+  --model-loader-extra-config '{"distributed": true}'
+```
+
+## Limitations
+
+- **Supported Formats**: Currently only supports `.safetensors` weight format (recommended format)
+- **Supported Device**: Distributed streaming is supported on cuda alike devices. Otherwise fallback to non distributed streaming
+
+## See Also
+
+- [Runai model streamer documentation](https://github.com/run-ai/runai-model-streamer)

docs/advanced_features/server_arguments.md

@@ -84,7 +84,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s
 | `--tokenizer-mode` | Tokenizer mode. 'auto' will use the fast tokenizer if available, and 'slow' will always use the slow tokenizer. | `auto` | `auto`, `slow` |
 | `--tokenizer-worker-num` | The worker num of the tokenizer manager. | `1` | Type: int |
 | `--skip-tokenizer-init` | If set, skip init tokenizer and pass input_ids in generate request. | `False` | bool flag (set to enable) |
-| `--load-format` | The format of the model weights to load. "auto" will try to load the weights in the safetensors format and fall back to the pytorch bin format if safetensors format is not available. "pt" will load the weights in the pytorch bin format. "safetensors" will load the weights in the safetensors format. "npcache" will load the weights in pytorch format and store a numpy cache to speed up the loading. "dummy" will initialize the weights with random values, which is mainly for profiling."gguf" will load the weights in the gguf format. "bitsandbytes" will load the weights using bitsandbytes quantization."layered" loads weights layer by layer so that one can quantize a layer before loading another to make the peak memory envelope smaller. "flash_rl" will load the weights in flash_rl format. "fastsafetensors" and "private" are also supported. | `auto` | `auto`, `pt`, `safetensors`, `npcache`, `dummy`, `sharded_state`, `gguf`, `bitsandbytes`, `layered`, `flash_rl`, `remote`, `remote_instance`, `fastsafetensors`, `private` |
+| `--load-format` | The format of the model weights to load. "auto" will try to load the weights in the safetensors format and fall back to the pytorch bin format if safetensors format is not available. "pt" will load the weights in the pytorch bin format. "safetensors" will load the weights in the safetensors format. "npcache" will load the weights in pytorch format and store a numpy cache to speed up the loading. "dummy" will initialize the weights with random values, which is mainly for profiling."gguf" will load the weights in the gguf format. "bitsandbytes" will load the weights using bitsandbytes quantization."layered" loads weights layer by layer so that one can quantize a layer before loading another to make the peak memory envelope smaller. "flash_rl" will load the weights in flash_rl format. "fastsafetensors" and "private" are also supported. "runai_streamer" enables direct model loading from object storage and shared file systems.| `auto` | `auto`, `pt`, `safetensors`, `npcache`, `dummy`, `sharded_state`, `gguf`, `bitsandbytes`, `layered`, `flash_rl`, `remote`, `remote_instance`, `fastsafetensors`, `private`, `runai_streamer` |
 | `--model-loader-extra-config` | Extra config for model loader. This will be passed to the model loader corresponding to the chosen load_format. | `{}` | Type: str |
 | `--trust-remote-code` | Whether or not to allow for custom models defined on the Hub in their own modeling files. | `False` | bool flag (set to enable) |
 | `--context-length` | The model's maximum context length. Defaults to None (will use the value from the model's config.json instead). | `None` | Type: int |

docs/index.rst

@@ -41,6 +41,7 @@ Its core features include:
    :caption: Advanced Features
 
    advanced_features/server_arguments.md
+   advanced_features/object_storage.md
    advanced_features/hyperparameter_tuning.md
    advanced_features/attention_backend.md
    advanced_features/speculative_decoding.ipynb

python/pyproject.toml

@@ -97,6 +97,7 @@ torch = [
 
 [project.optional-dependencies]
 checkpoint-engine = ["checkpoint-engine==0.1.2"]
+runai = ["runai-model-streamer[s3,gcs,azure]>=0.15.7"]
 diffusion = [
   "PyYAML==6.0.1",
   "cloudpickle==3.1.2",
@@ -108,7 +109,7 @@ diffusion = [
   "remote-pdb==2.1.0",
   "st_attn==0.0.7 ; platform_machine != 'aarch64' and platform_machine != 'arm64'",
   "vsa==0.0.4 ; platform_machine != 'aarch64' and platform_machine != 'arm64'",
-  "runai_model_streamer>=0.15.5",
+  "runai_model_streamer>=0.15.7",
   "cache-dit==1.3.0",
   "addict==2.4.0",
   "av==16.1.0",

python/sglang/srt/configs/load_config.py

@@ -31,6 +31,7 @@ class LoadFormat(str, enum.Enum):
     LOCAL_CACHED = "local_cached"
     FASTSAFETENSORS = "fastsafetensors"
     PRIVATE = "private"
+    RUNAI_STREAMER = "runai_streamer"
 
 
 @dataclass

python/sglang/srt/configs/model_config.py

@@ -34,6 +34,7 @@ from sglang.srt.utils.hf_transformers_utils import (
     get_hf_text_config,
     get_sparse_attention_config,
 )
+from sglang.srt.utils.runai_utils import ObjectStorageModel, is_runai_obj_uri
 from sglang.utils import is_in_ci
 
 logger = logging.getLogger(__name__)
@@ -129,6 +130,7 @@ class ModelConfig:
         self._validate_quantize_and_serve_config()
 
         # Get hf config
+        self._maybe_pull_model_for_runai(self.model_path)
         self._maybe_pull_model_tokenizer_from_remote()
         self.model_override_args = json.loads(model_override_args)
         kwargs = {}
@@ -1148,6 +1150,13 @@ class ModelConfig:
 
         return default_sampling_params
 
+    def _maybe_pull_model_for_runai(self, model: str) -> None:
+        if is_runai_obj_uri(model):
+            # local path for loading the config
+            self.model_path = ObjectStorageModel.get_path(model)
+            # remote path for loading the weights
+            self.model_weights = model
+
     def _maybe_pull_model_tokenizer_from_remote(self) -> None:
         """
         Pull the model config files to a temporary

python/sglang/srt/model_loader/loader.py

@@ -2863,6 +2863,245 @@ class ModelOptModelLoader(DefaultModelLoader):
         return model.eval()
 
 
+class RunaiModelStreamerLoader(BaseModelLoader):
+    """
+    Model loader that uses Runai Model Streamer to load a model.
+
+    Supports fast model loading from SSDs, shared filesystems and object storage (S3, GCS, Azure blob) with weight streaming.
+
+    Configuration (via load_config.model_loader_extra_config):
+        - distributed (bool): Enable distributed streaming - True by default for url paths (object storage)
+        - concurrency (int): Number of concurrent downloads
+        - memory_limit (int): Memory limit for streaming buffer
+
+    Note: Metadata files must be pre-downloaded via
+    ObjectStorageModel.download_and_get_path() before instantiation.
+    """
+
+    @dataclasses.dataclass
+    class Source:
+        """A source for weights."""
+
+        model_or_path: str
+        """The model ID or path."""
+
+        revision: Optional[str]
+        """The optional model revision."""
+
+        prefix: str = ""
+        """A prefix to prepend to all weights."""
+
+        fall_back_to_pt: bool = True
+        """Whether .pt weights can be used."""
+
+        model_config: Optional["ModelConfig"] = None
+        """The model configuration (for checking architecture, etc)."""
+
+        @classmethod
+        def init_new(cls, model_config: ModelConfig, model):
+            model_weights = model_config.model_path
+            if hasattr(model_config, "model_weights"):
+                model_weights = model_config.model_weights
+            return cls(
+                model_weights,
+                model_config.revision,
+                prefix="",
+                fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True),
+                model_config=model_config,
+            )
+
+    def __init__(self, load_config: LoadConfig):
+        super().__init__(load_config)
+        extra_config = load_config.model_loader_extra_config
+        allowed_keys = {"distributed", "concurrency", "memory_limit"}
+        unexpected_keys = set(extra_config.keys()) - allowed_keys
+
+        if unexpected_keys:
+            raise ValueError(
+                f"Unexpected extra config keys for load format "
+                f"{load_config.load_format}: "
+                f"{unexpected_keys}"
+            )
+
+        set_runai_streamer_env(load_config)
+
+        self._is_distributed = None
+        if load_config.model_loader_extra_config:
+            extra_config = load_config.model_loader_extra_config
+
+            if "distributed" in extra_config and isinstance(
+                extra_config.get("distributed"), bool
+            ):
+                self._is_distributed = extra_config.get("distributed")
+
+    def _prepare_weights(
+        self, model_name_or_path: str, revision: Optional[str]
+    ) -> Tuple[str, List[str]]:
+        """Prepare weights for the model.
+
+        If the model is not local, it will be downloaded."""
+        from sglang.srt.utils.runai_utils import is_runai_obj_uri, list_safetensors
+
+        is_object_storage_path = is_runai_obj_uri(model_name_or_path)
+        if self._is_distributed is None:
+            self._is_distributed = is_object_storage_path
+        is_local = os.path.isdir(model_name_or_path)
+        safetensors_pattern = "*.safetensors"
+        index_file = SAFE_WEIGHTS_INDEX_NAME
+
+        hf_folder = (
+            model_name_or_path
+            if (is_local or is_object_storage_path)
+            else download_weights_from_hf(
+                model_name_or_path,
+                self.load_config.download_dir,
+                [safetensors_pattern],
+                revision,
+                ignore_patterns=self.load_config.ignore_patterns,
+            )
+        )
+
+        server_args = get_global_server_args()
+        if server_args and server_args.model_checksum is not None:
+            from sglang.srt.utils.model_file_verifier import verify
+
+            checksums_source = server_args.model_checksum or model_name_or_path
+            verify(model_path=hf_folder, checksums_source=checksums_source)
+
+        hf_weights_files = list_safetensors(path=hf_folder)
+
+        # For models like Mistral-7B-Instruct-v0.3
+        # there are both sharded safetensors files and a consolidated
+        # safetensors file. Using both breaks.
+        # Here, we download the `model.safetensors.index.json` and filter
+        # any files not found in the index.
+        if not is_local and not is_object_storage_path:
+            download_safetensors_index_file_from_hf(
+                model_name_or_path,
+                index_file,
+                self.load_config.download_dir,
+                revision,
+            )
+        hf_weights_files = filter_duplicate_safetensors_files(
+            hf_weights_files, hf_folder, index_file
+        )
+
+        if len(hf_weights_files) == 0:
+            raise RuntimeError(
+                f"Cannot find any model weights with `{model_name_or_path}`"
+            )
+
+        return hf_folder, hf_weights_files
+
+    def _get_weights_iterator(
+        self, source: "Source"
+    ) -> Generator[Tuple[str, torch.Tensor], None, None]:
+        """Get an iterator for the model weights based on the load format."""
+        from sglang.srt.model_loader.weight_utils import (
+            runai_safetensors_weights_iterator,
+        )
+
+        hf_folder, hf_weights_files = self._prepare_weights(
+            source.model_or_path, source.revision
+        )
+
+        if source.model_config is not None:
+            hf_weights_files = maybe_add_mtp_safetensors(
+                hf_weights_files,
+                hf_folder,
+                "model.safetensors.index.json",
+                source.model_config.hf_config,
+            )
+
+        weights_iterator = runai_safetensors_weights_iterator(
+            hf_weights_files, self._is_distributed, self.target_device_str
+        )
+
+        if self.load_config.draft_model_idx is not None:
+            import re
+
+            def filter_weights(original_weights_iterator):
+                pattern = r"model.mtp.layers.(\d+)."
+                for name, tensor in original_weights_iterator:
+                    group = re.match(pattern, name)
+                    if group is not None:
+                        idx = int(group.group(1))
+                        if idx != self.load_config.draft_model_idx:
+                            continue
+                        new_name = name.replace(group.group(), "model.mtp.layers.0.")
+                    else:
+                        new_name = name
+                    yield (new_name, tensor)
+
+            weights_iterator = filter_weights(weights_iterator)
+
+        def apply_prefix(original_weights_iterator):
+            yield from (
+                (source.prefix + name, tensor)
+                for (name, tensor) in original_weights_iterator
+            )
+
+        return apply_prefix(weights_iterator)
+
+    def _get_all_weights(
+        self,
+        model_config: ModelConfig,
+        model: nn.Module,
+    ) -> Generator[Tuple[str, torch.Tensor], None, None]:
+
+        primary_weights = RunaiModelStreamerLoader.Source.init_new(model_config, model)
+        yield from self._get_weights_iterator(primary_weights)
+
+        secondary_weights = cast(
+            Iterable[RunaiModelStreamerLoader.Source],
+            getattr(model, "secondary_weights", ()),
+        )
+        for source in secondary_weights:
+            yield from self._get_weights_iterator(source)
+
+    def download_model(self, model_config: ModelConfig) -> None:
+        self._prepare_weights(model_config.model_path, model_config.revision)
+
+    def load_model(
+        self,
+        *,
+        model_config: ModelConfig,
+        device_config: DeviceConfig,
+    ) -> nn.Module:
+
+        if hasattr(model_config, "modelopt_quant") and model_config.modelopt_quant:
+            # Load base model using shared method
+            raise NotImplementedError(
+                "Runai Model Streamer Loader does not support ModelOpt quantization yet"
+            )
+
+        assert device_config.device_type in ("cuda", "cpu"), (
+            f"Runai Model Streamer only supports CUDA and CPU, "
+            f"got {device_config.device_type}"
+        )
+
+        if device_config.device_type == "cuda":
+            self.target_device_str = (
+                device_config.device_type + ":" + str(device_config.gpu_id)
+            )
+        else:
+            self.target_device_str = "cpu"
+
+        target_device = torch.device(device_config.device)
+        with set_default_torch_dtype(model_config.dtype):
+            with target_device:
+                model = _initialize_model(
+                    model_config,
+                    self.load_config,
+                )
+
+            DefaultModelLoader.load_weights_and_postprocess(
+                model, self._get_all_weights(model_config, model), target_device
+            )
+
+        return model.eval()
+
+
 def get_model_loader(
     load_config: LoadConfig, model_config: Optional[ModelConfig] = None
 ) -> BaseModelLoader:
@@ -2947,4 +3186,7 @@ def get_model_loader(
         except ImportError:
             raise ValueError("Failed to import sglang.private.private_model_loader")
 
+    if load_config.load_format == LoadFormat.RUNAI_STREAMER:
+        return RunaiModelStreamerLoader(load_config)
+
     return DefaultModelLoader(load_config)

python/sglang/srt/model_loader/weight_utils.py

@@ -58,6 +58,7 @@ from sglang.srt.utils import (
     log_info_on_rank0,
     print_warning_once,
 )
+from sglang.srt.utils.common import is_cuda_alike
 from sglang.utils import is_in_ci
 
 try:
@@ -1088,7 +1089,7 @@ def composed_weight_loader(
 
 
 def runai_safetensors_weights_iterator(
-    hf_weights_files: List[str],
+    hf_weights_files: List[str], is_distributed: bool = False, device: str = "cpu"
 ) -> Generator[Tuple[str, torch.Tensor], None, None]:
     """Iterate over the weights in the model safetensor files."""
     from runai_model_streamer import SafetensorsStreamer
@@ -1096,17 +1097,30 @@ def runai_safetensors_weights_iterator(
     enable_tqdm = (
         not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
     )
+    device = device if is_distributed and is_cuda_alike() else "cpu"
 
     with SafetensorsStreamer() as streamer:
-        for st_file in tqdm(
+
+        streamer.stream_files(
             hf_weights_files,
+            device=device,
+            is_distributed=is_distributed,
+        )
+        total_tensors = sum(
+            len(tensors_meta)
+            for tensors_meta in streamer.files_to_tensors_metadata.values()
+        )
+
+        tensor_iter = tqdm(
+            streamer.get_tensors(),
+            total=total_tensors,
             desc="Loading safetensors using Runai Model Streamer",
-            disable=not enable_tqdm,
             bar_format=BAR_FORMAT,
-            position=tqdm._get_free_pos(),
-        ):
-            streamer.stream_file(st_file)
-            yield from streamer.get_tensors()
+            disable=not enable_tqdm,
+            mininterval=2,
+        )
+
+        yield from tensor_iter
 
 
 def set_runai_streamer_env(load_config: LoadConfig):

python/sglang/srt/server_args.py

@@ -68,6 +68,7 @@ from sglang.srt.utils.common import (
 )
 from sglang.srt.utils.hf_transformers_utils import check_gguf_file
 from sglang.srt.utils.network import NetworkAddress, get_free_port, wait_port_available
+from sglang.srt.utils.runai_utils import ObjectStorageModel, is_runai_obj_uri
 from sglang.utils import is_in_ci
 
 logger = logging.getLogger(__name__)
@@ -91,6 +92,7 @@ LOAD_FORMAT_CHOICES = [
     "remote_instance",
     "fastsafetensors",
     "private",
+    "runai_streamer",
 ]
 
 QUANTIZATION_CHOICES = [
@@ -745,6 +747,8 @@ class ServerArgs:
         Orchestrates the handling of various server arguments, ensuring proper configuration and validation.
         """
 
+        self._maybe_download_model_for_runai()
+
         # Normalize load balancing defaults early (before dummy-model short-circuit).
         self._handle_load_balance_method()
 
@@ -846,6 +850,17 @@ class ServerArgs:
         # Handle any other necessary validations.
         self._handle_other_validations()
 
+    def _maybe_download_model_for_runai(self):
+        if is_runai_obj_uri(self.model_path):
+            ObjectStorageModel.download_and_get_path(self.model_path)
+
+        if (
+            self.tokenizer_path is not None
+            and is_runai_obj_uri(self.tokenizer_path)
+            and self.tokenizer_path != self.model_path
+        ):
+            ObjectStorageModel.download_and_get_path(self.tokenizer_path)
+
     def _handle_load_balance_method(self):
         if self.disaggregation_mode not in ("null", "prefill", "decode"):
             raise ValueError(
@@ -3130,7 +3145,9 @@ class ServerArgs:
                 "Detected Mistral native format checkpoint, setting load_format='mistral'"
             )
 
-        if is_remote_url(self.model_path):
+        if is_runai_obj_uri(self.model_path):
+            self.load_format = "runai_streamer"
+        elif is_remote_url(self.model_path):
             self.load_format = "remote"
 
         if self.custom_weight_loader is None:
@@ -6075,11 +6092,12 @@ class ServerArgs:
         }, "moe_dense_tp_size only support 1 and None currently"
 
         # Check served model name to not have colon as it is reserved for LoRA adapter syntax
-        assert ":" not in self.served_model_name, (
-            "served_model_name cannot contain a colon (':') character. "
-            "The colon is reserved for the 'model:adapter' syntax used in LoRA adapter specification. "
-            f"Invalid value: '{self.served_model_name}'"
-        )
+        if not is_runai_obj_uri(self.served_model_name):
+            assert ":" not in self.served_model_name, (
+                "served_model_name cannot contain a colon (':') character. "
+                "The colon is reserved for the 'model:adapter' syntax used in LoRA adapter specification. "
+                f"Invalid value: '{self.served_model_name}'"
+            )
 
         # Check LoRA
         self.check_lora_server_args()

python/sglang/srt/utils/hf_transformers_utils.py

@@ -27,6 +27,7 @@ import torch
 from huggingface_hub import snapshot_download
 
 from sglang.srt.utils import get_bool_env_var
+from sglang.srt.utils.runai_utils import ObjectStorageModel, is_runai_obj_uri
 
 # Compatibility shim: flash-attn-4 registers a bare ``flash_attn`` namespace
 # that makes ``is_flash_attn_2_available()`` return True, but lacks the v2 API
@@ -488,6 +489,9 @@ def get_config(
         kwargs["gguf_file"] = model
         model = Path(model).parent
 
+    if is_runai_obj_uri(model):
+        model = ObjectStorageModel.get_path(model)
+
     if is_remote_url(model):
         # BaseConnector implements __del__() to clean up the local dir.
         # Since config files need to exist all the time, so we DO NOT use
@@ -798,6 +802,9 @@ def get_tokenizer(
         kwargs["gguf_file"] = tokenizer_name
         tokenizer_name = Path(tokenizer_name).parent
 
+    if is_runai_obj_uri(tokenizer_name):
+        tokenizer_name = ObjectStorageModel.get_path(tokenizer_name)
+
     if is_remote_url(tokenizer_name):
         # BaseConnector implements __del__() to clean up the local dir.
         # Since config files need to exist all the time, so we DO NOT use

python/sglang/srt/utils/runai_utils.py

@@ -0,0 +1,134 @@
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/model_loader/runai_utils.py
+
+import hashlib
+import logging
+import os
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+SUPPORTED_SCHEMES = ["s3://", "gs://", "az://"]
+
+# Design Pattern: Single Metadata Download Before Process Launch
+
+#   1. Engine entrypoint (engine.py) or server arguments post init  (server_args.py):
+#     - Downloads config/tokenizer metadata ONCE before launching subprocesses
+#     - This happens in the main process, avoiding multi-process coordination
+#
+#   2. ModelConfig/HF Utils (model_config.py, hf_transformers_utils.py):
+#     - Use ObjectStorageModel.get_path() to retrieve the cached local path
+#     - NO re-download - just path resolution
+#
+#   3. RunaiModelStreamerLoader (loader.py):
+#     - Calls list_safetensors() which operates directly on the object storage URI
+#     - Streams weights lazily during model loading
+
+#   This avoids file locks, race conditions, and duplicate downloads
+
+
+def get_cache_dir() -> str:
+    # Expand user path (~) to ensure absolute paths for locking
+    path = os.getenv("SGLANG_CACHE_DIR", "~/.cache/sglang/")
+    return os.path.expanduser(path)
+
+
+def list_safetensors(path: str = "") -> list[str]:
+    """
+    List full file names from object path and filter by allow pattern.
+
+    Args:
+        path: The object storage path to list from.
+
+    Returns:
+        list[str]: List of full object storage paths allowed by the pattern
+    """
+    from runai_model_streamer import list_safetensors as runai_list_safetensors
+
+    return runai_list_safetensors(path)
+
+
+def is_runai_obj_uri(model_or_path: str | Path) -> bool:
+    # Cast to str to handle pathlib.Path inputs which lack string methods (like .lower)
+    return str(model_or_path).lower().startswith(tuple(SUPPORTED_SCHEMES))
+
+
+class ObjectStorageModel:
+    """
+    Model loader that uses Runai Model Streamer to load a model.
+
+      Supports object storage (S3, GCS) with lazy weight streaming.
+
+      Configuration (via load_config.model_loader_extra_config):
+          - distributed (bool): Enable distributed streaming
+          - concurrency (int): Number of concurrent downloads
+          - memory_limit (int): Memory limit for streaming buffer
+
+      Note: Metadata files must be pre-downloaded via
+      ObjectStorageModel.download_and_get_path() before instantiation.
+
+    Attributes:
+        dir: The temporary created directory.
+    """
+
+    def __init__(self, url: str) -> None:
+        self.dir = ObjectStorageModel.get_path(url)
+
+        from runai_model_streamer import ObjectStorageModel as RunaiObjectStorageModel
+
+        self._runai_obj = RunaiObjectStorageModel(model_path=url, dst=self.dir)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return self._runai_obj.__exit__(exc_type, exc_val, exc_tb)
+
+    def pull_files(
+        self,
+        allow_pattern: list[str] | None = None,
+        ignore_pattern: list[str] | None = None,
+    ) -> None:
+        """Pull files from object storage into the local cache directory.
+
+        Args:
+            allow_pattern: File patterns to include (e.g. ["*.json"]).
+            ignore_pattern: File patterns to exclude.
+        """
+        self._runai_obj.pull_files(allow_pattern, ignore_pattern)
+
+    @classmethod
+    def download_and_get_path(cls, model_path: str) -> str:
+        """
+        Downloads the model metadata (excluding heavy weights) and returns
+        the local directory path. Safe for concurrent usage by multiple processes
+        """
+        with cls(url=model_path) as downloader:
+            downloader.pull_files(
+                ignore_pattern=[
+                    "*.pt",
+                    "*.safetensors",
+                    "*.bin",
+                    "*.tensors",
+                    "*.pth",
+                ],
+            )
+            cache_dir = downloader.dir
+            logger.info(f"Runai Model : {cache_dir}, metadata ready.")
+        return cache_dir
+
+    @classmethod
+    def get_path(cls, model_path: str) -> str:
+        """
+        Returns the local directory path.
+        """
+        model_hash = hashlib.sha256(str(model_path).encode()).hexdigest()[:16]
+        base_dir = get_cache_dir()
+
+        # Ensure base cache dir exists
+        os.makedirs(os.path.join(base_dir, "model_streamer"), exist_ok=True)
+
+        return os.path.join(
+            base_dir,
+            "model_streamer",
+            model_hash,
+        )

scripts/ci/cuda/ci_install_dependency.sh

@@ -218,9 +218,9 @@ mark_step_done "Uninstall Flashinfer"
 # Install main package
 # ------------------------------------------------------------------------------
 # Install the main package
-EXTRAS="dev"
+EXTRAS="dev,runai"
 if [ -n "$OPTIONAL_DEPS" ]; then
-    EXTRAS="dev,${OPTIONAL_DEPS}"
+    EXTRAS="dev,runai,${OPTIONAL_DEPS}"
 fi
 echo "Installing python extras: [${EXTRAS}]"
 source "$(dirname "$0")/cache_nvidia_wheels.sh"

test/registered/model_loading/test_runai_model_loader.py

@@ -0,0 +1,50 @@
+import unittest
+
+import sglang as sgl
+from sglang.srt.environ import temp_set_env
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=120, suite="stage-b-test-1-gpu-small")
+
+TEST_GCS_MODEL = "gs://vertex-model-garden-public-us/codegemma/codegemma-2b/"
+
+PROMPTS = [
+    "Hello, my name is",
+    "The president of the United States is",
+    "The capital of France is",
+    "The future of AI is",
+]
+
+
+class TestRunaiModelLoader(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        with temp_set_env(
+            GOOGLE_CLOUD_PROJECT="fake-project",
+            RUNAI_STREAMER_GCS_USE_ANONYMOUS_CREDENTIALS="true",
+            CLOUD_STORAGE_EMULATOR_ENDPOINT="https://storage.googleapis.com",
+        ):
+            cls.engine = sgl.Engine(
+                model_path=TEST_GCS_MODEL,
+                load_format="runai_streamer",
+                cuda_graph_max_bs=1,
+                max_total_tokens=64,
+            )
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "engine") and cls.engine:
+            cls.engine.shutdown()
+
+    def test_generate_produces_output(self):
+        outputs = self.engine.generate(PROMPTS)
+        self.assertEqual(len(outputs), len(PROMPTS))
+        for i, output in enumerate(outputs):
+            text = output["text"]
+            self.assertIsInstance(text, str)
+            self.assertGreater(len(text), 0, f"Prompt {i} produced empty output")
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/test_runai_utils.py

@@ -0,0 +1,57 @@
+import unittest
+from pathlib import Path
+
+from sglang.srt.configs.load_config import LoadFormat
+from sglang.srt.utils.runai_utils import ObjectStorageModel, is_runai_obj_uri
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="stage-a-test-cpu")
+
+
+class TestRunaiUtils(CustomTestCase):
+    def test_is_runai_obj_uri_s3(self):
+        self.assertTrue(is_runai_obj_uri("s3://bucket/model/"))
+        self.assertTrue(is_runai_obj_uri("S3://Bucket/Model/"))
+
+    def test_is_runai_obj_uri_gs(self):
+        self.assertTrue(is_runai_obj_uri("gs://bucket/model/"))
+        self.assertTrue(is_runai_obj_uri("GS://Bucket/Model/"))
+
+    def test_is_runai_obj_uri_az(self):
+        self.assertTrue(is_runai_obj_uri("az://container/model/"))
+        self.assertTrue(is_runai_obj_uri("AZ://Container/Model/"))
+
+    def test_is_runai_obj_uri_local_paths(self):
+        self.assertFalse(is_runai_obj_uri("/path/to/model"))
+        self.assertFalse(is_runai_obj_uri("./relative/path"))
+        self.assertFalse(is_runai_obj_uri("meta-llama/Llama-3.2-1B"))
+
+    def test_is_runai_obj_uri_other_schemes(self):
+        self.assertFalse(is_runai_obj_uri("http://example.com/model"))
+        self.assertFalse(is_runai_obj_uri("https://example.com/model"))
+        self.assertFalse(is_runai_obj_uri("ftp://example.com/model"))
+
+    def test_is_runai_obj_uri_pathlib(self):
+        self.assertFalse(is_runai_obj_uri(Path("/local/model")))
+
+    def test_get_path_deterministic(self):
+        path1 = ObjectStorageModel.get_path("s3://bucket/model/")
+        path2 = ObjectStorageModel.get_path("s3://bucket/model/")
+        self.assertEqual(path1, path2)
+
+    def test_get_path_different_uris(self):
+        path1 = ObjectStorageModel.get_path("s3://bucket/model-a/")
+        path2 = ObjectStorageModel.get_path("s3://bucket/model-b/")
+        self.assertNotEqual(path1, path2)
+
+    def test_get_path_contains_model_streamer(self):
+        path = ObjectStorageModel.get_path("s3://bucket/model/")
+        self.assertIn("model_streamer", path)
+
+    def test_load_format_enum(self):
+        self.assertEqual(LoadFormat.RUNAI_STREAMER.value, "runai_streamer")
+
+
+if __name__ == "__main__":
+    unittest.main()