Переписывание UDF Spark для GPU с AI: ускорение в 22x

В эпоху, когда корпоративные приложения становятся всё более требовательными к вычислительным ресурсам, а агенты искусственного интеллекта (Agentic AI) генерируют огромные объемы данных для сложных запросов, традиционные подходы к обработке данных часто не справляются с нагрузкой. Ключевая задача для современных инженеров данных — обеспечить высокую пропускную способность без пропорционального роста затрат. Здесь на сцену выходит NVIDIA cuDF, технология, использующая параллелизм и высокую пропускную способность памяти графических процессоров (GPU) для ускорения аналитики данных. Особенно важно то, что плагин NVIDIA cuDF для Apache Spark позволяет интегрировать это ускорение в уже существующие конвейеры обработки больших данных, значительно улучшая соотношение стоимости и производительности (TCO).

Однако, несмотря на мощь GPU, многие конвейеры данных сталкиваются с серьезным узким местом: пользовательские функции (User-Defined Functions, UDF). Эти функции, позволяющие разработчикам добавлять собственную бизнес-логику, нарушают декларативную абстракцию Spark. Они заставляют систему переключаться с быстрого выполнения на GPU на медленное выполнение на CPU, создавая дополнительные накладные расходы на передачу данных. В этой статье мы подробно разберем, как с помощью AI-агентов и специализированных инструментов NVIDIA можно переписать эти UDF, устраняя узкие места и достигая ускорения более чем в 20 раз.

01Проблема пользовательских функций (UDF) в Spark

Чтобы понять суть решения, нужно сначала разобраться в природе проблемы. Apache Spark славится своей декларативной моделью программирования через SQL и DataFrames. Эта абстракция мощна тем, что разделяет описание того, что должен вернуть запрос, от того, как система должна его выполнить. Плагин cuDF использует эту особенность: он прозрачно распараллеливает операции по строкам, используя возможности SIMD (Single Instruction, Multiple Data) графических процессоров, не требуя изменений в коде пользователя.

Но UDF ломают эту идеальную картину. UDF — это произвольный пользовательский код, который система не может проанализировать или распараллелить общим способом. Когда в конвейере, который в остальном выполняется на GPU, встречается UDF, происходит следующее:

Данные, хранящиеся в столбцовом формате (columnar) на GPU, должны быть преобразованы в строковый формат (row-wise).
Эти данные передаются обратно на CPU.
UDF выполняется на процессоре построчно.
Результат снова преобразуется в столбцовый формат и возвращается на GPU.

Этот процесс, известный как "fallback" (возврат к CPU), создает два критических недостатка:

Накладные расходы на передачу данных: Постоянные циклы передачи данных между GPU и CPU (round-trip) съедают драгоценное время.
Потеря вычислительной мощности: UDF не использует преимущества параллельных вычислений GPU, оставаясь узким местом, которое ограничивает общую производительность всего запроса.

Сравнительная таблица подходов к переписыванию UDF: SQL, cuDF и CUDA, с указанием сложности и лучших сценариев применения.

02Три подхода к переписыванию UDF для GPU

К счастью, многие UDF можно переписать, используя правильные примитивы, чтобы они стали совместимыми с GPU. Это устраняет необходимость в fallback и открывает путь к ускорению более чем в 20 раз. NVIDIA выделяет три основных подхода, различающихся по сложности и области применения.

1. SQL (Низкая сложность)

Самый простой способ — заменить UDF эквивалентным выражением Spark SQL. Это работает для простой реляционной логики, например, для небольших или встроенных функций. Когда UDF заменяется на SQL-выражение, оно становится видимым в дереве выражений Spark. Во время планирования запроса плагин cuDF заменяет эти выражения их GPU-реализациями, удерживая данные и выполнение полностью на GPU.

Рассмотрим пример простой Scala UDF, которая преобразует температуру из Цельсия в Фаренгейт:

terminalscala

final class ToFahrenheit
    extends Function1[java.lang.Double, java.lang.Double]
    with Serializable {
  override def apply(celsius: java.lang.Double): java.lang.Double = {
    if (celsius == null) null
    else celsius * 9.0 / 5.0 + 32.0
  }
}
spark.udf.register("to_fahrenheit", new ToFahrenheit)
spark.sql("""
  SELECT
    *,
    to_fahrenheit(celsius) AS fahrenheit
  FROM temperatures
""")

