spawner/src/glicid_spawner/spawner.py

"""GLiCID spawner module."""

from batchspawner import SlurmSpawner
from jinja2 import Environment, PackageLoader, select_autoescape
from traitlets import Unicode

from .micromamba import get_envs
from .resources import CPU, GPU, RAM


class GlicidSpawner(SlurmSpawner):
    """Glicid SLURM Spawner."""

    batchspawner_singleuser_cmd = Unicode(
        'glicid-spawner-singleuser',
        help='Glicid spawner singleuser command.',
    ).tag(config=True)

    def _options_form_default(self) -> str:
        """JupyterHub rendered form template."""
        environment = Environment(
            loader=PackageLoader('glicid_spawner'),
            autoescape=select_autoescape(),
        )
        template = environment.get_template('interactive.html')

        return template.render(
            username=self.user.name,
            python_envs=get_envs(self.user.name),
            cpu_available=CPU,
            ram_available=RAM,
            gpu_available=GPU,
        )

    def options_from_form(self, formdata) -> dict:
        """Export options from form."""
        # Index of user resources choices
        i_cpu = int(formdata['cpu'][0])
        i_ram = int(formdata['ram'][0])
        i_gpu = int(formdata['gpu'][0])

        duration = min(
            CPU[i_cpu].max_duration,
            RAM[i_ram].max_duration,
            GPU[i_gpu].max_duration,
        )

        # Export options
        options = {}

        options['pyenv'] = formdata['python-env'][0]

        options['cpus-per-task'] = CPU[i_cpu].description
        options['mem'] = RAM[i_ram].description.replace(' ', '')

        if i_gpu:
            options['gres'] = 'gpu:' + GPU[i_gpu].description

        options['time'] = f'{duration:02d}:00:00'

        return options