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 8, 07:30 UTC
[Misc] Update tests and examples for Prithvi/Terratorch models (#34416)
Signed-off-by: Christian Pinto <christian.pinto@ibm.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Christian Pinto committed
Feb 14, 2026 at 07:03 UTC
342a7cda2d212205c4874f82a59559400bbec311
9 files changed
+162
-131
examples/pooling/plugin/prithvi_geospatial_mae_io_processor.py
+1
-1
@@ -28,7 +28,7 @@ def main():
28
)
29
30
llm = LLM(
31
- model="christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM",
31
+ model="ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
32
skip_tokenizer_init=True,
33
trust_remote_code=True,
34
enforce_eager=True,
examples/pooling/plugin/prithvi_geospatial_mae_offline.py
+1
-1
@@ -391,7 +391,7 @@ if __name__ == "__main__":
391
parser.add_argument(
392
"--model",
393
type=str,
394
- default="christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM",
394
+ default="ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
395
help="Path to a checkpoint file to load from.",
396
)
397
parser.add_argument(
examples/pooling/plugin/prithvi_geospatial_mae_online.py
+2
-4
@@ -14,9 +14,7 @@ import requests
14
# - install TerraTorch v1.1 (or later):
15
# pip install terratorch>=v1.1
16
# - start vllm in serving mode with the below args
17
-# --model='christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM'
18
-# --model-impl terratorch
19
-# --trust-remote-code
17
+# --model='ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11'
18
# --skip-tokenizer-init --enforce-eager
19
# --io-processor-plugin terratorch_segmentation
20
# --enable-mm-embeds
@@ -34,7 +32,7 @@ def main():
32
"out_data_format": "b64_json",
33
},
34
"priority": 0,
37
- "model": "christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM",
35
+ "model": "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
36
}
37
38
ret = requests.post(server_endpoint, json=request_payload_url)
requirements/test.in
+8
-1
@@ -56,7 +56,14 @@ runai-model-streamer[s3,gcs]==0.15.3
56
fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
57
pydantic>=2.12 # 2.11 leads to error on python 3.13
58
decord==0.6.0
59
-terratorch @ git+https://github.com/IBM/terratorch.git@1.1.rc3 # required for PrithviMAE test
59
+terratorch >= 1.2.2 # Required for Prithvi tests
60
+imagehash # Required for Prithvi tests
61
+segmentation-models-pytorch > 0.4.0 # Required for Prithvi tests
62
+
63
gpt-oss >= 0.0.7; python_version > '3.11'
64
65
perceptron # required for isaac test
66
+
67
+# Newer versions of datasets require torchcoded, that makes the tests fail in CI because of a missing library.
68
+# Older versions are in conflict with teerratorch requirements.
69
+datasets>=3.3.0,<=3.6.0
\ No newline at end of file
requirements/test.txt
+87
-94
@@ -1,7 +1,9 @@
1
# This file was autogenerated by uv via the following command:
2
# uv pip compile requirements/test.in -o requirements/test.txt --index-strategy unsafe-best-match --torch-backend cu129 --python-platform x86_64-manylinux_2_28 --python-version 3.12
3
absl-py==2.1.0
4
- # via rouge-score
4
+ # via
5
+ # rouge-score
6
+ # tensorboard
7
accelerate==1.0.1
8
# via
9
# lm-eval
@@ -31,9 +33,7 @@ albumentations==1.4.6
33
# -r requirements/test.in
34
# terratorch
35
alembic==1.16.4
34
- # via
35
- # mlflow
36
- # optuna
36
+ # via optuna
37
annotated-doc==0.0.4
38
# via fastapi
39
annotated-types==0.7.0
@@ -74,8 +74,6 @@ bitsandbytes==0.46.1
74
# lightning
75
black==24.10.0
76
# via datamodel-code-generator
77
-blinker==1.9.0
78
- # via flask
77
blobfile==3.0.0
78
# via -r requirements/test.in
79
bm25s==0.2.13
@@ -93,9 +91,7 @@ bounded-pool-executor==0.0.3
91
buildkite-test-collector==0.1.9
92
# via -r requirements/test.in
93
cachetools==5.5.2
96
- # via
97
- # google-auth
98
- # mlflow-skinny
94
+ # via google-auth
95
certifi==2024.8.30
96
# via
97
# fiona
@@ -106,6 +102,7 @@ certifi==2024.8.30
102
# pyproj
103
# rasterio
104
# requests
105
+ # sentry-sdk
106
cffi==1.17.1
107
# via soundfile
108
chardet==5.2.0
@@ -120,15 +117,14 @@ click==8.1.7
117
# click-plugins
118
# cligj
119
# fiona
123
- # flask
120
# jiwer
125
- # mlflow-skinny
121
# nltk
122
# rasterio
123
# ray
124
# schemathesis
125
# typer
126
# uvicorn
127
+ # wandb
128
click-plugins==1.1.1.2
129
# via
130
# fiona
@@ -137,8 +133,6 @@ cligj==0.7.2
133
# via
134
# fiona
135
# rasterio
140
-cloudpickle==3.1.1
141
- # via mlflow-skinny
136
colorama==0.4.6
137
# via
138
# perceptron
@@ -163,16 +157,15 @@ cupy-cuda12x==13.6.0
157
# via ray
158
cycler==0.12.1
159
# via matplotlib
166
-databricks-sdk==0.59.0
167
- # via mlflow-skinny
160
datamodel-code-generator==0.26.3
161
# via -r requirements/test.in
162
dataproperty==1.0.1
163
# via
164
# pytablewriter
165
# tabledata
174
-datasets==3.0.2
166
+datasets==3.3.0
167
# via
168
+ # -r requirements/test.in
169
# evaluate
170
# lm-eval
171
# mteb
@@ -180,6 +173,8 @@ decorator==5.1.1
173
# via librosa
174
decord==0.6.0
175
# via -r requirements/test.in
176
+diffusers==0.36.0
177
+ # via terratorch
178
dill==0.3.8
179
# via
180
# datasets
@@ -191,15 +186,11 @@ distlib==0.3.9
186
dnspython==2.7.0
187
# via email-validator
188
docker==7.1.0
194
- # via
195
- # gpt-oss
196
- # mlflow
189
+ # via gpt-oss
190
docopt==0.6.2
191
# via num2words
192
docstring-parser==0.17.0
193
# via jsonargparse
201
-efficientnet-pytorch==0.7.1
202
- # via segmentation-models-pytorch
194
einops==0.8.1
195
# via
196
# -r requirements/test.in
@@ -217,9 +208,7 @@ encodec==0.1.1
208
evaluate==0.4.3
209
# via lm-eval
210
fastapi==0.128.0
220
- # via
221
- # gpt-oss
222
- # mlflow-skinny
211
+ # via gpt-oss
212
fastparquet==2024.11.0
213
# via genai-perf
214
fastrlock==0.8.2
@@ -230,6 +219,7 @@ filelock==3.16.1
219
# via
220
# blobfile
221
# datasets
222
+ # diffusers
223
# huggingface-hub
224
# ray
225
# torch
@@ -237,8 +227,6 @@ filelock==3.16.1
227
# virtualenv
228
fiona==1.10.1
229
# via torchgeo
240
-flask==3.1.1
241
- # via mlflow
230
fonttools==4.55.0
231
# via matplotlib
232
fqdn==1.5.1
@@ -249,7 +237,7 @@ frozenlist==1.5.0
237
# via
238
# aiohttp
239
# aiosignal
252
-fsspec==2024.9.0
240
+fsspec==2024.12.0
241
# via
242
# datasets
243
# evaluate
@@ -257,6 +245,7 @@ fsspec==2024.9.0
245
# huggingface-hub
246
# lightning
247
# pytorch-lightning
248
+ # tacoreader
249
# torch
250
ftfy==6.3.1
251
# via open-clip-torch
@@ -269,7 +258,7 @@ geopandas==1.0.1
258
gitdb==4.0.12
259
# via gitpython
260
gitpython==3.1.44
272
- # via mlflow-skinny
261
+ # via wandb
262
google-api-core==2.24.2
263
# via
264
# google-cloud-core
@@ -277,7 +266,6 @@ google-api-core==2.24.2
266
# opencensus
267
google-auth==2.40.2
268
# via
280
- # databricks-sdk
269
# google-api-core
270
# google-cloud-core
271
# google-cloud-storage
@@ -296,25 +284,17 @@ googleapis-common-protos==1.70.0
284
# via google-api-core
285
gpt-oss==0.0.8
286
# via -r requirements/test.in
299
-graphene==3.4.3
300
- # via mlflow
287
graphql-core==3.2.6
302
- # via
303
- # graphene
304
- # graphql-relay
305
- # hypothesis-graphql
306
-graphql-relay==3.2.0
307
- # via graphene
288
+ # via hypothesis-graphql
289
greenlet==3.2.3
290
# via sqlalchemy
291
grpcio==1.78.0
292
# via
293
# grpcio-tools
294
# ray
295
+ # tensorboard
296
grpcio-tools==1.78.0
297
# via -r requirements/test.in
316
-gunicorn==23.0.0
317
- # via mlflow
298
h11==0.14.0
299
# via
300
# httpcore
@@ -338,12 +318,14 @@ httpcore==1.0.6
318
httpx==0.27.2
319
# via
320
# -r requirements/test.in
321
+ # diffusers
322
# perceptron
323
# schemathesis
324
huggingface-hub==0.36.2
325
# via
326
# accelerate
327
# datasets
328
+ # diffusers
329
# evaluate
330
# open-clip-torch
331
# peft
@@ -379,11 +361,13 @@ idna==3.10
361
# jsonschema
362
# requests
363
# yarl
364
+imagehash==4.3.2
365
+ # via -r requirements/test.in
366
imageio==2.37.0
367
# via scikit-image
368
importlib-metadata==8.7.0
369
# via
386
- # mlflow-skinny
370
+ # diffusers
371
# opentelemetry-api
372
importlib-resources==6.5.2
373
# via typeshed-client
@@ -395,14 +379,10 @@ isoduration==20.11.0
379
# via jsonschema
380
isort==5.13.2
381
# via datamodel-code-generator
398
-itsdangerous==2.2.0
399
- # via flask
382
jinja2==3.1.6
383
# via
384
# datamodel-code-generator
403
- # flask
385
# genai-perf
405
- # mlflow
386
# torch
387
jiwer==3.0.5
388
# via -r requirements/test.in
@@ -415,12 +395,14 @@ joblib==1.4.2
395
# librosa
396
# nltk
397
# scikit-learn
418
-jsonargparse==4.35.0
398
+jsonargparse==4.46.0
399
# via
400
# lightning
401
# terratorch
402
jsonlines==4.0.0
403
# via lm-eval
404
+jsonnet==0.21.0
405
+ # via jsonargparse
406
jsonpointer==3.0.0
407
# via jsonschema
408
jsonschema==4.23.0
@@ -449,13 +431,13 @@ libnacl==2.1.0
431
# via tensorizer
432
librosa==0.10.2.post1
433
# via -r requirements/test.in
452
-lightly==1.5.20
434
+lightly==1.5.22
435
# via
436
# terratorch
437
# torchgeo
438
lightly-utils==0.0.2
439
# via lightly
458
-lightning==2.5.1.post0
440
+lightning==2.6.1
441
# via
442
# terratorch
443
# torchgeo
@@ -476,12 +458,11 @@ lxml==5.3.0
458
mako==1.3.10
459
# via alembic
460
markdown==3.8.2
479
- # via mlflow
461
+ # via tensorboard
462
markdown-it-py==3.0.0
463
# via rich
464
markupsafe==3.0.1
465
# via
484
- # flask
466
# jinja2
467
# mako
468
# werkzeug
@@ -489,7 +470,6 @@ matplotlib==3.9.2
470
# via
471
# -r requirements/test.in
472
# lightning
492
- # mlflow
473
# pycocotools
474
# torchgeo
475
mbstrdecoder==1.1.3
@@ -501,10 +481,6 @@ mdurl==0.1.2
481
# via markdown-it-py
482
mistral-common==1.9.1
483
# via -r requirements/test.in
504
-mlflow==2.22.0
505
- # via terratorch
506
-mlflow-skinny==2.22.0
507
- # via mlflow
484
more-itertools==10.5.0
485
# via lm-eval
486
mpmath==1.3.0
@@ -523,8 +499,6 @@ multiprocess==0.70.16
499
# via
500
# datasets
501
# evaluate
526
-munch==4.0.0
527
- # via pretrainedmodels
502
mypy-extensions==1.0.0
503
# via black
504
networkx==3.2.1
@@ -553,6 +527,7 @@ numpy==2.2.6
527
# cupy-cuda12x
528
# datasets
529
# decord
530
+ # diffusers
531
# einx
532
# encodec
533
# evaluate
@@ -560,13 +535,13 @@ numpy==2.2.6
535
# genai-perf
536
# geopandas
537
# h5py
538
+ # imagehash
539
# imageio
540
# librosa
541
# lightly
542
# lightly-utils
543
# matplotlib
544
# mistral-common
569
- # mlflow
545
# mteb
546
# numba
547
# numexpr
@@ -578,6 +553,7 @@ numpy==2.2.6
553
# perceptron
554
# pycocotools
555
# pyogrio
556
+ # pywavelets
557
# rasterio
558
# rioxarray
559
# rouge-score
@@ -590,8 +566,10 @@ numpy==2.2.6
566
# shapely
567
# soxr
568
# statsmodels
569
+ # tensorboard
570
# tensorboardx
571
# tensorizer
572
+ # terratorch
573
# tifffile
574
# torchgeo
575
# torchmetrics
@@ -659,7 +637,6 @@ opencv-python-headless==4.13.0.90
637
# mistral-common
638
opentelemetry-api==1.35.0
639
# via
662
- # mlflow-skinny
640
# opentelemetry-exporter-prometheus
641
# opentelemetry-sdk
642
# opentelemetry-semantic-conventions
@@ -669,7 +646,6 @@ opentelemetry-proto==1.36.0
646
# via ray
647
opentelemetry-sdk==1.35.0
648
# via
672
- # mlflow-skinny
649
# opentelemetry-exporter-prometheus
650
# ray
651
opentelemetry-semantic-conventions==0.56b0
@@ -687,7 +663,6 @@ packaging==24.2
663
# evaluate
664
# fastparquet
665
# geopandas
690
- # gunicorn
666
# huggingface-hub
667
# hydra-core
668
# kornia
@@ -695,7 +670,6 @@ packaging==24.2
670
# lightning
671
# lightning-utilities
672
# matplotlib
698
- # mlflow-skinny
673
# optuna
674
# peft
675
# plotly
@@ -708,10 +682,12 @@ packaging==24.2
682
# rioxarray
683
# scikit-image
684
# statsmodels
685
+ # tensorboard
686
# tensorboardx
687
# torchmetrics
688
# transformers
689
# typepy
690
+ # wandb
691
# xarray
692
pandas==2.2.3
693
# via
@@ -720,8 +696,8 @@ pandas==2.2.3
696
# fastparquet
697
# genai-perf
698
# geopandas
723
- # mlflow
699
# statsmodels
700
+ # tacoreader
701
# torchgeo
702
# xarray
703
pathspec==0.12.1
@@ -740,7 +716,9 @@ perf-analyzer==0.1.0
716
# via genai-perf
717
pillow==10.4.0
718
# via
719
+ # diffusers
720
# genai-perf
721
+ # imagehash
722
# imageio
723
# lightly-utils
724
# matplotlib
@@ -748,6 +726,7 @@ pillow==10.4.0
726
# perceptron
727
# scikit-image
728
# segmentation-models-pytorch
729
+ # tensorboard
730
# torchgeo
731
# torchvision
732
platformdirs==4.3.6
@@ -755,6 +734,7 @@ platformdirs==4.3.6
734
# black
735
# pooch
736
# virtualenv
737
+ # wandb
738
plotly==5.24.1
739
# via genai-perf
740
pluggy==1.5.0
@@ -769,8 +749,6 @@ portalocker==2.10.1
749
# via sacrebleu
750
pqdm==0.2.0
751
# via -r requirements/test.in
772
-pretrainedmodels==0.7.4
773
- # via segmentation-models-pytorch
752
prometheus-client==0.22.0
753
# via
754
# opentelemetry-exporter-prometheus
@@ -786,12 +764,13 @@ protobuf==6.33.2
764
# google-api-core
765
# googleapis-common-protos
766
# grpcio-tools
789
- # mlflow-skinny
767
# opentelemetry-proto
768
# proto-plus
769
# ray
770
+ # tensorboard
771
# tensorboardx
772
# tensorizer
773
+ # wandb
774
psutil==6.1.0
775
# via
776
# accelerate
@@ -801,11 +780,12 @@ py==1.11.0
780
# via pytest-forked
781
py-spy==0.4.0
782
# via ray
804
-pyarrow==18.0.0
783
+pyarrow==23.0.0
784
# via
785
# datasets
786
# genai-perf
808
- # mlflow
787
+ # tacoreader
788
+ # terratorch
789
pyasn1==0.6.1
790
# via
791
# pyasn1-modules
@@ -831,11 +811,11 @@ pydantic==2.12.0
811
# gpt-oss
812
# lightly
813
# mistral-common
834
- # mlflow-skinny
814
# mteb
815
# openai-harmony
816
# pydantic-extra-types
817
# ray
818
+ # wandb
819
pydantic-core==2.41.1
820
# via pydantic
821
pydantic-extra-types==2.10.5
@@ -873,7 +853,6 @@ pytest==8.3.5
853
# pytest-subtests
854
# pytest-timeout
855
# schemathesis
876
- # terratorch
856
pytest-asyncio==0.24.0
857
# via -r requirements/test.in
858
pytest-cov==6.3.0
@@ -896,7 +875,6 @@ python-dateutil==2.9.0.post0
875
# via
876
# arrow
877
# botocore
899
- # graphene
878
# lightly
879
# matplotlib
880
# pandas
@@ -913,6 +891,8 @@ pytz==2024.2
891
# via
892
# pandas
893
# typepy
894
+pywavelets==1.9.0
895
+ # via imagehash
896
pyyaml==6.0.2
897
# via
898
# accelerate
@@ -923,7 +903,6 @@ pyyaml==6.0.2
903
# huggingface-hub
904
# jsonargparse
905
# lightning
926
- # mlflow-skinny
906
# omegaconf
907
# optuna
908
# peft
@@ -934,6 +913,7 @@ pyyaml==6.0.2
913
# timm
914
# transformers
915
# vocos
916
+ # wandb
917
rapidfuzz==3.12.1
918
# via jiwer
919
rasterio==1.4.3
@@ -951,6 +931,7 @@ referencing==0.35.1
931
# jsonschema-specifications
932
regex==2024.9.11
933
# via
934
+ # diffusers
935
# nltk
936
# open-clip-torch
937
# sacrebleu
@@ -959,8 +940,8 @@ regex==2024.9.11
940
requests==2.32.3
941
# via
942
# buildkite-test-collector
962
- # databricks-sdk
943
# datasets
944
+ # diffusers
945
# docker
946
# evaluate
947
# google-api-core
@@ -970,15 +951,16 @@ requests==2.32.3
951
# lightly
952
# lm-eval
953
# mistral-common
973
- # mlflow-skinny
954
# mteb
955
# pooch
956
# ray
957
# responses
958
# schemathesis
959
# starlette-testclient
960
+ # tacoreader
961
# tiktoken
962
# transformers
963
+ # wandb
964
responses==0.25.3
965
# via genai-perf
966
rfc3339-validator==0.1.4
@@ -991,6 +973,7 @@ rich==13.9.4
973
# lightning
974
# mteb
975
# perceptron
976
+ # terratorch
977
# typer
978
rioxarray==0.19.0
979
# via terratorch
@@ -1017,47 +1000,55 @@ sacrebleu==2.4.3
1000
safetensors==0.4.5
1001
# via
1002
# accelerate
1003
+ # diffusers
1004
# open-clip-torch
1005
# peft
1006
+ # segmentation-models-pytorch
1007
# timm
1008
# transformers
1009
schemathesis==3.39.15
1010
# via -r requirements/test.in
1011
scikit-image==0.25.2
1027
- # via albumentations
1012
+ # via
1013
+ # albumentations
1014
+ # terratorch
1015
scikit-learn==1.5.2
1016
# via
1017
# albumentations
1018
# librosa
1019
# lm-eval
1033
- # mlflow
1020
# mteb
1021
# sentence-transformers
1022
+ # terratorch
1023
scipy==1.13.1
1024
# via
1025
# albumentations
1026
# bm25s
1027
+ # imagehash
1028
# librosa
1041
- # mlflow
1029
# mteb
1030
# scikit-image
1031
# scikit-learn
1032
# sentence-transformers
1033
# statsmodels
1034
# vocos
1048
-segmentation-models-pytorch==0.4.0
1035
+segmentation-models-pytorch==0.5.0
1036
# via
1037
+ # -r requirements/test.in
1038
# terratorch
1039
# torchgeo
1040
sentence-transformers==5.2.0
1041
# via
1042
# -r requirements/test.in
1043
# mteb
1044
+sentry-sdk==2.52.0
1045
+ # via wandb
1046
setuptools==77.0.3
1047
# via
1048
# grpcio-tools
1049
# lightning-utilities
1050
# pytablewriter
1051
+ # tensorboard
1052
# torch
1053
shapely==2.1.1
1054
# via
@@ -1075,7 +1066,6 @@ six==1.16.0
1066
# python-dateutil
1067
# rfc3339-validator
1068
# rouge-score
1078
- # segmentation-models-pytorch
1069
smart-open==7.1.0
1070
# via ray
1071
smmap==5.0.2
@@ -1099,12 +1089,9 @@ soxr==0.5.0.post1
1089
sqlalchemy==2.0.41
1090
# via
1091
# alembic
1102
- # mlflow
1092
# optuna
1093
sqlitedict==2.1.0
1094
# via lm-eval
1106
-sqlparse==0.5.3
1107
- # via mlflow-skinny
1095
starlette==0.50.0
1096
# via
1097
# fastapi
@@ -1124,6 +1111,8 @@ tabledata==1.3.3
1111
# via pytablewriter
1112
tabulate==0.9.0
1113
# via sacrebleu
1114
+tacoreader==0.5.6
1115
+ # via terratorch
1116
tblib==3.1.0
1117
# via -r requirements/test.in
1118
tcolorpy==0.1.6
@@ -1133,13 +1122,19 @@ tenacity==9.1.2
1122
# gpt-oss
1123
# lm-eval
1124
# plotly
1125
+tensorboard==2.20.0
1126
+ # via terratorch
1127
+tensorboard-data-server==0.7.2
1128
+ # via tensorboard
1129
tensorboardx==2.6.4
1130
# via lightning
1131
tensorizer==2.10.1
1132
# via -r requirements/test.in
1133
termcolor==3.1.0
1141
- # via gpt-oss
1142
-terratorch @ git+https://github.com/IBM/terratorch.git@07184fcf91a1324f831ff521dd238d97fe350e3e
1134
+ # via
1135
+ # gpt-oss
1136
+ # terratorch
1137
+terratorch==1.2.2
1138
# via -r requirements/test.in
1139
threadpoolctl==3.5.0
1140
# via scikit-learn
@@ -1172,7 +1167,6 @@ torch==2.10.0+cu129
1167
# -r requirements/test.in
1168
# accelerate
1169
# bitsandbytes
1175
- # efficientnet-pytorch
1170
# encodec
1171
# kornia
1172
# lightly
@@ -1181,7 +1175,6 @@ torch==2.10.0+cu129
1175
# mteb
1176
# open-clip-torch
1177
# peft
1184
- # pretrainedmodels
1178
# pytorch-lightning
1179
# runai-model-streamer
1180
# segmentation-models-pytorch
@@ -1213,12 +1206,11 @@ torchvision==0.25.0+cu129
1206
# -r requirements/test.in
1207
# lightly
1208
# open-clip-torch
1216
- # pretrainedmodels
1209
# segmentation-models-pytorch
1210
# terratorch
1211
# timm
1212
# torchgeo
1221
-tqdm==4.66.6
1213
+tqdm==4.67.3
1214
# via
1215
# datasets
1216
# evaluate
@@ -1232,10 +1224,11 @@ tqdm==4.66.6
1224
# optuna
1225
# peft
1226
# pqdm
1235
- # pretrainedmodels
1227
# pytorch-lightning
1228
# segmentation-models-pytorch
1229
# sentence-transformers
1230
+ # tacoreader
1231
+ # terratorch
1232
# tqdm-multiprocess
1233
# transformers
1234
tqdm-multiprocess==0.0.11
@@ -1274,14 +1267,12 @@ typing-extensions==4.15.0
1267
# alembic
1268
# chz
1269
# fastapi
1277
- # graphene
1270
# grpcio
1271
# huggingface-hub
1272
# librosa
1273
# lightning
1274
# lightning-utilities
1275
# mistral-common
1284
- # mlflow-skinny
1276
# mteb
1277
# opentelemetry-api
1278
# opentelemetry-sdk
@@ -1299,6 +1290,7 @@ typing-extensions==4.15.0
1290
# typer
1291
# typeshed-client
1292
# typing-inspection
1293
+ # wandb
1294
typing-inspection==0.4.2
1295
# via pydantic
1296
tzdata==2024.2
@@ -1313,25 +1305,26 @@ urllib3==2.2.3
1305
# lightly
1306
# requests
1307
# responses
1308
+ # sentry-sdk
1309
# tritonclient
1310
uvicorn==0.35.0
1318
- # via
1319
- # gpt-oss
1320
- # mlflow-skinny
1311
+ # via gpt-oss
1312
vector-quantize-pytorch==1.21.2
1313
# via -r requirements/test.in
1314
virtualenv==20.31.2
1315
# via ray
1316
vocos==0.1.0
1317
# via -r requirements/test.in
1318
+wandb==0.24.2
1319
+ # via terratorch
1320
wcwidth==0.2.13
1321
# via ftfy
1322
webcolors==24.11.1
1323
# via jsonschema
1324
werkzeug==3.1.3
1325
# via
1333
- # flask
1326
# schemathesis
1327
+ # tensorboard
1328
word2number==1.1
1329
# via lm-eval
1330
wrapt==1.17.2
tests/models/multimodal/pooling/test_prithvi_mae.py
+1
-1
@@ -40,7 +40,7 @@ def _run_test(
40
vllm_model.llm.encode(prompt, pooling_task="plugin")
41
42
43
-MODELS = ["mgazz/Prithvi-EO-2.0-300M-TL-Sen1Floods11"]
43
+MODELS = ["ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11"]
44
45
46
@pytest.mark.core_model
tests/models/test_terratorch.py
+1
-1
@@ -13,7 +13,7 @@ from tests.utils import create_new_process_for_each_test
13
"model",
14
[
15
"ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
16
- "mgazz/Prithvi_v2_eo_300_tl_unet_agb",
16
+ "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-BurnScars",
17
],
18
)
19
def test_inference(
tests/plugins/prithvi_io_processor_plugin/prithvi_io_processor/prithvi_processor.py
+2
-6
@@ -44,12 +44,8 @@ datamodule_config: DataModuleConfig = {
44
"no_label_replace": -1,
45
"num_workers": 8,
46
"test_transform": [
47
- albumentations.Resize(
48
- always_apply=False, height=448, interpolation=1, p=1, width=448
49
- ),
50
- albumentations.pytorch.ToTensorV2(
51
- transpose_mask=False, always_apply=True, p=1.0
52
- ),
47
+ albumentations.Resize(height=448, interpolation=1, p=1, width=448),
48
+ albumentations.pytorch.ToTensorV2(transpose_mask=False, p=1.0),
49
],
50
}
51
tests/plugins_tests/test_io_processor_plugins.py
+59
-22
@@ -1,18 +1,39 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import base64
4
+import io
5
6
+import imagehash
7
import pytest
8
import requests
9
+from PIL import Image
10
11
from tests.utils import RemoteOpenAIServer
12
from vllm.config import VllmConfig
13
from vllm.entrypoints.pooling.pooling.protocol import IOProcessorResponse
14
from vllm.plugins.io_processors import get_io_processor
15
13
-MODEL_NAME = "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11"
16
+models_config = {
17
+ "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11": {
18
+ "image_url": "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/valencia_example_2024-10-26.tiff", # noqa: E501
19
+ "out_hash": "aa6d92ad25926a5e",
20
+ "plugin": "prithvi_to_tiff",
21
+ },
22
+ "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-BurnScars": {
23
+ "image_url": "https://huggingface.co/ibm-nasa-geospatial/Prithvi-EO-2.0-300M-BurnScars/resolve/main/examples/subsetted_512x512_HLS.S30.T10SEH.2018190.v1.4_merged.tif", # noqa: E501
24
+ "out_hash": "c07f4f602da73552",
25
+ "plugin": "prithvi_to_tiff",
26
+ },
27
+}
28
15
-image_url = "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/valencia_example_2024-10-26.tiff" # noqa: E501
29
+
30
+def _compute_image_hash(base64_data: str) -> str:
31
+ # Decode the base64 output and create image from byte stream
32
+ decoded_image = base64.b64decode(base64_data)
33
+ image = Image.open(io.BytesIO(decoded_image))
34
+
35
+ # Compute perceptual hash of the output image
36
+ return str(imagehash.phash(image))
37
38
39
def test_loading_missing_plugin():
@@ -22,33 +43,39 @@ def test_loading_missing_plugin():
43
44
45
@pytest.fixture(scope="function")
25
-def server():
46
+def server(model_name, plugin):
47
args = [
48
"--runner",
49
"pooling",
50
"--enforce-eager",
30
- "--trust-remote-code",
51
"--skip-tokenizer-init",
52
# Limit the maximum number of parallel requests
53
# to avoid the model going OOM in CI.
54
"--max-num-seqs",
55
"32",
56
"--io-processor-plugin",
37
- "prithvi_to_tiff",
38
- "--model-impl",
39
- "terratorch",
57
+ plugin,
58
"--enable-mm-embeds",
59
]
60
43
- with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
61
+ with RemoteOpenAIServer(model_name, args) as remote_server:
62
yield remote_server
63
64
65
@pytest.mark.asyncio
48
-@pytest.mark.parametrize("model_name", [MODEL_NAME])
66
+@pytest.mark.parametrize(
67
+ "model_name, image_url, plugin, expected_hash",
68
+ [
69
+ (model_name, config["image_url"], config["plugin"], config["out_hash"])
70
+ for model_name, config in models_config.items()
71
+ ],
72
+)
73
async def test_prithvi_mae_plugin_online(
74
server: RemoteOpenAIServer,
75
model_name: str,
76
+ image_url: str | dict,
77
+ plugin: str,
78
+ expected_hash: str,
79
):
80
request_payload_url = {
81
"data": {
@@ -74,16 +101,25 @@ async def test_prithvi_mae_plugin_online(
101
102
# verify the output is formatted as expected for this plugin
103
plugin_data = parsed_response.data
77
-
104
assert all(plugin_data.get(attr) for attr in ["type", "format", "data"])
105
80
- # We just check that the output is a valid base64 string.
81
- # Raises an exception and fails the test if the string is corrupted.
82
- base64.b64decode(plugin_data["data"])
106
+ # Compute the output image hash and compare it against the expected hash
107
+ image_hash = _compute_image_hash(plugin_data["data"])
108
+ assert image_hash == expected_hash, (
109
+ f"Image hash mismatch: expected {expected_hash}, got {image_hash}"
110
+ )
111
112
85
-@pytest.mark.parametrize("model_name", [MODEL_NAME])
86
-def test_prithvi_mae_plugin_offline(vllm_runner, model_name: str):
113
+@pytest.mark.parametrize(
114
+ "model_name, image_url, plugin, expected_hash",
115
+ [
116
+ (model_name, config["image_url"], config["plugin"], config["out_hash"])
117
+ for model_name, config in models_config.items()
118
+ ],
119
+)
120
+def test_prithvi_mae_plugin_offline(
121
+ vllm_runner, model_name: str, image_url: str | dict, plugin: str, expected_hash: str
122
+):
123
img_prompt = dict(
124
data=image_url,
125
data_format="url",
@@ -96,13 +132,12 @@ def test_prithvi_mae_plugin_offline(vllm_runner, model_name: str):
132
runner="pooling",
133
skip_tokenizer_init=True,
134
enable_mm_embeds=True,
99
- trust_remote_code=True,
135
enforce_eager=True,
136
# Limit the maximum number of parallel requests
137
# to avoid the model going OOM in CI.
103
- max_num_seqs=1,
104
- model_impl="terratorch",
105
- io_processor_plugin="prithvi_to_tiff",
138
+ max_num_seqs=32,
139
+ io_processor_plugin=plugin,
140
+ default_torch_num_threads=1,
141
) as llm_runner:
142
pooler_output = llm_runner.get_llm().encode(img_prompt, pooling_task="plugin")
143
output = pooler_output[0].outputs
@@ -110,6 +145,8 @@ def test_prithvi_mae_plugin_offline(vllm_runner, model_name: str):
145
# verify the output is formatted as expected for this plugin
146
assert all(hasattr(output, attr) for attr in ["type", "format", "data"])
147
113
- # We just check that the output is a valid base64 string.
114
- # Raises an exception and fails the test if the string is corrupted.
115
- base64.b64decode(output.data)
148
+ # Compute the output image hash and compare it against the expected hash
149
+ image_hash = _compute_image_hash(output.data)
150
+ assert image_hash == expected_hash, (
151
+ f"Image hash mismatch: expected {expected_hash}, got {image_hash}"
152
+ )