Refactor is_s3 to is_cloud_storage

Signed-off-by: Peter Schuurman <[email protected]>

Author: pwschuurman

vllm/engine/arg_utils.py

@@ -80,7 +80,7 @@
     is_interleaved,
     maybe_override_with_speculators,
 )
-from vllm.transformers_utils.utils import check_gguf_file, is_s3
+from vllm.transformers_utils.utils import check_gguf_file, is_cloud_storage
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 from vllm.utils.mem_constants import GiB_bytes
 from vllm.utils.network_utils import get_ip
@@ -1297,10 +1297,10 @@ def create_engine_config(
 
         # Check if the model is a speculator and override model/tokenizer/config
         # BEFORE creating ModelConfig, so the config is created with the target model
-        # Skip speculator detection for S3 models since HuggingFace cannot load
-        # configs directly from S3 URLs. S3 models can still use speculators with
-        # explicit --speculative-config.
-        if not is_s3(self.model):
+        # Skip speculator detection for cloud storage models (eg: S3, GCS) since
+        # HuggingFace cannot load configs directly from S3 URLs. S3 models can still
+        # use speculators with explicit --speculative-config.
+        if not is_cloud_storage(self.model):
             (self.model, self.tokenizer, self.speculative_config) = (
                 maybe_override_with_speculators(
                     model=self.model,

vllm/model_executor/model_loader/sharded_state_loader.py

@@ -19,7 +19,7 @@
     runai_safetensors_weights_iterator,
 )
 from vllm.transformers_utils.s3_utils import glob as s3_glob
-from vllm.transformers_utils.utils import is_s3
+from vllm.transformers_utils.utils import is_cloud_storage
 
 logger = init_logger(__name__)
 
@@ -90,7 +90,7 @@ def get_end_ptr(tensor: torch.Tensor) -> int:
         return result
 
     def _prepare_weights(self, model_name_or_path: str, revision: str | None):
-        if is_s3(model_name_or_path) or os.path.isdir(model_name_or_path):
+        if is_cloud_storage(model_name_or_path) or os.path.isdir(model_name_or_path):
             return model_name_or_path
         else:
             allow_patterns = ["*.safetensors"]
@@ -120,7 +120,7 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
         )
 
         filepaths = []
-        if is_s3(local_model_path):
+        if is_cloud_storage(local_model_path):
             file_pattern = f"*{self.pattern.format(rank=rank, part='*')}"
             filepaths = s3_glob(path=local_model_path, allow_pattern=[file_pattern])
         else:

vllm/transformers_utils/utils.py

@@ -15,8 +15,10 @@
 logger = init_logger(__name__)
 
 
-def is_s3(model_or_path: str) -> bool:
-    return model_or_path.lower().startswith("s3://")
+def is_cloud_storage(model_or_path: str) -> bool:
+    return model_or_path.lower().startswith(
+        "s3://"
+    ) or model_or_path.lower().startswith("gs://")
 
 
 def check_gguf_file(model: str | PathLike) -> bool: