Search code examples
pythonasynchronouspython-asyncio

Is there a way in Python asyncio to execute first bit of a for loop until a response is received?


Relatively new to asyncio and I need to know if I'm doing something fundamentally wrong. I have a general pattern I want to run in Python that looks like the following:

async def function(index):
    print(f'going to sleep: {index}')
    await asyncio.sleep(1) // some function that takes some time
    print(f'waking up: {index}')

async def main():
    await asyncio.wait([function(i) for i in range(10)])

I would like to call function 10 times, and while awaiting the response from asyncio.sleep(1) I would like to continue onto the next iteration of my loop. However, if a call to asyncio.sleep finishes while attempting to start another iteration of the loop I would like that response to be dealt with.

Currently, if I run this I get the following output:

going to sleep: 4
going to sleep: 8
going to sleep: 0
going to sleep: 5
going to sleep: 1
going to sleep: 2
going to sleep: 6
going to sleep: 9
going to sleep: 7
going to sleep: 3
waking up: 4
waking up: 8
waking up: 0
waking up: 5
waking up: 1
waking up: 2
waking up: 6
waking up: 9
waking up: 7
waking up: 3

I would like the result to be something similar to the following:

going to sleep: 4
going to sleep: 8
going to sleep: 0
going to sleep: 5
going to sleep: 1
going to sleep: 2
going to sleep: 6
waking up: 4
waking up: 8
waking up: 0
going to sleep: 9
going to sleep: 7
going to sleep: 3
waking up: 5
waking up: 1
waking up: 2
waking up: 6
waking up: 9
waking up: 7
waking up: 3

Is this possible with asyncio or am I completely off the mark?

Thanks


Solution

  • The asyncio.wait function can do in practical matters(*) what you want. By default, it will wait for all tasks to be completed before returning, but it can be changed to return the control to the caller after the first task is finished - so that you can take action, create new tasks based on its response, and so on - and then you can wait again with the remaining (and any newly created) tasks.

    (*) However, this does not mean you will see the output you typed in the question - every task will be entered and print "going to sleep ..." before any other task have a chance to do anything else - as soon as the first task yields to the event loop (through its await statement), the next ready task is executed (up to the point it also awaits, and then the next task is called, etc...). The first task will only be resumed when all the tasks had a chance to run their first part, up to the first await.

    Randomizing the sleeping times will make that easier to visualize:

    import random
    import asyncio
    
    
    async def function(index):
        interval = 0.1 + 0.1 * random.randint(0,10)
        print(f'going to sleep: {index} for {interval:.02f}s')
        await asyncio.sleep(interval) # some function that takes some time
        print(f'waking up: {index}')
        return index
    
    async def main():
        tasks = {asyncio.create_task(function(i), name=i) for i in range(10)}
        # All tasks created but none yet started execution, until
        # we run an `await` statement!
        print("starting")
        results = []
        i = 10
        while tasks:
            # Pending tasks are returned in a second set of the "wait" call:
            done, tasks = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
            for task in done:
                results.append(task.result())
                print(f"task {task.get_name()} done!")
            # one can create new tasks here, and just add then to the
            # "tasks" set:
            if len(results) < 5:
                tasks.add(asyncio.create_task(function(i)))
                i += 1
    
        print(f"results: {sorted(results)}")
    
    asyncio.run(main())
    
    

    (Please, when posting questions, do not skip the "boilerplate" lines needed to run the complete example (the import, and the asyncio.run call in this case). They were simple in this case, but in other questions that is often a barrier to tackle the problem.)

    Output:

    $ python  test.py 
    starting
    going to sleep: 0 for 1.10s
    going to sleep: 1 for 0.40s
    going to sleep: 2 for 0.80s
    going to sleep: 3 for 0.10s
    going to sleep: 4 for 0.30s
    going to sleep: 5 for 1.00s
    going to sleep: 6 for 0.60s
    going to sleep: 7 for 0.20s
    going to sleep: 8 for 0.80s
    going to sleep: 9 for 0.80s
    waking up: 3
    task 3 done!
    going to sleep: 10 for 0.60s
    waking up: 7
    task 7 done!
    going to sleep: 11 for 0.70s
    waking up: 4
    task 4 done!
    going to sleep: 12 for 0.70s
    waking up: 1
    task 1 done!
    going to sleep: 13 for 0.60s
    waking up: 6
    task 6 done!
    waking up: 10
    task Task-12 done!
    waking up: 2
    waking up: 8
    waking up: 9
    task 2 done!
    task 9 done!
    task 8 done!
    waking up: 11
    task Task-13 done!
    waking up: 5
    task 5 done!
    waking up: 13
    waking up: 12
    task Task-14 done!
    task Task-15 done!
    waking up: 0
    task 0 done!
    results: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]