chore: import upstream snapshot with attribution

python/sglang/srt/connector/__init__.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import enum
+import logging
+
+from sglang.srt.connector.base_connector import (
+    BaseConnector,
+    BaseFileConnector,
+    BaseKVConnector,
+)
+from sglang.srt.connector.redis import RedisConnector
+from sglang.srt.connector.remote_instance import RemoteInstanceConnector
+from sglang.srt.connector.s3 import S3Connector
+from sglang.srt.utils import parse_connector_type
+
+logger = logging.getLogger(__name__)
+
+
+class ConnectorType(str, enum.Enum):
+    FS = "filesystem"
+    KV = "KV"
+    INSTANCE = "instance"
+
+
+def _is_azure_blob_url(url: str, connector_type: str) -> bool:
+    """Detect Azure Blob Storage URLs.
+
+    Matches ``az://...`` URLs and ``https://<account>.blob.core.windows.net/...``
+    URLs, which are the two forms accepted by the ``blobfile`` library.
+    """
+    if connector_type == "az":
+        return True
+    return connector_type == "https" and ".blob.core.windows.net" in url
+
+
+def create_remote_connector(url, device=None, **kwargs) -> BaseConnector:
+    connector_type = parse_connector_type(url)
+    if connector_type == "redis":
+        return RedisConnector(url)
+    elif connector_type == "s3":
+        return S3Connector(url)
+    elif connector_type == "instance":
+        return RemoteInstanceConnector(url, device)
+    elif _is_azure_blob_url(url, connector_type):
+        # Imported lazily so the optional ``blobfile`` dependency is only
+        # required when an Azure URL is actually used.
+        from sglang.srt.connector.azure import AzureBlobConnector
+
+        return AzureBlobConnector(url)
+    else:
+        raise ValueError(f"Invalid connector type: {url}")
+
+
+def get_connector_type(client: BaseConnector) -> ConnectorType:
+    if isinstance(client, BaseKVConnector):
+        return ConnectorType.KV
+    if isinstance(client, BaseFileConnector):
+        return ConnectorType.FS
+    if isinstance(client, RemoteInstanceConnector):
+        return ConnectorType.INSTANCE
+
+    raise ValueError(f"Invalid connector type: {client}")
+
+
+__all__ = [
+    "BaseConnector",
+    "BaseFileConnector",
+    "BaseKVConnector",
+    "RedisConnector",
+    "RemoteInstanceConnector",
+    "S3Connector",
+    "ConnectorType",
+    "create_remote_connector",
+    "get_connector_type",
+]

python/sglang/srt/connector/azure.py

