RAPIDS cuDF: GPU acceleration, CUDA требования и профиль данных

NVIDIA GPU Driver (версия ≥ 525)
    │
CUDA Runtime (11.8 или 12.x)
    │
cuDF C++ library (собирается под конкретную CUDA)
    │
RAPIDS Accelerator JAR (rapids-4-spark_2.12-XX.YY.ZZ.jar)
    │
Apache Spark (3.3 / 3.4 / 3.5)
    │
JVM (JDK 8 / 11 / 17)

# Проверка версии драйвера и CUDA
nvidia-smi
# Вывод:
# Driver Version: 535.104.05   CUDA Version: 12.2

# Проверка Compute Capability конкретного GPU
nvidia-smi --query-gpu=name,compute_cap --format=csv
# Вывод:
# name, compute_cap
# NVIDIA A100-SXM4-80GB, 8.0

# Загрузка RAPIDS JAR (версия под Spark 3.5, CUDA 12.x)
RAPIDS_VERSION="24.02.0"
wget https://repo1.maven.org/maven2/com/nvidia/rapids-4-spark_2.12/${RAPIDS_VERSION}/rapids-4-spark_2.12-${RAPIDS_VERSION}-cuda12.jar

# Конфигурация GPU resources для Spark Executor Pod
apiVersion: v1
kind: Pod
spec:
  containers:
  - name: spark-executor
    resources:
      requests:
        cpu: "8"
        memory: "32Gi"
        nvidia.com/gpu: "1"        # запрашиваем 1 GPU на Executor
      limits:
        nvidia.com/gpu: "1"        # строгий лимит
    env:
    - name: NVIDIA_VISIBLE_DEVICES
      value: all
    - name: NVIDIA_DRIVER_CAPABILITIES
      value: compute,utility

# yarn-site.xml: включение GPU как ресурса
# yarn.resource-types = yarn.io/gpu
# yarn.nodemanager.resource-plugins = io.yarn.server.nodemanager.containermanager.resourceplugin.gpu.GpuResourcePlugin

# spark-submit для YARN с GPU
spark-submit \
  --master yarn \
  --deploy-mode cluster \
  --num-executors 4 \
  --executor-cores 8 \
  --executor-memory 32g \
  --conf spark.executor.resource.gpu.amount=1 \
  --conf spark.task.resource.gpu.amount=1 \
  --conf spark.executor.resource.gpu.discoveryScript=/usr/lib/spark/scripts/getGpusResources.sh \
  --jars rapids-4-spark_2.12-24.02.0-cuda12.jar \
  your_job.py

// CUDA kernel для предиката amount > threshold (упрощённо)
__global__ void filter_double_gt(
    const double* __restrict__ amounts,
    const uint8_t* __restrict__ validity,   // null-bitmap
    bool* __restrict__ result,
    int64_t n_rows,
    double threshold
) {
    // Каждый поток обрабатывает один элемент
    int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n_rows) return;

    // Проверяем validity bit (не null)
    bool is_valid = (validity[idx / 8] >> (idx % 8)) & 1;

    // Применяем предикат
    result[idx] = is_valid && (amounts[idx] > threshold);
}

// Запуск: 8M строк / 256 потоков на блок = 31250 блоков
// Все 31250 блоков запускаются параллельно → весь массив проверяется за ~5-10 мкс
filter_double_gt<<<31250, 256>>>(d_amounts, d_validity, d_result, 8000000, 1000.0);

GPU A (сервер 1) → PCIe → CPU RAM → сеть → CPU RAM → PCIe → GPU B (сервер 2)

GPU A (сервер 1) → NVLink/InfiniBand RDMA → GPU B (сервер 2)

.config("spark.shuffle.manager",
        "com.nvidia.spark.rapids.spark35.RapidsShuffleManager")
.config("spark.rapids.shuffle.mode", "UCX")
# UCX endpoint: InfiniBand или TCP fallback
.config("spark.rapids.shuffle.ucx.useWakeup", "true")

# Конфигурация RMM
.config("spark.rapids.memory.gpu.allocator", "ARENA")

# Какую долю VRAM зарезервировать под RMM-пул (0.0–1.0)
# Остаток оставляем под CUDA context, нативные буферы
.config("spark.rapids.memory.gpu.reserve", "0.1")   # 10% резерв

# Минимальный размер пула (байты). При старте RMM аллоцирует не менее этого
.config("spark.rapids.memory.gpu.minAllocFraction", "0.25")

# Максимальный размер пула (доля от VRAM)
.config("spark.rapids.memory.gpu.maxAllocFraction", "0.80")

# Пул host-памяти для staging буферов (PCIe transfer)
.config("spark.rapids.memory.host.spillStorageSize", "8589934592")  # 8 GB

# Размер батча в строках (по умолчанию ~1 млн строк для широких таблиц)
.config("spark.rapids.sql.reader.batchSizeRows", "1000000")

# Размер батча в байтах (альтернативное ограничение)
.config("spark.rapids.sql.reader.batchSizeBytes", "1073741824")  # 1 GB

# RAPIDS берёт минимум из двух ограничений

batchSizeBytes = VRAM_total × max_fraction × 0.5 / num_concurrent_tasks
              = 80 GB × 0.80 × 0.5 / 2 tasks = 16 GB на батч
batchSizeRows  = batchSizeBytes / avg_row_size_bytes
              = 16 GB / 200 bytes = 80 M строк

GPU VRAM (быстро, ~3.35 TB/s)
    ↓ overflow
Host RAM (средне, ~32 GB/s PCIe 4.0)
    ↓ overflow
Local NVMe SSD (медленно, ~7 GB/s)

# Настройка spill
.config("spark.rapids.memory.gpu.spillBatchSizeBytes", "536870912")  # 512 MB per spill
.config("spark.rapids.memory.host.spillStorageSize", "8589934592")   # 8 GB host spill

from pyspark.sql import SparkSession

def create_rapids_spark(
    app_name: str,
    gpu_vram_gb: int = 80,
) -> SparkSession:
    """
    Production-ready SparkSession с RAPIDS Accelerator.

    Параметры рассчитаны для A100 80GB:
    - VRAM: 80 GB (резерв 10% = 8 GB под CUDA, 72 GB под RMM)
    - Host RAM: 32 GB (heap 16 GB + offHeap 8 GB + overhead 8 GB)
    """
    rmm_max_bytes = int(gpu_vram_gb * 0.80 * (1024 ** 3))
    host_spill_bytes = 8 * (1024 ** 3)

    return (
        SparkSession.builder
        .appName(app_name)

        # RAPIDS plugin registration (через spark.plugins)
        .config("spark.plugins",
                "com.nvidia.spark.SQLPlugin")
        # Активация GPU execution
        .config("spark.rapids.sql.enabled", "true")

        # Ресурс GPU для Executor
        .config("spark.executor.resource.gpu.amount", "1")
        .config("spark.task.resource.gpu.amount", "1")

        # Память JVM Executor
        .config("spark.executor.memory", "16g")
        .config("spark.executor.memoryOverhead", "8g")
        .config("spark.memory.offHeap.enabled", "true")
        .config("spark.memory.offHeap.size", "8g")

        # RMM pool
        .config("spark.rapids.memory.gpu.allocator", "ARENA")
        .config("spark.rapids.memory.gpu.reserve", "0.10")
        .config("spark.rapids.memory.gpu.maxAllocFraction", "0.80")
        .config("spark.rapids.memory.host.spillStorageSize",
                str(host_spill_bytes))

        # Размер батча
        .config("spark.rapids.sql.reader.batchSizeRows", "1000000")
        .config("spark.rapids.sql.reader.batchSizeBytes",
                str(1 * 1024 ** 3))  # 1 GB

        # GPU Shuffle через UCX
        .config("spark.shuffle.manager",
                "com.nvidia.spark.rapids.spark35.RapidsShuffleManager")

        # AQE совместимость
        .config("spark.sql.adaptive.enabled", "true")

        # Диагностика (в production: WARNING)
        .config("spark.rapids.sql.explain", "NONE")

        .getOrCreate()
    )

spark = create_rapids_spark("rapids-demo")

# 1. Проверяем, что GPU обнаружен
gpu_info = spark.sparkContext._jvm.com.nvidia.spark.rapids.GpuDeviceManager \
    .getGpuResource(spark._jsparkSession)
print(f"GPU resource: {gpu_info}")

