chore: import upstream snapshot with attribution

tests/plugins/bge_m3_sparse_plugin/bge_m3_sparse_processor/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+def register_bge_m3_sparse_embeddings_processor():
+    return "bge_m3_sparse_processor.sparse_embeddings_processor.BgeM3SparseEmbeddingsProcessor"  # noqa: E501

tests/plugins/bge_m3_sparse_plugin/bge_m3_sparse_processor/sparse_embeddings_processor.py

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Sequence
+
+from vllm.config import PoolerConfig, VllmConfig
+from vllm.entrypoints.openai.engine.protocol import UsageInfo
+from vllm.inputs import PromptType
+from vllm.outputs import PoolingRequestOutput
+from vllm.plugins.io_processors.interface import IOProcessor
+from vllm.pooling_params import PoolingParams
+from vllm.renderers import BaseRenderer
+from vllm.tokenizers.detokenizer_utils import convert_ids_list_to_tokens
+
+from .types import (
+    SparseEmbeddingCompletionRequestMixin,
+    SparseEmbeddingResponse,
+    SparseEmbeddingResponseData,
+    SparseEmbeddingTokenWeight,
+)
+
+
+class BgeM3SparseEmbeddingsProcessor(
+    IOProcessor[SparseEmbeddingCompletionRequestMixin, SparseEmbeddingResponse]
+):
+    def __init__(self, vllm_config: VllmConfig, renderer: BaseRenderer):
+        super().__init__(vllm_config, renderer)
+        self.offline_requests: list[SparseEmbeddingCompletionRequestMixin] = []
+        self.online_requests: dict[str, SparseEmbeddingCompletionRequestMixin] = {}
+        self.renderer: BaseRenderer = renderer
+        self.default_pooling_params = {}
+        pooler_config: PoolerConfig = vllm_config.model_config.pooler_config
+        if pooler_config is not None:
+            for param in ["use_activation", "dimensions"]:
+                if getattr(pooler_config, param, None) is None:
+                    continue
+                self.default_pooling_params[param] = getattr(pooler_config, param)
+        self.embed_dimensions = vllm_config.model_config.embedding_size
+
+    def __repr__(self) -> str:
+        return (
+            f"BgeM3SparseEmbeddingsProcessor("
+            f"embed_dimensions={self.embed_dimensions}, "
+            f"default_pooling_params={self.default_pooling_params})"
+        )
+
+    def merge_pooling_params(
+        self,
+        params: PoolingParams | None = None,
+    ) -> PoolingParams:
+        if params is None:
+            params = PoolingParams()
+        # refer to PoolingCompletionRequest.to_pooling_params
+        # set and verify pooling params
+        params.skip_reading_prefix_cache = True
+        params.task = "embed&token_classify"
+        params.use_activation = True
+        params.dimensions = self.embed_dimensions
+        return params
+
+    def parse_request(
+        self, request_data: object
+    ) -> SparseEmbeddingCompletionRequestMixin:
+        # for vllm.entrypoints.llm.LLM, offline mode, calls `encode` directly.
+        if isinstance(request_data, dict):
+            return SparseEmbeddingCompletionRequestMixin(**request_data)
+        raise TypeError("request_data should be a dictionary")
+
+    def pre_process(
+        self,
+        prompt: SparseEmbeddingCompletionRequestMixin,
+        request_id: str | None = None,
+        **kwargs,
+    ) -> PromptType | Sequence[PromptType]:
+        if request_id is not None:
+            assert request_id not in self.online_requests, "request_id duplicated"
+            self.online_requests[request_id] = prompt
+        else:
+            self.offline_requests.append(prompt)
+        return prompt.input
+
+    def _get_sparse_embedding_request(self, request_id: str | None = None):
+        if request_id:
+            return self.online_requests.pop(request_id, None)
+        return self.offline_requests.pop(0)
+
+    def _build_sparse_embedding_token_weights(
+        self,
+        sparse_embedding: dict[int, float],
+        return_tokens: bool = False,
+    ) -> list[SparseEmbeddingTokenWeight]:
+        token_ids = sparse_embedding.keys()
+        token_weights = sparse_embedding.values()
+        tokens = [None] * len(token_ids)
+
+        if return_tokens and self.renderer is not None:
+            tokens = convert_ids_list_to_tokens(
+                self.renderer.get_tokenizer(), token_ids
+            )
+        sparse_embedding_output: list[SparseEmbeddingTokenWeight] = []
+        for token_id, weight, token in zip(token_ids, token_weights, tokens):
+            sparse_embedding_output.append(
+                SparseEmbeddingTokenWeight(
+                    token_id=token_id, weight=weight, token=token
+                )
+            )
+        return sparse_embedding_output
+
+    def post_process(
+        self,
+        model_output: Sequence[PoolingRequestOutput],
+        request_id: str | None = None,
+        **kwargs,
+    ) -> SparseEmbeddingResponse:
+        num_prompt_tokens = 0
+        response_data = []
+        raw_request = self._get_sparse_embedding_request(request_id)
+        has_dense_embed = raw_request.embed_task in ["dense", "dense&sparse"]
+        has_sparse_embed = raw_request.embed_task in ["sparse", "dense&sparse"]
+        embed_dimensions = self.embed_dimensions
+        for idx in range(len(model_output)):
+            mo = model_output[idx]
+            sparse_embedding_dict: dict[int, float] = {}
+            num_prompt_tokens += len(mo.prompt_token_ids)
+            dense_embedding: list[float] | None = None
+            sparse_embedding: list[SparseEmbeddingTokenWeight] | None = None
+            if has_dense_embed:
+                dense_embedding = mo.outputs.data[:embed_dimensions].tolist()
+            if has_sparse_embed:
+                sparse_weights = mo.outputs.data[embed_dimensions:].tolist()
+                if len(mo.prompt_token_ids) != len(sparse_weights):
+                    # this is the case that add_special_tokens is True,
+                    # which means first token and last token are special tokens
+                    mo.prompt_token_ids = mo.prompt_token_ids[1:]
+                for token_id, weight in zip(mo.prompt_token_ids, sparse_weights):
+                    sparse_embedding_dict[token_id] = max(
+                        weight, sparse_embedding_dict.get(token_id, 0.0)
+                    )
+                sparse_embedding = self._build_sparse_embedding_token_weights(
+                    sparse_embedding_dict,
+                    raw_request.return_tokens,
+                )
+
+            response_data.append(
+                SparseEmbeddingResponseData(
+                    index=idx,
+                    object=raw_request.embed_task,
+                    sparse_embedding=sparse_embedding,
+                    dense_embedding=dense_embedding,
+                )
+            )
+
+        usage = UsageInfo(
+            prompt_tokens=num_prompt_tokens,
+            total_tokens=num_prompt_tokens,
+        )
+        resp = SparseEmbeddingResponse(
+            data=response_data,
+            usage=usage,
+        )
+
+        return resp

