https://www.getdaft.io logo
Join Slack
Powered by
Hey all, I'm working to understand daft stateless ...
# general
s
Hey all, I'm working to understand daft stateless udfs. I've created this df:
Copy code
df = daft.from_pydict({"a": [1, 2, 3, 4], "b": [4, 3, 2 ,1]}).repartition(2)
with 2 non empty partitions and 2 columns. When I run a dummy udf that simply sleeps for 2 seconds, I observe what seems to be four sequential udf calls, potentially disregarding resource allocation. All four udf calls share the same pid. Here's the dummy udf I run:
Copy code
@daft.udf(return_dtype=daft.DataType.int64())
def f_1s(s: daft.Series):
    process_id = os.getpid()
    task_id = np.random.randint(1, 1000)

    print(f"start ---- pid {process_id} ---- task {task_id}")
    sleep(2)
    print(f"end ------ pid {process_id} ---- task {task_id}")
    return s
the output is:
Copy code
start ---- pid 25236 ---- task 705
end ------ pid 25236 ---- task 705
start ---- pid 25236 ---- task 70
end ------ pid 25236 ---- task 70
start ---- pid 25236 ---- task 166
end ------ pid 25236 ---- task 166
start ---- pid 25236 ---- task 888
end ------ pid 25236 ---- task 888
Can anyone explain why this happens, or suggest how I could enable parallel udf calls? Edit: It appears that this only applies when I retrieve the df with 
show
instead of 
collect
. Is it possible to initiate parallel calls across columns?
j
Hi @Sagi! I ran your example:
Copy code
from time import sleep
import numpy as np
import os
import daft

@daft.udf(return_dtype=daft.DataType.int64())
def f_1s(s: daft.Series):
    process_id = os.getpid()
    task_id = np.random.randint(1, 1000)

    print(f"start ---- pid {process_id} ---- task {task_id}")
    sleep(2)
    print(f"end ------ pid {process_id} ---- task {task_id}")
    return s

df = daft.from_pydict({"a": [1, 2, 3, 4], "b": [4, 3, 2 ,1]}).repartition(2)
df = df.with_column("c", f_1s(df["a"]))