@@ -0,0 +1,127 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import fnmatch
+import os
+from pathlib import Path
+from typing import Generator, Optional, Tuple
+
+import torch
+
+from sglang.srt.connector import BaseFileConnector
+
+
+def _filter_allow(paths: list[str], patterns: list[str]) -> list[str]:
+    return [
+        path
+        for path in paths
+        if any(fnmatch.fnmatch(path, pattern) for pattern in patterns)
+    ]
+
+
+def _filter_ignore(paths: list[str], patterns: list[str]) -> list[str]:
+    return [
+        path
+        for path in paths
+        if not any(fnmatch.fnmatch(path, pattern) for pattern in patterns)
+    ]
+
+
+def _normalize_url(url: str) -> str:
+    """Strip trailing slash so blobfile glob/listdir behave consistently."""
+    return url.rstrip("/")
+
+
+def list_files(
+    bf,
+    path: str,
+    allow_pattern: Optional[list[str]] = None,
+    ignore_pattern: Optional[list[str]] = None,
+) -> Tuple[str, list[str]]:
+    """List files from an Azure Blob Storage path and filter by pattern.
+
+    Args:
+        bf: The ``blobfile`` module.
+        path: An ``az://<account>/<container>/<prefix>`` or
+            ``https://<account>.blob.core.windows.net/<container>/<prefix>`` URL.
+        allow_pattern: A list of fnmatch patterns of which files to keep.
+        ignore_pattern: A list of fnmatch patterns of which files to drop.
+
+    Returns:
+        A tuple ``(base_dir, files)`` where ``base_dir`` is the normalized
+        prefix used as a directory anchor for relative paths, and ``files``
+        is the list of full URLs matched by the patterns.
+    """
+    base_dir = _normalize_url(path)
+    files = [p for p in bf.glob(base_dir + "/**") if not bf.isdir(p)]
+
+    files = _filter_ignore(files, ["*/"])
+    if allow_pattern is not None:
+        files = _filter_allow(files, allow_pattern)
+    if ignore_pattern is not None:
+        files = _filter_ignore(files, ignore_pattern)
+
+    return base_dir, files
+
+
+class AzureBlobConnector(BaseFileConnector):
+    """File connector for Azure Blob Storage.
+
+    Accepts both ``az://<account>/<container>/<path>`` URLs and HTTPS URLs of
+    the form ``https://<account>.blob.core.windows.net/<container>/<path>``.
+    Uses the third-party ``blobfile`` package, which handles authentication via
+    standard Azure credential chains (env vars, az CLI, managed identity).
+    """
+
+    def __init__(self, url: str) -> None:
+        try:
+            import blobfile as bf
+        except ImportError as e:
+            raise ImportError(
+                "AzureBlobConnector requires the 'blobfile' package. "
+                "Install it with `pip install blobfile`."
+            ) from e
+
+        super().__init__(url)
+        self.bf = bf
+
+    def glob(self, allow_pattern: Optional[list[str]] = None) -> list[str]:
+        _, files = list_files(self.bf, self.url, allow_pattern=allow_pattern)
+        return files
+
+    def pull_files(
+        self,
+        allow_pattern: Optional[list[str]] = None,
+        ignore_pattern: Optional[list[str]] = None,
+    ) -> None:
+        """Download files from Azure Blob Storage to ``self.local_dir``."""
+        base_dir, files = list_files(self.bf, self.url, allow_pattern, ignore_pattern)
+        if not files:
+            return
+
+        for file in files:
+            relative = file[len(base_dir) :].lstrip("/")
+            destination_file = os.path.join(self.local_dir, relative)
+            os.makedirs(Path(destination_file).parent, exist_ok=True)
+            self.bf.copy(file, destination_file, overwrite=True)
+
+    def weight_iterator(
+        self, rank: int = 0
+    ) -> Generator[Tuple[str, torch.Tensor], None, None]:
+        from sglang.srt.model_loader.weight_utils import (
+            runai_safetensors_weights_iterator,
+        )
+
+        # Pull *.safetensors locally first since runai_safetensors_weights_iterator
+        # expects local files. blobfile does not provide a streaming safetensors
+        # reader compatible with runai_model_streamer.
+        self.pull_files(allow_pattern=["*.safetensors"])
+        local_files = [
+            os.path.join(root, f)
+            for root, _, fs in os.walk(self.local_dir)
+            for f in fs
+            if f.endswith(".safetensors")
+        ]
+        return runai_safetensors_weights_iterator(local_files)
+
+    def close(self):
+        super().close()

python/sglang/srt/connector/base_connector.py

@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import shutil
+import signal
+import tempfile
+from abc import ABC, abstractmethod
+from typing import Generator, List, Optional, Tuple
+
+import torch
+
+
+class BaseConnector(ABC):
+    """
+    For fs connector such as s3:
+    <connector_type>://<path>/<filename>
+
+    For kv connector such as redis:
+    <connector_type>://<host>:<port>/<model_name>/keys/<key>
+    <connector_type://<host>:<port>/<model_name>/files/<filename>
+    """
+
+    def __init__(self, url: str):
+        self.url = url
+        self.closed = False
+        self.local_dir = tempfile.mkdtemp()
+        for sig in (signal.SIGINT, signal.SIGTERM):
+            existing_handler = signal.getsignal(sig)
+            signal.signal(sig, self._close_by_signal(existing_handler))
+
+    def get_local_dir(self):
+        return self.local_dir
+
+    @abstractmethod
+    def weight_iterator(
+        self, rank: int = 0
+    ) -> Generator[Tuple[str, torch.Tensor], None, None]:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def pull_files(
+        self,
+        allow_pattern: Optional[List[str]] = None,
+        ignore_pattern: Optional[List[str]] = None,
+    ) -> None:
+        raise NotImplementedError()
+
+    def close(self):
+        if self.closed:
+            return
+
+        self.closed = True
+        if os.path.exists(self.local_dir):
+            shutil.rmtree(self.local_dir)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
+    def __del__(self):
+        self.close()
+
+    def _close_by_signal(self, existing_handler=None):
+
+        def new_handler(signum, frame):
+            self.close()
+            if existing_handler:
+                existing_handler(signum, frame)
+
+        return new_handler
+
+
+class BaseKVConnector(BaseConnector):
+
+    @abstractmethod
+    def get(self, key: str) -> Optional[torch.Tensor]:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def getstr(self, key: str) -> Optional[str]:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def set(self, key: str, obj: torch.Tensor) -> None:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def setstr(self, key: str, obj: str) -> None:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def list(self, prefix: str) -> List[str]:
+        raise NotImplementedError()
+
+
+class BaseFileConnector(BaseConnector):
+    """
+    List full file names from remote fs path and filter by allow pattern.
+
+    Args:
+        allow_pattern: A list of patterns of which files to pull.
+
+    Returns:
+        list[str]: List of full paths allowed by the pattern
+    """
+
+    @abstractmethod
+    def glob(self, allow_pattern: str) -> List[str]:
+        raise NotImplementedError()

