Workflow State Session vs Task

[Tehsmash]

I've been playing around with llama deploy over the last couple of days and noticed a quirk and I don't know if it's a bug or intended behaviour.

When I issue two runs on the same session, it creates separate "tasks" for the session, but regardless of the input of the second task the output is always the same as the first session.run call.

>>> session = client.create_session()
>>> result = session.run("assistant", input="what colour are bananas?")
>>> result
'{"response":"Bananas are typically yellow when they\'re ripe, but they can also be green or have other colors depending on the stage of ripeness and variety."}'
>>> result = session.run("assistant", input="how big is the sun?")
>>> result
'{"response":"Bananas are typically yellow when they\'re ripe, but they can also be green or have other colors depending on the stage of ripeness and variety."}'

Looking at the logs of my workflow, on the second run it doesn't even execute the steps in the workflow it just returns the final result straight away.

Debugging this, it seems like every task received by the WorkflowService saves the workflow state and context using the session_id, which means a few things:

1. it's not possible to have concurrent tasks within a session (global context is shared)
2. the state of a completed task (completed workflow run) gets reused for new tasks which seems to result in the workflow returning instantly.

Can you confirm if there is a bug or misunderstanding here or not?

[masci]

@Tehsmash that looks definitely like a bug, the two consecutive runs are supposed to be independent and provide different results. I'm creating an issue out of this GH discussion, it would be super helpful if you could add the smallest snippet of code I can use to replicate the issue.

[Tehsmash]

Managed to recreate the issue with a limited amount of code:

core.py:

import asyncio

from llama_deploy import (
    ControlPlaneConfig,
    deploy_core,
)
from llama_deploy.message_queues.simple import SimpleMessageQueueConfig


def run():
    # initialize the message queue config
    message_queue_config = SimpleMessageQueueConfig()
    control_plane_config = ControlPlaneConfig()

    asyncio.run(
        deploy_core(
            control_plane_config=control_plane_config,
            message_queue_config=message_queue_config,
        )
    )

if __name__ == "__main__":
    run()

workflow.py:

import asyncio

from llama_deploy import (
    ControlPlaneConfig,
    WorkflowServiceConfig,
    deploy_workflow,
)
from llama_index.core.workflow import (
    Context,
    StartEvent,
    StopEvent,
    Workflow,
    step,
)


async def slow_get_string() -> str:
    await asyncio.sleep(5)
    return "slowstring"


class ExampleWorkflow(Workflow):

    @step
    async def first_step(self, ctx: Context, ev: StartEvent) -> StopEvent:
        i = ev.input

        response = await slow_get_string()

        return StopEvent(
                result=i
        )


def run():
    control_plane_config = ControlPlaneConfig()

    w = ExampleWorkflow(timeout=60)

    asyncio.run(
        deploy_workflow(
            workflow=w,
            workflow_config=WorkflowServiceConfig(
                host="127.0.0.1",
                port="8082",
                service_name="example",
            ),
            control_plane_config=control_plane_config,
        )
    )

if __name__ == "__main__":
    run()

broken.py:

from llama_deploy import ControlPlaneConfig, LlamaDeployClient

if __name__ == "__main__":
    client = LlamaDeployClient(ControlPlaneConfig())
    session = client.create_session()

    result_raw = session.run("example", input="input1")
    print(result_raw) # expect input1

    result_raw = session.run("example", input="input2")
    print(result_raw) # export input2 but get input1...

I run the control plane using python core.py and start the workflow with python workflow.py.

I exercise the session like:

 $ python broken.py
input1
input1

You can see that I get back the same output, but the workflow should always return the StartEvent input param.

It seems like having any kind of "await" inside of the workflow step causes this behaviour on the second call to the workflow, in my real use case this is a llm.achat call.

Further investigation shows that this has something to do with sharing the ctx between workflows runs. This happens today because the WorkflowService stores and loads the context from the session not from the task ID): https://github.com/run-llama/llama_deploy/blob/main/llama_deploy/services/workflow.py#L250-L274

Removing llama_deploy control plane from the mix I can recreate the same issue by copying the behaviour of the WorkflowService:

Reusing the context "normally":

w = ExampleWorkflow(timeout=60)
handler = w.run(input="input1")
result = await handler
print(result)

result = await w.run(ctx=handler.ctx, input="input2")
print(result)

Working Output:

input1
input2

Adding in the json pickle like the WorkflowService does:

w = ExampleWorkflow(timeout=60)
handler = w.run(input="input1")
result = await handler 
print(result)

s = handler.ctx.to_dict(serializer=JsonPickleSerializer())
ctx = Context.from_dict(w, s, serializer=JsonPickleSerializer())

result = await w.run(ctx=ctx, input="input2")
print(result)

Broken Output:

input1
input1