OlmoEarth平台:行星尺度的地理空间推理基础设施解析

0 阅读

从模型到平台:地理空间推理的规模化挑战

近年来,地球观测基础模型(如OlmoEarth)在森林砍伐监测、粮食安全评估和野火风险预测等领域展现出巨大潜力。然而,将这些模型从实验室原型转化为能够处理大陆级区域的实用工具,面临着远超传统机器学习推理的工程复杂性。一个典型的推理任务可能涉及数十TB的多模态卫星数据,需要跨多个数据提供商、多种投影和分辨率进行对齐,最终输出必须与地理坐标精确配准的地图。

OlmoEarth平台正是为应对这些挑战而设计的。它不仅仅是一个模型推理服务,而是一套完整的基础设施,涵盖了数据获取、预处理、模型推理、后处理以及结果可视化等全生命周期。其核心设计理念是:将计算任务与最合适的硬件匹配,通过大规模并行化压缩时间,并构建能够自动从分布式计算常见故障中恢复的健壮系统。

为什么卫星推理如此困难

与处理几兆字节文本或图像的典型ML模型不同,地球观测推理的输入数据量巨大,且高度异构。一个推理作业可能需要处理覆盖大片区域、跨越多个光谱波段、不同传感器类型和多个时间步的影像。这些数据来自不同提供商,采用不同的投影和分辨率,还可能包含被云层遮挡或缺失的观测。更关键的是,输出本身是一张地图,每个预测都必须与周围区域保持精确的投影和坐标网格对齐。

数据获取本身就是一个重大瓶颈。实际运行中,推理作业花费在下载和准备影像上的时间往往超过模型计算本身。因此,高效的数据管道至关重要,它必须处理高吞吐量的I/O,同时提供重投影和重采样所需的计算能力。

为任务匹配合适的硬件

由于数据获取和预处理通常主导推理作业的耗时,将这些工作分配给GPU会导致最昂贵的硬件执行更适合CPU的任务。因此,OlmoEarth平台将每个作业分为三个阶段,每个阶段匹配不同的硬件配置:

  • 数据获取与预处理(CPU,高I/O):负责获取、重投影、对齐和归一化影像,并将其写入为推理期间快速加载而优化的格式。
  • 推理(GPU):运行模型的前向传播,并将最小化处理的输出直接写入存储。
  • 后处理(CPU):将每个窗口的输出拼接在一起,应用掩码或重新缩放,并导出为用户友好的格式,如Zarr、GeoTIFF或GeoJSON。

Three-stage inference pipeline as cards: 01 Data acquisition on CPU (heavily I/O-bound), 02 Run inference on GPU (heavily GPU-bound), 03 Postprocessing on CPU

平台将这些阶段分布在多台机器上,同时保持GPU的充分利用。多进程数据加载器持续为每个GPU提供数据,而完成的输出则直接流式传输到对象存储。

一个请求,数百个工作节点,数千个进程

OlmoEarth Run是平台的大规模推理执行层。它将作业覆盖的地理区域划分为适合单个计算实例(工作节点)大小的分区,然后将这些分区进一步细分为OlmoEarth模型处理的小窗口。由于每个窗口可以在独立的前向传播中处理,地图某部分的工作无需等待其他部分完成。

OlmoEarth two-stage partitioning: Stage 1 splits one prediction request into machine-sized partitions fanned out across up to ~1,000 worker nodes; Stage 2 re-partitions each worker’s cell into model-sized windows processed in parallel, one GPU per node

在实践中,一个州大小的区域可能变成大约一百个分区,而大陆级运行则可能变成数千个。相邻分区略有重叠,在输出组装时会协调这些重叠,以确保最终栅格中不出现接缝。

由于分区相互独立,同一阶段可以同时在数千个计算实例上运行。最近,平台使用这种方法生成了覆盖整个北美的野火风险地图。在高峰期,该运行并行使用了约19,600个CPU和994个GPU,网络吞吐量超过168 GB/s。这种并行度将估计4,737小时的串行计算缩短到约30.5小时的墙钟时间,实现了155倍的加速。

Continent-scale inference run, by the numbers — a stat grid summarizing the North America wildfire-risk run

然而,扇出并非无限制。更多的工作节点会触及云配额,因此并行度是一个可调参数,可以根据单个作业进行调整。输出分辨率在数据量和计算量与细节之间权衡;模型大小在GPU时间与准确性之间权衡;缓存原始影像在存储与跨重复运行的速度之间权衡。正确的设置取决于手头的任务和预算。

查找和获取正确的像素

给定地理区域和时间范围,平台首先必须确定哪些卫星场景应输入模型。这意味着要跨不同提供商识别捕获了什么、在哪里、何时,这些提供商的目录、格式和发布延迟各不相同。选择标准还取决于数据源:对于光学影像(如Sentinel-2),通常需要云量最少的场景;而对于合成孔径雷达,可用的极化通道可能更重要。

