Add support for TerraTorch Geospatial models served via the vLLM /pooling endpoint

[mgazz]

Summary

This PR adds support for TerraTorch Geospatial models like Prithvi. Compared to traditional LLMs, Geospatial Models support image-to-image or tensor-to-tensor input/output, but they do not process text. Moreover they are served in vLLM at the /pooling endpoint.

Details

• Support for the /pooling endpoint
• A custom PoolingColumnMapper to map the prompt entry in datasets without creating conflict with the existing GenerativeColumnMapper
• A new request handler for the pooling endpoint: PoolingRequestHandler

Test Plan

Related Issues

• Resolves Add support for Geospatial models (served via vLLM) #587

---

• "I certify that all code in this PR is my own, except as noted below."

Use of AI

• Includes AI-assisted code completion
• Includes code generated by an AI application
• Includes AI-generated tests (NOTE: AI written tests should have a docstring that includes ## WRITTEN BY AI ##)

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @mgazz.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mgazz]

Please wait before merging. After the rebase we can benchmark the Prithvi model but Terramind returns a server error. I am currently investigating the reason behind it.

[mgazz]

@sjmonson we verified the problem was not a guidellm issue. The TerraMind deployment needed a bigger shared memory size.

[dbutenhof]

This looks fairly straightforward (even if the entire concept of geospatial inference is more than slightly mysterious to me).

I made a minor comment on your documentation, which might help readers to follow the sequence... but I think that's a minor issue and I'm approving anyway. (Not that it really matters as you need a rebase now.)

FYI, the unit test failure is a known issue, and will be resolved when you rebase.

[sjmonson]

@Mergifyio rebase

[mergify]

rebase

✅ Branch has been successfully rebased

[jaredoconnell]

I don't see anything blocking, but I do have a comment on a design decision. I'm not the most familiar with geospacial models, but I was able to confirm that it ran.

[sjmonson]

Couple small nits,

[dbutenhof]

The documentation now gives --request-format values that won't work since you removed the "legacy" alias.

[mgazz]

Sorry for the delay and thank you for the feedback. I will update the documentation to make it coherent. I will also provide a command to build the flood_detection_dataset.jsonl. This way we remove potential errors in designing a sample dataset.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @mgazz.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

@mgazz, this project requires a linear history on feature branches.
Your PR contains merge commits. Please rebase your branch against main
and remove them.

You can do this by running:
git pull --rebase upstream main

[sjmonson]

@Mergifyio rebase

[mergify]

rebase

✅ Branch has been successfully rebased