spawner/render/__main__.py

"""Web Server Gateway Interface with autoreload to render spawner template.

Usage: `python -m render`

"""
import sys
from argparse import ArgumentParser
from pathlib import Path
from traceback import format_exc

from flask import Flask, render_template, request
from glicid_spawner.form import options_from_form
from glicid_spawner.micromamba import MicromambaEnv
from glicid_spawner.resources import CPU, MEMORY, get_folders, gpu_max_duration
from glicid_spawner.slurm import gres, sinfo_from_file
from glicid_spawner.templates import render_template as glicid_template
from livereload import Server

# Dummy username and python environments
USERNAME = 'john-doe'
ENVS = [
    MicromambaEnv('USER', 'foo', f'/{USERNAME}/envs/foo'),
    MicromambaEnv('USER', 'bar', f'/{USERNAME}/envs/bar'),
    MicromambaEnv('GLOBAL', 'baz', '/global/envs/baz'),
]
FOLDERS = get_folders(USERNAME)

# Dummy SLURM config
DATA = Path(__file__).parent / '..' / 'tests' / 'data'
SINFO = sinfo_from_file(
    DATA / 'sinfo.txt', with_states=('idle', 'mixed', 'allocated', 'completing', 'planned')
)

# Single vs. multi-cluster implementation
SLURM_SINGLE_CLUSTER = {'cluster': SINFO.pop('cluster')}
SLURM_MULTI_CLUSTER = SINFO
GPU_SINGLE_CLUSTER = gpu_max_duration(gres(SLURM_SINGLE_CLUSTER))
GPU_MULTI_CLUSTER = gpu_max_duration(gres(SLURM_MULTI_CLUSTER))

# Format dummy options
OPTIONS = {
    'username': USERNAME,
    'folders': FOLDERS,
    'envs': ENVS,
    'cpus': CPU,
    'mems': MEMORY,
    # Multi-cluster by default. See `--single-cluster` flag in CLI for single cluster config.
    'gpus': GPU_MULTI_CLUSTER,
    'sinfo': SLURM_MULTI_CLUSTER,
}

# Flask app
app = Flask(__name__)
app.debug = True


@app.route('/')
def home():
    """Form spawner home page."""
    return render_template(
        'form.html',
        spawner_options_form=glicid_template('spawner_form.jinja', **OPTIONS),
        options=OPTIONS,
    )


@app.route('/submit', methods=['POST'])
def submit():
    """Reformat form data and extract spawner options.

    https://jupyterhub.readthedocs.io/en/stable/reference/spawners.html#spawner-options-from-form

    """
    formdata = dict(request.form.lists())

    # Trying to parse the options from the formdata
    try:
        return render_template(
            'options.html', formdata=formdata, options=options_from_form(formdata)
        )
    except Exception:
        return render_template('options.html', formdata=formdata, err=format_exc())


@app.route('/favicon.ico')
def favicon():
    """Dummy favicon."""
    return ''


def server_autoreload():
    """Start auto-reload server."""
    server = Server(app.wsgi_app)
    server.serve()


def cli(argv=None):
    """Command line interface."""
    parser = ArgumentParser('Spawner form render')
    parser.add_argument(
        '--single-cluster',
        action='store_true',
        help='Toggle SLURM config to single cluster configuration.',
    )

    args, _ = parser.parse_known_args(argv)

    if args.single_cluster:
        OPTIONS['gpus'] = GPU_SINGLE_CLUSTER
        OPTIONS['sinfo'] = SLURM_SINGLE_CLUSTER

    server_autoreload()


if __name__ == '__main__':
    cli(sys.argv)