A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Misc] Terratorch related fixes (#24337)
Signed-off-by: Christian Pinto <christian.pinto@ibm.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Christian Pinto committed
Sep 8, 2025 at 15:40 UTC
9cd76b71abf15b31878f8d9675546f809a6ba150
11 files changed
+18
-37
examples/offline_inference/prithvi_geospatial_mae_io_processor.py
+2
-2
@@ -18,7 +18,7 @@ from vllm.pooling_params import PoolingParams
18
19
def main():
20
torch.set_default_dtype(torch.float16)
21
- image_url = "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/India_900498_S2Hand.tif" # noqa: E501
21
+ image_url = "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/valencia_example_2024-10-26.tiff" # noqa: E501
22
23
img_prompt = dict(
24
data=image_url,
@@ -36,7 +36,7 @@ def main():
36
# to avoid the model going OOM.
37
# The maximum number depends on the available GPU memory
38
max_num_seqs=32,
39
- io_processor_plugin="prithvi_to_tiff_india",
39
+ io_processor_plugin="prithvi_to_tiff",
40
model_impl="terratorch",
41
)
42
examples/online_serving/prithvi_geospatial_mae.py
+2
-2
@@ -18,11 +18,11 @@ import requests
18
# --model-impl terratorch
19
# --task embed --trust-remote-code
20
# --skip-tokenizer-init --enforce-eager
21
-# --io-processor-plugin prithvi_to_tiff_india
21
+# --io-processor-plugin prithvi_to_tiff
22
23
24
def main():
25
- image_url = "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/India_900498_S2Hand.tif" # noqa: E501
25
+ image_url = "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/valencia_example_2024-10-26.tiff" # noqa: E501
26
server_endpoint = "http://localhost:8000/pooling"
27
28
request_payload_url = {
requirements/test.in
+1
-1
@@ -54,4 +54,4 @@ runai-model-streamer-s3==0.11.0
54
fastsafetensors>=0.1.10
55
pydantic>=2.10 # 2.9 leads to error on python 3.10
56
decord==0.6.0
57
-terratorch==1.1rc3 # required for PrithviMAE test
57
+terratorch @ git+https://github.com/IBM/terratorch.git@1.1.rc3 # required for PrithviMAE test
requirements/test.txt
+1
-1
@@ -1042,7 +1042,7 @@ tensorboardx==2.6.4
1042
# via lightning
1043
tensorizer==2.10.1
1044
# via -r requirements/test.in
1045
-terratorch==1.1rc3
1045
+terratorch @ git+https://github.com/IBM/terratorch.git@07184fcf91a1324f831ff521dd238d97fe350e3e
1046
# via -r requirements/test.in
1047
threadpoolctl==3.5.0
1048
# via scikit-learn
tests/entrypoints/openai/test_skip_tokenizer.py
+1
-1
@@ -11,7 +11,7 @@ import torch
11
12
from ...utils import RemoteOpenAIServer
13
14
-MODEL_NAME = "mgazz/Prithvi-EO-2.0-300M-TL-Sen1Floods11"
14
+MODEL_NAME = "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11"
15
DTYPE = "float16"
16
17
tests/models/registry.py
+2
-2
@@ -383,7 +383,7 @@ _EMBEDDING_EXAMPLE_MODELS = {
383
"Phi3VForCausalLM": _HfExamplesInfo("TIGER-Lab/VLM2Vec-Full",
384
trust_remote_code=True),
385
"Qwen2VLForConditionalGeneration": _HfExamplesInfo("MrLight/dse-qwen2-2b-mrl-v1"), # noqa: E501
386
- "PrithviGeoSpatialMAE": _HfExamplesInfo("mgazz/Prithvi-EO-2.0-300M-TL-Sen1Floods11", # noqa: E501
386
+ "PrithviGeoSpatialMAE": _HfExamplesInfo("ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11", # noqa: E501
387
dtype=torch.float16,
388
enforce_eager=True,
389
skip_tokenizer_init=True,
@@ -391,7 +391,7 @@ _EMBEDDING_EXAMPLE_MODELS = {
391
# going OOM in CI
392
max_num_seqs=32,
393
),
394
- "Terratorch": _HfExamplesInfo("mgazz/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
394
+ "Terratorch": _HfExamplesInfo("ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11", # noqa: E501
395
dtype=torch.float16,
396
enforce_eager=True,
397
skip_tokenizer_init=True,
tests/models/test_terratorch.py
+1
-1
@@ -11,7 +11,7 @@ from vllm.utils import set_default_torch_num_threads
11
@pytest.mark.parametrize(
12
"model",
13
[
14
- "mgazz/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
14
+ "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
15
"mgazz/Prithvi_v2_eo_300_tl_unet_agb"
16
],
17
)
tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/__init__.py
+2
-4
@@ -1,8 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-def register_prithvi_india():
4
- return "prithvi_io_processor.prithvi_processor.PrithviMultimodalDataProcessorIndia" # noqa: E501
3
4
7
-def register_prithvi_valencia():
8
- return "prithvi_io_processor.prithvi_processor.PrithviMultimodalDataProcessorValencia" # noqa: E501
5
+def register_prithvi():
6
+ return "prithvi_io_processor.prithvi_processor.PrithviMultimodalDataProcessor" # noqa: E501
tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py
+2
-18
@@ -234,6 +234,8 @@ def load_image(
234
235
class PrithviMultimodalDataProcessor(IOProcessor):
236
237
+ indices = [0, 1, 2, 3, 4, 5]
238
+
239
def __init__(self, vllm_config: VllmConfig):
240
241
super().__init__(vllm_config)
@@ -412,21 +414,3 @@ class PrithviMultimodalDataProcessor(IOProcessor):
414
format="tiff",
415
data=out_data,
416
request_id=request_id)
415
-
416
-
417
-class PrithviMultimodalDataProcessorIndia(PrithviMultimodalDataProcessor):
418
-
419
- def __init__(self, vllm_config: VllmConfig):
420
-
421
- super().__init__(vllm_config)
422
-
423
- self.indices = [1, 2, 3, 8, 11, 12]
424
-
425
-
426
-class PrithviMultimodalDataProcessorValencia(PrithviMultimodalDataProcessor):
427
-
428
- def __init__(self, vllm_config: VllmConfig):
429
-
430
- super().__init__(vllm_config)
431
-
432
- self.indices = [0, 1, 2, 3, 4, 5]
tests/plugins/prithvi_io_processor_plugin/setup.py
+1
-2
@@ -9,8 +9,7 @@ setup(
9
packages=["prithvi_io_processor"],
10
entry_points={
11
"vllm.io_processor_plugins": [
12
- "prithvi_to_tiff_india = prithvi_io_processor:register_prithvi_india", # noqa: E501
13
- "prithvi_to_tiff_valencia = prithvi_io_processor:register_prithvi_valencia", # noqa: E501
12
+ "prithvi_to_tiff = prithvi_io_processor:register_prithvi", # noqa: E501
13
]
14
},
15
)
tests/plugins_tests/test_io_processor_plugins.py
+3
-3
@@ -11,7 +11,7 @@ from vllm.entrypoints.openai.protocol import IOProcessorResponse
11
from vllm.plugins.io_processors import get_io_processor
12
from vllm.pooling_params import PoolingParams
13
14
-MODEL_NAME = "mgazz/Prithvi-EO-2.0-300M-TL-Sen1Floods11"
14
+MODEL_NAME = "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11"
15
16
image_url = "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/valencia_example_2024-10-26.tiff" # noqa: E501
17
@@ -35,7 +35,7 @@ def server():
35
"--max-num-seqs",
36
"32",
37
"--io-processor-plugin",
38
- "prithvi_to_tiff_valencia",
38
+ "prithvi_to_tiff",
39
"--model-impl",
40
"terratorch",
41
]
@@ -107,7 +107,7 @@ def test_prithvi_mae_plugin_offline(vllm_runner, model_name: str):
107
# to avoid the model going OOM in CI.
108
max_num_seqs=1,
109
model_impl="terratorch",
110
- io_processor_plugin="prithvi_to_tiff_valencia",
110
+ io_processor_plugin="prithvi_to_tiff",
111
) as llm_runner:
112
pooler_output = llm_runner.get_llm().encode(
113
img_prompt,