pythonasync-awaitpython-asynciopython-contextvars

Using contextvar to keep track of async loop in Python


I am trying to work with a simple async example in Python, largely following this excellent answer here.

My goal is to set up a context variable and keep track of the series of calls by continuously appending to it. I know that context variables can be accessed with the .get() method and their values altered with the .set() method. In the below case however, the variable doesn't get modified despite the series of calls to the function sum() that is apparent from the console.

Edit: Based on Michael Butscher's comment below I replaced the original context variable (which was a string) with a list: output_list and modified the list iteratively using .append(). This now does enable me to view the final output but not the intermediate ones in the individual sum() methods.

Full code:

import asyncio
import contextvars
import time

output_list = contextvars.ContextVar('output_list', default=list())

async def sleep():
    print(f'Time: {time.time() - start:.2f}')
    await asyncio.sleep(1)

async def sum(name, numbers):
    total = 0
    for number in numbers:
        print(f'Task {name}: Computing {total}+{number}')
        await sleep()
        total += number
    output_list.set(output_list.get().append(f"{name}"))
    print(f'Task {name}: Sum = {total}\n')
    print(f'Partial output from task {name}:', output_list.get())

start = time.time()

loop = asyncio.get_event_loop()
tasks = [
    loop.create_task(sum("A", [1, 2])),
    loop.create_task(sum("B", [1, 2, 3])),
]
loop.run_until_complete(asyncio.wait(tasks))
loop.close()

end = time.time()
print(f'Time: {end-start:.2f} sec')
print("Final output_str =", output_list.get())

How can I iteratively follow the expansion of the context variable list output_list?

My desired console output is:

Task A: Computing 0+1
Time: 0.00
Task B: Computing 0+1
Time: 0.00
Task A: Computing 1+2
Time: 1.02
Task B: Computing 1+2
Time: 1.02
Task A: Sum = 3

Partial output from task A: ['A']
Task B: Computing 3+3
Time: 2.02
Task B: Sum = 6

Partial output from task B: ['A', 'B']
Time: 3.03 sec
Final output_str = ['A', 'B']

Instead, I am getting:

Task A: Computing 0+1
Time: 0.00
Task B: Computing 0+1
Time: 0.00
Task A: Computing 1+2
Time: 1.02
Task B: Computing 1+2
Time: 1.02
Task A: Sum = 3

Partial output from task A: None
Task B: Computing 3+3
Time: 2.02
Task B: Sum = 6

Partial output from task B: None
Time: 3.03 sec
Final output_str = ['A', 'B']

Solution

  • According to the asyncio documentation:

    Tasks support the contextvars module. When a Task is created it copies the current context and later runs its coroutine in the copied context.

    So, if you declare at the top of your program cvar = contextvars.ContextVar('cvar', default='x'), when you create a task, this will copy the current context, and if you modify cvar whithin it will just affect the copy but no the original context. That's the main reason why you got ''(empty string) at your final output.

    To achieve the "track" you want you must use a global variable in order to modify it anywhere. But if you want to play around with asyncio and contextvars to see how it works, see the example below:

    import asyncio
    import contextvars
    import time
    
    output = contextvars.ContextVar('output', default='No changes at all') 
    
    async def sleep():
        print(f'Time: {time.time() - start:.2f}')
        await asyncio.sleep(1)
    
    async def sum(name, numbers):
        total = 0
        for number in numbers:
            print(f'Task {name}: Computing {total}+{number}')
            await sleep()
            total += number
            output.set(output.get()+name) #Here we modify the respective context
        print(f'Task {name}: Sum = {total}\n')
        print(f'Partial output from task {name}:', output.get())
        return output.get() #Here we return the variable modified
    start = time.time()
    
    # main() will have its own copy of the context
    async def main():
        output.set('Changed - ') # Change output var in this function context
        # task1 and task2 will copy this context (In this contect output=='Changed - ')
        task1 = asyncio.create_task(sum("A", [1, 2])) #This task has its own copy of the context of main()
        task2 = asyncio.create_task(sum("B", [1, 2, 3])) #This task also has its own copy of the context of main()
        done, pending = await asyncio.wait({task1,task2})
        resultTask1 = task1.result() # get the value of return of task1
        resultTask2 = task2.result() # get the value of return of task1
        print('Result1: ', resultTask1)
        print('Result2: ', resultTask2)
        print('Variable output in main(): ',output.get()) # However, output in main() is sitill 'Changed - '
        output.set(output.get()+'/'+resultTask1+'/'+resultTask2) #Modify the var in this context
        print('Variable modified in main(): ', output.get())
        return output.get() #Return modified value
    
    x = asyncio.run(main()) # Assign the return value to x
    
    end = time.time()
    print(f'Time: {end-start:.2f} sec')
    print("Final output (without changes) =", output.get())
    output.set(x)
    print("Final output (changed) =", output.get())
    
    ##### OUTPUT #####
    # Time: 0.00
    # Task B: Computing 0+1
    # Time: 0.00
    # Task A: Computing 1+2
    # Time: 1.01
    # Task B: Computing 1+2
    # Time: 1.01
    # Task A: Sum = 3
    
    # Partial output from task A: Changed - AA
    # Task B: Computing 3+3
    # Time: 2.02
    # Task B: Sum = 6
    
    # Partial output from task B: Changed - BBB
    # Result1:  Changed - AA
    # Result2:  Changed - BBB
    # Variable output in main():  Changed -
    # Variable modified in main():  Changed - /Changed - AA/Changed - BBB
    # Time: 3.03 sec
    # Final output (without changes) = No changes at all
    # Final output (changed) = Changed - /Changed - AA/Changed - BBB
    
    

    As you can see, it is impossible to modify the same variable at the same time. While task1 is modifying its copy, task2 is modifying its copy too.