stack/notebooks/gpu_test.py

import marimo

__generated_with = "0.19.7"
app = marimo.App(width="medium")

with app.setup:
    import os
    import platform
    import subprocess
    import sys
    import time

    import marimo as mo
    import numpy as np
    import pandas as pd
    import polars as pl


@app.cell(hide_code=True)
def gpu_diagnostics():
    try:
        smi_result = subprocess.run(
            ["nvidia-smi"], capture_output=True, text=True, timeout=10
        )
        nvidia_smi_output = smi_result.stdout
        gpu_detected = smi_result.returncode == 0
    except Exception as exc:
        nvidia_smi_output = str(exc)
        gpu_detected = False

    visible_devices = os.environ.get("NVIDIA_VISIBLE_DEVICES", "not set")
    driver_capabilities = os.environ.get("NVIDIA_DRIVER_CAPABILITIES", "not set")

    mo.md(f"""
    # GPU & System Check

    **GPU Available:** `{gpu_detected}`
    **NVIDIA_VISIBLE_DEVICES:** `{visible_devices}`
    **NVIDIA_DRIVER_CAPABILITIES:** `{driver_capabilities}`

    ```
    {nvidia_smi_output}
    ```
    """)
    return


@app.cell(hide_code=True)
def polars_gpu_benchmark():
    benchmark_rows = 50_000_000
    bench_rng = np.random.default_rng(42)

    bench_gen_start = time.perf_counter()
    bench_df = pl.DataFrame(
        {
            "id": np.arange(benchmark_rows),
            "group": bench_rng.choice(["A", "B", "C", "D", "E"], size=benchmark_rows),
            "value_1": bench_rng.standard_normal(benchmark_rows),
            "value_2": bench_rng.uniform(0, 1000, size=benchmark_rows),
            "value_3": bench_rng.integers(0, 100, size=benchmark_rows),
        }
    )
    bench_gen_elapsed = time.perf_counter() - bench_gen_start

    # Pre-create lazy frame to exclude setup from timing
    bench_lazy = bench_df.lazy()

    # GPU-supported aggregations only (no quantile/median which cause fallback)
    bench_agg_expr = [
        pl.col("value_1").mean().alias("mean_v1"),
        pl.col("value_1").std().alias("std_v1"),
        pl.col("value_2").sum().alias("sum_v2"),
        pl.col("value_2").min().alias("min_v2"),
        pl.col("value_2").max().alias("max_v2"),
        pl.col("value_3").mean().alias("mean_v3"),
        pl.len().alias("count"),
    ]

    # --- GPU collect ---
    bench_gpu_agg_start = time.perf_counter()
    bench_gpu_agg_result = (
        bench_lazy.group_by("group")
        .agg(*bench_agg_expr)
        .sort("group")
        .collect(engine="gpu")
    )
    bench_gpu_agg_elapsed = time.perf_counter() - bench_gpu_agg_start

    # --- CPU collect ---
    bench_cpu_agg_start = time.perf_counter()
    bench_cpu_agg_result = (
        bench_lazy.group_by("group").agg(*bench_agg_expr).sort("group").collect()
    )
    bench_cpu_agg_elapsed = time.perf_counter() - bench_cpu_agg_start

    bench_agg_speedup = (
        bench_cpu_agg_elapsed / bench_gpu_agg_elapsed
        if bench_gpu_agg_elapsed > 0
        else float("inf")
    )

    # GPU-supported window functions only (no rank which causes fallback)
    bench_window_expr = [
        pl.col("value_1").mean().over("group").alias("group_mean"),
        pl.col("value_2").sum().over("group").alias("group_sum"),
    ]

    # --- GPU window ---
    bench_gpu_window_start = time.perf_counter()
    bench_gpu_window_result = (
        bench_lazy.with_columns(*bench_window_expr).head(5).collect(engine="gpu")
    )
    bench_gpu_window_elapsed = time.perf_counter() - bench_gpu_window_start

    # --- CPU window ---
    bench_cpu_window_start = time.perf_counter()
    bench_cpu_window_result = (
        bench_lazy.with_columns(*bench_window_expr).head(5).collect()
    )
    bench_cpu_window_elapsed = time.perf_counter() - bench_cpu_window_start

    bench_window_speedup = (
        bench_cpu_window_elapsed / bench_gpu_window_elapsed
        if bench_gpu_window_elapsed > 0
        else float("inf")
    )

    mo.md(f"""
    # Polars GPU vs CPU — {benchmark_rows:,} rows

    | Operation | GPU | CPU | Speedup |
    |---|---|---|---|
    | Data generation | `{bench_gen_elapsed:.3f}s` | — | — |
    | GroupBy aggregation | `{bench_gpu_agg_elapsed:.3f}s` | `{bench_cpu_agg_elapsed:.3f}s` | **{bench_agg_speedup:.1f}x** |
    | Window functions | `{bench_gpu_window_elapsed:.3f}s` | `{bench_cpu_window_elapsed:.3f}s` | **{bench_window_speedup:.1f}x** |
    """)

    mo.hstack(
        [
            mo.ui.table(bench_gpu_agg_result, label="GPU Aggregation Results"),
            mo.ui.table(bench_gpu_window_result, label="GPU Window Functions (head 5)"),
        ]
    )
    return


