minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Noa Neria <noa@run.ai>
This commit is contained in:
Noa Neria
2026-04-01 18:41:22 -07:00
committed by GitHub
parent ae3b207dfd
commit 8d9145d97e
14 changed files with 659 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/pyproject.toml
+2 -1
View File
@@ -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
+1
View File
@@ -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
+9
View File
@@ -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
+242
View File
@@ -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
+21 -7
View File
@@ -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
+24 -6
View File
@@ -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
+7
View File
@@ -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
+134
View File
@@ -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,
)