Unable to use Prefect with stateful flows #15806
-
|
I am currently bumping against an apparent limitation with stateful flows. Consider the following example: import numpy as np
import threading
from prefect import flow, task
from prefect.futures import as_completed
from prefect_dask import DaskTaskRunner
from prefect.task_runners import ThreadPoolTaskRunner
@task
def work(arg):
print(f"Starting task in a function with argument: {arg}")
return arg
class MainObj:
def __init__(self):
# This object should not be pickled
self.large_object = np.random.random((10000, 10000))
self.lock = threading.Lock()
self.counter = 0
#@flow(task_runner=ThreadPoolTaskRunner)
@flow(task_runner=DaskTaskRunner())
def run(self):
futures = [work.submit(i) for i in range(3)]
for fut in as_completed(futures):
res = fut.result()
self.process_result(res)
def process_result(self, res):
print(f"Task finished with {res}")
self.counter += res
if __name__ == "__main__":
obj = MainObj()
print(f"Counter is {obj.counter}")
obj.run()
print(f"Counter is {obj.counter}")Here, the flow needs to update the The code works with However, it fails with This is because If so, then it's reasonable that the code doesn't work: the task runner cannot update the
Is there any way to use Prefect in this case? Any insight would be very much appreciated! |
Beta Was this translation helpful? Give feedback.
Replies: 2 comments 6 replies
-
|
would something like this work for your case? example# /// script
# dependencies = [
# "numpy",
# "prefect-dask",
# ]
# ///
import time
import dask.array as da
from prefect_dask import DaskTaskRunner
from prefect import flow, task
from prefect.futures import as_completed
@task
def process_chunk(chunk_data, chunk_id: int):
print(f"Processing chunk {chunk_id} with shape: {chunk_data.shape}")
time.sleep(1) # Simulate work
# Calculate mean absolute value - naturally some will be above 0.5
result = abs(chunk_data).mean().compute()
print(f"Chunk {chunk_id} result: {result:.4f}")
return {"chunk_id": chunk_id, "result": result}
@task
def process_followup(previous_result: float):
print(f"Starting followup for result: {previous_result}")
time.sleep(3) # Longer sleep to make it obvious
print(f"Completed followup for result: {previous_result}")
return previous_result * 0.5
@flow(task_runner=DaskTaskRunner(), log_prints=True)
def process_large_dataset():
large_array = da.random.random((3000, 3000), chunks=(1000, 1000)) * 2 - 1
print(
f"Created array with shape {large_array.shape} and chunks {large_array.chunks}"
)
futures = []
for i in range(large_array.numblocks[0]):
chunk = large_array[i * 1000 : (i + 1) * 1000, :]
futures.append(process_chunk.submit(chunk, i))
# Track all results including followups
all_results = {}
followup_futures = [] # Separate list for followups
print("Starting to process results...")
for future in as_completed(futures):
result = future.result()
all_results[result["chunk_id"]] = result["result"]
print(f"Processed chunk {result['chunk_id']}")
if result["result"] > 0.5:
print(f"Spawning followup task for chunk {result['chunk_id']}")
followup_future = process_followup.submit(result["result"])
followup_futures.append(followup_future)
print(
f"Main processing complete. Waiting for {len(followup_futures)} followup tasks..."
)
# Explicitly wait for followups
for future in as_completed(followup_futures):
result = future.result()
print(f"Completed followup with result: {result:.4f}")
return all_results
if __name__ == "__main__":
print(sorted(process_large_dataset().items()))
here we are:
feel free to let me know if there's some nuance where prefect specifically is getting in the way |
Beta Was this translation helpful? Give feedback.
-
|
Thanks a lot @zzstoatzz for the suggestion! This approach could probably work in principle, but would require a significant refactor. Specifically, it would require moving all the class methods to the inside of The |
Beta Was this translation helpful? Give feedback.
-
|
I recently stumbled upon Covalent, which works in the desired manner: import numpy as np
import threading
import time
import covalent as ct
@ct.electron
def work(arg):
print(f"Starting task in a function with argument: {arg}")
time.sleep(1)
return arg
class MainObj:
def __init__(self):
# This object should not be pickled
self.large_object = np.random.random((10000, 10000))
self.lock = threading.Lock()
self.counter = 0
def run(self):
dask_executor = (
ct.executor.DaskExecutor()
) # Not actually necessary, since Dask is also the default executor
futures = [
ct.dispatch(ct.lattice(work, executor=dask_executor))(i) for i in range(3)
]
for fut in futures:
wresult = ct.get_result(dispatch_id=fut, wait=True)
res = wresult.get_node_result(node_id=0)["output"].get_deserialized()
self.process_result(res)
def process_result(self, res):
print(f"Task finished with {res}")
self.counter += res
if __name__ == "__main__":
obj = MainObj()
print(f"Counter is {obj.counter}")
obj.run()
print(f"Counter is {obj.counter}")The main difference here is that |
Beta Was this translation helpful? Give feedback.
-
|
hmm, I'm not sure I understand your intended use of if the above represents what you're trying to do, what about this? i.e just call your dask work (instance method or not) from a flow docs on this # /// script
# dependencies = [
# "numpy",
# "prefect-dask",
# ]
# ///
import numpy as np
from prefect_dask import DaskTaskRunner
from prefect import flow, task
@task(task_run_name="dask work on input: {arg}")
def work(arg):
print(f"Starting task in a function with argument: {arg}")
return arg
@flow(task_runner=DaskTaskRunner())
def submit_to_dask(n: int = 100):
return [work.submit(1) for i in range(n)]
class MainObj:
def __init__(self, x: int, y: int, initial_counter: int = 0):
self.large_object = np.random.random((x, y))
self.counter = initial_counter
def process_result(self, res):
self.counter += res
# @flow # (if you really need to wrap this method)
def run(self):
for future in submit_to_dask():
self.process_result(future.result())
@flow(log_prints=True) # but instead i'd just define a main function
def setup_experiment(x: int, y: int, initial_counter: int = 0) -> int:
obj = MainObj(x, y, initial_counter)
print(f"Initial counter: {obj.counter}")
obj.run()
print(f"Final counter: {obj.counter}")
return obj.counter
if __name__ == "__main__":
print(setup_experiment(10000, 10000, 0))» PREFECT_LOGGING_LEVEL=CRITICAL uv run proofs/oop_dask.py
100
let me know if thats helpful! |
Beta Was this translation helpful? Give feedback.
-
|
Thanks again for the help! The Your example is interesting and nearly works for what I want to do. Though, in my case, I also need to receive and process the task results as they complete. I didn't have that part in the covalent example since In the end, the limitation of just running one flow at the time seems to be the core limitation for my very specific use-case. Otherwise, this approach would work for similar workflows that don't require so much interactivity. I will look more into |
Beta Was this translation helpful? Give feedback.
-
|
fwiw, in my example the flow is returning the futures (before resolving them), so it's not clear to me why you couldn't do as_completed on what you get from calling the flow, like in the first example i shared in this discussion |
Beta Was this translation helpful? Give feedback.
-
|
Well, the following crashes for some reason: # /// script
# dependencies = [
# "numpy",
# "prefect-dask",
# ]
# ///
import numpy as np
from prefect.futures import as_completed
from prefect_dask import DaskTaskRunner
from prefect import flow, task
@task(task_run_name="dask work on input: {arg}")
def work(arg):
print(f"Starting task in a function with argument: {arg}")
return arg
@flow(task_runner=DaskTaskRunner())
def submit_to_dask(n: int = 100):
return [work.submit(1) for i in range(n)]
class MainObj:
def __init__(self, x: int, y: int, initial_counter: int = 0):
self.large_object = np.random.random((x, y))
self.counter = initial_counter
def process_result(self, res):
print(f"Got result: {res}")
self.counter += res
def run(self):
for future in as_completed(submit_to_dask()):
self.process_result(future.result())
@flow(log_prints=True)
def setup_experiment(x: int, y: int, initial_counter: int = 0) -> int:
obj = MainObj(x, y, initial_counter)
print(f"Initial counter: {obj.counter}")
obj.run()
print(f"Final counter: {obj.counter}")
return obj.counter
if __name__ == "__main__":
print(setup_experiment(10000, 10000, 0))with this message: 18:20:28.143 | ERROR | Flow run 'viridian-platypus' - Encountered exception during execution: AttributeError("'State' object has no attribute '_final_state'")
Traceback (most recent call last):
File "/home/raphael/.../python3.12/site-packages/prefect/flow_engine.py", line 655, in run_context
yield self
File "/home/raphael/.../python3.12/site-packages/prefect/flow_engine.py", line 699, in run_flow_sync
engine.call_flow_fn()
File "/home/raphael/.../python3.12/site-packages/prefect/flow_engine.py", line 678, in call_flow_fn
result = call_with_parameters(self.flow.fn, self.parameters)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/raphael/.../python3.12/site-packages/prefect/utilities/callables.py", line 206, in call_with_parameters
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/home/raphael/4QuantumChemistry/4GOL_GoldRush/4GOL.01_src/workflows/prefect_alt/alt2.py", line 42, in setup_experiment
obj.run()
File "/home/raphael/4QuantumChemistry/4GOL_GoldRush/4GOL.01_src/workflows/prefect_alt/alt2.py", line 34, in run
for future in as_completed(submit_to_dask()):
File "/home/raphael/.../python3.12/site-packages/prefect/futures.py", line 340, in as_completed
done = {f for f in unique_futures if f._final_state}
^^^^^^^^^^^^^^
File "/home/raphael/.../python3.12/site-packages/pydantic/main.py", line 856, in __getattr__
raise AttributeError(f'{type(self).__name__!r} object has no attribute {item!r}')
AttributeError: 'State' object has no attribute '_final_state'and there is not message "Got result..." from Defining the flow as such instead: @flow(task_runner=DaskTaskRunner())
def submit_to_dask(n: int = 100):
return as_completed([work.submit(1) for i in range(n)])and def run(self):
for future in submit_to_dask():
self.process_result(future.result())the execution instead finishes with this message: 18:24:01.404 | ERROR | asyncio - Future exception was never retrieved
future: <Future finished exception=TimeoutError('Timeout')>
TimeoutError: Timeout
18:24:01.874 | INFO | Flow run 'armored-ibis' - Finished in state Completed()
18:24:01.875 | INFO | Flow run 'dark-platypus' - Final counter: 0The counter also indicates that the results haven't been processed. |
Beta Was this translation helpful? Give feedback.
-
|
interesting! that seems like a shortcoming of the dask task runner implementation - we'll take a look at that! thanks for the example |
Beta Was this translation helpful? Give feedback.
hmm, I'm not sure I understand your intended use of
threading.Lockin this contextif the above represents what you're trying to do, what about this? i.e just call your dask work (instance method or not) from a flow
docs on this