AI导哈

GPU加速MCP数据服务器

一个GPU加速的模型上下文协议服务器，用于查询存储在S3上的地理空间数据集，作为mcp-data-server的替代品，通过Polars + RAPIDS cuDF实现GPU加速。

  • by boettiger-lab
  • Python
  • 本地部署
  • 1 星标
  • 更新于
在 GitHub 查看 boettiger-lab/mcp-gpu-data-server

使用说明

GPU加速MCP数据服务器

GPU加速 模型上下文协议 用于查询S3上存储为蜂巢分区拼花游戏的地理空间数据集的服务器。直接替换 mcp数据服务器 通过Polars+RAPIDS cuDF实现GPU加速。

为什么这么复杂？（请先阅读此内容）

此服务器中的设计决策并不明显。每个都增加了真正的复杂性，并且每个都有特定的性能发现。本节解释了推理过程，以便未来的维护人员了解权衡。

问题：GPU计算速度很快，但数据路径很慢

在NRP Nautilus（RTX 4000 Ada、100G InfiniBand、Ceph S3内部端点）上，初始基准测试显示 CPU（DuckDB）比GPU（Polars/cuDF）快2-4倍 用于S3支持的H3加入查询。瓶颈不是计算，而是数据路径：

gpu mode:      S3 → Polars Rust object_store → CPU RAM → GPU SQL (GPUEngine)
gpu-cudf mode: S3 → kvikio pread (6 Gbps) → CPU RAM → Polars parse → GPU SQL (GPUEngine)
cpu mode:      S3 → Polars Rust object_store → CPU RAM → CPU SQL

大型数据集的瓶颈是S3下载速度，而不是计算速度。kvikio的并行分块HTTP实现了6.25 Gbps，而Polars Rust object_store实现了0.97 Gbps，对于大于5 MB的文件，这是6.5倍的改进。

为什么是kvikio（以及为什么是pread，而不是read）

kvikio 是NVIDIA的高性能I/O库。对于HTTP远程文件，它使用并发分块范围请求来饱和高带宽网络。

NRP 100G InfiniBand上的实测吞吐量 （碳美洲，28个文件，3.22 GiB）：

运输 时间 吞吐量
kvikio pread() （64个线程，16 MiB块） 4.1秒 6.25 Gbps
Polars生锈 object_store 26.6s 0.97 Gbps

速度提高6.5倍。到达那里需要两个不明显的细节：

  1. pread() 不 read() — RemoteFile.read() 是单线程HTTP GET。 RemoteFile.pread() 为并行分块范围请求激活kvikio的内部线程池。API的名称没有给出任何提示。

  2. KVIKIO_NTHREADS env变量，不是 set_num_threads() — kvikio.defaults.set_num_threads(64) 默默地接受调用，但kvikio 25.02中的值不变。在库初始化之前，必须通过环境变量设置线程数。

为什么仍然使用s3fs（仅用于全局分辨率）

HTTP没有列出API的目录。查找存在哪些分区文件 s3://public-carbon/.../hex/**，我们需要S3的ListObjects API，这需要一个S3客户端。 s3fs.glob() 被使用 仅 对于这个文件发现步骤（每个查询每个数据集一个API调用）。

实际的数据传输使用kvikio纯HTTP，完全绕过S3 SDK：

s3://public-carbon/.../hex/h0=576.../data_0.parquet   ← s3fs discovers this path
http://rook-ceph-rgw.../public-carbon/.../data_0.parquet  ← kvikio downloads at 6 Gbps

这就是为什么 cudf.read_parquet(storage_options=...) 是 未使用 尽管RAPIDS文档暗示它应该在内部使用kvikio。在实践中，它通过PyArrow的S3文件系统路由，而不是kvikio。看 问题#3 为了进行全面调查。

为什么kvikio有利于大文件而不是小文件

kvikio的并行分块下载在许多并发范围请求中分摊了每个连接的开销。当单个文件足够大以容纳多个块时，它会有所帮助：

数据集 文件 平均文件大小 kvikio效益
IUCN十六进制 548 0.1 MB 无-开销占主导地位
WDPA十六进制 116 2 MB 中等
碳十六进制 94 78 MB（最大768 MB） 6.5倍更快
GBIF十六进制 419 307 MB（最大522 MB） 最佳情况

基准查询Q3a-Q5a（碳×IUCN/WDPA）和Q6a（GBIF×IUCN，美洲子集）是有意义的GPU测试。Q1/Q2（仅IUCN）太小，无法显示任何S3运输差异。

