Tools for making Prefect work better for typical data science workflows.
$ pip install prefect-ds
prefect_ds
is a lightweight wrapper around Prefect
, designed
to make it easier to run workflows I typically encounter when doing data science — especially
tasks related to analyzing large datasets and building models. Specifically, it implements the
following:
A result handler that reads to and writes from Pandas DataFrames. It should be able to handle
any file type Pandas supports, and unlike built-in handlers like LocalResultHandler
requires
the full specification of the file path — this makes it easy to inspect task results, or
use those results in other analysis. It also has support for templating, so task arguments
can be injected into the filenames (useful for things like map
).
>>> import os
>>> os.environ["PREFECT__LOGGING__LEVEL"] = "ERROR"
>>> import pandas as pd
>>> import time
>>> from prefect import task
>>> from prefect_ds.pandas_result_handler import PandasResultHandler
>>> # Note the use of the task argument {id} as a template in the filename
>>> @task(result_handler=PandasResultHandler("data_{id}.csv", "csv"))
... def demo_task(id):
... time.sleep(5)
... return pd.DataFrame({"one": [1, 2, 3], "two": [4, 5, 6]})
Note that in order to use the templating functionality of PandasResultHandler
, you will need
to run your flow using the DSTaskRunner
(see below for more details).
A state handler that implements filename-based checkpointing, in concert with the specialty
result handlers in prefect_ds. It intercepts the state change from Pending
to Running
, runs the read
method of the result handler, and if successful loads
the result of that method as the result of the task, then sets the task to the Success
state. Conversely, if the read
method fails, the task is run as normal and instead the
checkpoint_handler
runs the write
method of the result handler afterwards. Using the
checkpoint_handler
makes it much easier to cache data across Prefect runs — you don't have to
explicitly persist the final flow state between runs, and you don't have to have the cache expire
after a certain amount of time.
This handler combines with DSTaskRunner
, an extension to Prefect's TaskRunner
that implements
the necessary hacks to allow for the templating of task arguments. This templating is required
to handle cases like map
, where without the templating the checkpoint_handler
will read from/write
to the same file for every iteration of the map
.
>>> import contextlib
>>> from prefect import Flow
>>> from prefect.engine import FlowRunner
>>> from prefect_ds.checkpoint_handler import checkpoint_handler
>>> from prefect_ds.task_runner import DSTaskRunner
>>> with Flow("test") as flow:
... output = demo_task(1)
>>> # First, clean up any existing task results
>>> with contextlib.suppress(FileNotFoundError):
... os.remove("data_1.csv")
>>> # Run the flow. Since the result file doesn't exist, will run the task
>>> start = time.time()
>>> state = FlowRunner(flow=flow, task_runner_cls=DSTaskRunner).run(
... task_runner_state_handlers=[checkpoint_handler]
... )
>>> print(f"Took more than 5 seconds: {(time.time() - start) > 5}")
Took more than 5 seconds: True
>>> # Run the flow again. Now that the result file exists, the task is short-circuited
>>> start = time.time()
>>> state = FlowRunner(flow=flow, task_runner_cls=DSTaskRunner).run(
... task_runner_state_handlers=[checkpoint_handler]
... )
>>> print(f"Took less than 1 second: {(time.time() - start) < 1}")
Took less than 1 second: True
An extension to Prefect's FlowRunner
, which will automatically purge the results of upstream
tasks once all of their downstream tasks have been run. This can be useful if your task outputs are
large datasets; by default Prefect stores the results of every task for the duration of the
flow, which can overwhelm your RAM if your results are all things like multi-GB Pandas DataFrames.
While PandasResultHandler
, checkpoint_handler
, and DSTaskRunner
are all designed to be used
together, DSFlowRunner
can have value on its own.
>>> from prefect_ds.flow_runner import DSFlowRunner
>>> @task()
... def generate_data():
... return pd.DataFrame({"one": [1, 2, 3], "two": [4, 5, 6]})
>>> @task()
... def double_data(input_data):
... return input_data * 2
>>> with Flow("test") as flow:
... initial_data = generate_data()
... two_x_data = double_data(initial_data)
... four_x_data = double_data(two_x_data)
>>> state = DSFlowRunner(flow=flow).run(return_tasks=flow.tasks)
>>> state.result[initial_data].result # actual type is prefect_ds.result.PurgedResult
>>> state.result[two_x_data].result # another PurgedResult
>>> state.result[four_x_data].result # final result is correct, even though upstream results are gone
one two
0 4 16
1 8 20
2 12 24
While these components have unit tests covering what I consider to be typical use cases, I have not attempted to comprehensively test every possible interaction with Prefect. As my understanding of Prefect is still relatively immature, I expect there are edge cases where the added functionality of prefect_ds breaks something in Prefect. I think this is especially likely with Prefect Cloud, which I have not done any testing on. If you find such a situation, please feel free to post an issue describing the problem and (ideally) including a minimum reproducible example of the bug.