https://www.getdaft.io logo
Join Slack
Powered by
Hello! I have a question about flotilla runner mem...
# general
b
Hello! I have a question about flotilla runner memory usage. I'm running a daft query against a ray cluster running on k8s (azure). I'm attempting to read a large parquet table (~1.8T, 50k files) from azure blob store, performing a group by with a count aggregation, and then writing the resulting dataframe to another directory. The script and query looks like:
Copy code
df = daft.read_parquet(
    "<abfs://my_container@my_service_account.dfs.core.windows.net/in_table_path/>"
)
df = df.groupby("group_col").agg(col("count_col").count())
df.write_parquet("<abfs://my_container@my_service_account.dfs.core.windows.net/out_table_path/>")
(it's worth noting that I have daft pinned to 0.5.19 since I currently get an error using 
write_parquet
to azure blob store on later versions. It is the same issue as here: https://github.com/Eventual-Inc/Daft/issues/5336) When I submit a job running this script to my ray cluster, it eventually fails with this error:
Copy code
ray.exceptions.OutOfMemoryError: Task was killed due to the node running low on memory.
Memory on the node (IP: 10.224.2.36, ID: 637c809d2d9dfed8dd5c109e7ff9b81e4c10a024322086029a8f6def) where the task (actor ID: 53ff539b6349c1bee7efc01802000000, name=flotilla-plan-runner:RemoteFlotillaRunner.__init__, pid=1631, memory used=24.82GB) was running was 30.46GB / 32.00GB (0.951874), which exceeds the memory usage threshold of 0.95.
Is this unexpected that the flotilla runner on the head node is consuming so much memory?
Running explain, the physical plan is like:
Copy code
== Physical Plan ==

* Project: col(0: group_col), col(1: count_col) as count_col
|   Resource request = None
|
* Group-By Aggregate: col(0: group_col)
|   Aggregations = sum(col(1: count_col))
|
* Repartition: Hash
|   Num partitions = Some(200)
|   By = col(0: group_col)
|   Num partitions = 200
|
* Group-By Aggregate: col(1: group_col)
|   Aggregations = count(col(0: count_col), Valid)
|
* Project: col(1: count_col), col(0: group_col)
|   Resource request = None
|
* ScanTaskSource:
|   Num Scan Tasks = 50000
|   Estimated Scan Bytes = 1992341647247
|   Pushdowns: {projection: [count_col, group_col]}
...
My head pod is provisioned with 4cpu / 32G, and I have 3 worker pods with 64cpu / 256G
also these are the top 10 memory users on the head node:
Copy code
PID	MEM(GB)	COMMAND
1645	24.48	ray::RemoteFlotillaRunner.get_next_partition
217	0.50	ray-dashboard-DataHead-0 (/home/ray/anaconda3/bin/python3.9 -c from multiprocessing.spawn import spa...
222	0.50	ray-dashboard-ServeHead-0 (/home/ray/anaconda3/bin/python3.9 -c from multiprocessing.spawn import sp...
42	0.40	/home/ray/anaconda3/lib/python3.9/site-packages/ray/core/src/ray/gcs/gcs_server --log_dir=/tmp/ray/s...
134	0.30	/home/ray/anaconda3/bin/python3.9 /home/ray/anaconda3/lib/python3.9/site-packages/ray/dashboard/dash...
221	0.30	ray-dashboard-ReportHead-0 (/home/ray/anaconda3/bin/python3.9 -c from multiprocessing.spawn import s...
218	0.29	ray-dashboard-EventHead-0 (/home/ray/anaconda3/bin/python3.9 -c from multiprocessing.spawn import sp...
223	0.28	ray-dashboard-StateHead-0 (/home/ray/anaconda3/bin/python3.9 -c from multiprocessing.spawn import sp...
705	0.28	/home/ray/anaconda3/lib/python3.9/site-packages/ray/core/src/ray/raylet/raylet --raylet_socket_name=...
219	0.26	ray-dashboard-JobHead-0 (/home/ray/anaconda3/bin/python3.9 -c from multiprocessing.spawn import spaw...
c
I think it's because of the large amounts of intermediate object refs during the repartition.
Copy code
* Repartition: Hash
|   Num partitions = Some(200)
|   By = col(0: group_col)
|   Num partitions = 200
This will basically split each input into 200 partitions, and repartition the data by the group_col. This means that there will be 50000 * 200 = 10M intermediate objects
Each object ref takes up around 3kb of memory on the head node, which means 30gb of memory in total
which kinda checks out
in the later versions of daft we added an optimization called pre shuffle merge to coalesce the inputs before the repartition, which should help with this.
I see in the issue someone proposed a fix https://github.com/Eventual-Inc/Daft/issues/5336#issuecomment-3406517249, does that work for u?
b
gotcha, so the number of input files is a big factor as well. So during a repartition, all object refs are held in memory on the head node, and the pre-shuffle merge reduces the number of object refs? That proposed fix would work for me, yes. I haven't tried the monkey patch yet (I will now) but it sounds like I can also increase the head node memory for now.
c
So during a repartition, all object refs are held in memory on the head node, and the pre-shuffle merge reduces the number of object refs?
Exactly. The preshuffle merge will coalesce partitions up till a 1gb threshold, and only then do the repartition on this one partition. Since you have ~1.8tb of data across 50k files, assuming each file is roughly the same size, that means roughly 40mb per file, and so 24-25 files per gb. pre shuffle merge should be able to do some aggressive merging then, and should help a bit.
Also looks like the azure fs fix is pretty simple so i have it ready here https://github.com/Eventual-Inc/Daft/pull/5412
b
That's good to know about the 1gb threshold - thanks for explaining all that. And thank you for writing the fix!
👍 1
Open in Slack
PreviousNext
Powered by