-
Notifications
You must be signed in to change notification settings - Fork 184
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Re-introduce Proper State Management (#306)
- Loading branch information
1 parent
27244ea
commit 36adfa3
Showing
17 changed files
with
895 additions
and
553 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,14 @@ | ||
from llama_deploy import deploy_core, ControlPlaneConfig, SimpleMessageQueueConfig | ||
|
||
|
||
async def main(): | ||
await deploy_core( | ||
ControlPlaneConfig(), | ||
SimpleMessageQueueConfig(), | ||
) | ||
|
||
|
||
if __name__ == "__main__": | ||
import asyncio | ||
|
||
asyncio.run(main()) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,45 @@ | ||
import asyncio | ||
from llama_index.core.workflow import ( | ||
Context, | ||
Workflow, | ||
StartEvent, | ||
StopEvent, | ||
step, | ||
) | ||
from llama_deploy import deploy_workflow, ControlPlaneConfig, WorkflowServiceConfig | ||
|
||
|
||
class SessionWorkflow(Workflow): | ||
@step() | ||
async def step_1(self, ctx: Context, ev: StartEvent) -> StopEvent: | ||
cur_val = await ctx.get("count", default=0) | ||
await ctx.set("count", cur_val + 1) | ||
|
||
return StopEvent(result=cur_val + 1) | ||
|
||
|
||
session_workflow = SessionWorkflow(timeout=10) | ||
|
||
|
||
async def main(): | ||
# sanity check | ||
result = await session_workflow.run(arg1="hello_world") | ||
assert result == 1, "Sanity check failed" | ||
|
||
outer_task = asyncio.create_task( | ||
deploy_workflow( | ||
session_workflow, | ||
WorkflowServiceConfig( | ||
host="127.0.0.1", | ||
port=8002, | ||
service_name="session_workflow", | ||
), | ||
ControlPlaneConfig(), | ||
) | ||
) | ||
|
||
await asyncio.gather(outer_task) | ||
|
||
|
||
if __name__ == "__main__": | ||
asyncio.run(main()) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,25 @@ | ||
#!/bin/bash | ||
|
||
# Kill any previously running scripts | ||
echo "Killing any previously running scripts" | ||
pkill -f "launch_" | ||
|
||
# Wait for processes to terminate | ||
sleep 2 | ||
|
||
set -e | ||
|
||
echo "Launching core" | ||
python ./launch_core.py & | ||
sleep 5 | ||
|
||
echo "Launching workflow" | ||
python ./launch_workflow.py & | ||
sleep 5 | ||
|
||
echo "Running client tests" | ||
python ./test_run_client.py | ||
|
||
# Kill any previously running scripts | ||
echo "Killing any previously running scripts" | ||
pkill -f "launch_" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,67 @@ | ||
from llama_deploy import AsyncLlamaDeployClient, ControlPlaneConfig, LlamaDeployClient | ||
|
||
|
||
def test_run_client(): | ||
client = LlamaDeployClient(ControlPlaneConfig(), timeout=10) | ||
|
||
# sanity check | ||
sessions = client.list_sessions() | ||
assert len(sessions) == 0, "Sessions list is not empty" | ||
|
||
# create a session | ||
session = client.create_session() | ||
|
||
# kick off run | ||
_ = session.run("session_workflow", arg1="hello_world") | ||
|
||
# kick off another run | ||
res = session.run("session_workflow", arg1="hello_world") | ||
|
||
# if the session state is working across runs, | ||
# the count should be 2 | ||
assert res == "2", f"Session state is not working across runs, result was {res}" | ||
|
||
# delete the session | ||
client.delete_session(session.session_id) | ||
|
||
# sanity check | ||
sessions = client.list_sessions() | ||
assert len(sessions) == 0, "Sessions list is not empty" | ||
|
||
|
||
async def test_run_client_async(): | ||
client = AsyncLlamaDeployClient(ControlPlaneConfig(), timeout=10) | ||
|
||
# sanity check | ||
sessions = await client.list_sessions() | ||
assert len(sessions) == 0, "Sessions list is not empty" | ||
|
||
# create a session | ||
session = await client.create_session() | ||
|
||
# kick off run | ||
_ = await session.run("session_workflow", arg1="hello_world") | ||
|
||
# kick off another run | ||
res = await session.run("session_workflow", arg1="hello_world") | ||
|
||
# if the session state is working across runs, | ||
# the count should be 2 | ||
assert res == "2", f"Session state is not working across runs, result was {res}" | ||
|
||
# delete the session | ||
await client.delete_session(session.session_id) | ||
|
||
# sanity check | ||
sessions = await client.list_sessions() | ||
assert len(sessions) == 0, "Sessions list is not empty" | ||
|
||
|
||
if __name__ == "__main__": | ||
import asyncio | ||
|
||
print("Running async test") | ||
asyncio.run(test_run_client_async()) | ||
|
||
print("Running sync test") | ||
test_run_client() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.