From 4bcd3149de6718d85aab63b63542e0d199c1f4a3 Mon Sep 17 00:00:00 2001 From: Eric Vo Date: Fri, 14 Aug 2026 10:20:32 -0700 Subject: [PATCH 1/7] Add 1TRC benchmark skeleton --- benchmarks/1trc.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) create mode 100644 benchmarks/1trc.py diff --git a/benchmarks/1trc.py b/benchmarks/1trc.py new file mode 100644 index 00000000000..363a961e3fc --- /dev/null +++ b/benchmarks/1trc.py @@ -0,0 +1,12 @@ +import argparse + + +def time_dask_1trc(): + pass + + +def time_ak_1trc(): + pass + +def create_parser(): + parser = argparse.ArgumentParser(description="Measure performance of the 1 trillion row challenge.") \ No newline at end of file From 4c6957c47d276a057b70102363f896ac91f9ab59 Mon Sep 17 00:00:00 2001 From: Eric Vo Date: Tue, 18 Aug 2026 01:32:19 -0700 Subject: [PATCH 2/7] Add 1TRC lookup df and Arkouda + Dask benchmark + graphs --- benchmarks/1trc.py | 278 ++++++++++++++- benchmarks/graph_infra/1trc.perfkeys | 4 + benchmarks/graph_infra/GRAPHLIST | 4 +- benchmarks/graph_infra/arkouda-1trc.graph | 11 + benchmarks/run_benchmarks.py | 82 +++-- resources/1trc-testing/lookup.csv | 414 ++++++++++++++++++++++ 6 files changed, 756 insertions(+), 37 deletions(-) create mode 100644 benchmarks/graph_infra/1trc.perfkeys create mode 100644 benchmarks/graph_infra/arkouda-1trc.graph create mode 100644 resources/1trc-testing/lookup.csv diff --git a/benchmarks/1trc.py b/benchmarks/1trc.py index 363a961e3fc..55bd96a2175 100644 --- a/benchmarks/1trc.py +++ b/benchmarks/1trc.py @@ -1,12 +1,280 @@ +#!/usr/bin/env python3 + import argparse +import os +import time + +from glob import glob +from math import ceil + +import numpy as np +import pandas as pd + +import arkouda as ak + +from server_util.test.server_test_util import get_default_temp_directory + + +CHUNK_SIZE = 10_000 +STD = 10.0 +FILE_PATTERN = "measurements-*.parquet" +LOOKUP_PATH = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "resources", + "1trc-testing", + "lookup.csv", +) + + +def load_lookup(path=LOOKUP_PATH): + """Load the table of station names and their mean temperatures.""" + # generate_chunk maps positional station ids, so the index must stay a RangeIndex + return pd.read_csv(path) + + +def generate_chunk(partition_idx, chunksize, std, lookup_df, out_dir="."): + """Generate some sample data based on the lookup table.""" + rng = np.random.default_rng(partition_idx) # Deterministic data generation + df = pd.DataFrame( + { + # Choose a random station from the lookup table for each row in our output + "station": rng.integers(0, len(lookup_df) - 1, int(chunksize)), + # Generate a normal distibution around zero for each row in our output + # Because the std is the same for every station we can adjust the mean for each row afterwards + "measure": rng.normal(0, std, int(chunksize)), + } + ) + + # Offset each measurement by the station's mean value + df.measure += df.station.map(lookup_df.mean_temp) + # Round the temprature to one decimal place + df.measure = df.measure.round(decimals=1) + # Convert the station index to the station name + df.station = df.station.map(lookup_df.station) + + # Save this chunk to the output file + filename = os.path.join(out_dir, f"measurements-{partition_idx}.parquet") + df.to_parquet(filename, engine="pyarrow") + + +def generate_data(size, out_dir, chunksize=CHUNK_SIZE): + """Write ``size`` rows of measurements to parquet files and return their directory.""" + os.makedirs(out_dir, exist_ok=True) + lookup_df = load_lookup() + for i in range(ceil(size / chunksize)): + generate_chunk(i, chunksize, STD, lookup_df, out_dir) + return out_dir + + +def measurement_files(data): + files = sorted(glob(os.path.join(data, FILE_PATTERN))) + if not files: + raise ValueError(f"No files matching {FILE_PATTERN} found in {data}") + return files + + +def remove_files(path): + for f in glob(os.path.join(path, FILE_PATTERN)): + os.remove(f) + + +def start_dask_cluster(args, num_jobs): + """Start a PBS backed dask cluster; ``(None, None)`` means use the local scheduler.""" + if args.dask_cluster != "pbs": + return None, None + + import dask + + from dask.distributed import Client + from dask_jobqueue import PBSCluster + + dask.config.set( + { + "distributed.scheduler.worker-ttl": "1h", + "distributed.comm.timeouts.connect": "120s", + "distributed.comm.timeouts.tcp": "120s", + "distributed.worker.memory.target": 0.6, + "distributed.worker.memory.spill": 0.7, + "distributed.worker.memory.pause": 0.85, + "distributed.worker.memory.terminate": False, + "temporary-directory": args.dask_scratch, + } + ) + cluster_args = { + "cores": args.dask_cores, + "memory": args.dask_memory, + "walltime": args.dask_walltime, + "local_directory": args.dask_scratch, + } + if args.dask_queue: + cluster_args["queue"] = args.dask_queue + if args.dask_account: + cluster_args["account"] = args.dask_account -def time_dask_1trc(): - pass + cluster = PBSCluster(**cluster_args) + client = Client(cluster) + print("scaling dask to {} PBS jobs".format(num_jobs)) + cluster.scale(jobs=num_jobs) + client.wait_for_workers(n_workers=num_jobs) + return client, cluster -def time_ak_1trc(): - pass +def time_dask_1trc(size, trials, data): + """Time the reference dask implementation of the challenge.""" + print(">>> dask 1trc") + try: + import dask.dataframe as dd + except ImportError: + print("dask is not installed, skipping the dask comparison") + return None + + glob_path = os.path.join(data, FILE_PATTERN) + totalbytes = sum(os.path.getsize(f) for f in measurement_files(data)) + + timings = [] + result = None + for _ in range(trials): + start = time.time() + df = dd.read_parquet(glob_path, dtype_backend="pyarrow") + # split_out=1 forces a tree reduction + result = df.groupby("station").agg(["min", "max", "mean"], split_out=1).compute() + result = result.sort_values("station") + end = time.time() + timings.append(end - start) + + tavg = sum(timings) / trials + print("dask Average time = {:.4f} sec".format(tavg)) + print("dask Average rate = {:.4f} GiB/sec".format(totalbytes / tavg / 2**30)) + print(result.head()) + return tavg + + +def materialize_result_df(station_keys, mins, maxs, means): + result_df = pd.DataFrame( + { + ("measure", "min"): mins.to_ndarray(), + ("measure", "max"): maxs.to_ndarray(), + ("measure", "mean"): means.to_ndarray(), + }, + index=pd.Index(station_keys.to_ndarray(), name="station"), + ) + result_df.columns = pd.MultiIndex.from_tuples(result_df.columns) + return result_df + + +def time_ak_1trc(size, trials, data): + """Time the arkouda implementation of the challenge.""" + print(">>> arkouda 1trc") + cfg = ak.get_config() + file_paths = measurement_files(data) + print( + "numLocales = {}, numNodes = {}, files = {:,}".format( + cfg["numLocales"], cfg["numNodes"], len(file_paths) + ) + ) + + timings = [] + totalbytes = 0 + result = None + for _ in range(trials): + start = time.time() + columns = ak.read(file_paths) + stations = columns["station"] + measures = columns["measure"] + + order = ak.argsort(stations) + stations = stations[order] + measures = measures[order] + + grouped = ak.GroupBy(stations, assume_sorted=True) + station_keys, mins = grouped.min(measures) + _, maxs = grouped.max(measures) + _, means = grouped.mean(measures) + end = time.time() + + timings.append(end - start) + totalbytes = stations.nbytes + measures.size * measures.itemsize + result = (station_keys, mins, maxs, means) + + tavg = sum(timings) / trials + print("arkouda Average time = {:.4f} sec".format(tavg)) + print("arkouda Average rate = {:.4f} GiB/sec".format(totalbytes / tavg / 2**30)) + # Pulling the results back to the client is not part of the measured time + print(materialize_result_df(*result).head()) + return tavg + def create_parser(): - parser = argparse.ArgumentParser(description="Measure performance of the 1 trillion row challenge.") \ No newline at end of file + parser = argparse.ArgumentParser(description="Measure performance of the 1 trillion row challenge.") + parser.add_argument("hostname", help="Hostname of arkouda server") + parser.add_argument("port", type=int, help="Port of arkouda server") + parser.add_argument("-n", "--size", type=int, default=10**6, help="Number of rows to compute with") + parser.add_argument( + "-t", "--trials", type=int, default=6, help="Number of times to run the benchmark" + ) + parser.add_argument( + "-d", + "--data", + required=False, + type=str, + help="Optional dataset directory to use, otherwise a dataset will be generated", + ) + parser.add_argument( + "-p", + "--path", + default=os.path.join(get_default_temp_directory(), "ak-1trc-test"), + help="Target path for the generated dataset", + ) + parser.add_argument( + "--dask-cluster", + default="local", + choices=("local", "pbs"), + help="Scheduler for the dask comparison. 'pbs' scales a PBSCluster to the same number " + "of jobs as the arkouda server has nodes", + ) + parser.add_argument("--dask-cores", type=int, default=128, help="Cores per PBS job") + parser.add_argument("--dask-memory", default="500GB", help="Memory per PBS job") + parser.add_argument("--dask-walltime", default="5-00:00:00", help="Walltime per PBS job") + parser.add_argument("--dask-queue", default="", help="PBS queue to submit to") + parser.add_argument("--dask-account", default="", help="PBS account to charge") + parser.add_argument( + "--dask-scratch", + default=os.path.join(get_default_temp_directory(), "dask-scratch"), + help="Worker scratch directory", + ) + return parser + + +if __name__ == "__main__": + import sys + + args = create_parser().parse_args() + setattr(ak, "verbose", False) + ak.connect(args.hostname, args.port) + + # Use the supplied dataset if there is one, otherwise generate it + data = args.data if args.data else generate_data(args.size, args.path) + + print("number of trials = ", args.trials) + print("number of rows = ", args.size) + + # Give dask the same number of jobs as the server has nodes so the two series line up + num_nodes = ak.get_config()["numNodes"] + dask_client, dask_cluster = start_dask_cluster(args, num_nodes) + try: + dask_time = time_dask_1trc(args.size, args.trials, data) + finally: + if dask_client is not None: + dask_client.close() + dask_cluster.close() + + arkouda_time = time_ak_1trc(args.size, args.trials, data) + + if dask_time and arkouda_time: + print("arkouda/dask ratio = {:.2f}x".format(arkouda_time / dask_time)) + + if not args.data: + remove_files(args.path) + + sys.exit(0) diff --git a/benchmarks/graph_infra/1trc.perfkeys b/benchmarks/graph_infra/1trc.perfkeys new file mode 100644 index 00000000000..592e6102b5a --- /dev/null +++ b/benchmarks/graph_infra/1trc.perfkeys @@ -0,0 +1,4 @@ +arkouda Average time = +arkouda Average rate = +dask Average time = +dask Average rate = diff --git a/benchmarks/graph_infra/GRAPHLIST b/benchmarks/graph_infra/GRAPHLIST index 1265d259291..db12c69cc43 100644 --- a/benchmarks/graph_infra/GRAPHLIST +++ b/benchmarks/graph_infra/GRAPHLIST @@ -7,4 +7,6 @@ arkouda-comp.graph # suite: Sort Cases arkouda-sort-cases.graph # suite: Bigint Benchmarks -arkouda-bigint.graph \ No newline at end of file +arkouda-bigint.graph +# suite: 1TRC +arkouda-1trc.graph diff --git a/benchmarks/graph_infra/arkouda-1trc.graph b/benchmarks/graph_infra/arkouda-1trc.graph new file mode 100644 index 00000000000..15a7708274c --- /dev/null +++ b/benchmarks/graph_infra/arkouda-1trc.graph @@ -0,0 +1,11 @@ +perfkeys: arkouda Average rate = +graphkeys: 1TRC GiB/s +files: 1trc.dat +graphtitle: 1TRC Arkouda Performance +ylabel: Performance (GiB/s) + +perfkeys: dask Average rate = +graphkeys: 1TRC GiB/s +files: 1trc.dat +graphtitle: 1TRC Dask Performance +ylabel: Performance (GiB/s) diff --git a/benchmarks/run_benchmarks.py b/benchmarks/run_benchmarks.py index 1b5cf4d2b9c..b54dc73a960 100755 --- a/benchmarks/run_benchmarks.py +++ b/benchmarks/run_benchmarks.py @@ -27,6 +27,7 @@ logging.basicConfig(level=logging.INFO) BENCHMARKS = [ + "1trc", "stream", "argsort", "coargsort", @@ -157,6 +158,12 @@ def create_parser(): ) parser.add_argument("-sp", "--server-port", default="5555", help="Port number to use for the server") parser.add_argument("--server-args", action="append", help="Additional server arguments") + parser.add_argument( + "--node-counts", + default="", + help="Comma separated list of node counts to sweep, e.g. '16,32,64,128'. Each count is " + "recorded as its own configuration so that all counts share a graph", + ) parser.add_argument("--numtrials", default=1, type=int, help="Number of trials to run") parser.add_argument( "benchmarks", @@ -209,46 +216,59 @@ def main(): parser = create_parser() args, client_args = parser.parse_known_args() args.graph_dir = args.graph_dir or os.path.join(args.dat_dir, "html") - config_dat_dir = os.path.join(args.dat_dir, args.description) run_isolated = bool(args.isolated) - if args.save_data or args.gen_graphs: - os.makedirs(config_dat_dir, exist_ok=True) - - if not run_isolated: - _my_start_server(args) + node_counts = [n.strip() for n in args.node_counts.split(",") if n.strip()] + if node_counts: + # One configuration per node count puts every count on the same graph as its own series + sweep = [(n, "{}-nodes".format(n)) for n in node_counts] + args.configs = args.configs or ",".join("{}:v".format(desc) for _, desc in sweep) + else: + sweep = [(args.num_locales, args.description)] args.benchmarks = args.benchmarks or BENCHMARKS - for benchmark in args.benchmarks: - if run_isolated: + + for num_locales, description in sweep: + args.num_locales = num_locales + config_dat_dir = os.path.join(args.dat_dir, description) + + if args.save_data or args.gen_graphs: + os.makedirs(config_dat_dir, exist_ok=True) + + if not run_isolated: _my_start_server(args) - for trial in range(args.numtrials): - benchmark_py = os.path.join(benchmark_dir, "{}.py".format(benchmark)) - out = run_client(benchmark_py, client_args) - if args.save_data or args.gen_graphs: - add_to_dat(benchmark, out, config_dat_dir, args.graph_infra) - print(out) + for benchmark in args.benchmarks: + if run_isolated: + _my_start_server(args) + + for trial in range(args.numtrials): + benchmark_py = os.path.join(benchmark_dir, "{}.py".format(benchmark)) + out = run_client(benchmark_py, client_args) + if args.save_data or args.gen_graphs: + add_to_dat(benchmark, out, config_dat_dir, args.graph_infra) + print(out) + + if run_isolated: + stop_arkouda_server() - if run_isolated: + if not run_isolated: stop_arkouda_server() - if not run_isolated: - stop_arkouda_server() - - if args.save_data or args.gen_graphs: - comp_file = os.getenv("ARKOUDA_PRINT_PASSES_FILE", "") - if os.path.isfile(comp_file): - with open(comp_file, "r") as f: - out = f.read() - add_to_dat("comp-time", out, config_dat_dir, args.graph_infra) - emitted_code_file = os.getenv("ARKOUDA_EMITTED_CODE_SIZE_FILE", "") - if os.path.isfile(emitted_code_file): - with open(emitted_code_file, "r") as f: - out = f.read() - add_to_dat("emitted-code-size", out, config_dat_dir, args.graph_infra) - if args.gen_graphs: - generate_graphs(args) + if args.save_data or args.gen_graphs: + comp_file = os.getenv("ARKOUDA_PRINT_PASSES_FILE", "") + if os.path.isfile(comp_file): + with open(comp_file, "r") as f: + out = f.read() + add_to_dat("comp-time", out, config_dat_dir, args.graph_infra) + emitted_code_file = os.getenv("ARKOUDA_EMITTED_CODE_SIZE_FILE", "") + if os.path.isfile(emitted_code_file): + with open(emitted_code_file, "r") as f: + out = f.read() + add_to_dat("emitted-code-size", out, config_dat_dir, args.graph_infra) + + if args.gen_graphs: + generate_graphs(args) if __name__ == "__main__": diff --git a/resources/1trc-testing/lookup.csv b/resources/1trc-testing/lookup.csv new file mode 100644 index 00000000000..cbde9f632c4 --- /dev/null +++ b/resources/1trc-testing/lookup.csv @@ -0,0 +1,414 @@ +station,mean_temp +Abha,18.0 +Abidjan,26.0 +Abéché,29.4 +Accra,26.4 +Addis Ababa,16.0 +Adelaide,17.3 +Aden,29.1 +Ahvaz,25.4 +Albuquerque,14.0 +Alexandra,11.0 +Alexandria,20.0 +Algiers,18.2 +Alice Springs,21.0 +Almaty,10.0 +Amsterdam,10.2 +Anadyr,-6.9 +Anchorage,2.8 +Andorra la Vella,9.8 +Ankara,12.0 +Antananarivo,17.9 +Antsiranana,25.2 +Arkhangelsk,1.3 +Ashgabat,17.1 +Asmara,15.6 +Assab,30.5 +Astana,3.5 +Athens,19.2 +Atlanta,17.0 +Auckland,15.2 +Austin,20.7 +Baghdad,22.77 +Baguio,19.5 +Baku,15.1 +Baltimore,13.1 +Bamako,27.8 +Bangkok,28.6 +Bangui,26.0 +Banjul,26.0 +Barcelona,18.2 +Bata,25.1 +Batumi,14.0 +Beijing,12.9 +Beirut,20.9 +Belgrade,12.5 +Belize City,26.7 +Benghazi,19.9 +Bergen,7.7 +Berlin,10.3 +Bilbao,14.7 +Birao,26.5 +Bishkek,11.3 +Bissau,27.0 +Blantyre,22.2 +Bloemfontein,15.6 +Boise,11.4 +Bordeaux,14.2 +Bosaso,30.0 +Boston,10.9 +Bouaké,26.0 +Bratislava,10.5 +Brazzaville,25.0 +Bridgetown,27.0 +Brisbane,21.4 +Brussels,10.5 +Bucharest,10.8 +Budapest,11.3 +Bujumbura,23.8 +Bulawayo,18.9 +Burnie,13.1 +Busan,15.0 +Cabo San Lucas,23.9 +Cairns,25.0 +Cairo,21.4 +Calgary,4.4 +Canberra,13.1 +Cape Town,16.2 +Changsha,17.4 +Charlotte,16.1 +Chiang Mai,25.8 +Chicago,9.8 +Chihuahua,18.6 +Chișinău,10.2 +Chittagong,25.9 +Chongqing,18.6 +Christchurch,12.2 +City of San Marino,11.8 +Colombo,27.4 +Columbus,11.7 +Conakry,26.4 +Copenhagen,9.1 +Cotonou,27.2 +Cracow,9.3 +Da Lat,17.9 +Da Nang,25.8 +Dakar,24.0 +Dallas,19.0 +Damascus,17.0 +Dampier,26.4 +Dar es Salaam,25.8 +Darwin,27.6 +Denpasar,23.7 +Denver,10.4 +Detroit,10.0 +Dhaka,25.9 +Dikson,-11.1 +Dili,26.6 +Djibouti,29.9 +Dodoma,22.7 +Dolisie,24.0 +Douala,26.7 +Dubai,26.9 +Dublin,9.8 +Dunedin,11.1 +Durban,20.6 +Dushanbe,14.7 +Edinburgh,9.3 +Edmonton,4.2 +El Paso,18.1 +Entebbe,21.0 +Erbil,19.5 +Erzurum,5.1 +Fairbanks,-2.3 +Fianarantsoa,17.9 +"Flores, Petén",26.4 +Frankfurt,10.6 +Fresno,17.9 +Fukuoka,17.0 +Gabès,19.5 +Gaborone,21.0 +Gagnoa,26.0 +Gangtok,15.2 +Garissa,29.3 +Garoua,28.3 +George Town,27.9 +Ghanzi,21.4 +Gjoa Haven,-14.4 +Guadalajara,20.9 +Guangzhou,22.4 +Guatemala City,20.4 +Halifax,7.5 +Hamburg,9.7 +Hamilton,13.8 +Hanga Roa,20.5 +Hanoi,23.6 +Harare,18.4 +Harbin,5.0 +Hargeisa,21.7 +Hat Yai,27.0 +Havana,25.2 +Helsinki,5.9 +Heraklion,18.9 +Hiroshima,16.3 +Ho Chi Minh City,27.4 +Hobart,12.7 +Hong Kong,23.3 +Honiara,26.5 +Honolulu,25.4 +Houston,20.8 +Ifrane,11.4 +Indianapolis,11.8 +Iqaluit,-9.3 +Irkutsk,1.0 +Istanbul,13.9 +İzmir,17.9 +Jacksonville,20.3 +Jakarta,26.7 +Jayapura,27.0 +Jerusalem,18.3 +Johannesburg,15.5 +Jos,22.8 +Juba,27.8 +Kabul,12.1 +Kampala,20.0 +Kandi,27.7 +Kankan,26.5 +Kano,26.4 +Kansas City,12.5 +Karachi,26.0 +Karonga,24.4 +Kathmandu,18.3 +Khartoum,29.9 +Kingston,27.4 +Kinshasa,25.3 +Kolkata,26.7 +Kuala Lumpur,27.3 +Kumasi,26.0 +Kunming,15.7 +Kuopio,3.4 +Kuwait City,25.7 +Kyiv,8.4 +Kyoto,15.8 +La Ceiba,26.2 +La Paz,23.7 +Lagos,26.8 +Lahore,24.3 +Lake Havasu City,23.7 +Lake Tekapo,8.7 +Las Palmas de Gran Canaria,21.2 +Las Vegas,20.3 +Launceston,13.1 +Lhasa,7.6 +Libreville,25.9 +Lisbon,17.5 +Livingstone,21.8 +Ljubljana,10.9 +Lodwar,29.3 +Lomé,26.9 +London,11.3 +Los Angeles,18.6 +Louisville,13.9 +Luanda,25.8 +Lubumbashi,20.8 +Lusaka,19.9 +Luxembourg City,9.3 +Lviv,7.8 +Lyon,12.5 +Madrid,15.0 +Mahajanga,26.3 +Makassar,26.7 +Makurdi,26.0 +Malabo,26.3 +Malé,28.0 +Managua,27.3 +Manama,26.5 +Mandalay,28.0 +Mango,28.1 +Manila,28.4 +Maputo,22.8 +Marrakesh,19.6 +Marseille,15.8 +Maun,22.4 +Medan,26.5 +Mek'ele,22.7 +Melbourne,15.1 +Memphis,17.2 +Mexicali,23.1 +Mexico City,17.5 +Miami,24.9 +Milan,13.0 +Milwaukee,8.9 +Minneapolis,7.8 +Minsk,6.7 +Mogadishu,27.1 +Mombasa,26.3 +Monaco,16.4 +Moncton,6.1 +Monterrey,22.3 +Montreal,6.8 +Moscow,5.8 +Mumbai,27.1 +Murmansk,0.6 +Muscat,28.0 +Mzuzu,17.7 +N'Djamena,28.3 +Naha,23.1 +Nairobi,17.8 +Nakhon Ratchasima,27.3 +Napier,14.6 +Napoli,15.9 +Nashville,15.4 +Nassau,24.6 +Ndola,20.3 +New Delhi,25.0 +New Orleans,20.7 +New York City,12.9 +Ngaoundéré,22.0 +Niamey,29.3 +Nicosia,19.7 +Niigata,13.9 +Nouadhibou,21.3 +Nouakchott,25.7 +Novosibirsk,1.7 +Nuuk,-1.4 +Odesa,10.7 +Odienné,26.0 +Oklahoma City,15.9 +Omaha,10.6 +Oranjestad,28.1 +Oslo,5.7 +Ottawa,6.6 +Ouagadougou,28.3 +Ouahigouya,28.6 +Ouarzazate,18.9 +Oulu,2.7 +Palembang,27.3 +Palermo,18.5 +Palm Springs,24.5 +Palmerston North,13.2 +Panama City,28.0 +Parakou,26.8 +Paris,12.3 +Perth,18.7 +Petropavlovsk-Kamchatsky,1.9 +Philadelphia,13.2 +Phnom Penh,28.3 +Phoenix,23.9 +Pittsburgh,10.8 +Podgorica,15.3 +Pointe-Noire,26.1 +Pontianak,27.7 +Port Moresby,26.9 +Port Sudan,28.4 +Port Vila,24.3 +Port-Gentil,26.0 +Portland (OR),12.4 +Porto,15.7 +Prague,8.4 +Praia,24.4 +Pretoria,18.2 +Pyongyang,10.8 +Rabat,17.2 +Rangpur,24.4 +Reggane,28.3 +Reykjavík,4.3 +Riga,6.2 +Riyadh,26.0 +Rome,15.2 +Roseau,26.2 +Rostov-on-Don,9.9 +Sacramento,16.3 +Saint Petersburg,5.8 +Saint-Pierre,5.7 +Salt Lake City,11.6 +San Antonio,20.8 +San Diego,17.8 +San Francisco,14.6 +San Jose,16.4 +San José,22.6 +San Juan,27.2 +San Salvador,23.1 +Sana'a,20.0 +Santo Domingo,25.9 +Sapporo,8.9 +Sarajevo,10.1 +Saskatoon,3.3 +Seattle,11.3 +Ségou,28.0 +Seoul,12.5 +Seville,19.2 +Shanghai,16.7 +Singapore,27.0 +Skopje,12.4 +Sochi,14.2 +Sofia,10.6 +Sokoto,28.0 +Split,16.1 +St. John's,5.0 +St. Louis,13.9 +Stockholm,6.6 +Surabaya,27.1 +Suva,25.6 +Suwałki,7.2 +Sydney,17.7 +Tabora,23.0 +Tabriz,12.6 +Taipei,23.0 +Tallinn,6.4 +Tamale,27.9 +Tamanrasset,21.7 +Tampa,22.9 +Tashkent,14.8 +Tauranga,14.8 +Tbilisi,12.9 +Tegucigalpa,21.7 +Tehran,17.0 +Tel Aviv,20.0 +Thessaloniki,16.0 +Thiès,24.0 +Tijuana,17.8 +Timbuktu,28.0 +Tirana,15.2 +Toamasina,23.4 +Tokyo,15.4 +Toliara,24.1 +Toluca,12.4 +Toronto,9.4 +Tripoli,20.0 +Tromsø,2.9 +Tucson,20.9 +Tunis,18.4 +Ulaanbaatar,-0.4 +Upington,20.4 +Ürümqi,7.4 +Vaduz,10.1 +Valencia,18.3 +Valletta,18.8 +Vancouver,10.4 +Veracruz,25.4 +Vienna,10.4 +Vientiane,25.9 +Villahermosa,27.1 +Vilnius,6.0 +Virginia Beach,15.8 +Vladivostok,4.9 +Warsaw,8.5 +"Washington, D.C.",14.6 +Wau,27.8 +Wellington,12.9 +Whitehorse,-0.1 +Wichita,13.9 +Willemstad,28.0 +Winnipeg,3.0 +Wrocław,9.6 +Xi'an,14.1 +Yakutsk,-8.8 +Yangon,27.5 +Yaoundé,23.8 +Yellowknife,-4.3 +Yerevan,12.4 +Yinchuan,9.0 +Zagreb,10.7 +Zanzibar City,26.0 +Zürich,9.3 From f86120f90a3c428fab58fda532e2ee5f84045568 Mon Sep 17 00:00:00 2001 From: Eric Vo Date: Tue, 18 Aug 2026 02:04:26 -0700 Subject: [PATCH 3/7] Simplify code; edit arguments for clarity and correct values --- benchmarks/1trc.py | 63 ++++++++++++++++++---------------------------- 1 file changed, 25 insertions(+), 38 deletions(-) diff --git a/benchmarks/1trc.py b/benchmarks/1trc.py index 55bd96a2175..3ce703ae1ee 100644 --- a/benchmarks/1trc.py +++ b/benchmarks/1trc.py @@ -10,6 +10,11 @@ import numpy as np import pandas as pd +import dask +import dask.dataframe as dd +from dask.distributed import Client +from dask_jobqueue import PBSCluster + import arkouda as ak from server_util.test.server_test_util import get_default_temp_directory @@ -73,21 +78,17 @@ def measurement_files(data): return files +def dataset_bytes(files): + """On-disk size of the dataset, used as a common rate denominator for both engines.""" + return sum(os.path.getsize(f) for f in files) + + def remove_files(path): for f in glob(os.path.join(path, FILE_PATTERN)): os.remove(f) def start_dask_cluster(args, num_jobs): - """Start a PBS backed dask cluster; ``(None, None)`` means use the local scheduler.""" - if args.dask_cluster != "pbs": - return None, None - - import dask - - from dask.distributed import Client - from dask_jobqueue import PBSCluster - dask.config.set( { "distributed.scheduler.worker-ttl": "1h", @@ -120,23 +121,15 @@ def start_dask_cluster(args, num_jobs): return client, cluster -def time_dask_1trc(size, trials, data): +def time_dask_1trc(trials, file_paths, totalbytes): """Time the reference dask implementation of the challenge.""" print(">>> dask 1trc") - try: - import dask.dataframe as dd - except ImportError: - print("dask is not installed, skipping the dask comparison") - return None - - glob_path = os.path.join(data, FILE_PATTERN) - totalbytes = sum(os.path.getsize(f) for f in measurement_files(data)) timings = [] result = None for _ in range(trials): start = time.time() - df = dd.read_parquet(glob_path, dtype_backend="pyarrow") + df = dd.read_parquet(file_paths, dtype_backend="pyarrow") # split_out=1 forces a tree reduction result = df.groupby("station").agg(["min", "max", "mean"], split_out=1).compute() result = result.sort_values("station") @@ -159,15 +152,13 @@ def materialize_result_df(station_keys, mins, maxs, means): }, index=pd.Index(station_keys.to_ndarray(), name="station"), ) - result_df.columns = pd.MultiIndex.from_tuples(result_df.columns) return result_df -def time_ak_1trc(size, trials, data): +def time_ak_1trc(trials, file_paths, totalbytes): """Time the arkouda implementation of the challenge.""" print(">>> arkouda 1trc") cfg = ak.get_config() - file_paths = measurement_files(data) print( "numLocales = {}, numNodes = {}, files = {:,}".format( cfg["numLocales"], cfg["numNodes"], len(file_paths) @@ -175,7 +166,6 @@ def time_ak_1trc(size, trials, data): ) timings = [] - totalbytes = 0 result = None for _ in range(trials): start = time.time() @@ -194,8 +184,9 @@ def time_ak_1trc(size, trials, data): end = time.time() timings.append(end - start) - totalbytes = stations.nbytes + measures.size * measures.itemsize result = (station_keys, mins, maxs, means) + # Release the per-trial server arrays outside the timed region + del columns, stations, measures, order, grouped tavg = sum(timings) / trials print("arkouda Average time = {:.4f} sec".format(tavg)) @@ -226,15 +217,8 @@ def create_parser(): default=os.path.join(get_default_temp_directory(), "ak-1trc-test"), help="Target path for the generated dataset", ) - parser.add_argument( - "--dask-cluster", - default="local", - choices=("local", "pbs"), - help="Scheduler for the dask comparison. 'pbs' scales a PBSCluster to the same number " - "of jobs as the arkouda server has nodes", - ) - parser.add_argument("--dask-cores", type=int, default=128, help="Cores per PBS job") - parser.add_argument("--dask-memory", default="500GB", help="Memory per PBS job") + parser.add_argument("--dask-cores", type=int, default=32, help="Cores per PBS job") + parser.add_argument("--dask-memory", default="400GB", help="Memory per PBS job") parser.add_argument("--dask-walltime", default="5-00:00:00", help="Walltime per PBS job") parser.add_argument("--dask-queue", default="", help="PBS queue to submit to") parser.add_argument("--dask-account", default="", help="PBS account to charge") @@ -256,6 +240,10 @@ def create_parser(): # Use the supplied dataset if there is one, otherwise generate it data = args.data if args.data else generate_data(args.size, args.path) + # Both engines read the same files, so resolve them and their size once + file_paths = measurement_files(data) + totalbytes = dataset_bytes(file_paths) + print("number of trials = ", args.trials) print("number of rows = ", args.size) @@ -263,13 +251,12 @@ def create_parser(): num_nodes = ak.get_config()["numNodes"] dask_client, dask_cluster = start_dask_cluster(args, num_nodes) try: - dask_time = time_dask_1trc(args.size, args.trials, data) + dask_time = time_dask_1trc(args.trials, file_paths, totalbytes) finally: - if dask_client is not None: - dask_client.close() - dask_cluster.close() + dask_client.close() + dask_cluster.close() - arkouda_time = time_ak_1trc(args.size, args.trials, data) + arkouda_time = time_ak_1trc(args.trials, file_paths, totalbytes) if dask_time and arkouda_time: print("arkouda/dask ratio = {:.2f}x".format(arkouda_time / dask_time)) From 609f8c6b5bd24e71abae63f327926a76e90eb929 Mon Sep 17 00:00:00 2001 From: Eric Vo Date: Wed, 19 Aug 2026 03:24:07 -0500 Subject: [PATCH 4/7] Comment out dask code and adjust defaults --- benchmarks/1trc.py | 166 +++++++++++----------- benchmarks/graph_infra/1trc.perfkeys | 4 +- benchmarks/graph_infra/arkouda-1trc.graph | 12 +- 3 files changed, 91 insertions(+), 91 deletions(-) diff --git a/benchmarks/1trc.py b/benchmarks/1trc.py index 3ce703ae1ee..6da7f8c9c73 100644 --- a/benchmarks/1trc.py +++ b/benchmarks/1trc.py @@ -10,17 +10,17 @@ import numpy as np import pandas as pd -import dask -import dask.dataframe as dd -from dask.distributed import Client -from dask_jobqueue import PBSCluster +# import dask +# import dask.dataframe as dd +# from dask.distributed import Client +# from dask_jobqueue import PBSCluster import arkouda as ak from server_util.test.server_test_util import get_default_temp_directory -CHUNK_SIZE = 10_000 +CHUNK_SIZE = 1_000_000 STD = 10.0 FILE_PATTERN = "measurements-*.parquet" LOOKUP_PATH = os.path.join( @@ -79,7 +79,7 @@ def measurement_files(data): def dataset_bytes(files): - """On-disk size of the dataset, used as a common rate denominator for both engines.""" + """On-disk size of the dataset, used as the rate denominator.""" return sum(os.path.getsize(f) for f in files) @@ -88,59 +88,59 @@ def remove_files(path): os.remove(f) -def start_dask_cluster(args, num_jobs): - dask.config.set( - { - "distributed.scheduler.worker-ttl": "1h", - "distributed.comm.timeouts.connect": "120s", - "distributed.comm.timeouts.tcp": "120s", - "distributed.worker.memory.target": 0.6, - "distributed.worker.memory.spill": 0.7, - "distributed.worker.memory.pause": 0.85, - "distributed.worker.memory.terminate": False, - "temporary-directory": args.dask_scratch, - } - ) - - cluster_args = { - "cores": args.dask_cores, - "memory": args.dask_memory, - "walltime": args.dask_walltime, - "local_directory": args.dask_scratch, - } - if args.dask_queue: - cluster_args["queue"] = args.dask_queue - if args.dask_account: - cluster_args["account"] = args.dask_account - - cluster = PBSCluster(**cluster_args) - client = Client(cluster) - print("scaling dask to {} PBS jobs".format(num_jobs)) - cluster.scale(jobs=num_jobs) - client.wait_for_workers(n_workers=num_jobs) - return client, cluster - - -def time_dask_1trc(trials, file_paths, totalbytes): - """Time the reference dask implementation of the challenge.""" - print(">>> dask 1trc") - - timings = [] - result = None - for _ in range(trials): - start = time.time() - df = dd.read_parquet(file_paths, dtype_backend="pyarrow") - # split_out=1 forces a tree reduction - result = df.groupby("station").agg(["min", "max", "mean"], split_out=1).compute() - result = result.sort_values("station") - end = time.time() - timings.append(end - start) - - tavg = sum(timings) / trials - print("dask Average time = {:.4f} sec".format(tavg)) - print("dask Average rate = {:.4f} GiB/sec".format(totalbytes / tavg / 2**30)) - print(result.head()) - return tavg +# def start_dask_cluster(args, num_jobs): +# dask.config.set( +# { +# "distributed.scheduler.worker-ttl": "1h", +# "distributed.comm.timeouts.connect": "120s", +# "distributed.comm.timeouts.tcp": "120s", +# "distributed.worker.memory.target": 0.6, +# "distributed.worker.memory.spill": 0.7, +# "distributed.worker.memory.pause": 0.85, +# "distributed.worker.memory.terminate": False, +# "temporary-directory": args.dask_scratch, +# } +# ) +# +# cluster_args = { +# "cores": args.dask_cores, +# "memory": args.dask_memory, +# "walltime": args.dask_walltime, +# "local_directory": args.dask_scratch, +# } +# if args.dask_queue: +# cluster_args["queue"] = args.dask_queue +# if args.dask_account: +# cluster_args["account"] = args.dask_account +# +# cluster = PBSCluster(**cluster_args) +# client = Client(cluster) +# print("scaling dask to {} PBS jobs".format(num_jobs)) +# cluster.scale(jobs=num_jobs) +# client.wait_for_workers(n_workers=num_jobs) +# return client, cluster + + +# def time_dask_1trc(trials, file_paths, totalbytes): +# """Time the reference dask implementation of the challenge.""" +# print(">>> dask 1trc") +# +# timings = [] +# result = None +# for _ in range(trials): +# start = time.time() +# df = dd.read_parquet(file_paths, dtype_backend="pyarrow") +# # split_out=1 forces a tree reduction +# result = df.groupby("station").agg(["min", "max", "mean"], split_out=1).compute() +# result = result.sort_values("station") +# end = time.time() +# timings.append(end - start) +# +# tavg = sum(timings) / trials +# print("dask Average time = {:.4f} sec".format(tavg)) +# print("dask Average rate = {:.4f} GiB/sec".format(totalbytes / tavg / 2**30)) +# print(result.head()) +# return tavg def materialize_result_df(station_keys, mins, maxs, means): @@ -191,6 +191,7 @@ def time_ak_1trc(trials, file_paths, totalbytes): tavg = sum(timings) / trials print("arkouda Average time = {:.4f} sec".format(tavg)) print("arkouda Average rate = {:.4f} GiB/sec".format(totalbytes / tavg / 2**30)) + # Pulling the results back to the client is not part of the measured time print(materialize_result_df(*result).head()) return tavg @@ -200,7 +201,7 @@ def create_parser(): parser = argparse.ArgumentParser(description="Measure performance of the 1 trillion row challenge.") parser.add_argument("hostname", help="Hostname of arkouda server") parser.add_argument("port", type=int, help="Port of arkouda server") - parser.add_argument("-n", "--size", type=int, default=10**6, help="Number of rows to compute with") + parser.add_argument("-n", "--size", type=int, default=10**8, help="Number of rows to compute with") parser.add_argument( "-t", "--trials", type=int, default=6, help="Number of times to run the benchmark" ) @@ -214,19 +215,19 @@ def create_parser(): parser.add_argument( "-p", "--path", - default=os.path.join(get_default_temp_directory(), "ak-1trc-test"), + default=os.path.join(get_default_temp_directory(), "1trc-test"), help="Target path for the generated dataset", ) - parser.add_argument("--dask-cores", type=int, default=32, help="Cores per PBS job") - parser.add_argument("--dask-memory", default="400GB", help="Memory per PBS job") - parser.add_argument("--dask-walltime", default="5-00:00:00", help="Walltime per PBS job") - parser.add_argument("--dask-queue", default="", help="PBS queue to submit to") - parser.add_argument("--dask-account", default="", help="PBS account to charge") - parser.add_argument( - "--dask-scratch", - default=os.path.join(get_default_temp_directory(), "dask-scratch"), - help="Worker scratch directory", - ) + # parser.add_argument("--dask-cores", type=int, default=32, help="Cores per PBS job") + # parser.add_argument("--dask-memory", default="400GB", help="Memory per PBS job") + # parser.add_argument("--dask-walltime", default="5-00:00:00", help="Walltime per PBS job") + # parser.add_argument("--dask-queue", default="", help="PBS queue to submit to") + # parser.add_argument("--dask-account", default="", help="PBS account to charge") + # parser.add_argument( + # "--dask-scratch", + # default=os.path.join(get_default_temp_directory(), "dask-scratch"), + # help="Worker scratch directory", + # ) return parser @@ -238,28 +239,27 @@ def create_parser(): ak.connect(args.hostname, args.port) # Use the supplied dataset if there is one, otherwise generate it - data = args.data if args.data else generate_data(args.size, args.path) + data = args.data if args.data else generate_data(args.size, args.path, args.chunk_size) - # Both engines read the same files, so resolve them and their size once file_paths = measurement_files(data) totalbytes = dataset_bytes(file_paths) print("number of trials = ", args.trials) print("number of rows = ", args.size) - # Give dask the same number of jobs as the server has nodes so the two series line up - num_nodes = ak.get_config()["numNodes"] - dask_client, dask_cluster = start_dask_cluster(args, num_nodes) - try: - dask_time = time_dask_1trc(args.trials, file_paths, totalbytes) - finally: - dask_client.close() - dask_cluster.close() + # # Give dask the same number of jobs as the server has nodes so the two series line up + # num_nodes = ak.get_config()["numNodes"] + # dask_client, dask_cluster = start_dask_cluster(args, num_nodes) + # try: + # dask_time = time_dask_1trc(args.trials, file_paths, totalbytes) + # finally: + # dask_client.close() + # dask_cluster.close() arkouda_time = time_ak_1trc(args.trials, file_paths, totalbytes) - if dask_time and arkouda_time: - print("arkouda/dask ratio = {:.2f}x".format(arkouda_time / dask_time)) + # if dask_time and arkouda_time: + # print("arkouda/dask ratio = {:.2f}x".format(arkouda_time / dask_time)) if not args.data: remove_files(args.path) diff --git a/benchmarks/graph_infra/1trc.perfkeys b/benchmarks/graph_infra/1trc.perfkeys index 592e6102b5a..5f16b8ebab8 100644 --- a/benchmarks/graph_infra/1trc.perfkeys +++ b/benchmarks/graph_infra/1trc.perfkeys @@ -1,4 +1,4 @@ arkouda Average time = arkouda Average rate = -dask Average time = -dask Average rate = +# dask Average time = +# dask Average rate = diff --git a/benchmarks/graph_infra/arkouda-1trc.graph b/benchmarks/graph_infra/arkouda-1trc.graph index 15a7708274c..e4b8d53589f 100644 --- a/benchmarks/graph_infra/arkouda-1trc.graph +++ b/benchmarks/graph_infra/arkouda-1trc.graph @@ -1,11 +1,11 @@ perfkeys: arkouda Average rate = -graphkeys: 1TRC GiB/s +graphkeys: Arkouda (GiB/s) files: 1trc.dat -graphtitle: 1TRC Arkouda Performance +graphtitle: 1TRC Performance ylabel: Performance (GiB/s) -perfkeys: dask Average rate = -graphkeys: 1TRC GiB/s +perfkeys: arkouda Average time = +graphkeys: Arkouda (sec) files: 1trc.dat -graphtitle: 1TRC Dask Performance -ylabel: Performance (GiB/s) +graphtitle: 1TRC Time +ylabel: Time (seconds) From 7a6c18c802b08c144219252815052d5ce18d3189 Mon Sep 17 00:00:00 2001 From: Eric Vo Date: Wed, 19 Aug 2026 03:44:18 -0500 Subject: [PATCH 5/7] Remove unrelated benchmark changes --- benchmarks/run_benchmarks.py | 81 ++++++++++++++---------------------- 1 file changed, 31 insertions(+), 50 deletions(-) diff --git a/benchmarks/run_benchmarks.py b/benchmarks/run_benchmarks.py index b54dc73a960..e121e160b6c 100755 --- a/benchmarks/run_benchmarks.py +++ b/benchmarks/run_benchmarks.py @@ -158,12 +158,6 @@ def create_parser(): ) parser.add_argument("-sp", "--server-port", default="5555", help="Port number to use for the server") parser.add_argument("--server-args", action="append", help="Additional server arguments") - parser.add_argument( - "--node-counts", - default="", - help="Comma separated list of node counts to sweep, e.g. '16,32,64,128'. Each count is " - "recorded as its own configuration so that all counts share a graph", - ) parser.add_argument("--numtrials", default=1, type=int, help="Number of trials to run") parser.add_argument( "benchmarks", @@ -216,59 +210,46 @@ def main(): parser = create_parser() args, client_args = parser.parse_known_args() args.graph_dir = args.graph_dir or os.path.join(args.dat_dir, "html") + config_dat_dir = os.path.join(args.dat_dir, args.description) run_isolated = bool(args.isolated) - node_counts = [n.strip() for n in args.node_counts.split(",") if n.strip()] - if node_counts: - # One configuration per node count puts every count on the same graph as its own series - sweep = [(n, "{}-nodes".format(n)) for n in node_counts] - args.configs = args.configs or ",".join("{}:v".format(desc) for _, desc in sweep) - else: - sweep = [(args.num_locales, args.description)] - - args.benchmarks = args.benchmarks or BENCHMARKS - - for num_locales, description in sweep: - args.num_locales = num_locales - config_dat_dir = os.path.join(args.dat_dir, description) + if args.save_data or args.gen_graphs: + os.makedirs(config_dat_dir, exist_ok=True) - if args.save_data or args.gen_graphs: - os.makedirs(config_dat_dir, exist_ok=True) + if not run_isolated: + _my_start_server(args) - if not run_isolated: + args.benchmarks = args.benchmarks or BENCHMARKS + for benchmark in args.benchmarks: + if run_isolated: _my_start_server(args) - for benchmark in args.benchmarks: - if run_isolated: - _my_start_server(args) - - for trial in range(args.numtrials): - benchmark_py = os.path.join(benchmark_dir, "{}.py".format(benchmark)) - out = run_client(benchmark_py, client_args) - if args.save_data or args.gen_graphs: - add_to_dat(benchmark, out, config_dat_dir, args.graph_infra) - print(out) - - if run_isolated: - stop_arkouda_server() + for trial in range(args.numtrials): + benchmark_py = os.path.join(benchmark_dir, "{}.py".format(benchmark)) + out = run_client(benchmark_py, client_args) + if args.save_data or args.gen_graphs: + add_to_dat(benchmark, out, config_dat_dir, args.graph_infra) + print(out) - if not run_isolated: + if run_isolated: stop_arkouda_server() - if args.save_data or args.gen_graphs: - comp_file = os.getenv("ARKOUDA_PRINT_PASSES_FILE", "") - if os.path.isfile(comp_file): - with open(comp_file, "r") as f: - out = f.read() - add_to_dat("comp-time", out, config_dat_dir, args.graph_infra) - emitted_code_file = os.getenv("ARKOUDA_EMITTED_CODE_SIZE_FILE", "") - if os.path.isfile(emitted_code_file): - with open(emitted_code_file, "r") as f: - out = f.read() - add_to_dat("emitted-code-size", out, config_dat_dir, args.graph_infra) - - if args.gen_graphs: - generate_graphs(args) + if not run_isolated: + stop_arkouda_server() + + if args.save_data or args.gen_graphs: + comp_file = os.getenv("ARKOUDA_PRINT_PASSES_FILE", "") + if os.path.isfile(comp_file): + with open(comp_file, "r") as f: + out = f.read() + add_to_dat("comp-time", out, config_dat_dir, args.graph_infra) + emitted_code_file = os.getenv("ARKOUDA_EMITTED_CODE_SIZE_FILE", "") + if os.path.isfile(emitted_code_file): + with open(emitted_code_file, "r") as f: + out = f.read() + add_to_dat("emitted-code-size", out, config_dat_dir, args.graph_infra) + if args.gen_graphs: + generate_graphs(args) if __name__ == "__main__": From 0271e35a742d662fa90c84a6706f972a0aef92f1 Mon Sep 17 00:00:00 2001 From: Eric Vo Date: Wed, 19 Aug 2026 10:53:28 -0700 Subject: [PATCH 6/7] Add correctness test + cleanup files before and after generation --- benchmarks/1trc.py | 131 ++++++++++++++++++++++++++++++++------------- 1 file changed, 94 insertions(+), 37 deletions(-) diff --git a/benchmarks/1trc.py b/benchmarks/1trc.py index 6da7f8c9c73..3edc7f7e5ba 100644 --- a/benchmarks/1trc.py +++ b/benchmarks/1trc.py @@ -4,6 +4,7 @@ import os import time +from contextlib import contextmanager, nullcontext from glob import glob from math import ceil @@ -20,7 +21,7 @@ from server_util.test.server_test_util import get_default_temp_directory -CHUNK_SIZE = 1_000_000 +CORRECTNESS_SIZE = 10**4 STD = 10.0 FILE_PATTERN = "measurements-*.parquet" LOOKUP_PATH = os.path.join( @@ -62,13 +63,19 @@ def generate_chunk(partition_idx, chunksize, std, lookup_df, out_dir="."): df.to_parquet(filename, engine="pyarrow") -def generate_data(size, out_dir, chunksize=CHUNK_SIZE): - """Write ``size`` rows of measurements to parquet files and return their directory.""" +@contextmanager +def generate_data(size, out_dir, chunksize): + """Write ``size`` rows of measurements to parquet files and yield their directory. + """ os.makedirs(out_dir, exist_ok=True) + remove_files(out_dir) lookup_df = load_lookup() for i in range(ceil(size / chunksize)): generate_chunk(i, chunksize, STD, lookup_df, out_dir) - return out_dir + try: + yield out_dir + finally: + remove_files(out_dir) def measurement_files(data): @@ -155,6 +162,27 @@ def materialize_result_df(station_keys, mins, maxs, means): return result_df +def ak_1trc(file_paths): + """Compute the per-station min/max/mean with arkouda. + + Returns the result and the intermediates, so a caller can free the server-side arrays + outside of a timed region. + """ + columns = ak.read(file_paths) + stations = columns["station"] + measures = columns["measure"] + + order = ak.argsort(stations) + stations = stations[order] + measures = measures[order] + + grouped = ak.GroupBy(stations, assume_sorted=True) + station_keys, mins = grouped.min(measures) + _, maxs = grouped.max(measures) + _, means = grouped.mean(measures) + return (station_keys, mins, maxs, means), (columns, stations, measures, order, grouped) + + def time_ak_1trc(trials, file_paths, totalbytes): """Time the arkouda implementation of the challenge.""" print(">>> arkouda 1trc") @@ -169,24 +197,12 @@ def time_ak_1trc(trials, file_paths, totalbytes): result = None for _ in range(trials): start = time.time() - columns = ak.read(file_paths) - stations = columns["station"] - measures = columns["measure"] - - order = ak.argsort(stations) - stations = stations[order] - measures = measures[order] - - grouped = ak.GroupBy(stations, assume_sorted=True) - station_keys, mins = grouped.min(measures) - _, maxs = grouped.max(measures) - _, means = grouped.mean(measures) + result, intermediates = ak_1trc(file_paths) end = time.time() timings.append(end - start) - result = (station_keys, mins, maxs, means) # Release the per-trial server arrays outside the timed region - del columns, stations, measures, order, grouped + del intermediates tavg = sum(timings) / trials print("arkouda Average time = {:.4f} sec".format(tavg)) @@ -197,6 +213,30 @@ def time_ak_1trc(trials, file_paths, totalbytes): return tavg +def check_correctness(path): + """Run the challenge on a small generated dataset and compare against pandas.""" + data_dir = os.path.join(path, "correctness") + # Two chunks so the multi-file read path is exercised + with generate_data(CORRECTNESS_SIZE, data_dir, chunksize=CORRECTNESS_SIZE // 2) as data: + file_paths = measurement_files(data) + + result, _ = ak_1trc(file_paths) + ak_result = materialize_result_df(*result) + + expected = ( + pd.concat([pd.read_parquet(f) for f in file_paths]) + .groupby("station")["measure"] + .agg(["min", "max", "mean"]) + .sort_index() + ) + + assert list(ak_result.index) == list(expected.index), "station keys do not match pandas" + for stat in ("min", "max", "mean"): + assert np.allclose( + ak_result[("measure", stat)].to_numpy(), expected[stat].to_numpy() + ), "{} does not match pandas".format(stat) + + def create_parser(): parser = argparse.ArgumentParser(description="Measure performance of the 1 trillion row challenge.") parser.add_argument("hostname", help="Hostname of arkouda server") @@ -218,6 +258,19 @@ def create_parser(): default=os.path.join(get_default_temp_directory(), "1trc-test"), help="Target path for the generated dataset", ) + parser.add_argument( + "-c", + "--chunk-size", + type=int, + default=1_000_000, + help="Number of rows to write to each generated parquet file", + ) + parser.add_argument( + "--correctness-only", + default=False, + action="store_true", + help="Only check correctness, not performance.", + ) # parser.add_argument("--dask-cores", type=int, default=32, help="Cores per PBS job") # parser.add_argument("--dask-memory", default="400GB", help="Memory per PBS job") # parser.add_argument("--dask-walltime", default="5-00:00:00", help="Walltime per PBS job") @@ -238,30 +291,34 @@ def create_parser(): setattr(ak, "verbose", False) ak.connect(args.hostname, args.port) - # Use the supplied dataset if there is one, otherwise generate it - data = args.data if args.data else generate_data(args.size, args.path, args.chunk_size) + if args.correctness_only: + check_correctness(args.path) + sys.exit(0) - file_paths = measurement_files(data) - totalbytes = dataset_bytes(file_paths) + # Use the supplied dataset if there is one, otherwise generate one that is cleaned up after + dataset = ( + nullcontext(args.data) if args.data else generate_data(args.size, args.path, args.chunk_size) + ) - print("number of trials = ", args.trials) - print("number of rows = ", args.size) + with dataset as data: + file_paths = measurement_files(data) + totalbytes = dataset_bytes(file_paths) - # # Give dask the same number of jobs as the server has nodes so the two series line up - # num_nodes = ak.get_config()["numNodes"] - # dask_client, dask_cluster = start_dask_cluster(args, num_nodes) - # try: - # dask_time = time_dask_1trc(args.trials, file_paths, totalbytes) - # finally: - # dask_client.close() - # dask_cluster.close() + print("number of trials = ", args.trials) + print("number of rows = ", args.size) - arkouda_time = time_ak_1trc(args.trials, file_paths, totalbytes) + # # Give dask the same number of jobs as the server has nodes so the two series line up + # num_nodes = ak.get_config()["numNodes"] + # dask_client, dask_cluster = start_dask_cluster(args, num_nodes) + # try: + # dask_time = time_dask_1trc(args.trials, file_paths, totalbytes) + # finally: + # dask_client.close() + # dask_cluster.close() - # if dask_time and arkouda_time: - # print("arkouda/dask ratio = {:.2f}x".format(arkouda_time / dask_time)) + arkouda_time = time_ak_1trc(args.trials, file_paths, totalbytes) - if not args.data: - remove_files(args.path) + # if dask_time and arkouda_time: + # print("arkouda/dask ratio = {:.2f}x".format(arkouda_time / dask_time)) sys.exit(0) From bbe6440c223eaea28284d340145b0372a51fa8a1 Mon Sep 17 00:00:00 2001 From: Eric Vo Date: Wed, 19 Aug 2026 11:38:20 -0700 Subject: [PATCH 7/7] Format --- benchmarks/1trc.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/benchmarks/1trc.py b/benchmarks/1trc.py index 3edc7f7e5ba..0cfa832f5d8 100644 --- a/benchmarks/1trc.py +++ b/benchmarks/1trc.py @@ -65,8 +65,7 @@ def generate_chunk(partition_idx, chunksize, std, lookup_df, out_dir="."): @contextmanager def generate_data(size, out_dir, chunksize): - """Write ``size`` rows of measurements to parquet files and yield their directory. - """ + """Write ``size`` rows of measurements to parquet files and yield their directory.""" os.makedirs(out_dir, exist_ok=True) remove_files(out_dir) lookup_df = load_lookup() @@ -232,9 +231,9 @@ def check_correctness(path): assert list(ak_result.index) == list(expected.index), "station keys do not match pandas" for stat in ("min", "max", "mean"): - assert np.allclose( - ak_result[("measure", stat)].to_numpy(), expected[stat].to_numpy() - ), "{} does not match pandas".format(stat) + assert np.allclose(ak_result[("measure", stat)].to_numpy(), expected[stat].to_numpy()), ( + "{} does not match pandas".format(stat) + ) def create_parser():