tests/plugins/bge_m3_sparse_plugin/bge_m3_sparse_processor/types.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Literal, get_args
+
+from pydantic import BaseModel, Field
+
+from vllm.entrypoints.openai.engine.protocol import UsageInfo
+from vllm.entrypoints.pooling.base.protocol import (
+    CompletionRequestMixin,
+    EmbedRequestMixin,
+)
+
+EmbedTask = Literal[
+    "sparse",
+    "dense",
+    "dense&sparse",
+]
+
+EMBED_TASKS: tuple[EmbedTask, ...] = get_args(EmbedTask)
+
+
+class SparseEmbeddingCompletionRequestMixin(CompletionRequestMixin, EmbedRequestMixin):
+    return_tokens: bool | None = Field(
+        default=None,
+        description="Whether to return dict shows the mapping of token_id to text."
+        "`None` or False means not return.",
+    )
+    embed_task: EmbedTask = Field(
+        default="dense&sparse",
+        description="embed task, can be one of 'sparse', 'dense' , 'dense&sparse', "
+        "default to 'dense&sparse'",
+    )
+
+    def to_embed_requests_offline(self) -> list[EmbedRequestMixin]:
+        if isinstance(self.input, list):
+            return [self] * len(self.input)
+        return [self]
+
+    def to_embed_requests_online(self) -> list[EmbedRequestMixin]:
+        return [self]
+
+
+class SparseEmbeddingTokenWeight(BaseModel):
+    token_id: int
+    weight: float
+    token: str | None
+
+
+class SparseEmbeddingResponseData(BaseModel):
+    index: int
+    object: str = "dense&sparse"
+    sparse_embedding: list[SparseEmbeddingTokenWeight] | None
+    dense_embedding: list[float] | None
+
+
+class SparseEmbeddingResponse(BaseModel):
+    data: list[SparseEmbeddingResponseData]
+    usage: UsageInfo

tests/plugins/bge_m3_sparse_plugin/setup.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from setuptools import setup
+
+setup(
+    name="bge-m3-sparse-plugin",
+    version="0.1",
+    packages=["bge_m3_sparse_processor"],
+    entry_points={
+        "vllm.io_processor_plugins": [
+            "bge_m3_sparse_plugin = bge_m3_sparse_processor:register_bge_m3_sparse_embeddings_processor",  # noqa: E501
+        ]
+    },
+)

tests/plugins/colbert_query_plugin/colbert_query_processor/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+def register_colbert_query_embedding_processor():
+    return "colbert_query_processor.query_embedding_processor.ColBERTQueryEmbeddingProcessor"  # noqa: E501

tests/plugins/colbert_query_plugin/colbert_query_processor/query_embedding_processor.py