# 2. Читаем данные и проверяем план
df = spark.read.parquet("/data/orders_large/")
df.filter("o_totalprice > 50000").groupBy("o_orderstatus").count().explain()

# В плане должны быть строки вида:
# GpuHashAggregateExec(keys=[o_orderstatus], funcs=[count(1)])
# +- GpuFilterExec(o_totalprice > 50000.0)
#    +- GpuBatchScanExec[o_orderstatus, o_totalprice] GpuParquetScan

# 3. Простейшая проверка GPU
import subprocess
result = subprocess.run(["nvidia-smi", "--query-gpu=name,memory.used,memory.free",
                        "--format=csv,noheader"], capture_output=True, text=True)
print(result.stdout)

from pyspark.sql import functions as F, SparkSession

spark = create_rapids_spark("ecommerce-gold")

# Чтение партиционированных Parquet таблиц
df_transactions = spark.read.parquet("/data/transactions/")
df_products = spark.read.parquet("/data/products/")
df_customers = spark.read.parquet("/data/customers/")

# Gold-витрина: продажи по категориям и регионам
gold = (
    df_transactions
    .filter(
        F.col("transaction_date") >= F.date_sub(F.current_date(), 90)
    )
    .join(
        df_products.select("product_id", "category_l1", "category_l2", "base_price"),
        "product_id"
    )
    .join(
        df_customers.select("customer_id", "region", "customer_segment"),
        "customer_id"
    )
    .groupBy("category_l1", "category_l2", "region", "customer_segment")
    .agg(
        F.count("*").alias("transaction_count"),
        F.sum("amount").alias("total_revenue"),
        F.avg("amount").alias("avg_transaction"),
        F.sum(F.col("amount") * (1 - F.col("discount_rate"))).alias("net_revenue"),
        F.countDistinct("customer_id").alias("unique_customers"),
        F.max("amount").alias("max_transaction")
    )
    .orderBy(F.col("total_revenue").desc())
)

gold.explain()

== Physical Plan ==
GpuColumnarToRowExec
+- GpuSortExec[total_revenue DESC]
   +- GpuHashAggregateExec(keys=[category_l1, category_l2, region, customer_segment],
   |     funcs=[count(1), sum(amount), avg(amount), sum(net), count(distinct cust_id), max(amount)])
   +- GpuColumnarExchange hashpartitioning(category_l1,..., 200)
      +- GpuHashAggregateExec(keys=[...], funcs=[partial_count, partial_sum, ...])
         +- GpuProjectExec[..., amount*(1-discount_rate) AS net]
            +- GpuBroadcastHashJoinExec[customer_id]
               :- GpuBroadcastHashJoinExec[product_id]
               :  :- GpuFilterExec(transaction_date >= date_sub(current_date(), 90))
               :  :  +- GpuBatchScanExec[...] GpuParquetScan /data/transactions/
               :  +- GpuBroadcastExchangeExec HashedRelation
               :     +- GpuBatchScanExec[product_id, category_l1, ...] GpuParquetScan
               +- GpuBroadcastExchangeExec HashedRelation
                  +- GpuBatchScanExec[customer_id, region, ...] GpuParquetScan

from pyspark.sql.types import StringType

@udf(returnType=StringType())
def classify_revenue(amount):
    """Плохой паттерн: Python UDF прерывает GPU-выполнение"""
    if amount > 1_000_000: return "platinum"
    elif amount > 100_000: return "gold"
    elif amount > 10_000: return "silver"
    else: return "bronze"

gold_with_tier = gold.withColumn("revenue_tier",
                                  classify_revenue(F.col("total_revenue")))
gold_with_tier.explain()

== Physical Plan ==
GpuColumnarToRowExec
+- GpuRowToColumnarExec                    <-- возврат в GPU
   +- BatchEvalPython [classify_revenue]   <-- Python UDF на CPU (строковый режим)
      +- GpuColumnarToRowExec              <-- выход из GPU
         +- GpuSortExec[total_revenue DESC]
            +- GpuHashAggregateExec(...)
               ...

gold_with_tier = gold.withColumn(
    "revenue_tier",
    F.when(F.col("total_revenue") > 1_000_000, "platinum")
     .when(F.col("total_revenue") > 100_000, "gold")
     .when(F.col("total_revenue") > 10_000, "silver")
     .otherwise("bronze")
)
# F.when → GpuProjectExec: полностью нативный CUDA-kernel

