Search code examples
pythonlinuxsshgnu-screen

How to have a python script restart itself in a screen session?


I have a long running python script that runs on a remote machine. I'd like to have it check if it is running in a screen session and restart itself in an attached one if not. That way I will still see the output if I am watching it, but it will continue to run if I get disconnected. So far I've tried several variations of subprocess.run(["screen", "-m", *sys.argv], shell=True) and os.execv("/usr/bin/screen", ["screen", python + " " + sys.argv[0]] + sys.argv[1:]) but none of them seem to work.

What is the right way to do this?


Solution

  • I was finally able to come with the right combination. It turns out that I needed to use the -q option and not pass shell=True.

    import os
    import sys
    import time
    import subprocess
    
    
    def in_screen_session() -> bool:
        return os.environ.get("STY") is not None
    
    
    def in_ssh_session() -> bool:
        return os.environ.get("SSH_CLIENT") is not None
    
    
    def restart_in_screen():
        python = sys.executable
        # This version launches screen, but doesn't run the script.
        # os.execv("/usr/bin/screen", ["-q " + python + " " + sys.argv[0]] + sys.argv[1:])
    
        # This version works!
        subprocess.run(["screen", "-q", python, *sys.argv])
    
    
    print(f"Args: {sys.argv}")
    
    print("Checking for ssh session.")
    if in_ssh_session():
        print("In ssh session. Checking for screen session.")
        time.sleep(2)
        if in_screen_session():
            print("In screen session. Exiting.")
            time.sleep(2)
            sys.exit(0)
        else:
            print("Not in screen session. Restarting in screen session.")
            restart_in_screen()
    
    print("Not in ssh session.")
    if in_screen_session():
        print("In screen session. Exiting.")
    else:
        print("Not in screen session. Exiting.")