is it possible to define resource request on Expre...
# general
y
is it possible to define resource request on Expressions such as:
Copy code
df = df.with_column(
            "thumbnail",
            df["image_url"]
                .str.replace("conductor", "s3")
                .url.download()
                .image.decode()
                .image.resize(32, 32)
                .image.encode(image_format=daft.ImageFormat.JPEG),
        )
If not, how to best control memory usage of such operations?
j
Are you running into OOM issues because of the memory usage?
Also are you using the local runner or the Ray runner?
y
I’m using this on ray, and yes the nodes are running out of memory
I assume due to the extra memory required for loading and operating on the images
j
Yes indeed. We are working on releasing our new streaming execution engine which will provide a much better experience around memory pressure (especially for expensive operations such as these which does big URL downloads and image decoding) At the moment, the best solution is to make finer partitions: https://www.getdaft.io/projects/docs/en/stable/user_guide/poweruser/memory.html Could you also share the plan you’re running? You can access this with
df.explain(True)
. I’m wondering if this is a purely map-style workload without a shuffle
y
Copy code
It should be yes, and I am already making finer partitions indeed. I guess I can try to go even finer.

== Unoptimized Logical Plan ==

* Project: col(id), col(image_width), col(image_height), col(embedding_mofi), col(image_url), col(annotations_url), image_encode(image_resize(image_decode(download(replace(col(image_url),
|     lit("xxx"), lit("s3")))))) as thumbnail
|
* Repartition: Scheme = IntoPartitions
|   Num partitions = 40000
|
* GlobScanOperator
|   Glob paths = [<s3://vedaq-datacuration/datasets/xxx_sa1b/dataset_index/20240214/*.parquet>]
|   Coerce int96 timestamp unit = Nanoseconds
|   IO config = S3 config = { Region name = us-west-2, Endpoint URL = <https://xxx.data.xxx.com>, Max connections = 8, Retry initial backoff ms = 1000, Connect timeout ms = 300000, Read
|     timeout ms = 30000, Max retries = 25, Retry mode = adaptive, Anonymous = false, Use SSL = true, Verify SSL = true, Check hostname SSL = true, Requester pays = false, Force Virtual Addressing
|     = false, Profile Name = xxx-notary }, Azure config = { Use Fabric Endpoint = false, Anonymous = false, Use SSL = true }, GCS config = { Anonymous = false }, HTTP config = { user_agent
|     = daft/0.0.1 }
|   Use multithreading = false
|   File schema = id#Binary, image_width#Int32, image_height#Int32, embedding_mofi#List(Float32), image_url#Utf8, annotations_url#Utf8
|   Partitioning keys = []
|   Output schema = id#Binary, image_width#Int32, image_height#Int32, embedding_mofi#List(Float32), image_url#Utf8, annotations_url#Utf8


== Optimized Logical Plan ==

* Project: col(id), col(image_width), col(image_height), col(embedding_mofi), col(image_url), col(annotations_url), image_encode(image_resize(image_decode(download(replace(col(image_url),
|     lit("xxx"), lit("s3")))))) as thumbnail
|
* Repartition: Scheme = IntoPartitions
|   Num partitions = 40000
|
* GlobScanOperator
|   Glob paths = [<s3://vedaq-datacuration/datasets/xxx_sa1b/dataset_index/20240214/*.parquet>]
|   Coerce int96 timestamp unit = Nanoseconds
|   IO config = S3 config = { Region name = us-west-2, Endpoint URL = <https://xxx.data.xxx.com>, Max connections = 8, Retry initial backoff ms = 1000, Connect timeout ms = 300000, Read
|     timeout ms = 30000, Max retries = 25, Retry mode = adaptive, Anonymous = false, Use SSL = true, Verify SSL = true, Check hostname SSL = true, Requester pays = false, Force Virtual Addressing
|     = false, Profile Name = xxx-notary }, Azure config = { Use Fabric Endpoint = false, Anonymous = false, Use SSL = true }, GCS config = { Anonymous = false }, HTTP config = { user_agent
|     = daft/0.0.1 }
|   Use multithreading = false
|   File schema = id#Binary, image_width#Int32, image_height#Int32, embedding_mofi#List(Float32), image_url#Utf8, annotations_url#Utf8
|   Partitioning keys = []
|   Output schema = id#Binary, image_width#Int32, image_height#Int32, embedding_mofi#List(Float32), image_url#Utf8, annotations_url#Utf8


== Physical Plan ==

* Project: col(id), col(image_width), col(image_height), col(embedding_mofi), col(image_url), col(annotations_url), image_encode(image_resize(image_decode(download(replace(col(image_url),
|     lit("xxx"), lit("s3")))))) as thumbnail
|   Clustering spec = { Num partitions = 40000 }
|
* Flatten
|
* Split: Input num partitions = 200
|   Output num partitions = 40000
|
* TabularScan:
|   Num Scan Tasks = 200
|   Estimated Scan Bytes = 18624301464
|   Clustering spec = { Num partitions = 200 }
|   Schema: {id#Binary, image_width#Int32, image_height#Int32, embedding_mofi#List(Float32), image_url#Utf8, annotations_url#Utf8}
|   Scan Tasks: [
|   {File {<s3://vedaq-datacuration/datasets/xxx_sa1b/dataset_index/20240214/part-00000-25ba56a2-afdc-4530-9cb1-5d54d7f025d5-c000.snappy.parquet}}>
|   {File {<s3://vedaq-datacuration/datasets/xxx_sa1b/dataset_index/20240214/part-00001-25ba56a2-afdc-4530-9cb1-5d54d7f025d5-c000.snappy.parquet}}>
|   {File {<s3://vedaq-datacuration/datasets/xxx_sa1b/dataset_index/20240214/part-00002-25ba56a2-afdc-4530-9cb1-5d54d7f025d5-c000.snappy.parquet}}>
|   ...
|   {File {<s3://vedaq-datacuration/datasets/xxx_sa1b/dataset_index/20240214/part-00197-25ba56a2-afdc-4530-9cb1-5d54d7f025d5-c000.snappy.parquet}}>
|   {File {<s3://vedaq-datacuration/datasets/xxx_sa1b/dataset_index/20240214/part-00198-25ba56a2-afdc-4530-9cb1-5d54d7f025d5-c000.snappy.parquet}}>
|   {File {<s3://vedaq-datacuration/datasets/xxx_sa1b/dataset_index/20240214/part-00199-25ba56a2-afdc-4530-9cb1-5d54d7f025d5-c000.snappy.parquet}}>
|   ]
j
Thanks! You can also use instances that have a higher memory:CPU ratio which will help as well. Cc @Sammy Sidhu @Colin Ho as well for thoughts on streaming execution
This is super helpful, let me talk to the team tomorrow on the progress of our streaming executor. Really appreciate the feedback, keep us posted!
y
got it, thanks @jay. Am I to understand that today Daft doesn't support streaming, meaning each phase will fully finish before moving on to the next phase?
Or is the current execution model also streaming based?
c
Am I to understand that today Daft doesn't support streaming, meaning each phase will fully finish before moving on to the next phase?
Yes, the current executor does not support streaming, but the new one will!
j
Yup! What Colin just said. You can actually try out our streaming engine today already (it runs on a single machine only right now). I’ll let Colin elaborate on how you can try it We’re going to work on integrating this with our distributed Runner as well after releasing it end-of-month for local execution. That should help with your workload significantly, since instead of materializing an entire partition’s worth of URLs, we will instead be streaming that into image decoding/resizing/JPEG encoding.
y
Got it, so that implies that all data in each stage must fit into the cluster's memory (or spill to disk).
j
Got it, so that implies that all data must fit into the cluster’s memory (or spill to disk).
Not quite, if your workload for example has 1000 partitions, but your cluster only has capacity for 8 partitions at a time, we can pipeline that work.
(e.g. if your pipeline ends with a write to Parquet, then we can perform 8 units of work at a time that end with writing to Parquet)
y
ah ok that's great. I guess that's what I meant by streaming.
j
Oh! Yeah we think of that more as pipelining (we already do that today). When we talk about streaming, we are referring more to the work that gets performed within each task: 1. Read a single partition of data 2. Download from URLs in that partition 3. Decode as images 4. Resize the images 5. Encode as JPEG 6. Write the data out Today, these steps happen in a batch synchronous fashion (materializing the entire partition’s worth of data after each step). In the future we will be integrating our streaming engine here so that we can instead process this morsel-at-a-time, so that you get much more stable memory (i.e. fewer URLs are materialized all at once).
The reason why you might be seeing OOMs today is because the materialization of URLs (step 2) and decoding of images (step 3) are very big inflations of memory (potentially 1000x), especially if you do them on entire partitions of data.
Hope that clarifies!
y
yes, thanks @jay
Follow up question. I have a very simple pipeline which reads parquet files, computes a new column, saves parquet files:
Copy code
df = daft.read_parquet(
    "s3://..."
)
# df = df.into_partitions(df.num_partitions() * 2)
df = df.with_column("total_pixels", col("image_width") * col("image_height"))
df.write_parquet("s3://...")
When I uncomment the
into_partitions()
things work great. When I double the partitions, it fails with OOM error. Which is unexpected.
Copy code
== Physical Plan ==

* Project: col(id), col(image_width), col(image_height), col(embedding_mofi), col(image_url), col(annotations_url), col(image_width) * col(image_height) as total_pixels
|   Clustering spec = { Num partitions = 400 }
|
* Flatten
|
* Split: Input num partitions = 200
|   Output num partitions = 400
|
* TabularScan:
|   Num Scan Tasks = 200
|   Estimated Scan Bytes = 18624301464
|   Clustering spec = { Num partitions = 200 }
|   Schema: {id#Binary, image_width#Int32, image_height#Int32, embedding_mofi#List(Float32), image_url#Utf8, annotations_url#Utf8}
|   Scan Tasks: [
|   ...
That's the physical plan. with re-partitioning.
Copy code
== Physical Plan ==

* Project: col(id), col(image_width), col(image_height), col(embedding_mofi), col(image_url), col(annotations_url), col(image_width) * col(image_height) as total_pixels
|   Clustering spec = { Num partitions = 200 }
|
* TabularScan:
|   Num Scan Tasks = 200
|   Estimated Scan Bytes = 18624301464
|   Clustering spec = { Num partitions = 200 }
|   Schema: {id#Binary, image_width#Int32, image_height#Int32, embedding_mofi#List(Float32), image_url#Utf8, annotations_url#Utf8}
|   Scan Tasks: [
|   ...
|   ]
This is witout ^
is there any reason the first would use up more memory?
(aside from minor bookkeeping of more partitions)
Does
Copy code
|
* Flatten
|
* Split: Input num partitions = 200
|   Output num partitions = 400
|
force a materialization of the full dataset?
c
Yeah
Split
will force a materialization of the full dataset. I think what you could try instead to manipulate the number of partitions is set the
scan_tasks_min_size_bytes
and
scan_tasks_max_size_bytes
parameters. Example:
Copy code
from daft.context import set_execution_config
set_execution_config(scan_tasks_min_size_bytes=48 * 1024 * 1024, scan_tasks_max_size_bytes=192 * 1024 * 1024)
These two parameters essentially set a target range of size bytes per scan task (for load balancing reasons), and do merging/splitting of the parquet files accordingly. The default range is 96 - 384mb. If you set them lower you should in theory end up with more partitions
y
ah! ok. good to know.
I think it may be worth to call this out in the docs for
into_partitions
As currently it makes it seem as almost a no-op.
I thought it would just split my partitions as it goes, partition by partition.
c
Yes, that makes sense, will add it to the docs, thanks!
j
I think the intention for
into_partitions
was actually more to do the “split my partitions as it goes”… Would that be useful in this situation?
y
In my case yes, as I’m using it to control my memory usage by making the partitions smaller.
But it kind of blows up with the current behavior
controlling partition sizes with
scan_tasks_min_size_bytes
scan_tasks_max_size_bytes
seems not straightforward at first try, but I need to play more with it.
For context, what I am trying to do end-to-end is:
Copy code
df = daft.read_parquet(
    "s3://..."
)
df = df.with_column(
            "thumbnail",
            df["image_url"]
                .str.replace("conductor", "s3")
                .url.download()
                .image.decode(on_error='null')
                .image.resize(32, 32)
                .image.encode(image_format=daft.ImageFormat.JPEG),
        )
df.write_parquet("s3://...")
I need smaller partitions to control the memory of the image downloading/decoding.
But ideally, I don't load the whole dataframe into memory either (to support
into_partitions
) as it's huge (100M+ rows).
The equivalent of the following in ray data:
Copy code
ray.data.read_parquet("s3://...", override_num_blocks=100_000) \
            .map(get_image_thumbnail_or_null) \
            .write_parquet("s3://...")
override_num_blocks
seems to do the right thing
j
Let me check with the team tomorrow to get their thoughts. I think the current behavior of
.into_partitions
is not ideal, but also in the longer term we really shouldn’t need our users to have to worry about partitioning.
y
👍 thanks!
I tried a little workaround:
Copy code
ds = ray.data.read_parquet("s3://...", override_num_blocks=100_000)
        df = daft.from_ray_dataset(ds)
And it seems to work.
Probably there is some innefficiency in doing so
j
BTW had a deeper conversation with the team today. We’ll work on some fixes here: 1. Fix
into_partitions
behavior 2. Better memory estimations for each task That should give us much better memory stability already. More updates to come in ~1 week!
y
very cool, thank you jay!
j
into_partitions
fix coming through! https://github.com/Eventual-Inc/Daft/pull/3080
🙌 1
y
very cool, thanks! cc @Henry T
🙌 1