Package: prefect-ds
prefect DSTools for making Prefect work better for typical data science workflows.Install$pipinstallprefect-dsUsageprefect_dsis a lightweight wrapper aroundPrefect, 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:PandasResultHandlerA 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 likeLocalResultHandlerrequires
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 likemap).>>>importos>>>os.environ["PREFECT__LOGGING__LEVEL"]="ERROR">>>importpandasaspd>>>importtime>>>fromprefectimporttask>>>fromprefect_ds.pandas_result_handlerimportPandasResultHandler>>># Note the use of the task argument {id} as a template in the filename>>>@task(result_handler=PandasResultHandler("data_{id}.csv","csv"))...defdemo_task(id):...time.sleep(5)...returnpd.DataFrame({"one":[1,2,3],"two":[4,5,6]})Note that in order to use the templating functionality ofPandasResultHandler, you will need
to run your flow using theDSTaskRunner(see below for more details).checkpoint_handlerandDSTaskRunnerA state handler that implements filename-based checkpointing, in concert with the specialty
result handlers in prefect_ds. It intercepts the state change fromPendingtoRunning, runs thereadmethod of the result handler, and if successful loads
the result of that method as the result of the task, then sets the task to theSuccessstate. Conversely, if thereadmethod fails, the task is run as normal and instead thecheckpoint_handlerruns thewritemethod of the result handler afterwards. Using thecheckpoint_handlermakes 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 withDSTaskRunner, an extension to Prefect'sTaskRunnerthat implements
the necessary hacks to allow for the templating of task arguments. This templating is required
to handle cases likemap, where without the templating thecheckpoint_handlerwill read from/write
to the same file for every iteration of themap.>>>importcontextlib>>>fromprefectimportFlow>>>fromprefect.engineimportFlowRunner>>>fromprefect_ds.checkpoint_handlerimportcheckpoint_handler>>>fromprefect_ds.task_runnerimportDSTaskRunner>>>withFlow("test")asflow:...output=demo_task(1)>>># First, clean up any existing task results>>>withcontextlib.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}")Tookmorethan5seconds: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}")Tooklessthan1second:TrueDSFlowRunnerAn extension to Prefect'sFlowRunner, 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.
WhilePandasResultHandler,checkpoint_handler, andDSTaskRunnerare all designed to be used
together,DSFlowRunnercan have value on its own.>>>fromprefect_ds.flow_runnerimportDSFlowRunner>>>@task()...defgenerate_data():...returnpd.DataFrame({"one":[1,2,3],"two":[4,5,6]})>>>@task()...defdouble_data(input_data):...returninput_data*2>>>withFlow("test")asflow:...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 goneonetwo0416182021224CaveatWhile 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.