how to display parallel output results within Jupyter Notebook, using AsyncSSH + IPyWidget?

Question

ChatGPT is running in circles now and it keeps failing at the task: having multiple boxes monitoring in real time the output of a remote python script.

So far, here is the notebook code:

import asyncssh
import asyncio
from ipywidgets import Output, HBox
import traceback

class MySSHClientSession(asyncssh.SSHClientSession):
    def __init__(self, output_widget):
        super().__init__()
        self._output_widget = output_widget

    def data_received(self, data, datatype):
        if datatype == asyncssh.EXTENDED_DATA_STDERR:
            self._output_widget.append_stderr(data)
        else:
            self._output_widget.append_stdout(data)

    def connection_lost(self, exc):
        if exc:
            self._output_widget.append_stderr(f"SSH session error: {exc}")

async def run_remote_command(host, username, password, command, output_widget):
    try:
        async with asyncssh.connect(host, username=username, password=password, known_hosts=None) as conn:
            
            chan,session = await conn.create_session(lambda: MySSHClientSession(output_widget), command)
            await chan.wait_closed()
    except Exception as e:
        output_widget.append_stderr(f"Error connecting to {host}: {str(e)}\n")


async def main():

    host_infos=[parse_creds(i) for i in range(6)]
    cmds=[f"python /scripts/print_hostname.py {P}" for P in range(1,7)]
    outputs = [Output(layout={'border': '1px solid white', 'width': '200px'}) for _ in host_infos]
    
    tasks = [run_remote_command(host_info[0], host_info[1], host_info[2], command, out) for host_info, command, out in zip(host_infos, cmds, outputs)]
    
    display(HBox(outputs))
    await asyncio.gather(*tasks)

# Run the asynchronous function
asyncio.create_task(main())

while troubleshooting, we simplified the code of print_hostname.py to the following:

import time

print("Début du script sur la machine.")

for i in range(5):
    print(f"Étape {i} sur la machine.")
    time.sleep(4)

print("Fin du script sur la machine.")

I don't know what to try anymore. We went from Threads, to pure asyncio, to managing the output in a while True loop.

And I think the while True loop is the key, but I can't figure out how to implement it in the above code?

Answer 1

The real trick is to add flush=True to all the print commands in the remote python script.

For the sake of answering the question, here is the final remote script that displays P times the hostname:

import sys
import socket
import time

def main(P):
    hostname = socket.gethostname()

    print(f"Starting {P} prints...", flush=True)

    for i in range(P):
        print(hostname, flush=True)
        time.sleep(2)

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python print_hostname.py <P>")
        sys.exit(1)

    P = int(sys.argv[1])
    
    main(P)

And actually, the async function does not even need a special factory in that case, so I was able to go back to a simpler code like the following.

def print_message(message, output_index):
    outputs[output_index].append_stdout(message)

async def run_remote_command(host, username, password, command, output_index):
    try:
        async with asyncssh.connect(host, username=username, password=password, known_hosts=None) as conn:
            #async with conn.create_process(f"ping {host}") as process:
            async with conn.create_process(command) as process:

                print_message(f'{host}: CONNECTION OPEN\n', output_index)

                while True: 
                    line = await process.stdout.readline()
                    if not line:
                        break
                        
                    print_message(line, output_index)
                    #await asyncio.sleep(0.5)

            await conn.wait_closed()
               
    except Exception as e:
        print_message(f"Error connecting to {host}: {str(e)}\n", output_index)