def assert_gpu_fallback_write(write_func, read_func, base_path, cpu_fallback_class_name, conf={}): conf = _prep_incompat_conf(conf) print('### CPU RUN ###') cpu_start = time.time() cpu_path = base_path + '/CPU' with_cpu_session(lambda spark : write_func(spark, cpu_path), conf=conf) cpu_end = time.time() print('### GPU RUN ###') jvm = spark_jvm() jvm.com.nvidia.spark.rapids.ExecutionPlanCaptureCallback.startCapture() gpu_start = time.time() gpu_path = base_path + '/GPU' with_gpu_session(lambda spark : write_func(spark, gpu_path), conf=conf) gpu_end = time.time() jvm.com.nvidia.spark.rapids.ExecutionPlanCaptureCallback.assertCapturedAndGpuFellBack(cpu_fallback_class_name, 2000) print('### WRITE: GPU TOOK {} CPU TOOK {} ###'.format( gpu_end - gpu_start, cpu_end - cpu_start)) (cpu_bring_back, cpu_collect_type) = _prep_func_for_compare( lambda spark: read_func(spark, cpu_path), 'COLLECT') (gpu_bring_back, gpu_collect_type) = _prep_func_for_compare( lambda spark: read_func(spark, gpu_path), 'COLLECT') from_cpu = with_cpu_session(cpu_bring_back, conf=conf) from_gpu = with_cpu_session(gpu_bring_back, conf=conf) if should_sort_locally(): from_cpu.sort(key=_RowCmp) from_gpu.sort(key=_RowCmp) assert_equal(from_cpu, from_gpu)
def assert_gpu_fallback_collect(func, cpu_fallback_class_name, conf={}): (bring_back, collect_type) = _prep_func_for_compare(func, 'COLLECT') conf = _prep_incompat_conf(conf) print('### CPU RUN ###') cpu_start = time.time() from_cpu = with_cpu_session(bring_back, conf=conf) cpu_end = time.time() print('### GPU RUN ###') jvm = spark_jvm() jvm.com.nvidia.spark.rapids.ExecutionPlanCaptureCallback.startCapture() gpu_start = time.time() from_gpu = with_gpu_session(bring_back, conf=conf) gpu_end = time.time() jvm.com.nvidia.spark.rapids.ExecutionPlanCaptureCallback.assertCapturedAndGpuFellBack(cpu_fallback_class_name, 2000) print('### {}: GPU TOOK {} CPU TOOK {} ###'.format(collect_type, gpu_end - gpu_start, cpu_end - cpu_start)) if should_sort_locally(): from_cpu.sort(key=_RowCmp) from_gpu.sort(key=_RowCmp) assert_equal(from_cpu, from_gpu)
def assert_cpu_and_gpu_are_equal_collect_with_capture(func, exist_classes='', non_exist_classes='', conf={}): (bring_back, collect_type) = _prep_func_for_compare(func, 'COLLECT_WITH_DATAFRAME') conf = _prep_incompat_conf(conf) print('### CPU RUN ###') cpu_start = time.time() from_cpu, cpu_df = with_cpu_session(bring_back, conf=conf) cpu_end = time.time() print('### GPU RUN ###') gpu_start = time.time() from_gpu, gpu_df = with_gpu_session(bring_back, conf=conf) gpu_end = time.time() jvm = spark_jvm() if exist_classes: for clz in exist_classes.split(','): jvm.com.nvidia.spark.rapids.ExecutionPlanCaptureCallback.assertContains(gpu_df._jdf, clz) if non_exist_classes: for clz in non_exist_classes.split(','): jvm.com.nvidia.spark.rapids.ExecutionPlanCaptureCallback.assertNotContain(gpu_df._jdf, clz) print('### {}: GPU TOOK {} CPU TOOK {} ###'.format(collect_type, gpu_end - gpu_start, cpu_end - cpu_start)) if should_sort_locally(): from_cpu.sort(key=_RowCmp) from_gpu.sort(key=_RowCmp) assert_equal(from_cpu, from_gpu)