Tmux Job Spawner

spawner.py

#!/usr/bin/env python3
"""
Run a list of commands in tmux.
"""

from subprocess import call
from os.path import realpath, isfile
from uuid import uuid4
from datetime import datetime
import os, sys

SESSION_BASE = 'spawner'


def tmux_has_session(name):
    """
    Return true if tmux has a session with the given name.
    """
    return 0 == call(['tmux', 'has-session', '-t', name])


def tmux_new_unique_session():
    """
    Create a completely new session to work in.
    """
    number = 1
    while tmux_has_session(SESSION_BASE + str(number)):
        number += 1
    name = SESSION_BASE + str(number)
    call(['tmux', 'new-session', '-d', '-s', name])
    return name


def tmux_set_remain(session):
    """
    Tell tmux that all windows should remain open even after they die.
    """
    return call(['tmux', 'set-option', '-t', session, 'set-remain-on-exit',
                 'on'])


def tmux_spawn(session, command, output):
    """
    Run this script in tmux, so we can spawn the command.
    """
    filename = realpath(__file__)
    window = '%s:%s' % (session, output)
    uuid = str(uuid4()).replace('-', '')
    return call([
        'tmux',
        'new-window',
        '-t', session,
        '-n', uuid, # name the window after the filename we will output
        'python', filename, command, output, session, uuid
    ])


def run_spawner(filename):
    """
    Load the commands from a file.  Start up a tmux session, and run each
    command in a window within that session.  When done, write the output to
    specified files.
    """
    commands = [c.strip() for c in open(filename)]
    commands = filter(lambda v: v, commands)
    commands = [c.split(':', 1) for c in commands]
    session = tmux_new_unique_session()
    tmux_set_remain(session)
    results = [tmux_spawn(session, c, o) for o, c in commands]
    print('spawned all commands in session "%s"' % session)


def tmux_save_window(filename, session, window):
    """
    This takes the contents of your window, and saves them.

    Only works if this is running within a tmux session.  If not, who knows
    which window it will save, or if it will find one at all???
    """
    buffername = window
    windowtitle = '%s:%s' % (session, window)
    return call([
        'tmux',
        'select-window', '-t', windowtitle, ';',
        'capture-pane', '-S', '-E', '-b', buffername, ';',
        'save-buffer', '-b', buffername, filename, ';',
        'delete-buffer', '-b', buffername,
    ])

def create_filename(filename):
    if not isfile(filename):
        return filename
    else:
        template = filename + '-%d'
        number = 2
        while isfile(template % number):
            number += 1
        return template % number


def spawn(command, output, session, uuid):
    """
    This is called once we're running in tmux.  Just output the command, run it,
    and then save this buffer into the filename we were supposed to use.
    """
    print('$ ' + command)
    start = datetime.now()
    sys.stdout.flush()
    os.system(command)
    runtime = datetime.now() - start
    print('Ran for %s (H:MM:SS.fraction)' % runtime)
    sys.stdout.flush()
    output = create_filename(output)
    tmux_save_window(output, session, uuid)
    input()


if __name__ == '__main__':
    from sys import argv
    if len(argv) == 5:
        spawn(argv[1], argv[2], argv[3], argv[4])
    else:
        run_spawner(argv[1])

Have you ever wished you could run a whole bunch of programs quickly, see their progress in tmux live, but also have their progress written to a file automatically when they complete? Say, for a class like Machine Learning, where you need to run lots of long running programs? Never fear! This script has you covered from now on.

When run with one argument (./spawner.py program-list.txt), this nifty program creates a new tmux session (spawnerN for some number N) and then runs all of the programs specified, each in their own window. When they're finished, it saves their output (being careful never to destroy old output).

Here is an example. Imagine you need to run prog with three different arguments. Here is your list.txt:

ls.txt:ls
echo_hello.txt:echo hello
echo_star.txt:echo *

As you can see, on each line there is a filename, a colon, and then a command. When you run ./spawner.py list.txt, it will create a tmux session (probably spawner1, but whatever it uses, it will tell you). It will run each program in its own window within that session. When each program exits, it will write the output to the filename specified (along with how long it took to run). For example, here is echo_hello.txt:

$ echo hello
hello
Ran for 0:00:00.002664 (H:MM:SS.fraction)

The spawner takes special care (actually, a huge amount of care, check out the code for details) to makes sure that your program arguments are at the top, and the runtime is at the bottom. This won't work perfectly with long output - however many lines of buffering you have Tmux configured for is all you'll be able to get out of this program.

When run with two arguments, this program assumes it's being run in tmux, and that it is being instructed to spawn one of your programs. So please only use one argument with this script!