python/sglang/srt/connector/redis.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from typing import Generator, List, Optional, Tuple
+from urllib.parse import urlparse
+
+import torch
+
+from sglang.srt.connector import BaseKVConnector
+from sglang.srt.connector.serde import create_serde
+from sglang.srt.connector.utils import pull_files_from_db
+
+logger = logging.getLogger(__name__)
+
+
+class RedisConnector(BaseKVConnector):
+
+    def __init__(self, url: str):
+        import redis
+
+        super().__init__(url)
+        parsed_url = urlparse(url)
+        self.connection = redis.Redis(host=parsed_url.hostname, port=parsed_url.port)
+        self.model_name = parsed_url.path.lstrip("/")
+        # TODO: more serde options
+        self.s, self.d = create_serde("safe")
+
+    def get(self, key: str) -> Optional[torch.Tensor]:
+        val = self.connection.get(key)
+
+        if val is None:
+            logger.error("Key %s not found", key)
+            return None
+
+        return self.d.from_bytes(val)
+
+    def getstr(self, key: str) -> Optional[str]:
+        val = self.connection.get(key)
+        if val is None:
+            logger.error("Key %s not found", key)
+            return None
+
+        return val.decode("utf-8")
+
+    def set(self, key: str, tensor: torch.Tensor) -> None:
+        assert tensor is not None
+        self.connection.set(key, self.s.to_bytes(tensor))
+
+    def setstr(self, key: str, obj: str) -> None:
+        self.connection.set(key, obj)
+
+    def list(self, prefix: str) -> List[str]:
+        cursor = 0
+        all_keys: List[bytes] = []
+
+        while True:
+            ret: Tuple[int, List[bytes]] = self.connection.scan(
+                cursor=cursor, match=f"{prefix}*"
+            )  # type: ignore
+            cursor, keys = ret
+            all_keys.extend(keys)
+            if cursor == 0:
+                break
+
+        return [key.decode("utf-8") for key in all_keys]
+
+    def weight_iterator(
+        self, rank: int = 0
+    ) -> Generator[Tuple[str, bytes], None, None]:
+        keys = self.list(f"{self.model_name}/keys/rank_{rank}/")
+        for key in keys:
+            val = self.get(key)
+            key = key.removeprefix(f"{self.model_name}/keys/rank_{rank}/")
+            yield key, val
+
+    def pull_files(
+        self,
+        allow_pattern: Optional[List[str]] = None,
+        ignore_pattern: Optional[List[str]] = None,
+    ) -> None:
+        pull_files_from_db(self, self.model_name, allow_pattern, ignore_pattern)
+
+    def close(self):
+        self.connection.close()
+        super().close()

python/sglang/srt/connector/remote_instance.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from typing import Generator, Optional, Tuple
+from urllib.parse import urlparse
+
+import torch
+import torch.distributed as dist
+
+from sglang.srt.connector import BaseConnector
+from sglang.srt.utils import init_custom_process_group
+
+logger = logging.getLogger(__name__)
+
+
+class RemoteInstanceConnector(BaseConnector):
+
+    def __init__(self, url: str, device: torch.device = "cpu"):
+        assert (
+            device.type == "cuda" or device.type == "npu"
+        ), "RemoteInstanceConnector only supports cuda device."
+        super().__init__(url)
+        self.url = url
+        self.device = device
+
+    def build_group(
+        self,
+        gpu_id: int = -1,
+        tp_rank: int = -1,
+        instance_ip: str = None,
+        group_rank: int = 1,
+        world_size: int = 2,
+    ):
+        assert (
+            self.device.type == "cuda" or self.device.type == "npu"
+        ), "RemoteInstanceConnector only supports cuda device."
+        assert (
+            gpu_id != -1 and tp_rank != -1
+        ), "gpu_id and tp_rank must be specified for RemoteInstanceConnector. "
+
+        self.device_id = torch.device(self.device.type, gpu_id)
+
+        parsed_url = urlparse(self.url)
+        master_address = parsed_url.hostname
+        master_port = parsed_url.port
+        group_name = f"send_weights_{instance_ip}_{master_port}_{tp_rank}"
+        backend = "nccl"
+
+        logger.info(
+            f"init custom process group: master_address={master_address}, master_port={master_port}, "
+            f"rank_offset={group_rank}, world_size={world_size}, group_name={group_name}, backend={backend}"
+        )
+
+        try:
+            self._model_update_group = init_custom_process_group(
+                backend=backend,
+                init_method=f"tcp://{master_address}:{master_port}",
+                world_size=world_size,
+                rank=group_rank,
+                group_name=group_name,
+                device_id=self.device_id,
+            )
+            dist.barrier(group=self._model_update_group)
+            return True, "Succeeded to initialize custom process group."
+        except Exception as e:
+            message = f"Failed to initialize custom process group: {e}."
+            logger.error(message)
+            return False, message
+
+    # Implemented as a no-op to make BaseConnector interface consistent.
+    def pull_files(
+        self,
+        allow_pattern: Optional[list[str]] = None,
+        ignore_pattern: Optional[list[str]] = None,
+    ) -> None:
+        return
+
+    # Implemented as a no-op to make BaseConnector interface consistent.
+    def weight_iterator(
+        self, rank: int = 0
+    ) -> Generator[Tuple[str, torch.Tensor], None, None]:
+        return

