Fix model loading from S3 buckets

docker/Dockerfile

@@ -312,6 +312,11 @@ RUN --mount=type=cache,target=/root/.cache/pip curl --proto '=https' --tlsv1.2 -
     && python3 -m pip install --force-reinstall dist/*.whl \
     && rm -rf /root/.cargo /root/.rustup target dist ~/.cargo
 
+# Install S3 dependencies
+RUN python3 -m pip install --no-cache-dir --break-system-packages \
+    boto3 \
+    runai-model-streamer \
+    runai-model-streamer-s3
 
 # Add yank script
 COPY --chown=root:root --chmod=755 docker/configs/yank /usr/local/bin/yank

python/sglang/srt/configs/model_config.py

@@ -97,6 +97,7 @@ def __init__(
         model_impl: Union[str, ModelImpl] = ModelImpl.AUTO,
         sampling_defaults: str = "openai",
         quantize_and_serve: bool = False,
+        device: Optional[str] = None,
     ) -> None:
         # Parse args
         self.model_path = model_path
@@ -106,6 +107,7 @@ def __init__(
         self.model_impl = model_impl
         self.sampling_defaults = sampling_defaults
         self.quantize_and_serve = quantize_and_serve
+        self.device = device
 
         # Validate quantize_and_serve configuration
         self._validate_quantize_and_serve_config()
@@ -121,6 +123,7 @@ def __init__(
             trust_remote_code=trust_remote_code,
             revision=revision,
             model_override_args=self.model_override_args,
+            device=self.device,
             **kwargs,
         )
         self.hf_text_config = get_hf_text_config(self.hf_config)
@@ -234,6 +237,7 @@ def from_server_args(
             model_impl=server_args.model_impl,
             sampling_defaults=server_args.sampling_defaults,
             quantize_and_serve=server_args.quantize_and_serve,
+            device=server_args.device,
             override_config_file=server_args.decrypted_config_file,
             **kwargs,
         )
@@ -805,7 +809,7 @@ def _maybe_pull_model_tokenizer_from_remote(self) -> None:
             # BaseConnector implements __del__() to clean up the local dir.
             # Since config files need to exist all the time, so we DO NOT use
             # with statement to avoid closing the client.
-            client = create_remote_connector(self.model_path)
+            client = create_remote_connector(self.model_path, self.device)
             if is_remote_url(self.model_path):
                 client.pull_files(allow_pattern=["*config.json"])
                 self.model_weights = self.model_path

python/sglang/srt/connector/s3.py

@@ -101,7 +101,9 @@ def pull_files(
             return
 
         for file in files:
-            destination_file = os.path.join(self.local_dir, file.removeprefix(base_dir))
+            destination_file = os.path.join(
+                self.local_dir, file.removeprefix(f"{base_dir}/")
+            )
             local_dir = Path(destination_file).parent
             os.makedirs(local_dir, exist_ok=True)
             self.client.download_file(bucket_name, file, destination_file)

python/sglang/srt/managers/tokenizer_manager.py

@@ -263,6 +263,7 @@ def __init__(
                     tokenizer_mode=server_args.tokenizer_mode,
                     trust_remote_code=server_args.trust_remote_code,
                     revision=server_args.revision,
+                    device=server_args.device,
                 )
                 self._initialize_multi_item_delimiter_text()
 

python/sglang/srt/utils/hf_transformers_utils.py

@@ -192,6 +192,7 @@ def get_config(
     trust_remote_code: bool,
     revision: Optional[str] = None,
     model_override_args: Optional[dict] = None,
+    device: Optional[str] = None,
     **kwargs,
 ):
     is_gguf = check_gguf_file(model)
@@ -203,7 +204,7 @@ def get_config(
         # BaseConnector implements __del__() to clean up the local dir.
         # Since config files need to exist all the time, so we DO NOT use
         # with statement to avoid closing the client.
-        client = create_remote_connector(model)
+        client = create_remote_connector(model, device)
         client.pull_files(ignore_pattern=["*.pt", "*.safetensors", "*.bin"])
         model = client.get_local_dir()
 
@@ -361,6 +362,7 @@ def get_tokenizer(
     tokenizer_mode: str = "auto",
     trust_remote_code: bool = False,
     tokenizer_revision: Optional[str] = None,
+    device: Optional[str] = None,
     **kwargs,
 ) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]:
     """Gets a tokenizer for the given model name via Huggingface."""
@@ -387,7 +389,7 @@ def get_tokenizer(
         # BaseConnector implements __del__() to clean up the local dir.
         # Since config files need to exist all the time, so we DO NOT use
         # with statement to avoid closing the client.
-        client = create_remote_connector(tokenizer_name)
+        client = create_remote_connector(tokenizer_name, device)
         client.pull_files(ignore_pattern=["*.pt", "*.safetensors", "*.bin"])
         tokenizer_name = client.get_local_dir()