Переписав это через SQL, мы получаем:

terminalsql

spark.sql("""
  SELECT
    *,
    celsius * 9.0 / 5.0 + 32.0 AS fahrenheit
  FROM temperatures
""")

Теперь функция видна оптимизатору, и данные не покидают GPU.

2. cuDF / RapidsUDF (Средняя сложность)

Когда SQL-примитивов недостаточно или логика слишком сложна (например, парсинг строк или работа с вложенными типами), можно использовать интерфейс RapidsUDF. Этот интерфейс предоставляет контракт, требующий реализации метода evaluateColumnar. Вместо работы со скалярными значениями, этот метод принимает ColumnVector — объект, представляющий столбец данных, уже находящийся в памяти GPU.

Вот как выглядит переписанная версия нашей функции температуры с использованием Java API для cuDF:

terminalscala

import ai.rapids.cudf.{ColumnVector, Scalar}
import com.nvidia.spark.RapidsUDF
import scala.util.Using

final class ToFahrenheit
    extends Function1[java.lang.Double, java.lang.Double]
    with Serializable
    with RapidsUDF {
  override def apply(celsius: java.lang.Double): java.lang.Double = {
    if (celsius == null) null
    else celsius * 9.0 / 5.0 + 32.0
  }
  override def evaluateColumnar(
      numRows: Int,
      args: ColumnVector*): ColumnVector = {
    Using.Manager { use =>
      val scale = use(Scalar.fromDouble(9.0 / 5.0))
      val offset = use(Scalar.fromDouble(32.0))
      val scaled = use(args.head.mul(scale))
      scaled.add(offset)
    }.get
  }
}

Плагин cuDF распознает этот интерфейс и передает данные напрямую на GPU. Метод evaluateColumnar использует операции из Java API для cuDF (такие как mul и add) для параллельной обработки всего столбца. Это дает больше гибкости и контроля над выполнением.

3. CUDA (Высокая сложность)

Для самой сложной логики, которую невозможно эффективно собрать из существующих операций cuDF, можно использовать пользовательские ядра CUDA. Это достигается через вызов кастомного CUDA-ядра из метода RapidsUDF с помощью Java Native Interface (JNI). Этот подход требует максимальной экспертизы, но предоставляет полный контроль на уровне ядра, что критично для высокопроизводительных вычислений.

Пример кода Scala для простой UDF преобразования температуры Цельсия в Фаренгейт.

💡

Совет по выбору подхода. Начинайте с SQL. Если логика слишком сложна для SQL, переходите к RapidsUDF. Используйте кастомные CUDA-ядра только тогда, когда первые два варианта не обеспечивают нужной производительности или выразительности. Это позволит минимизировать сложность поддержки кода.

03Роль AI-агентов в оптимизации

Несмотря на очевидные преимущества, переписывание UDF — задача трудоемкая. Она требует тщательного выбора API, семантической валидации (чтобы убедиться, что результат не изменился) и измерения производительности. Для инженеров, не имеющих глубокого опыта в векторном программировании, работа с столбцовыми API может быть запутанной.

Здесь на помощь приходят AI-агенты. При наличии правильных ограничений (guardrails), агент может:

Исследовать незнакомые API cuDF.
Реализовать кандидаты на переписывание.
Быстро итерировать, используя исполняемые тесты и бенчмарки для проверки корректности и производительности каждого изменения.

NVIDIA представила cuDF skills for Apache Spark — набор навыков (skills), которые декомпозируют процесс переписывания на четыре сфокусированных цикла обратной связи. Каждый навык определяет не только работу, но и доказательства, необходимые для перехода к следующему этапу.

Четыре навыка (skills), которые направляют AI-агентов через тестирование, конвертацию, бенчмаркинг и оптимизацию с циклами обратной связи.

1. Каркас и проектирование (Scaffolding and design)

Надежное выполнение агента начинается с надежной среды разработки. Навыки предоставляют полный шаблон проекта, содержащий специфичный для cuDF и Spark каркас. В него входят:

Тестовые и бенчмарк-хаусы.
Кодированные процедуры для проверки выходных данных.
Конфигурация сборки.

Инкапсулируя эту инфраструктуру в код, агент может сосредоточиться на самой логике UDF, а не на повторном создании специализированной инфраструктуры. Кроме того, навыки включают лучшие практики программирования на GPU с cuDF, такие как управление памятью, служа справочником для передовых моделей.

2. Тестирование и валидация

Агент использует предоставленные тестовые наборы для сравнения результатов оригинальной CPU UDF и новой GPU-реализации. Это гарантирует, что переписанный код ведет себя идентично исходному, сохраняя корректность бизнес-логики.

3. Бенчмаркинг

После подтверждения корректности агент запускает бенчмарки для измерения производительности. Это позволяет количественно оценить ускорение и убедиться, что переписывание действительно решает проблему узкого места.

4. Оптимизация

На основе результатов бенчмарков агент может предложить дальнейшие оптимизации, например, выбор более эффективных операций cuDF или настройку параметров выполнения.

График, показывающий сокращение времени выполнения с 274,55 до 12,78 часов (ускорение в 22 раза) для рабочей нагрузки в 600 миллиардов строк.

⚠️

Важно: AI-агенты не заменяют инженера, а усиливают его. Всегда проверяйте результаты, особенно в критически важных производственных средах. Используйте агентов для ускорения рутинной работы и поиска оптимальных решений, но финальное утверждение должно оставаться за человеком.

04Реальный кейс: ускорение в 22 раза

Теория подтверждается практикой. Крупная потребительская технологическая компания использовала этот рабочий процесс для конвертации более 30 производственных UDF, используемых в их конвейерах данных. Переписывание сохранило поведение оригинальных CPU UDF, но устранило дорогостоящие fallback-операции.

Один из примеров демонстрирует впечатляющий результат: UDF, влиявшая более чем на 3000 рабочих задач (jobs), показала ускорение в 22 раза при реализации как RapidsUDF, совместимая с cuDF. В рамках рабочей нагрузки объемом 600 миллиардов строк время выполнения оператором сократилось с 274,55 часов до 12,78 часов.

Благодаря помощи агентов, такие приросты производительности стали гораздо более достижимыми. Инженеры могут итеративно улучшать переписывание за часы, а не дни или недели, не жертвуя корректностью. Это не только ускоряет обработку данных, но и значительно снижает затраты на инфраструктуру, так как GPU-кластеры обрабатывают больше данных за то же время.

05Что это значит на практике

Для инженеров данных и разработчиков, работающих с Apache Spark, внедрение GPU-ускорения через cuDF и AI-агентов означает переход от ручного, трудоемкого оптимизирования к автоматизированному, масштабируемому процессу. Вот ключевые шаги для начала работы:

Установка навыков: Установите cuDF skills for Apache Spark на GitHub. Это можно сделать напрямую или с помощью установщика навыков npx:

terminalbash

npx skills add NVIDIA/cudf-spark --skill '*'

Изучение материалов: Для более глубокого понимания рабочего процесса и практических уроков посмотрите сессию Data + AI Summit: Coding with AI Agents to Supercharge Your Spark UDFs with GPU Acceleration.
Начало работы: Следуйте руководству cuDF plugin for Apache Spark getting-started guide, чтобы включить GPU-ускорение в ваших конвейерах Spark. Это откроет путь к лучшей производительности и снижению TCO.

В заключение, переписывание UDF для GPU с помощью AI-агентов — это не просто техническая оптимизация, это стратегическое преимущество. Оно позволяет компаниям обрабатывать растущие объемы данных быстрее и дешевле, освобождая инженеров от рутины и позволяя им сосредоточиться на инновациях. В мире, где скорость принятия решений на основе данных становится критической, такие инструменты становятся незаменимыми.

Если вы хотите узнать больше о GPU-ускоренной обработке структурированных данных, посетите веб-страницу NVIDIA cuDF или ознакомьтесь с документацией по плагину cuDF для Apache Spark. Начните с малого, переписав одну простую UDF, и посмотрите, как это изменит производительность вашего конвейера.

Источник: NVIDIA Developer ↗

Как переписать UDF в Apache Spark для GPU: руководство с AI-агентами