python/sglang/srt/connector/s3.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import fnmatch
+import os
+from pathlib import Path
+from typing import Generator, Optional, Tuple
+
+import torch
+
+from sglang.srt.connector import BaseFileConnector
+
+
+def _filter_allow(paths: list[str], patterns: list[str]) -> list[str]:
+    return [
+        path
+        for path in paths
+        if any(fnmatch.fnmatch(path, pattern) for pattern in patterns)
+    ]
+
+
+def _filter_ignore(paths: list[str], patterns: list[str]) -> list[str]:
+    return [
+        path
+        for path in paths
+        if not any(fnmatch.fnmatch(path, pattern) for pattern in patterns)
+    ]
+
+
+def list_files(
+    s3,
+    path: str,
+    allow_pattern: Optional[list[str]] = None,
+    ignore_pattern: Optional[list[str]] = None,
+) -> tuple[str, str, list[str]]:
+    """
+    List files from S3 path and filter by pattern.
+
+    Args:
+        s3: S3 client to use.
+        path: The S3 path to list from.
+        allow_pattern: A list of patterns of which files to pull.
+        ignore_pattern: A list of patterns of which files not to pull.
+
+    Returns:
+        tuple[str, str, list[str]]: A tuple where:
+            - The first element is the bucket name
+            - The second element is string represent the bucket
+              and the prefix as a dir like string
+            - The third element is a list of files allowed or
+              disallowed by pattern
+    """
+    parts = path.removeprefix("s3://").split("/")
+    prefix = "/".join(parts[1:])
+    bucket_name = parts[0]
+
+    objects = s3.list_objects_v2(Bucket=bucket_name, Prefix=prefix)
+    paths = [obj["Key"] for obj in objects.get("Contents", [])]
+
+    paths = _filter_ignore(paths, ["*/"])
+    if allow_pattern is not None:
+        paths = _filter_allow(paths, allow_pattern)
+
+    if ignore_pattern is not None:
+        paths = _filter_ignore(paths, ignore_pattern)
+
+    return bucket_name, prefix, paths
+
+
+class S3Connector(BaseFileConnector):
+
+    def __init__(self, url: str) -> None:
+        import boto3
+
+        super().__init__(url)
+        self.client = boto3.client("s3")
+
+    def glob(self, allow_pattern: Optional[list[str]] = None) -> list[str]:
+        bucket_name, _, paths = list_files(
+            self.client, path=self.url, allow_pattern=allow_pattern
+        )
+        return [f"s3://{bucket_name}/{path}" for path in paths]
+
+    def pull_files(
+        self,
+        allow_pattern: Optional[list[str]] = None,
+        ignore_pattern: Optional[list[str]] = None,
+    ) -> None:
+        """
+        Pull files from S3 storage into the temporary directory.
+
+        Args:
+            s3_model_path: The S3 path of the model.
+            allow_pattern: A list of patterns of which files to pull.
+            ignore_pattern: A list of patterns of which files not to pull.
+
+        """
+        bucket_name, base_dir, files = list_files(
+            self.client, self.url, allow_pattern, ignore_pattern
+        )
+        if len(files) == 0:
+            return
+
+        for file in files:
+            destination_file = os.path.join(self.local_dir, file.removeprefix(base_dir))
+            local_dir = Path(destination_file).parent
+            os.makedirs(local_dir, exist_ok=True)
+            self.client.download_file(bucket_name, file, destination_file)
+
+    def weight_iterator(
+        self, rank: int = 0
+    ) -> Generator[Tuple[str, torch.Tensor], None, None]:
+        from sglang.srt.model_loader.weight_utils import (
+            runai_safetensors_weights_iterator,
+        )
+
+        # only support safetensor files now
+        hf_weights_files = self.glob(allow_pattern=["*.safetensors"])
+        return runai_safetensors_weights_iterator(hf_weights_files)
+
+    def close(self):
+        self.client.close()
+        super().close()