为什么在gpu-cudf模式下需要显式分区修剪（DPP）

所有数据集均按以下方式分区 h0 （H3分辨率-0细胞）。完整的碳数据集是94个文件×7.3 GiB，对于RTX 4000 Ada的20 GB VRAM来说太大了。当查询包括 WHERE h0 IN (...)，只需要读取匹配的分区文件。

波拉斯的懒惰 gpu 模式免费获得此功能：查询优化器将过滤器向下推送到 scan_parquet 蜂箱隔板修剪。这 gpu-cudf 模式急切地读取文件（kvikio然后是cudf），因此必须执行DPP 阅读前明确：SQL被解析为 h0 IN (...) 谓词，只有匹配的文件才会传递给kvikio。

没有DPP，Q3（全球碳，没有过滤器）会损坏吊舱。有了它，Q3a（美洲，94个分区中的28个）正常完成。

为什么不使用RDMA

GPU直接存储 可以完全绕过CPU RAM（NIC→ GPU VRAM直接），消除了PCIe瓶颈。然而，NRP Nautilus Ceph S3使用纯HTTP，没有公开支持RDMA的端点。用于远程文件的GPUDirect还需要 compat_mode=False kvikio和特殊内核驱动程序(nvidia-fs.ko).在这个集群上，kvikio运行 compat_mode=2 （兼容模式）。在将PCIe传输到GPU之前，数据始终会存储在CPU RAM中。

发动机模式

通过设置 QUERY_ENGINE 环境变量：

模式 S3传输 Parquet解析 GPU计算 何时使用
gpu （默认） Polars-Rust-object_store CPU（懒惰） 是（GPUEngine） 通用；可靠
gpu-cudf kvikio pread（6 Gbps） CPU（Polars） 是（GPUEngine） 大文件（carbon、GBIF）
cpu Polars Rust object_store CPU（懒惰） 否 没有可用的GPU

gpu-cudf 部署在NRP上(k8s/deployment.yaml)与 KVIKIO_NTHREADS=64 和 KVIKIO_TASK_SIZE=16777216.

MCP工具

工具 说明
list_datasets() 列出可用的STAC集合
get_dataset(id) 获取元数据、S3路径、列模式
query(sql) 执行SQL，返回markdown表

SQL方言

LLM使用内联编写DuckDB风格的SQL read_parquet('s3://...')该引擎提取拼花源，将其注册为Polars LazyFrames（或通过cudf热切加载），重写SQL以使用表别名，然后执行。

不支持 （Polars SQL中没有等效项）：

  • h3_cell_to_parent(), h3_h3_to_string() --使用预先计算 h0–h11 列直接
  • CAST(x AS TYPE) 在JOIN ON子句中--改为在CTE中预铸

自动重写：

  • APPROX_COUNT_DISTINCT(x) → COUNT(DISTINCT x)
  • COPY (...) TO 's3://...' → 通过s3fs写入

与CPU（DuckDB）版本的差异

看 基准测试结果 和 问题#5 以进行详细比较。摘要：小数据集（<1 GiB）的CPU获胜；GPU在kvikio管道的大型数据集（碳，GBIF）方面具有竞争力。

配置

变量 默认值 描述
QUERY_ENGINE gpu gpu, gpu-cudf，或 cpu
KVIKIO_NTHREADS 1 （于25.02年破裂） 吃起来 64 通过env var。
KVIKIO_TASK_SIZE 4194304 吃起来 16777216 （16 MiB）通过环境变量
S3_ENDPOINT_URL http://rook-ceph-rgw-nautiluss3.rook Ceph内部端点
AWS_ACCESS_KEY_ID （空→ 匿名） S3凭据
AWS_SECRET_ACCESS_KEY （空→ 匿名） S3凭据
STAC_CATALOG_URL NRP公共目录 数据集目录

本地运行（CPU模式）

pip install -r requirements.txt
python server.py

测试仅运行CPU（不需要GPU）：

pytest tests/ -v

Kubernetes部署

kubectl apply -f k8s/

要暂停并释放GPU（NRP策略：不要在GPU节点上空闲）：

kubectl -n biodiversity scale deployment/gpu-mcp --replicas=0
# Resume:
kubectl -n biodiversity scale deployment/gpu-mcp --replicas=1

基准测试

uv run --with mcp benchmarks/benchmark.py --queries Q1,Q2,Q3a,Q4a,Q5a,Q6a --runs 3

看 benchmarks/ 用于查询定义和结果CSV。

相关 MCP