OlmoEarth Run: как Ai2 считает спутниковые карты на целый континент

Ai2 выкатила OlmoEarth Platform — пайплайн от дообучения геомоделей до inference на масштабе континента. Семейство OlmoEarth претренировано на примерно 10 ТБ мультимодальных спутниковых снимков; платформа укладывает такой прогон примерно в сутки, квадратный километр обрабатывается за доли цента.

Три этапа под разное железо

Спутниковый inference не укладывается в шаблон «несколько мегабайт — ответ за секунду». Входы — терабайты снимков с разными проекциями, спектральными каналами и пропусками из-за облаков. Часто дольше всего занимает не модель, а скачивание и подготовка картинок — GPU на это не отдают.

Каждый job делят на три стадии: CPU с высоким I/O ищет снимки, выравнивает проекции и нормализует данные; GPU гоняет только прямой проход модели; CPU на финале сшивает окна, применяет маски и экспортирует в GeoTIFF, Zarr или GeoJSON.

Карта режется на тысячи окон

OlmoEarth Run делит регион на партиции по размеру воркера, затем на окна для модели. Область размера штата — около сотни партиций, континент — тысячи. Соседние клетки слегка перекрываются; overlap сглаживают при сборке, чтобы на финальном растре не осталось швов.

Недавний прогон карты пожарной опасности по всей Северной Америке на пике занял около 19 600 CPU и 994 GPU, сеть держала более 168 GB/s. Серийные 4737 часов вычислений сжались до 30,5 часов по часам — ускорение 155×. Параллелизм упирается в облачные квоты: разрешение выхода, размер модели и кэш сырых снимков — ручки бюджета и скорости на конкретном job.

Свой индекс вместо лавины запросов

Для каждого региона и периода платформа выбирает сцены для модели — у Sentinel-2 важна облачность, у радара поляризация. Публичные STAC-каталоги используют, но один большой job генерирует тысячи запросов, и внешние API ESA или Microsoft Planetary Computer под такой burst не заточены.

Решение — собственный индекс метаданных: для AWS Open Data приходят SNS-уведомления о новых сценах, остальные провайдеры поллят раз в несколько минут. В рантайме читают только нужные байты из cloud-optimized форматов COG или Zarr, а не целые сцены.

Упавший воркер — не катастрофа

Каждая задача крутится в Docker-контейнере runner'а: параметры, работа, результат, shutdown. Задачи перезапускаемые и безопасные при повторе — упавший воркер просто запускают снова. Провайдер может тормозить, облака закрыть снимок, задача падает — платформа отслеживает статус, ретраит и переключается на альтернативный источник, если есть.

Что в roadmap

В планах — автозапуск job'ов по расписанию или при появлении новой сцены, алерты на изменения ландшафта и агентные интерфейсы: от курации данных до подсказок, как улучшить дообученную модель. Также хотят глобальные embeddings — часть задач считать без полного прохода по сырым снимкам, а OlmoEarth Run перенести на любой облак с Docker и blob storage.

Источник: The OlmoEarth Platform: Geospatial inference at planetary scale.