# Скачиваем Qualification Tool (входит в RAPIDS SDK)
wget https://repo1.maven.org/maven2/com/nvidia/rapids-4-spark-tools_2.12/24.02.0/rapids-4-spark-tools_2.12-24.02.0.jar

# Запускаем анализ логов (event logs со Spark History Server)
java -jar rapids-4-spark-tools_2.12-24.02.0.jar qualification \
  --output-directory /tmp/rapids-qual-output \
  --event-log /var/spark-history/app-20240115-103022

# Отчёт появится в /tmp/rapids-qual-output/qualification_report.csv

App ID: app-20240115-103022
App Duration: 3847 seconds
Estimated GPU Duration: 412 seconds
Estimated Speedup: 9.3x
GPU Opportunity Score: 87/100

Stage Analysis:
  Stage 2 (SortMergeJoin):   3120s → 280s  (11.1x)  GPU: SUPPORTED
  Stage 5 (HashAggregate):    490s →  68s   (7.2x)   GPU: SUPPORTED
  Stage 8 (PythonUDF):        237s → 237s   (1.0x)   GPU: NOT_SUPPORTED (Python UDF)
  Stage 11 (WindowSpec):       -  →   -             GPU: PARTIAL (some functions)

# Намеренно занижаем batchSizeRows - слишком большой батч для демонстрации
spark_oom = (
    SparkSession.builder
    .config("spark.rapids.sql.enabled", "true")
    .config("spark.rapids.memory.gpu.maxAllocFraction", "0.80")
    # Батч 200M строк × средний размер строки 200 байт = 40 GB → > VRAM!
    .config("spark.rapids.sql.reader.batchSizeRows", "200000000")
    .config("spark.rapids.sql.reader.batchSizeBytes", str(40 * 1024**3))
    .getOrCreate()
)

try:
    # Большой join: обе таблицы > 1M строк → SortMergeJoin (не broadcast)
    result = (
        spark_oom.read.parquet("/data/large_table_a/")  # 500M строк
        .join(spark_oom.read.parquet("/data/large_table_b/"), "key")  # 800M строк
        .count()
    )
except Exception as e:
    print(f"GPU OOM: {e}")
    # RAPIDS Memory error: Insufficient GPU memory
    # RMM pool exhausted: requested 38.2 GB, available 12.1 GB
    # Falling back... (если spill включён - вместо crash начнётся spill)

watch -n 0.5 nvidia-smi --query-gpu=name,memory.used,memory.free,utilization.gpu \
  --format=csv,noheader
# Name                          | Mem Used | Mem Free | GPU%
# NVIDIA A100-SXM4-80GB         | 72348MiB | 9844MiB  | 98%  ← норма при полной нагрузке
# NVIDIA A100-SXM4-80GB         | 79200MiB |  992MiB  | 45%  ← почти OOM, GPU ждёт spill

# Правильные параметры для JOIN 500M × 800M строк на A100 80GB
.config("spark.rapids.sql.reader.batchSizeRows", "2000000")   # 2M строк
.config("spark.rapids.sql.reader.batchSizeBytes", str(2 * 1024**3))  # 2 GB

# Убеждаемся что spill включён как safety net
.config("spark.rapids.memory.gpu.spillBatchSizeBytes", str(512 * 1024**2))

import time

