Remote tmux session on GCP using a Python wrapper

When experimenting with ML models, I use Google Cloud Platform (GCP) with 24G VRAM GPU. Not great, but better than at home. I set up the GCP such that I can restart from cold to a working model within minuutes. Most of this is scripted via a Makefile in https://github.com/andygill/cloud-control.

However, Sometimes you run out of quotes. I often want to:

• ssh into the remote machine; and

• start session using a login shell (so that the correct PATH is set, etc); and

• change into the desired working directory; and

• start a tmux so you can monitor the session, connect if you drop the ssh; and

• set the tmux such that if the command finishes, the session remains.

All reasonable, and quite involved.

So I use the following. In the makefile I have:

REMOTE=ssh ${INSTANCE}
CONDA_BIN=/opt/conda/bin
REMOTE_PYTHON=${CONDA_BIN}/python3.10
PYTHON_REMOTE = ${REMOTE} -t ${REMOTE_PYTHON} scripts/remote.py

boot::
    ${PYTHON_REMOTE} --pwd ${COMFY_UI} --tmux ${COMFY_UI} "python main.py"

This means that we call the remote python, which does the needed remote setup, using this program:

import argparse
import os

def main():
    parser = argparse.ArgumentParser(
        description="Parse arguments for password, tmux session, and commands."
    )

    # Define required arguments
    parser.add_argument("--pwd", required=False, help="Password input")
    parser.add_argument("--tmux", required=False, help="Tmux session name")

    # Capture all remaining arguments
    parser.add_argument(
        "commands", nargs=argparse.REMAINDER, help="Remaining command-line arguments"
    )

    args = parser.parse_args()

    if args.pwd:
        os.chdir(args.pwd)

    formatted_commands = " ".join([f'"{c}"' if " " in c else c for c in args.commands])

    if args.tmux:
        formatted_commands = (
            f"tmux new-session -s {args.tmux} '{formatted_commands} ; bash'"
        )

    os.execvp("bash", ["bash", "--login", "-i", "-c", formatted_commands])

if __name__ == "__main__":
    main()

Quite a bit of layering, but this means that There is one line needed to be in the correct environment, and get things done.

Solve a problem once, script it up, document and amortize!