Example items-search API call: a POST to /api/v1/items/search querying the Sentinel-2 L2A collection over a date range and polygon, sorted by cloud cover, with the JSON response record and cloud-optimized asset pointers

平台尽可能依赖公共STAC目录和开放标准。但大型推理作业可能同时生成数千个元数据查询,这远远超出了外部服务(如ESA或Microsoft Planetary Computer的STAC API)设计的并发处理能力。

为避免压垮这些服务,OlmoEarth平台维护自己的元数据索引,并随着新影像的发布而更新。对于通过AWS开放数据托管的数据集,平台会为每个新场景接收SNS通知。当提供商不提供变更流时,平台会每隔几分钟轮询其上游索引。因此,平台对外部服务的请求遵循新发布的稳定节奏,而不是大型推理作业产生的突发流量。

OlmoEarth Datasets architecture: one metadata index over many satellite providers (AWS Open Data, Google Cloud, USGS, Copernicus, NASA/ASF, Microsoft Planetary Computer) feeding an ingestion queue, worker pods, and an Elasticsearch index that serves search and tile APIs

每个索引条目存储场景元数据以及指向底层像素可用位置的指针。在运行时,平台选择最佳来源,并针对云优化格式(如COG或Zarr)执行窗口化读取,仅检索给定分区所需的字节,而不是下载整个场景。

该索引还支持平台的标注工具。由于它维护指向Sentinel-1、Sentinel-2、Landsat和NISAR影像的指针,这些影像采用云优化格式,因此可以通过相同的窗口化读取系统从任何索引场景提供瓦片,而无需构建单独的摄取管道。

使这一工作流程最顺畅的提供商具有三个共同特征,平台建议将其作为发布地球观测数据的最佳实践:新影像可用时提供基于队列的通知;存储在主要云平台上,没有特殊的速率限制或可用性瓶颈;以及支持范围读取的云优化格式。

Close-up of a wildfire-risk heatmap over terrain, from blue (lower risk) to red (higher risk)

大规模故障处理

OlmoEarth平台设计为自动从故障中恢复。对于每个阶段和地理分区内的每个任务,它动态配置一台运行runner Docker容器的虚拟机。runner检索任务参数,执行工作,返回结果,然后关闭。由于每个任务都是可重入且幂等的,间歇性故障可以通过重新运行来安全处理。

在此规模下,故障是预期内的:提供商可能缓慢或暂时不可用;元数据可能表明影像存在,但缺少所需的波段或窗口;云层覆盖可能导致可用观测不足;或者任务可能直接崩溃。平台通过任务跟踪、自动重试、在可用时回退到备用提供商,以及明确区分可重试和致命错误来应对。一个单独的监控进程检测停滞或停止的runner,并重新启动其任务。

未来方向

平台的路线图由合作伙伴确定的差距和他们最需要的能力所塑造。正在研究的领域包括:

  • 自动化模型运行:提前安排推理作业,或在影像索引注册到感兴趣区域的新场景时触发。
  • 变化检测和警报:当监测的景观发生变化时通知用户,使森林砍伐或洪水等事件以警报形式出现,而不是需要手动查找和检查的栅格。
  • 智能体工具和界面:智能体可以降低使用地理空间模型的门槛,从数据整理和特征工程到识别改进微调模型的方法。平台希望让任何技术水平的用户都能完成以前需要经验丰富的ML研究员才能完成的工作。
  • 更快的模型:与研究团队合作,开发更高效的架构,减少每个窗口的GPU时间。
  • 更多模态:将新的卫星传感器和数据源添加到OlmoEarth模型和支持它们的影像索引中。目前正专注于整合天气数据(ERA-5)和提供更多环境因素细节的卫星。
  • 嵌入:开发专门的嵌入模型,并在全球范围内预计算嵌入。对于许多任务,对这些嵌入运行推理可以取代对原始影像的完整前向传播,使工作负载大大加快且成本降低。微调和直接推理在挑战性任务上仍将保持最大性能,但嵌入可能开启更广泛的效率应用。
  • 随处运行:OlmoEarth Run仅需要能够运行Docker镜像的虚拟机和对对象存储的访问。目前平台在Google Cloud上运行,但架构设计为支持多云以及合作伙伴自己的账户和计算环境中的部署。

Wildfire-risk map of North America generated on the OlmoEarth Platform, shown as a blue-to-red heatmap over a satellite basemap

结语

地理空间基础模型,尤其是将其投入运营,仍是一项新兴技术。许多最有条件使用这些模型的组织——那些从事保护、粮食安全、灾害响应和气候工作的组织——从未接触过这样的基础设施。这些团队对了解地球的需求与其预算和技术资源所能提供的支持之间的差距仍然很大。OlmoEarth平台正是为了缩小这一差距而构建的。