@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Iterator, Sequence
+from typing import cast
+
+from vllm.config import VllmConfig
+from vllm.entrypoints.openai.engine.protocol import UsageInfo
+from vllm.inputs import PromptType, TokensPrompt
+from vllm.outputs import PoolingRequestOutput
+from vllm.plugins.io_processors.interface import IOProcessor
+from vllm.pooling_params import PoolingParams
+from vllm.renderers import BaseRenderer
+from vllm.utils.collection_utils import is_list_of
+
+from .types import (
+    QUERY_MAXLEN,
+    ColBERTEmbeddingCompletionRequestMixin,
+    ColBERTEmbeddingResponse,
+    ColBERTEmbeddingResponseData,
+)
+
+QUERY_MARKER_TOKEN = "[QueryMarker]"
+DOCUMENT_MARKER_TOKEN = "[DocumentMarker]"
+
+
+class ColBERTQueryEmbeddingProcessor(
+    IOProcessor[ColBERTEmbeddingCompletionRequestMixin, ColBERTEmbeddingResponse]
+):
+    """This IO processor only supports the ColBERT-style model jinaai/jina-colbert-v2.
+    It does not support all ColBERT-style variants (e.g. colbert-ir/colbertv2.0).
+    """
+
+    def __init__(self, vllm_config: VllmConfig, renderer: BaseRenderer):
+        super().__init__(vllm_config, renderer)
+        self.requests_cache: dict[str, ColBERTEmbeddingCompletionRequestMixin] = {}
+        self.renderer: BaseRenderer = renderer
+        # Context window (8192 for jinaai/jina-colbert-v2); caps document
+        # content length minus the 3 special-token slots.
+        self.max_model_len = vllm_config.model_config.max_model_len
+        self._query_marker_id: int | None = None
+        self._document_marker_id: int | None = None
+
+    def __repr__(self) -> str:
+        return (
+            f"ColBERTQueryEmbeddingProcessor("
+            f"query_maxlen={QUERY_MAXLEN}, "
+            f"doc_maxlen={self.max_model_len}, "
+            f"query_marker_token={QUERY_MARKER_TOKEN!r}, "
+            f"document_marker_token={DOCUMENT_MARKER_TOKEN!r})"
+        )
+
+    def _resolve_marker_ids(self, tokenizer) -> tuple[int, int]:
+        if self._query_marker_id is not None and self._document_marker_id is not None:
+            return self._query_marker_id, self._document_marker_id
+
+        unk_id = getattr(tokenizer, "unk_token_id", None)
+        marker_ids: list[int] = []
+        for marker in (QUERY_MARKER_TOKEN, DOCUMENT_MARKER_TOKEN):
+            marker_id = tokenizer.convert_tokens_to_ids(marker)
+            if marker_id is None or marker_id == unk_id:
+                raise ValueError(
+                    f"Marker token {marker!r} not found in the tokenizer "
+                    "vocabulary. This plugin requires a ColBERT model whose "
+                    "tokenizer defines both "
+                    f"{QUERY_MARKER_TOKEN!r} and {DOCUMENT_MARKER_TOKEN!r} "
+                    "(e.g. jinaai/jina-colbert-v2)."
+                )
+            marker_ids.append(marker_id)
+
+        self._query_marker_id, self._document_marker_id = marker_ids
+        return self._query_marker_id, self._document_marker_id
+
+    def _iter_content_token_ids(
+        self,
+        tokenizer,
+        request_input: list[int] | list[list[int]] | str | list[str],
+    ) -> Iterator[list[int]]:
+        if isinstance(request_input, str):
+            yield tokenizer.encode(request_input, add_special_tokens=False)
+            return
+
+        if not isinstance(request_input, list) or not request_input:
+            raise ValueError("input must be a non-empty string or list")
+
+        if is_list_of(request_input, int):
+            yield list(cast(list[int], request_input))
+            return
+
+        for item in request_input:
+            if isinstance(item, str):
+                yield tokenizer.encode(item, add_special_tokens=False)
+            else:
+                yield list(cast(list[int], item))
+
+    def _build_query_prompt(
+        self,
+        tokenizer,
+        content_ids: list[int],
+    ) -> TokensPrompt:
+        """[CLS] [QueryMarker] <tokens> [SEP] [MASK]... up to QUERY_MAXLEN."""
+        query_marker_id, _ = self._resolve_marker_ids(tokenizer)
+        mask_token_id = tokenizer.mask_token_id
+        if mask_token_id is None:
+            raise ValueError(
+                "Tokenizer has no mask token; cannot perform query expansion."
+            )
+
+        # [CLS], marker and [SEP] take 3 slots.
+        content_ids = content_ids[: QUERY_MAXLEN - 3]
+        token_ids = [
+            tokenizer.cls_token_id,
+            query_marker_id,
+            *content_ids,
+            tokenizer.sep_token_id,
+        ]
+        token_ids += [mask_token_id] * (QUERY_MAXLEN - len(token_ids))
+        return TokensPrompt(prompt_token_ids=token_ids)
+
+    def _build_document_prompt(
+        self,
+        tokenizer,
+        content_ids: list[int],
+    ) -> TokensPrompt:
+        """[CLS] [DocumentMarker] <tokens> [SEP]"""
+        _, document_marker_id = self._resolve_marker_ids(tokenizer)
+
+        content_ids = content_ids[: self.max_model_len - 3]
+        token_ids = [
+            tokenizer.cls_token_id,
+            document_marker_id,
+            *content_ids,
+            tokenizer.sep_token_id,
+        ]
+        return TokensPrompt(prompt_token_ids=token_ids)
+
+    def parse_data(self, data: object) -> ColBERTEmbeddingCompletionRequestMixin:
+        if isinstance(data, dict):
+            return ColBERTEmbeddingCompletionRequestMixin(**data)
+        raise TypeError("request data should be a dictionary")
+
+    def pre_process(
+        self,
+        prompt: ColBERTEmbeddingCompletionRequestMixin,
+        request_id: str | None = None,
+        **kwargs,
+    ) -> PromptType | Sequence[PromptType]:
+        cache_key = request_id or "offline"
+        assert cache_key not in self.requests_cache, "request_id duplicated"
+        self.requests_cache[cache_key] = prompt
+
+        tokenizer = self.renderer.get_tokenizer()
+        prompts: list[TokensPrompt] = []
+        for content_ids in self._iter_content_token_ids(tokenizer, prompt.input):
+            if prompt.input_type == "query":
+                prompts.append(self._build_query_prompt(tokenizer, content_ids))
+            else:
+                prompts.append(self._build_document_prompt(tokenizer, content_ids))
+        return prompts
+
+    def merge_pooling_params(
+        self,
+        params: PoolingParams | None = None,
+    ) -> PoolingParams:
+        if params is None:
+            params = PoolingParams()
+        params.task = "token_embed"
+        params.skip_reading_prefix_cache = True
+        return params
+
+    def post_process(
+        self,
+        model_output: Sequence[PoolingRequestOutput],
+        request_id: str | None = None,
+        **kwargs,
+    ) -> ColBERTEmbeddingResponse:
+        raw_request = self.requests_cache.pop(request_id or "offline")
+
+        num_prompt_tokens = 0
+        response_data: list[ColBERTEmbeddingResponseData] = []
+        for idx, output in enumerate(model_output):
+            num_prompt_tokens += len(output.prompt_token_ids)
+            response_data.append(
+                ColBERTEmbeddingResponseData(
+                    index=idx,
+                    input_type=raw_request.input_type,
+                    embedding=output.outputs.data.tolist(),
+                )
+            )
+
+        usage = UsageInfo(
+            prompt_tokens=num_prompt_tokens,
+            total_tokens=num_prompt_tokens,
+        )
+        return ColBERTEmbeddingResponse(data=response_data, usage=usage)

