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

Question

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?

Answer 1

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.")