Trying Out Spark Jobs on Amazon EMR Serverless with the New 32 vCPU Worker Size

"""
CPU-bound均等分散ベンチマーク for EMR Serverless

設計:
- 各パーティションで60秒間ひたすら桁数計算（整数の累乗）を実行
- 60秒経過後に処理できた桁数を返し、全パーティションの合計を出す
- I/O・シャッフルなし → CPU使用率が均等に分散することが期待される

引数:
  sys.argv[1] = duration_seconds (デフォルト: 60)
  sys.argv[2] = num_partitions (デフォルト: 4)
"""
import subprocess
import sys
import time
from datetime import datetime

from pyspark.sql import SparkSession

def log(msg: str) -> None:
    print(f"[{datetime.now().isoformat()}] {msg}")

def print_cpu_info() -> None:
    """Driver側でCPU情報をログに出す。"""
    try:
        out = subprocess.run(["lscpu"], capture_output=True, text=True, timeout=10)
        log("=== lscpu (driver) ===")
        print(out.stdout)
    except Exception as e:
        log(f"cpu info failed: {e}")

def cpu_burn_digits(partition_index, iterator, duration_seconds):
    import time as t
    import subprocess as sp
    import math

    if partition_index == 0:
        try:
            out = sp.run(["lscpu"], capture_output=True, text=True, timeout=10)
            print(f"[partition-{partition_index}] === lscpu (executor) ===")
            print(out.stdout)
        except Exception:
            pass

    start = t.time()
    total_digits = 0
    iterations = 0
    base = 2
    exp = 1000

    while True:
        elapsed = t.time() - start
        if elapsed >= duration_seconds:
            break

        result = base ** exp
        digits = int(math.log10(result)) + 1 if result > 0 else 1
        total_digits += digits
        iterations += 1

        exp += 100
        if exp > 50000:
            exp = 1000
            base += 1

    actual_elapsed = t.time() - start
    yield (partition_index, total_digits, iterations, actual_elapsed)

def run(duration_seconds: int, num_partitions: int) -> None:
    spark = SparkSession.builder.appName("CPUBoundBench").getOrCreate()
    sc = spark.sparkContext

    log("=== Job start ===")
    log(f"duration_seconds={duration_seconds} num_partitions={num_partitions}")
    log(f"defaultParallelism={sc.defaultParallelism}")

    print_cpu_info()

    start = time.time()

    rdd = sc.parallelize(range(num_partitions), num_partitions)
    results = rdd.mapPartitionsWithIndex(
        lambda idx, it: cpu_burn_digits(idx, it, duration_seconds)
    ).collect()

    total_wall_time = time.time() - start

    log("=== Per-partition results ===")
    grand_total_digits = 0
    for partition_idx, digits, iters, elapsed in results:
        log(f"  partition[{partition_idx}]: digits={digits:,}, iterations={iters}, time={elapsed:.2f}s")
        grand_total_digits += digits

    log(f"=== TOTAL DIGITS (all partitions): {grand_total_digits:,} ===")
    log(f"total_wall_time={total_wall_time:.2f}s")
    log("=== Job end ===")

    spark.stop()

if __name__ == "__main__":
    duration = int(sys.argv[1]) if len(sys.argv) > 1 else 60
    parts = int(sys.argv[2]) if len(sys.argv) > 2 else 4
    run(duration, parts)

partition[0]:  digits=1,295,403,119, iterations=34815, time=60.00s
partition[1]:  digits=1,294,568,243, iterations=34805, time=60.00s
partition[2]:  digits=1,296,170,388, iterations=34824, time=60.00s
partition[3]:  digits=1,294,816,756, iterations=34808, time=60.00s
partition[4]:  digits=1,292,021,453, iterations=34773, time=60.00s
partition[5]:  digits=1,293,033,148, iterations=34786, time=60.00s
partition[6]:  digits=1,284,219,297, iterations=34654, time=60.00s
partition[7]:  digits=1,294,733,733, iterations=34807, time=60.00s
partition[8]:  digits=1,294,816,756, iterations=34808, time=60.00s
partition[9]:  digits=1,291,944,930, iterations=34772, time=60.00s
partition[10]: digits=1,294,321,402, iterations=34802, time=60.00s
partition[11]: digits=1,296,256,569, iterations=34825, time=60.00s
partition[12]: digits=1,295,066,940, iterations=34811, time=60.00s
partition[13]: digits=1,293,751,941, iterations=34795, time=60.00s
partition[14]: digits=1,294,485,777, iterations=34804, time=60.00s
partition[15]: digits=1,295,912,960, iterations=34821, time=60.00s
partition[16]: digits=1,295,912,960, iterations=34821, time=60.00s
partition[17]: digits=1,295,066,940, iterations=34811, time=60.00s
partition[18]: digits=1,294,076,233, iterations=34799, time=60.00s
partition[19]: digits=1,292,796,894, iterations=34783, time=60.00s
partition[20]: digits=1,295,657,204, iterations=34818, time=60.00s
partition[21]: digits=1,293,430,620, iterations=34791, time=60.00s
partition[22]: digits=1,293,350,754, iterations=34790, time=60.00s
partition[23]: digits=1,289,806,569, iterations=34743, time=60.00s
partition[24]: digits=1,294,983,360, iterations=34810, time=60.00s
partition[25]: digits=1,289,806,569, iterations=34743, time=60.00s
partition[26]: digits=1,291,189,920, iterations=34762, time=60.00s
partition[27]: digits=1,289,244,165, iterations=34735, time=60.00s
partition[28]: digits=1,296,777,553, iterations=34831, time=60.00s
partition[29]: digits=1,284,494,185, iterations=34659, time=60.00s
partition[30]: digits=1,290,380,860, iterations=34751, time=60.00s
partition[31]: digits=1,293,590,909, iterations=34793, time=60.00s
TOTAL DIGITS: 41,382,089,107