tests/plugins/colbert_query_plugin/colbert_query_processor/types.py

@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Literal, get_args
+
+from pydantic import BaseModel, Field
+
+from vllm.entrypoints.openai.engine.protocol import UsageInfo
+from vllm.entrypoints.pooling.base.protocol import CompletionRequestMixin
+
+InputType = Literal["query", "document"]
+INPUT_TYPES: tuple[InputType, ...] = get_args(InputType)
+QUERY_MAXLEN = 32
+
+
+class ColBERTEmbeddingCompletionRequestMixin(CompletionRequestMixin):
+    input_type: InputType = Field(
+        description="Whether to encode the input as a ColBERT 'query' "
+        f"(query marker + [mask] expansion to {QUERY_MAXLEN} tokens) or as a "
+        "'document' (document marker only). Required.",
+    )
+
+
+class ColBERTEmbeddingResponseData(BaseModel):
+    index: int
+    object: str = "embedding"
+    input_type: InputType
+    embedding: list[list[float]]
+
+
+class ColBERTEmbeddingResponse(BaseModel):
+    data: list[ColBERTEmbeddingResponseData]
+    usage: UsageInfo

tests/plugins/colbert_query_plugin/setup.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from setuptools import setup
+
+setup(
+    name="colbert-query-plugin",
+    version="0.1",
+    packages=["colbert_query_processor"],
+    entry_points={
+        "vllm.io_processor_plugins": [
+            "colbert_query_plugin = colbert_query_processor:register_colbert_query_embedding_processor",  # noqa: E501
+        ]
+    },
+)

tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+def register_prithvi():
+    return "prithvi_io_processor.prithvi_processor.PrithviMultimodalDataProcessor"  # noqa: E501

tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py