python/sglang/srt/connector/serde/__init__.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# inspired by LMCache
+from typing import Optional, Tuple
+
+import torch
+
+from sglang.srt.connector.serde.safe_serde import SafeDeserializer, SafeSerializer
+from sglang.srt.connector.serde.serde import Deserializer, Serializer
+
+
+def create_serde(serde_type: str) -> Tuple[Serializer, Deserializer]:
+    s: Optional[Serializer] = None
+    d: Optional[Deserializer] = None
+
+    if serde_type == "safe":
+        s = SafeSerializer()
+        d = SafeDeserializer()
+    else:
+        raise ValueError(f"Unknown serde type: {serde_type}")
+
+    return s, d
+
+
+__all__ = [
+    "Serializer",
+    "Deserializer",
+    "SafeSerializer",
+    "SafeDeserializer",
+    "create_serde",
+]

python/sglang/srt/connector/serde/safe_serde.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Union
+
+import torch
+from safetensors.torch import load, save
+
+from sglang.srt.connector.serde.serde import Deserializer, Serializer
+
+
+class SafeSerializer(Serializer):
+
+    def __init__(self):
+        super().__init__()
+
+    def to_bytes(self, t: torch.Tensor) -> bytes:
+        return save({"tensor_bytes": t.cpu().contiguous()})
+
+
+class SafeDeserializer(Deserializer):
+
+    def __init__(self):
+        # TODO: dtype options
+        super().__init__(torch.float32)
+
+    def from_bytes_normal(self, b: Union[bytearray, bytes]) -> torch.Tensor:
+        return load(bytes(b))["tensor_bytes"]
+
+    def from_bytes(self, b: Union[bytearray, bytes]) -> torch.Tensor:
+        return self.from_bytes_normal(b)

python/sglang/srt/connector/serde/serde.py

@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import abc
+from abc import ABC, abstractmethod
+
+import torch
+
+
+class Serializer(ABC):
+
+    @abstractmethod
+    def to_bytes(self, t: torch.Tensor) -> bytes:
+        """
+        Serialize a pytorch tensor to bytes. The serialized bytes should contain
+        both the data and the metadata (shape, dtype, etc.) of the tensor.
+
+        Input:
+            t: the input pytorch tensor, can be on any device, in any shape,
+               with any dtype
+
+        Returns:
+            bytes: the serialized bytes
+        """
+        raise NotImplementedError
+
+
+class Deserializer(metaclass=abc.ABCMeta):
+
+    def __init__(self, dtype):
+        self.dtype = dtype
+
+    @abstractmethod
+    def from_bytes(self, bs: bytes) -> torch.Tensor:
+        """
+        Deserialize a pytorch tensor from bytes.
+
+        Input:
+            bytes: a stream of bytes
+
+        Output:
+            torch.Tensor: the deserialized pytorch tensor
+        """
+        raise NotImplementedError

python/sglang/srt/connector/utils.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from pathlib import Path
+from typing import Optional
+from urllib.parse import urlparse
+
+from sglang.srt.connector import BaseConnector
+
+
+def parse_model_name(url: str) -> str:
+    """
+    Parse the model name from the url.
+    Only used for db connector
+    """
+    parsed_url = urlparse(url)
+    return parsed_url.path.lstrip("/")
+
+
+def pull_files_from_db(
+    connector: BaseConnector,
+    model_name: str,
+    allow_pattern: Optional[list[str]] = None,
+    ignore_pattern: Optional[list[str]] = None,
+) -> None:
+    prefix = f"{model_name}/files/"
+    local_dir = connector.get_local_dir()
+    files = connector.list(prefix)
+
+    for file in files:
+        destination_file = os.path.join(local_dir, file.removeprefix(prefix))
+        local_dir = Path(destination_file).parent
+        os.makedirs(local_dir, exist_ok=True)
+        with open(destination_file, "wb") as f:
+            f.write(connector.getstr(file).encode("utf-8"))