Running WebSocket server and separate while loop simultaneously

Question

The following code is for a RC flight controller based around an Raspberry Pi (Zero 2W). The purpose of the code is to receive input controls from my Iphone while also managing a few PID controllers to control the RC planes control surfaces. The script is written in python.

import smbus
from time import sleep
import time
import math
import RPi.GPIO as GPIO
from gpiozero import Servo
from gpiozero.pins.pigpio import PiGPIOFactory
import asyncio
import websockets
import json

#Initialize sensors

async def handler(websockets):
    async for message in websockets:
        print("Raw message recieved:" + str(message))
        try:
            data = json.loads(message)
            pitch = data.get("pitch", 0)
            roll = data.get("roll", 0)
            yaw = data.get("yaw", 0)
            throttle = data.get("throttle", 0)
            
            print(f"Pitch: {pitch}, Roll: {roll}, Yaw: {yaw}, Throttle: {throttle}")
        
        except json.JSONDecodeError as e:
            print("Invalid JSON recieved:" , e)

    
async def WebSocketInitialization():
    async with websockets.serve(handler, "0.0.0.0", ####):
        print("Websocket server started on port ####")
        await asyncio.Future()



async def main():
    #Initialize Servos & Sensors


    #Declearing variables for PID Controllers


        #PID controllers
    while True:
        #Reading Sensors
        
        #Code for all PID Controllers

        #Code for executing output of PID controllers
            

async def AsyncStart():
    asyncio.gather(WebSocketInitialization(), main())
    
asyncio.run(AsyncStart())

If the last line of the code is exchanged for "asyncio.run(WebSocketInitialization())", the WebSocket server and sending information form my iphone works perfectly. If the last line of the code is exchanged for "asyncio.run(main())", the PID controllers works wonderfully. However, I can't get these two to work on the same time.

If I try to run the code pasted above, the PID controllers initializes and runs but the WebSocket server isn't initialized (I never get the "Websocket server started on port ####") message.

I have tried TaskGroup() instead of .gather without success. I also tried to move the WebSocket code inside the PID-controllers while True: loop, but without success.

Answer 1

TL;DR:

you almost got it - all you missed is an await instruction in your AsyncStart function.

Just change it to:

async def AsyncStart():
    await asyncio.gather(WebSocketInitialization(), main())

(of course, I am assuming your main code contains pauses to let the asyncio loop run - if itself doesn't ever need to "await" on anything, be sure to include some lines reading await asyncio.sleep(0) so that the task can switch. On the websockets side, the async for... will fulfill this role)

What is going on:

By default an asyncio Loop will start with one main task, which can create any number of other tasks - but when that main task is done, the still running tasks are cancelled, and the loop stops.

What creates the concurrent tasks is not the the .gather call - it is the call to the WebSocketIntialization and main functions (technically, they are co-routines, but both are awaitables that can be executed in the running async loop). What the asyncio.gather call does is wrap all the arguments into proper tasks, and then become itself an awaitable, which will be complete when all its passed arguments are complete.

Without awaiting the gather call itself, all that happens is that your AsyncStart coroutine ends imediatelly, and both the gather coroutine, and the calls to WebSocketInitialization and main are not awaited - and either cancelled or "fogotten".

In order to get a cleaner pattern than this, Python 3.11 introduced the async.TaskGroup class - it works with the "asynchronous context manager" protocol - the async with statement block, and will just leave the block when all spawned tasks as complete. (It feels more natural, but if any of the spawned tasks raise an exception, it will immediately cancel all the others - while other calls to manage concurrent tasks like asyncio.gather and asyncio.wait give you more control. I suspect this is the reason TaskGroups didn't get more traction.

Anyway, if you feel like, you can rewrite your code as:

async def WebSocketInitialization():
     # leave as it is 
     ...

async def PIDControl():
     # bring your code currently in  `main` here, 
     # along with all other things you want...
     ...

async def main():
    # prefer to use the name `main`  to the co-routine which
    # will orchestrate the others - instead of creating an arbitray
    # name as `AsyncStart` 
    async with asyncio.TaskGroup() as tg:
         tg.create_task(PIDControl())
         tg.create_task(WebSocketInitialization())

And that is it: both tasks will be awaited at the end of the async with block. If any of them is terminated with an error (the WebSocketIn... could choose to raise an exception if it receives a command to stop everything, for example), the other will be cancelled, and the asyncio loop stopped.