@@ -0,0 +1,385 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import datetime
+import os
+import tempfile
+import urllib.request
+from collections.abc import Sequence
+from typing import Any
+
+import albumentations
+import numpy as np
+import pybase64 as base64
+import rasterio
+import regex as re
+import torch
+from einops import rearrange
+from terratorch.datamodules import Sen1Floods11NonGeoDataModule
+
+from vllm.config import VllmConfig
+from vllm.inputs import PromptType
+from vllm.logger import init_logger
+from vllm.outputs import PoolingRequestOutput
+from vllm.plugins.io_processors.interface import IOProcessor
+from vllm.renderers import BaseRenderer
+
+from .types import DataModuleConfig, ImagePrompt, ImageRequestOutput
+
+logger = init_logger(__name__)
+
+NO_DATA = -9999
+NO_DATA_FLOAT = 0.0001
+OFFSET = 0
+PERCENTILE = 99
+
+DEFAULT_INPUT_INDICES = [0, 1, 2, 3, 4, 5]
+
+datamodule_config: DataModuleConfig = {
+    "bands": ["BLUE", "GREEN", "RED", "NIR_NARROW", "SWIR_1", "SWIR_2"],
+    "batch_size": 16,
+    "constant_scale": 0.0001,
+    "data_root": "/dccstor/geofm-finetuning/datasets/sen1floods11",
+    "drop_last": True,
+    "no_data_replace": 0.0,
+    "no_label_replace": -1,
+    "num_workers": 8,
+    "test_transform": [
+        albumentations.Resize(height=448, interpolation=1, p=1, width=448),
+        albumentations.pytorch.ToTensorV2(transpose_mask=False, p=1.0),
+    ],
+}
+
+
+def save_geotiff(image: torch.Tensor, meta: dict, out_format: str) -> str | bytes:
+    """Save multi-band image in Geotiff file.
+
+    Args:
+        image: np.ndarray with shape (bands, height, width)
+        output_path: path where to save the image
+        meta: dict with meta info.
+    """
+    if out_format == "path":
+        # create temp file
+        file_path = os.path.join(os.getcwd(), "prediction.tiff")
+        with rasterio.open(file_path, "w", **meta) as dest:
+            for i in range(image.shape[0]):
+                dest.write(image[i, :, :], i + 1)
+
+        return file_path
+    elif out_format == "b64_json":
+        with tempfile.NamedTemporaryFile() as tmpfile:
+            with rasterio.open(tmpfile.name, "w", **meta) as dest:
+                for i in range(image.shape[0]):
+                    dest.write(image[i, :, :], i + 1)
+
+            file_data = tmpfile.read()
+            return base64.b64encode(file_data)
+
+    else:
+        raise ValueError("Unknown output format")
+
+
+def _convert_np_uint8(float_image: torch.Tensor):
+    image = float_image.numpy() * 255.0
+    image = image.astype(dtype=np.uint8)
+
+    return image
+
+
+def read_geotiff(
+    file_path: str | None = None,
+    path_type: str | None = None,
+    file_data: bytes | None = None,
+) -> tuple[torch.Tensor, dict, tuple[float, float] | None]:
+    """Read all bands from *file_path* and return image + meta info.
+
+    Args:
+        file_path: path to image file.
+
+    Returns:
+        np.ndarray with shape (bands, height, width)
+        meta info dict
+    """
+
+    if all([x is None for x in [file_path, path_type, file_data]]):
+        raise Exception("All input fields to read_geotiff are None")
+    write_to_file: bytes | None = None
+    path: str | None = None
+    if file_data is not None:
+        # with tempfile.NamedTemporaryFile() as tmpfile:
+        #     tmpfile.write(file_data)
+        #     path = tmpfile.name
+
+        write_to_file = file_data
+    elif file_path is not None and path_type == "url":
+        resp = urllib.request.urlopen(file_path)
+        # with tempfile.NamedTemporaryFile() as tmpfile:
+        #     tmpfile.write(resp.read())
+        #     path = tmpfile.name
+        write_to_file = resp.read()
+    elif file_path is not None and path_type == "path":
+        path = file_path
+    elif file_path is not None and path_type == "b64_json":
+        image_data = base64.b64decode(file_path)
+        # with tempfile.NamedTemporaryFile() as tmpfile:
+        #     tmpfile.write(image_data)
+        #     path = tmpfile.name
+        write_to_file = image_data
+    else:
+        raise Exception("Wrong combination of parameters to read_geotiff")
+
+    with tempfile.NamedTemporaryFile() as tmpfile:
+        path_to_use = None
+        if write_to_file:
+            tmpfile.write(write_to_file)
+            path_to_use = tmpfile.name
+        elif path:
+            path_to_use = path
+
+        with rasterio.open(path_to_use) as src:
+            img = src.read()
+            meta = src.meta
+            try:
+                coords = src.lnglat()
+            except Exception:
+                # Cannot read coords
+                coords = None
+
+    return img, meta, coords
+
+
+def load_image(
+    data: list[str],
+    path_type: str,
+    mean: list[float] | None = None,
+    std: list[float] | None = None,
+    indices: list[int] | None | None = None,
+):
+    """Build an input example by loading images in *file_paths*.
+
+    Args:
+        file_paths: list of file paths .
+        mean: list containing mean values for each band in the
+              images in *file_paths*.
+        std: list containing std values for each band in the
+             images in *file_paths*.
+
+    Returns:
+        np.array containing created example
+        list of meta info for each image in *file_paths*
+    """
+
+    imgs = []
+    metas = []
+    temporal_coords = []
+    location_coords = []
+
+    for file in data:
+        # if isinstance(file, bytes):
+        #     img, meta, coords = read_geotiff(file_data=file)
+        # else:
+        img, meta, coords = read_geotiff(file_path=file, path_type=path_type)
+        # Rescaling (don't normalize on nodata)
+        img = np.moveaxis(img, 0, -1)  # channels last for rescaling
+        if indices is not None:
+            img = img[..., indices]
+        if mean is not None and std is not None:
+            img = np.where(img == NO_DATA, NO_DATA_FLOAT, (img - mean) / std)
+
+        imgs.append(img)
+        metas.append(meta)
+        if coords is not None:
+            location_coords.append(coords)
+
+        try:
+            match = re.search(r"(\d{7,8}T\d{6})", file)
+            if match:
+                year = int(match.group(1)[:4])
+                julian_day = match.group(1).split("T")[0][4:]
+                if len(julian_day) == 3:
+                    julian_day = int(julian_day)
+                else:
+                    julian_day = (
+                        datetime.datetime.strptime(julian_day, "%m%d")
+                        .timetuple()
+                        .tm_yday
+                    )
+                temporal_coords.append([year, julian_day])
+        except Exception:
+            logger.exception("Could not extract timestamp for %s", file)
+
+    imgs = np.stack(imgs, axis=0)  # num_frames, H, W, C
+    imgs = np.moveaxis(imgs, -1, 0).astype("float32")  # C, num_frames, H, W
+    imgs = np.expand_dims(imgs, axis=0)  # add batch di
+
+    return imgs, temporal_coords, location_coords, metas
+
+
+class PrithviMultimodalDataProcessor(IOProcessor[ImagePrompt, ImageRequestOutput]):
+    indices = [0, 1, 2, 3, 4, 5]
+
+    def __init__(self, vllm_config: VllmConfig, renderer: BaseRenderer):
+        super().__init__(vllm_config, renderer)
+
+        self.datamodule = Sen1Floods11NonGeoDataModule(
+            data_root=datamodule_config["data_root"],
+            batch_size=datamodule_config["batch_size"],
+            num_workers=datamodule_config["num_workers"],
+            bands=datamodule_config["bands"],
+            drop_last=datamodule_config["drop_last"],
+            test_transform=datamodule_config["test_transform"],
+        )
+        self.img_size = 512
+        self.h1 = 1
+        self.w1 = 1
+        self.original_h = 512
+        self.original_w = 512
+        self.batch_size = 1
+        self.meta_data = None
+        self.requests_cache: dict[str, dict[str, Any]] = {}
+        self.indices = DEFAULT_INPUT_INDICES
+
+    def parse_data(self, data: object) -> ImagePrompt:
+        if isinstance(data, dict):
+            return ImagePrompt(**data)
+
+        raise ValueError("Prompt data should be an `ImagePrompt`")
+
+    def pre_process(
+        self,
+        prompt: ImagePrompt,
+        request_id: str | None = None,
+        **kwargs,
+    ) -> PromptType | Sequence[PromptType]:
+        image_data = dict(prompt)
+
+        if request_id:
+            self.requests_cache[request_id] = {
+                "out_format": image_data["out_data_format"],
+            }
+
+        input_data, temporal_coords, location_coords, meta_data = load_image(
+            data=[image_data["data"]],
+            indices=self.indices,
+            path_type=image_data["data_format"],
+        )
+
+        self.meta_data = meta_data[0]
+
+        if input_data.mean() > 1:
+            input_data = input_data / 10000  # Convert to range 0-1
+
+        self.original_h, self.original_w = input_data.shape[-2:]
+        pad_h = (self.img_size - (self.original_h % self.img_size)) % self.img_size
+        pad_w = (self.img_size - (self.original_w % self.img_size)) % self.img_size
+        input_data = np.pad(
+            input_data,
+            ((0, 0), (0, 0), (0, 0), (0, pad_h), (0, pad_w)),
+            mode="reflect",
+        )
+
+        batch = torch.tensor(input_data)
+        windows = batch.unfold(3, self.img_size, self.img_size).unfold(
+            4, self.img_size, self.img_size
+        )
+        self.h1, self.w1 = windows.shape[3:5]
+        windows = rearrange(
+            windows,
+            "b c t h1 w1 h w -> (b h1 w1) c t h w",
+            h=self.img_size,
+            w=self.img_size,
+        )
+
+        # Split into batches if number of windows > batch_size
+        num_batches = (
+            windows.shape[0] // self.batch_size
+            if windows.shape[0] > self.batch_size
+            else 1
+        )
+        windows = torch.tensor_split(windows, num_batches, dim=0)
+
+        if temporal_coords:
+            temporal_coords = torch.tensor(temporal_coords).unsqueeze(0)
+        else:
+            temporal_coords = None
+        if location_coords:
+            location_coords = torch.tensor(location_coords[0]).unsqueeze(0)
+        else:
+            location_coords = None
+
+        prompts = []
+        for window in windows:
+            # Apply standardization
+            window = self.datamodule.test_transform(
+                image=window.squeeze().numpy().transpose(1, 2, 0)
+            )
+            window = self.datamodule.aug(window)["image"]
+            prompts.append(
+                {
+                    "prompt_token_ids": [1],
+                    "multi_modal_data": {
+                        "image": {
+                            "pixel_values": window.to(torch.float16)[0],
+                            "location_coords": location_coords.to(torch.float16),
+                        }
+                    },
+                }
+            )
+
+        return prompts
+
+    def post_process(
+        self,
+        model_output: Sequence[PoolingRequestOutput],
+        request_id: str | None = None,
+        **kwargs,
+    ) -> ImageRequestOutput:
+        pred_imgs_list = []
+
+        if request_id and (request_id in self.requests_cache):
+            out_format = self.requests_cache[request_id]["out_format"]
+        else:
+            out_format = "b64_json"
+
+        for output in model_output:
+            y_hat = output.outputs.data.argmax(dim=0)
+            pred = torch.nn.functional.interpolate(
+                y_hat[None, None, ...].float(),
+                size=self.img_size,
+                mode="nearest",
+            )
+            pred_imgs_list.append(pred)
+
+        pred_imgs: torch.Tensor = torch.concat(pred_imgs_list, dim=0)
+
+        # Build images from patches
+        pred_imgs = rearrange(
+            pred_imgs,
+            "(b h1 w1) c h w -> b c (h1 h) (w1 w)",
+            h=self.img_size,
+            w=self.img_size,
+            b=1,
+            c=1,
+            h1=self.h1,
+            w1=self.w1,
+        )
+
+        # Cut padded area back to original size
+        pred_imgs = pred_imgs[..., : self.original_h, : self.original_w]
+
+        # Squeeze (batch size 1)
+        pred_imgs = pred_imgs[0]
+
+        if not self.meta_data:
+            raise ValueError("No metadata available for the current task")
+        self.meta_data.update(count=1, dtype="uint8", compress="lzw", nodata=0)
+        out_data = save_geotiff(
+            _convert_np_uint8(pred_imgs), self.meta_data, out_format
+        )
+
+        return ImageRequestOutput(
+            type=out_format,
+            format="tiff",
+            data=out_data,
+        )

tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/types.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any, Literal, TypedDict
+
+import albumentations
+from pydantic import BaseModel
+
+
+class DataModuleConfig(TypedDict):
+    bands: list[str]
+    batch_size: int
+    constant_scale: float
+    data_root: str
+    drop_last: bool
+    no_data_replace: float
+    no_label_replace: int
+    num_workers: int
+    test_transform: list[albumentations.core.transforms_interface.BasicTransform]
+
+
+class ImagePrompt(BaseModel):
+    data_format: Literal["b64_json", "bytes", "url", "path"]
+    """
+    This is the data type for the input image
+    """
+
+    image_format: str
+    """
+    This is the image format (e.g., jpeg, png, etc.)
+    """
+
+    out_data_format: Literal["b64_json", "url"]
+
+    data: Any
+    """
+    Input image data
+    """
+
+
+class ImageRequestOutput(BaseModel):
+    """
+    The output data of an image request to vLLM.
+
+    Args:
+        type (str): The data content type [path, object]
+        format (str): The image format (e.g., jpeg, png, etc.)
+        data (Any): The resulting data.
+    """
+
+    type: Literal["path", "b64_json"]
+    format: str
+    data: str

tests/plugins/prithvi_io_processor_plugin/setup.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from setuptools import setup
+
+setup(
+    name="prithvi_io_processor_plugin",
+    version="0.1",
+    packages=["prithvi_io_processor"],
+    entry_points={
+        "vllm.io_processor_plugins": [
+            "prithvi_to_tiff = prithvi_io_processor:register_prithvi",  # noqa: E501
+        ]
+    },
+)

