Apache Spark: переписывание UDF под GPU с помощью ИИ-агентов

UDF в Apache Spark ломают декларативную модель: система не может разобрать произвольный пользовательский код и параллелить его, поэтому GPU-ускоритель вынужден переводить данные из колоночного формата в построчный и возвращаться к CPU. Именно из-за этого в остальном ускоренный конвейер теряет производительность. NVIDIA предлагает решать проблему переписыванием UDF под GPU — и делать это с помощью ИИ-агентов, которые подбирают API, проверяют корректность и замеряют ускорение.

Почему UDF тормозят GPU-конвейер

Плагин NVIDIA cuDF для Apache Spark раскрывает декларативный интерфейс SQL и DataFrame: он отделяет то, что должен вернуть запрос, от того, как система будет его исполнять. Благодаря этому плагин прозрачно параллелит построчные операции на массивном SIMD-параллелизме GPU без изменений в пользовательском коде.

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

Отсюда две потери в ускоренном запросе: накладные расходы на круговую передачу данных из-за fallback и упущенное ускорение вычислений, которое было бы доступно при использовании GPU-примитивов.

Три способа переписать UDF

Небольшое точечное изменение кода переводит многие UDF в форму, пригодную для GPU. Это устраняет накладные расходы на возврат к CPU и даёт ускорение более чем в 20 раз. NVIDIA описывает три подхода.

Подход	Сложность	Для чего подходит
SQL	Низкая, самый высокий уровень абстракции	Простая реляционная логика, например небольшие или встраиваемые функции
cuDF	Средняя, больше гибкости и контроля над исполнением	Большинство случаев: умеренно сложная логика вроде разбора строк, вложенные типы
CUDA	Высокая, максимум гибкости и контроль на уровне ядра	Самая сложная или критичная по производительности логика

Простейший пример — UDF на Scala, переводящая температуру из Цельсия в Фаренгейт:

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
""")

Самый прямой путь для простых случаев — заменить UDF эквивалентным выражением Spark SQL:

SELECT
  *,
  celsius * 9.0 / 5.0 + 32.0 AS fahrenheit
FROM temperatures

Функция становится видна в дереве выражений Spark. При планировании запроса плагин cuDF подменяет эти выражения их GPU-реализациями, и данные с исполнением остаются целиком на GPU.

RapidsUDF: работа прямо с GPU-колонками

Когда примитивов SQL не хватает, UDF может работать напрямую с GPU-колонками через интерфейс RapidsUDF — контракт, который предоставляет плагин cuDF. Контракт требует единственный метод evaluateColumnar:

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 распознаёт экземпляры RapidsUDF и передаёт в evaluateColumnar объект ColumnVector, лежащий в памяти GPU. Функция собирает операцию из вызовов Java API для cuDF и обрабатывает колонку параллельно. Этот путь даёт более богатый API для сложной обработки.

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

Четыре навыка и циклы проверки

Переписывание UDF требует значительных инженерных усилий: подбора API, семантической валидации и замера производительности, а колоночные API незнакомы тем, у кого меньше опыта векторного программирования. Именно эту работу хорошо выполняют кодинг-агенты — при правильных ограничителях они ищут незнакомые API, реализуют варианты переписывания и быстро итерируют, а исполняемые тесты и замеры производительности решают, приемлемо ли каждое изменение.

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

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

Результат: 30+ UDF и ускорение в 22 раза

Крупная компания потребительских технологий перевела по этому процессу более 30 UDF, используемых в производственных конвейерах данных, в GPU-ускоренные реализации. Переписывание сохранило поведение исходных CPU-функций и убрало дорогой fallback из остальных GPU-ускоренных нагрузок. Одна UDF затрагивала более 3000 задач, и в этой задаче ускорение при реализации через RapidsUDF с поддержкой cuDF достигло 22x. С агентами такие результаты становятся куда доступнее: итерации по переписыванию занимают часы, а не дни или недели, и без потери корректности. В одном из примеров переписывание сократило время работы оператора с 274,55 до 12,78 часа на нагрузке в 600 миллиардов строк.

Как начать

Набор cuDF skills for Apache Spark ставится с GitHub — напрямую или через установщик npx:

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

Для более подробного разбора процесса и практических уроков есть сессия Data + AI Summit «Coding with AI Agents to Supercharge Your Spark UDFs with GPU Acceleration». Чтобы включить GPU-ускорение в конвейерах Spark, следует идти по getting-started guide для плагина cuDF. Оценить подход к агентам в других задачах помогает разбор оценки ИИ-агентов, а для локального запуска агентов пригодится NVIDIA DGX Spark.

Частые вопросы

Кого касается эта новость?

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

Что именно теряется из-за UDF?

Две вещи: накладные расходы на круговую передачу данных при fallback и потенциальное ускорение вычислений, которое было бы при использовании GPU-примитивов.

Как выбрать способ переписывания?

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

Как работает RapidsUDF?

Плагин cuDF распознаёт экземпляры интерфейса RapidsUDF и передаёт в метод evaluateColumnar объект ColumnVector с данными в памяти GPU. Функция собирает обработку из операций Java API для cuDF и параллельно обрабатывает колонку.

Что делать, если логику нельзя собрать из операций cuDF?

RapidsUDF может вызывать собственное ядро CUDA через Java Native Interface (JNI). Примеры конвертаций приведены в cuDF examples for Spark UDFs.

Зачем здесь ИИ-агенты?

Переписывание требует подбора API, семантической валидации и замеров производительности, а колоночные API незнакомы части практиков. Агенты ищут незнакомые API, реализуют варианты и быстро итерируют, а тесты и бенчмарки определяют, приемлемо ли изменение.

Сколько циклов в наборе навыков?

Четыре: тестирование, конвертация, бенчмаркинг и оптимизация. У каждого — свои требуемые доказательства до перехода к следующему этапу.

Какое ускорение получилось на практике?

Крупная компания потребительских технологий переписала более 30 производственных UDF. Одна из них затрагивала свыше 3000 задач и дала ускорение 22x при реализации через RapidsUDF с поддержкой cuDF.

Как установить набор навыков?

С GitHub — напрямую или командой npx skills add NVIDIA/cudf-spark --skill '*'.

Источник: developer.nvidia.com