@app.cell(hide_code=True)
def pandas_vs_polars_gpu():
    cmp_rows = 50_000_000
    cmp_rng = np.random.default_rng(99)

    # Generate data once, outside timing
    cmp_categories = cmp_rng.choice(["X", "Y", "Z"], size=cmp_rows)
    cmp_amounts = cmp_rng.standard_normal(cmp_rows).astype(np.float64)

    # --- Pandas (CPU only) - time only the compute, not DataFrame creation ---
    cmp_pandas_df = pd.DataFrame({"category": cmp_categories, "amount": cmp_amounts})
    cmp_pandas_start = time.perf_counter()
    cmp_pandas_agg = cmp_pandas_df.groupby("category")["amount"].agg(
        ["mean", "std", "sum"]
    )
    cmp_pandas_elapsed = time.perf_counter() - cmp_pandas_start

    # --- Polars GPU - time only the compute ---
    cmp_polars_df = pl.DataFrame({"category": cmp_categories, "amount": cmp_amounts})
    cmp_lazy = cmp_polars_df.lazy()

    cmp_gpu_start = time.perf_counter()
    cmp_gpu_agg = (
        cmp_lazy.group_by("category")
        .agg(
            pl.col("amount").mean().alias("mean"),
            pl.col("amount").std().alias("std"),
            pl.col("amount").sum().alias("sum"),
        )
        .sort("category")
        .collect(engine="gpu")
    )
    cmp_gpu_elapsed = time.perf_counter() - cmp_gpu_start

    # --- Polars CPU ---
    cmp_cpu_start = time.perf_counter()
    cmp_cpu_agg = (
        cmp_lazy.group_by("category")
        .agg(
            pl.col("amount").mean().alias("mean"),
            pl.col("amount").std().alias("std"),
            pl.col("amount").sum().alias("sum"),
        )
        .sort("category")
        .collect()
    )
    cmp_cpu_elapsed = time.perf_counter() - cmp_cpu_start

    cmp_gpu_vs_pandas = (
        cmp_pandas_elapsed / cmp_gpu_elapsed if cmp_gpu_elapsed > 0 else float("inf")
    )
    cmp_cpu_vs_pandas = (
        cmp_pandas_elapsed / cmp_cpu_elapsed if cmp_cpu_elapsed > 0 else float("inf")
    )

    mo.md(f"""
    # Three-Way Comparison — {cmp_rows:,} rows

    | Engine | GroupBy Time | vs Pandas |
    |---|---|---|
    | Pandas (CPU) | `{cmp_pandas_elapsed:.3f}s` | 1.0x |
    | Polars (CPU) | `{cmp_cpu_elapsed:.3f}s` | **{cmp_cpu_vs_pandas:.1f}x** |
    | Polars (GPU) | `{cmp_gpu_elapsed:.3f}s` | **{cmp_gpu_vs_pandas:.1f}x** |
    """)
    return


@app.cell(hide_code=True)
def environment_info():
    python_version = sys.version
    platform_info = platform.platform()
    cpu_cores = os.cpu_count()
    marimo_version = mo.__version__
    polars_version = pl.__version__
    pandas_version = pd.__version__
    numpy_version = np.__version__

    mo.md(f"""
    # Environment

    | Component | Version |
    |---|---|
    | Python | `{python_version}` |
    | Platform | `{platform_info}` |
    | CPU cores | `{cpu_cores}` |
    | Marimo | `{marimo_version}` |
    | Polars | `{polars_version}` |
    | Pandas | `{pandas_version}` |
    | NumPy | `{numpy_version}` |
    """)
    return


if __name__ == "__main__":
    app.run()