tests/plugins/vllm_add_dummy_endpoint_plugin/setup.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from setuptools import setup
+
+setup(
+    name="vllm_add_dummy_endpoint_plugin",
+    version="0.1",
+    packages=["vllm_add_dummy_endpoint_plugin"],
+    entry_points={
+        "vllm.endpoint_plugins": [
+            "dummy_admin_endpoint_plugin = vllm_add_dummy_endpoint_plugin:DummyAdminEndpointPlugin"  # noqa
+        ]
+    },
+)

tests/plugins/vllm_add_dummy_endpoint_plugin/vllm_add_dummy_endpoint_plugin/__init__.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Worked example `vllm.endpoint_plugins` entry point.
+
+Reports scheduler config via `collective_rpc`. Demonstrates the full
+contract: `attach_router` registers the route at Phase A (`build_app`) and
+`init_state` stashes the `EngineClient` the route handler needs at Phase B
+(`init_app_state`).
+
+`required_tasks` is `None`, so this plugin is also eligible on the CPU only
+render server which has no `EngineClient`. `init_state` is called with
+`engine_client=None` in that case and the route handler returns 503 rather
+than reaching for a client that doesn't exist.
+"""
+
+from fastapi import FastAPI, HTTPException, Request
+
+
+class DummyAdminEndpointPlugin:
+    name = "dummy_admin_endpoint_plugin"
+    required_tasks: tuple[str, ...] | None = None
+
+    def attach_router(self, app: FastAPI) -> None:
+        @app.get("/v1/admin/scheduler_config")
+        async def scheduler_config(raw_request: Request):
+            engine_client = raw_request.app.state.dummy_engine_client
+            if engine_client is None:
+                raise HTTPException(
+                    status_code=503,
+                    detail="scheduler_config requires an engine, which this "
+                    "server does not have",
+                )
+            results = await engine_client.collective_rpc("get_scheduler_config")
+            return {"scheduler_config": results}
+
+    async def init_state(self, engine_client, state, args) -> None:
+        state.dummy_engine_client = engine_client

tests/plugins/vllm_add_dummy_model/setup.py

@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from setuptools import setup
+
+setup(
+    name="vllm_add_dummy_model",
+    version="0.1",
+    packages=["vllm_add_dummy_model"],
+    entry_points={
+        "vllm.general_plugins": ["register_dummy_model = vllm_add_dummy_model:register"]
+    },
+)

tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/__init__.py

@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm import ModelRegistry
+
+
+def register():
+    # Test directly passing the model
+    from .my_opt import MyOPTForCausalLM
+
+    if "MyOPTForCausalLM" not in ModelRegistry.get_supported_archs():
+        ModelRegistry.register_model("MyOPTForCausalLM", MyOPTForCausalLM)
+
+    # Test passing lazy model
+    if "MyGemma2Embedding" not in ModelRegistry.get_supported_archs():
+        ModelRegistry.register_model(
+            "MyGemma2Embedding",
+            "vllm_add_dummy_model.my_gemma_embedding:MyGemma2Embedding",
+        )
+
+    if "MyLlava" not in ModelRegistry.get_supported_archs():
+        ModelRegistry.register_model("MyLlava", "vllm_add_dummy_model.my_llava:MyLlava")

tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_gemma_embedding.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+
+from vllm.config import VllmConfig
+from vllm.model_executor.layers.pooler import DispatchPooler
+from vllm.model_executor.models.gemma2 import Gemma2Model
+from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix
+from vllm.sequence import IntermediateTensors
+
+
+class MyGemma2Embedding(nn.Module):
+    is_pooling_model = True
+
+    hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""})
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+
+        self.model = Gemma2Model(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
+
+        pooler_config = vllm_config.model_config.pooler_config
+        assert pooler_config is not None
+
+        self.pooler = DispatchPooler.for_embedding(pooler_config)
+
+        self.make_empty_intermediate_tensors = (
+            self.model.make_empty_intermediate_tensors
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor | IntermediateTensors:
+        hidden_states = self.model(
+            input_ids,
+            positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+        )
+
+        if isinstance(hidden_states, IntermediateTensors):
+            return hidden_states
+
+        # Return all-zero embeddings
+        return torch.zeros_like(hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
+        weights = self.hf_to_vllm_mapper.apply(weights)
+        weights = (
+            (name, data) for name, data in weights if not name.startswith("lm_head.")
+        )
+        return self.model.load_weights(weights)

tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_llava.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import torch
+
+from vllm.model_executor.models.llava import (
+    LlavaDummyInputsBuilder,
+    LlavaForConditionalGeneration,
+    LlavaMultiModalProcessor,
+    LlavaProcessingInfo,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    LlavaMultiModalProcessor,
+    info=LlavaProcessingInfo,
+    dummy_inputs=LlavaDummyInputsBuilder,
+)
+class MyLlava(LlavaForConditionalGeneration):
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
+        # this dummy model always predicts the first token
+        logits = super().compute_logits(hidden_states)
+        if logits is not None:
+            logits.zero_()
+            logits[:, 0] += 1.0
+        return logits

tests/plugins/vllm_add_dummy_model/vllm_add_dummy_model/my_opt.py

@@ -0,0 +1,17 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import torch
+
+from vllm.model_executor.models.opt import OPTForCausalLM
+
+
+class MyOPTForCausalLM(OPTForCausalLM):
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
+        # this dummy model always predicts the first token
+        logits = super().compute_logits(hidden_states)
+        if logits is not None:
+            logits.zero_()
+            logits[:, 0] += 1.0
+        return logits

tests/plugins/vllm_add_dummy_platform/setup.py

@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from setuptools import setup
+
+setup(
+    name="vllm_add_dummy_platform",
+    version="0.1",
+    packages=["vllm_add_dummy_platform"],
+    entry_points={
+        "vllm.platform_plugins": [
+            "dummy_platform_plugin = vllm_add_dummy_platform:dummy_platform_plugin"  # noqa
+        ],
+        "vllm.general_plugins": [
+            "dummy_custom_ops = vllm_add_dummy_platform:register_ops"
+        ],
+    },
+)

tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/__init__.py

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+def dummy_platform_plugin() -> str | None:
+    return "vllm_add_dummy_platform.dummy_platform.DummyPlatform"
+
+
+def register_ops():
+    import vllm_add_dummy_platform.dummy_custom_ops  # noqa

tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/dummy_attention_backend.py

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.attention.backends.placeholder_attn import PlaceholderAttentionBackend
+
+
+class DummyAttentionBackend(PlaceholderAttentionBackend):
+    @staticmethod
+    def get_name() -> str:
+        return "Dummy_Backend"

tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/dummy_custom_ops.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
+
+
+# Register CustomRotaryEmbedding to CustomOP.
+@RotaryEmbedding.register_oot
+class DummyRotaryEmbedding(RotaryEmbedding):
+    """Original rotary positional embedding."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.addition_config = True
+
+    def forward_oot(self, *args, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
+        return super().forward_oot(*args, **kwargs)

tests/plugins/vllm_add_dummy_platform/vllm_add_dummy_platform/dummy_platform.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import TYPE_CHECKING
+
+from vllm.platforms.interface import Platform, PlatformEnum
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+else:
+    VllmConfig = None
+
+
+class DummyPlatform(Platform):
+    _enum = PlatformEnum.OOT
+    device_name = "DummyDevice"
+    device_type: str = "privateuseone"
+    dispatch_key: str = "PrivateUse1"
+
+    @classmethod
+    def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
+        vllm_config.compilation_config.custom_ops = ["all"]
+
+    def get_attn_backend_cls(
+        self,
+        backend_name,
+        head_size,
+        dtype,
+        kv_cache_dtype,
+        block_size,
+        use_mla,
+        has_sink,
+        use_sparse,
+        use_mm_prefix,
+    ):
+        return "vllm_add_dummy_platform.dummy_attention_backend.DummyAttentionBackend"  # noqa E501

tests/plugins/vllm_add_dummy_stat_logger/dummy_stat_logger/dummy_stat_logger.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.v1.metrics.loggers import StatLoggerBase
+
+
+class DummyStatLogger(StatLoggerBase):
+    """
+    A dummy stat logger for testing purposes.
+    Implements the minimal interface expected by StatLoggerManager.
+    """
+
+    def __init__(self, vllm_config, engine_idx=0):
+        self.vllm_config = vllm_config
+        self.engine_idx = engine_idx
+        self.recorded = []
+        self.logged = False
+        self.engine_initialized = False
+
+    def record(self, scheduler_stats, iteration_stats, mm_cache_stats, engine_idx):
+        self.recorded.append(
+            (scheduler_stats, iteration_stats, mm_cache_stats, engine_idx)
+        )
+
+    def log(self):
+        self.logged = True
+
+    def log_engine_initialized(self):
+        self.engine_initialized = True

tests/plugins/vllm_add_dummy_stat_logger/setup.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from setuptools import setup
+
+setup(
+    name="dummy_stat_logger",
+    version="0.1",
+    packages=["dummy_stat_logger"],
+    entry_points={
+        "vllm.stat_logger_plugins": [
+            "dummy_stat_logger = dummy_stat_logger.dummy_stat_logger:DummyStatLogger"  # noqa
+        ]
+    },
+)