Files
wehub-resource-sync a0c8464e58
Build Package / build (ubuntu-latest) (push) Failing after 1s
CodeQL / Analyze (python) (push) Failing after 1s
Core Typecheck / core-typecheck (push) Failing after 1s
Linting / lint (push) Failing after 1s
llama-dev tests / test-llama-dev (push) Failing after 1s
Publish Sub-Package to PyPI if Needed / publish_subpackage_if_needed (push) Has been skipped
Sync Docs to Developer Hub / sync-docs (push) Failing after 0s
Build Package / build (windows-latest) (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:26:52 +08:00

554 lines
17 KiB
Plaintext

{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Parallel Execution of Same Event Example\n",
"\n",
"In this example, we'll demonstrate how to use the workflow functionality to achieve similar capabilities while allowing parallel execution of multiple events of the same type. \n",
"By setting the `num_workers` parameter in `@step` decorator, we can control the number of steps executed simultaneously, enabling efficient parallel processing.\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
"# Installing Dependencies\n",
"\n",
"First, we need to install the necessary dependencies:\n",
"\n",
"* LlamaIndex core for most functionalities\n",
"* llama-index-utils-workflow for workflow capabilities"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# %pip install llama-index-core llama-index-utils-workflow -q"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Importing Required Libraries\n",
"After installing the dependencies, we can import the required libraries:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import asyncio\n",
"from llama_index.core.workflow import (\n",
" step,\n",
" Context,\n",
" Workflow,\n",
" Event,\n",
" StartEvent,\n",
" StopEvent,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We will create two workflows: one that can process multiple data items in parallel by using the `@step(num_workers=N)` decorator, and another without setting num_workers, for comparison. \n",
"By using the `num_workers` parameter in the `@step` decorator, we can limit the number of steps executed simultaneously, thus controlling the level of parallelism. This approach is particularly suitable for scenarios that require processing similar tasks while managing resource usage. \n",
"For example, you can execute multiple sub-queries at once, but please note that num_workers cannot be set without limits. It depends on your workload or token limits.\n",
"# Defining Event Types\n",
"We'll define two event types: one for input events to be processed, and another for processing results:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"class ProcessEvent(Event):\n",
" data: str\n",
"\n",
"\n",
"class ResultEvent(Event):\n",
" result: str"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Creating Sequential and Parallel Workflows\n",
"Now, we'll create a SequentialWorkflow and a ParallelWorkflow class that includes three main steps:\n",
"\n",
"- start: Initialize and send multiple parallel events\n",
"- process_data: Process data\n",
"- combine_results: Collect and merge all processing results"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import random\n",
"\n",
"\n",
"class SequentialWorkflow(Workflow):\n",
" @step\n",
" async def start(self, ctx: Context, ev: StartEvent) -> ProcessEvent:\n",
" data_list = [\"A\", \"B\", \"C\"]\n",
" await ctx.store.set(\"num_to_collect\", len(data_list))\n",
" for item in data_list:\n",
" ctx.send_event(ProcessEvent(data=item))\n",
" return None\n",
"\n",
" @step(num_workers=1)\n",
" async def process_data(self, ev: ProcessEvent) -> ResultEvent:\n",
" # Simulate some time-consuming processing\n",
" processing_time = 2 + random.random()\n",
" await asyncio.sleep(processing_time)\n",
" result = f\"Processed: {ev.data}\"\n",
" print(f\"Completed processing: {ev.data}\")\n",
" return ResultEvent(result=result)\n",
"\n",
" @step\n",
" async def combine_results(\n",
" self, ctx: Context, ev: ResultEvent\n",
" ) -> StopEvent | None:\n",
" num_to_collect = await ctx.store.get(\"num_to_collect\")\n",
" results = ctx.collect_events(ev, [ResultEvent] * num_to_collect)\n",
" if results is None:\n",
" return None\n",
"\n",
" combined_result = \", \".join([event.result for event in results])\n",
" return StopEvent(result=combined_result)\n",
"\n",
"\n",
"class ParallelWorkflow(Workflow):\n",
" @step\n",
" async def start(self, ctx: Context, ev: StartEvent) -> ProcessEvent:\n",
" data_list = [\"A\", \"B\", \"C\"]\n",
" await ctx.store.set(\"num_to_collect\", len(data_list))\n",
" for item in data_list:\n",
" ctx.send_event(ProcessEvent(data=item))\n",
" return None\n",
"\n",
" @step(num_workers=3)\n",
" async def process_data(self, ev: ProcessEvent) -> ResultEvent:\n",
" # Simulate some time-consuming processing\n",
" processing_time = 2 + random.random()\n",
" await asyncio.sleep(processing_time)\n",
" result = f\"Processed: {ev.data}\"\n",
" print(f\"Completed processing: {ev.data}\")\n",
" return ResultEvent(result=result)\n",
"\n",
" @step\n",
" async def combine_results(\n",
" self, ctx: Context, ev: ResultEvent\n",
" ) -> StopEvent | None:\n",
" num_to_collect = await ctx.store.get(\"num_to_collect\")\n",
" results = ctx.collect_events(ev, [ResultEvent] * num_to_collect)\n",
" if results is None:\n",
" return None\n",
"\n",
" combined_result = \", \".join([event.result for event in results])\n",
" return StopEvent(result=combined_result)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In these two workflows:\n",
"\n",
"- The start method initializes and sends multiple ProcessEvent.\n",
"- The process_data method uses\n",
" - only the `@step` decorator in SequentialWorkflow\n",
" - uses the `@step(num_workers=3)` decorator in ParallelWorkflow to limit the number of simultaneously executing workers to 3.\n",
"- The combine_results method collects all processing results and merges them.\n",
"\n",
"# Running the Workflow\n",
"Finally, we can create a main function to run our workflow:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Start a sequential workflow without setting num_workers in the step of process_data\n",
"Completed processing: A\n",
"Completed processing: B\n",
"Completed processing: C\n",
"Workflow result: Processed: A, Processed: B, Processed: C\n",
"Time taken: 7.439495086669922 seconds\n",
"------------------------------\n",
"Start a parallel workflow with setting num_workers in the step of process_data\n",
"Completed processing: C\n",
"Completed processing: A\n",
"Completed processing: B\n",
"Workflow result: Processed: C, Processed: A, Processed: B\n",
"Time taken: 2.5881590843200684 seconds\n"
]
}
],
"source": [
"import time\n",
"\n",
"sequential_workflow = SequentialWorkflow()\n",
"\n",
"print(\n",
" \"Start a sequential workflow without setting num_workers in the step of process_data\"\n",
")\n",
"start_time = time.time()\n",
"result = await sequential_workflow.run()\n",
"end_time = time.time()\n",
"print(f\"Workflow result: {result}\")\n",
"print(f\"Time taken: {end_time - start_time} seconds\")\n",
"print(\"-\" * 30)\n",
"\n",
"parallel_workflow = ParallelWorkflow()\n",
"\n",
"print(\n",
" \"Start a parallel workflow with setting num_workers in the step of process_data\"\n",
")\n",
"start_time = time.time()\n",
"result = await parallel_workflow.run()\n",
"end_time = time.time()\n",
"print(f\"Workflow result: {result}\")\n",
"print(f\"Time taken: {end_time - start_time} seconds\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Note\n",
"\n",
"- Without setting `num_workers=1`, it might take a total of 6-9 seconds. By setting `num_workers=3`, the processing occurs in parallel, handling 3 items at a time, and only takes 2-3 seconds total.\n",
"- In ParallelWorkflow, the order of the completed results may differ from the input order, depending on the completion time of the tasks.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This example demonstrates the execution speed with and without using num_workers, and how to implement parallel processing in a workflow. By setting num_workers, we can control the degree of parallelism, which is very useful for scenarios that need to balance performance and resource usage."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Checkpointing"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Checkpointing a parallel execution Workflow like the one defined above is also possible. To do so, we must wrap the `Workflow` with a `WorkflowCheckpointer` object and perfrom the runs with these instances. During the execution of the workflow, checkpoints are stored in this wrapper object and can be used for inspection and as starting points for run executions."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from llama_index.core.workflow.checkpointer import WorkflowCheckpointer"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Completed processing: C\n",
"Completed processing: A\n",
"Completed processing: B\n"
]
},
{
"data": {
"text/plain": [
"'Processed: C, Processed: A, Processed: B'"
]
},
"execution_count": null,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"wflow_ckptr = WorkflowCheckpointer(workflow=parallel_workflow)\n",
"handler = wflow_ckptr.run()\n",
"await handler"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Checkpoints for the above run are stored in the `WorkflowCheckpointer.checkpoints` Dict attribute."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Run: 90812bec-b571-4513-8ad5-aa957ad7d4fb has ['process_data', 'process_data', 'process_data', 'combine_results']\n"
]
}
],
"source": [
"for run_id, ckpts in wflow_ckptr.checkpoints.items():\n",
" print(f\"Run: {run_id} has {[c.last_completed_step for c in ckpts]}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can run from any of the checkpoints stored, using `WorkflowCheckpointer.run_from(checkpoint=...)` method. Let's take the first checkpoint that was stored after the first completion of \"process_data\" and run from it."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Completed processing: B\n",
"Completed processing: A\n"
]
},
{
"data": {
"text/plain": [
"'Processed: C, Processed: B, Processed: A'"
]
},
"execution_count": null,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"ckpt = wflow_ckptr.checkpoints[run_id][0]\n",
"handler = wflow_ckptr.run_from(ckpt)\n",
"await handler"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Invoking a `run_from` or `run` will create a new run entry in the `checkpoints` attribute. In the latest run from the specified checkpoint, we can see that only two more \"process_data\" steps and the final \"combine_results\" steps were left to be completed."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Run: 90812bec-b571-4513-8ad5-aa957ad7d4fb has ['process_data', 'process_data', 'process_data', 'combine_results']\n",
"Run: 4e1d24cd-c672-4ed1-bb5b-b9f1a252abed has ['process_data', 'process_data', 'combine_results']\n"
]
}
],
"source": [
"for run_id, ckpts in wflow_ckptr.checkpoints.items():\n",
" print(f\"Run: {run_id} has {[c.last_completed_step for c in ckpts]}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now, if we use the checkpoint associated with the second completion of \"process_data\" of the same initial run as the starting point, then we should see a new entry that only has two steps: \"process_data\" and \"combine_results\"."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'90812bec-b571-4513-8ad5-aa957ad7d4fb'"
]
},
"execution_count": null,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# get the run_id of the first initial run\n",
"first_run_id = next(iter(wflow_ckptr.checkpoints.keys()))\n",
"first_run_id"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Completed processing: B\n"
]
},
{
"data": {
"text/plain": [
"'Processed: C, Processed: A, Processed: B'"
]
},
"execution_count": null,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"ckpt = wflow_ckptr.checkpoints[first_run_id][\n",
" 1\n",
"] # checkpoint after the second \"process_data\" step\n",
"handler = wflow_ckptr.run_from(ckpt)\n",
"await handler"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Run: 90812bec-b571-4513-8ad5-aa957ad7d4fb has ['process_data', 'process_data', 'process_data', 'combine_results']\n",
"Run: 4e1d24cd-c672-4ed1-bb5b-b9f1a252abed has ['process_data', 'process_data', 'combine_results']\n",
"Run: e4f94fcd-9b78-4e28-8981-e0232d068f6e has ['process_data', 'combine_results']\n"
]
}
],
"source": [
"for run_id, ckpts in wflow_ckptr.checkpoints.items():\n",
" print(f\"Run: {run_id} has {[c.last_completed_step for c in ckpts]}\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Similarly, if we start with the checkpoint for the third completion of \"process_data\" of the initial run, then we should only see the final \"combine_results\" step."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"'Processed: C, Processed: A, Processed: B'"
]
},
"execution_count": null,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"ckpt = wflow_ckptr.checkpoints[first_run_id][\n",
" 2\n",
"] # checkpoint after the third \"process_data\" step\n",
"handler = wflow_ckptr.run_from(ckpt)\n",
"await handler"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Run: 90812bec-b571-4513-8ad5-aa957ad7d4fb has ['process_data', 'process_data', 'process_data', 'combine_results']\n",
"Run: 4e1d24cd-c672-4ed1-bb5b-b9f1a252abed has ['process_data', 'process_data', 'combine_results']\n",
"Run: e4f94fcd-9b78-4e28-8981-e0232d068f6e has ['process_data', 'combine_results']\n",
"Run: c498a1a0-cf4c-4d80-a1e2-a175bb90b66d has ['combine_results']\n"
]
}
],
"source": [
"for run_id, ckpts in wflow_ckptr.checkpoints.items():\n",
" print(f\"Run: {run_id} has {[c.last_completed_step for c in ckpts]}\")"
]
}
],
"metadata": {
"colab": {
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "llama-index-core",
"language": "python",
"name": "llama-index-core"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3"
}
},
"nbformat": 4,
"nbformat_minor": 4
}