def benchmark_query(spark, query: str, name: str, runs: int = 3) -> float:
    times = []
    for i in range(runs):
        t0 = time.time()
        spark.sql(query).count()
        times.append(time.time() - t0)
    median = sorted(times)[runs // 2]
    print(f"{name}: {median:.2f}s (median of {runs})")
    return median

# TPC-H Q1: чистая агрегация - идеальный кейс для GPU
Q1 = """
SELECT l_returnflag, l_linestatus,
       SUM(l_quantity) AS sum_qty,
       SUM(l_extendedprice) AS sum_base_price,
       SUM(l_extendedprice * (1 - l_discount)) AS sum_disc_price,
       SUM(l_extendedprice * (1 - l_discount) * (1 + l_tax)) AS sum_charge,
       AVG(l_quantity) AS avg_qty,
       AVG(l_extendedprice) AS avg_price,
       AVG(l_discount) AS avg_disc,
       COUNT(*) AS count_order
FROM lineitem
WHERE l_shipdate <= date_sub(date('1998-12-01'), 90)
GROUP BY l_returnflag, l_linestatus
ORDER BY l_returnflag, l_linestatus
"""

# TPC-H Q5: 5-way join - показывает силу GPU hash joins
Q5 = """
SELECT n_name, SUM(l_extendedprice * (1 - l_discount)) AS revenue
FROM customer, orders, lineitem, supplier, nation, region
WHERE c_custkey = o_custkey
  AND l_orderkey = o_orderkey
  AND l_suppkey = s_suppkey
  AND c_nationkey = s_nationkey
  AND s_nationkey = n_nationkey
  AND n_regionkey = r_regionkey
  AND r_name = 'ASIA'
  AND o_orderdate >= date('1994-01-01')
  AND o_orderdate < date('1995-01-01')
GROUP BY n_name
ORDER BY revenue DESC
"""

# Создаём сессии для сравнения
spark_cpu = SparkSession.builder.appName("baseline").getOrCreate()
spark_gpu = create_rapids_spark("rapids-bench")

data_path = "/data/tpch/sf100"

for table in ["customer", "orders", "lineitem", "supplier", "nation", "region"]:
    spark_cpu.read.parquet(f"{data_path}/{table}").createOrReplaceTempView(table)
    spark_gpu.read.parquet(f"{data_path}/{table}").createOrReplaceTempView(table)

print("=== TPC-H SF=100 Benchmark ===")
print(f"\n{'Query':<8} {'CPU JVM':>12} {'RAPIDS GPU':>12} {'Speedup':>10}")
print("-" * 45)

for qid, query in [("Q1", Q1), ("Q5", Q5)]:
    cpu_t = benchmark_query(spark_cpu, query, f"CPU {qid}")
    gpu_t = benchmark_query(spark_gpu, query, f"GPU {qid}")
    print(f"{qid:<8} {cpu_t:>10.2f}s {gpu_t:>11.2f}s {cpu_t/gpu_t:>9.2f}x")

# ПЛОХО: UDF до JOIN разрывает GPU-путь на самом нагруженном месте
df_flagged = df.withColumn("flag", my_python_udf(F.col("amount")))  # 2B строк!
df_flagged.join(df_other, "key").groupBy("category").agg(F.sum("amount"))
# UDF → GpuColumnarToRow (2B × 200 байт = 400 GB через PCIe)!

# ХОРОШО: UDF вынесен после агрегации (результат < 1000 строк)
df.join(df_other, "key").groupBy("category").agg(F.sum("amount")) \
  .withColumn("flag", my_python_udf(F.col("total")))
# UDF работает на 1000 строках - overhead ничтожен

# Добавьте в cron или monitoring:
nvidia-smi --query-gpu=utilization.gpu,memory.used --format=csv,noheader
# Если GPU utilization < 30% - ищите причину: fallback, IO-ожидание, skew

# Проблемный скрипт - найдите источники fallback
from pyspark.sql import functions as F
from pyspark.sql.types import StringType, ArrayType

spark = create_rapids_spark("homework-job")

# UDF 1: регулярные выражения со сложным паттерном
@udf(returnType=StringType())
def extract_sku(product_name):
    import re
    match = re.search(r'SKU-(\d{4,8})', product_name)
    return match.group(1) if match else None

# UDF 2: агрегация в Python-объект
@udf(returnType=ArrayType(StringType()))
def top_categories(categories_str):
    cats = categories_str.split(",")
    return sorted(set(cats))[:3]

df = (
    spark.read.parquet("/data/ecommerce_large/")
    # Сложный REGEXP → fallback
    .withColumn("sku", extract_sku(F.col("product_name")))
    # Агрегирующий UDF → fallback
    .withColumn("top_cats", top_categories(F.col("category_path")))
    # Locale-зависимая функция → fallback
    .withColumn("formatted_price",
                F.format_number(F.col("price"), 2))
    .groupBy("region", "sku")
    .agg(F.sum("revenue").alias("total"))
)

df.explain()
df.write.mode("overwrite").parquet("/data/output/")