df.show()
And observed only 2 tasks (which is expected):
Copy code
/Users/jaychia/code/Daft/daft/dataframe/dataframe.py:1665: UserWarning: No columns specified for repartition, so doing a random shuffle. If you do not require rebalancing of partitions, you may instead prefer using `df.into_partitions(N)` which is a cheaper operation that avoids shuffling data.
  warnings.warn(
LocalLimit-LocalLimit-FanoutRandom [Stage:2]:   0%|                                                                                                                                     | 0/1 [00:00<?, ?it/sstart ---- pid 65652 ---- task 7500%|                                                                                                                                                    | 0/1 [00:00<?, ?it/s]
end ------ pid 65652 ---- task 750
                                                                                                                                                                                                             start ---- pid 65652 ---- task 7550%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:02<00:00,  2.01s/it]
end ------ pid 65652 ---- task 755
╭───────┬───────┬───────╮
│ a     ┆ b     ┆ c     │                                                                                                                                                                                     
│ ---   ┆ ---   ┆ ---   │
│ Int64 ┆ Int64 ┆ Int64 │
╞═══════╪═══════╪═══════╡
│ 3     ┆ 2     ┆ 3     │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 1     ┆ 4     ┆ 1     │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 2     ┆ 3     ┆ 2     │
├╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 4     ┆ 1     ┆ 4     │
╰───────┴───────┴───────╯

(Showing first 4 of 4 rows)
1. The UDFs run on the same PID because our local runner runs on multithreading. By default we run the same number of threads as there are CPUs. a. This does mean that our stateless UDFs are susceptible to GIL contention. b. In v0.4, our “stateful UDFs” will allow a mode where users can specify a 
.with_concurrency
. This will instead spin up N number of processes (each representing one replica of an instantiated instance of the UDF’s class) and we won’t be limited by the GIL. i. The downside here is of course the overhead of passing data back-and-forth between processes, and also process initialization time. 2. Interestingly, if I run this on swordfish, there is only 1 “task”. This is because Swordfish figures out its own notion of parallelism based on the morsels I believe.
--- On further reading of your examples, I think I understand your question a bit more. You’re probably calling the UDF twice on each column, and asking if Daft can perform these operations in parallel.
Copy code
df = df.with_column("a_prime", f_1s(df["a"]))
df = df.with_column("b_prime", f_1s(df["b"]))
df.show()
Like this? We don’t do this today, and even our normal expressions don’t get called in parallel either I believe (e.g. if you call 
.str.lower()
on both columns, they get executed serially) Moreover, with the GIL and multithreading we’re likely to see contention here as well which would prevent “true” parallelism. We expect stateless UDFs to not be performance-sensitive and more of a convenient way to run Python on the dataframes. --- Things get a little more interesting when you get to our new stateful UDFs. We actually spin up two different process pools, one for each UDF. We can then get pipeline parallelism across these (they can be executing on different partitions in parallel). Here’s an example (you have to run this with the environment variable 
DAFT_ENABLE_ACTOR_POOL_PROJECTIONS=1
)
Copy code
from time import sleep
import numpy as np
import os
import daft

@daft.udf(return_dtype=daft.DataType.int64())
class F:

    def __init__(self, udf_id: str):
        self.udf_id = udf_id

    def __call__(self, s: daft.Series):
        process_id = os.getpid()
        task_id = np.random.randint(1, 1000)

        print(f"start ---- udf_id {self.udf_id} ---- pid {process_id} ---- task {task_id}")
        sleep(2)
        print(f"end ------ udf_id {self.udf_id} ---- pid {process_id} ---- task {task_id}")
        return s

# Set the number of replicas per F
F = F.with_concurrency(1)

F1 = F.with_init_args("F1")
F2 = F.with_init_args("F2")

df = daft.from_pydict({"a": [1, 2, 3, 4], "b": [4, 3, 2 ,1]}).repartition(2)
df = df.with_column("a_prime", F1(df["a"]))
df = df.with_column("b_prime", F2(df["b"]))

df.show()
You’ll notice here when executing that we get pipelined parallelism between 
F1
and 
F2
! Also note the PID is different, because we’re running them on separate process pools.
Copy code
start ---- udf_id F1 ---- pid 75060 ---- task 796
end ------ udf_id F1 ---- pid 75060 ---- task 796
start ---- udf_id F1 ---- pid 75060 ---- task 499
start ---- udf_id F2 ---- pid 75079 ---- task 4
end ------ udf_id F1 ---- pid 75060 ---- task 499
end ------ udf_id F2 ---- pid 75079 ---- task 4
start ---- udf_id F2 ---- pid 75079 ---- task 688
end ------ udf_id F2 ---- pid 75079 ---- task 688
The physical plan also tells a better story of how this works:
Copy code
== Physical Plan ==

* ActorPoolProject:
|   Projection = [col(a), col(b), col(a_prime), pyclass_udf(col(b)) as b_prime]
|   UDFs = [__main__.F]
|   Concurrency = 1
|   Clustering spec = { Num partitions = 2 }
|   Resource request = None
|
* ActorPoolProject:
|   Projection = [col(a), col(b), pyclass_udf(col(a)) as a_prime]
|   UDFs = [__main__.F]
|   Concurrency = 1
|   Clustering spec = { Num partitions = 2 }
|   Resource request = None
|
* ShuffleExchange:
|     Strategy: NaiveFullyMaterializingMapReduce
|     Target Spec: Random(RandomClusteringConfig { num_partitions: 2 })
|     Number of Partitions: 1 → 2
|
* InMemoryScan:
|   Schema = a#Int64, b#Int64,
|   Size bytes = 64,
|   Clustering spec = { Num partitions = 1 }
Note that there are 2 ActorPoolProjects, each will spin up its own process pool and run partitions in a pipelined manner!
s
@jay Amazing, that really clarified some things for me—thanks! I’ll just check to ensure that changing the multiprocessing start method from 'fork' to 'spawn' won’t cause materialized partitions to be duplicated in every process pool’s memory. Do you think daft will support elementwise mapping across multiple columns in the near future to allow parallel execution across columns?
j
I’ll just check to ensure that changing the multiprocessing start method from ‘fork’ to ‘spawn’ won’t cause materialized partitions to be duplicated in every process pool’s memory.
This is a good point… @Kevin Wang we might want to check up on that. We don’t currently have plans for elementwise mapping!
5 Views
Open in Slack
PreviousNext
Powered by