spawner/src/glicid_spawner/spawner.py

"""GLiCID spawner module."""

import re

from batchspawner import SlurmSpawner
from traitlets import Unicode, default

from .form import options_form, options_from_form


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

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

    req_qos = Unicode(
        'short',
        help='QoS name to submit job to resource manager',
    ).tag(config=True)

    slurm_job_id_re = Unicode(r'(\d+)(?:;(\w+))?').tag(config=True)

    def parse_job_id(self, output):
        """Parse job id with cluster name support.

        If cluster name is present, `job_id` will be a string
        and suffix with `-M job_cluster` name.

        """
        for job_id, job_cluster in re.findall(self.slurm_job_id_re, output):
            return f'{job_id} -M {job_cluster}' if job_cluster else int(job_id)

        self.log.error(f'GlicidSpawner unable to parse job ID from text: {output}')
        return None

    @default('options_form')
    def _options_form_default(self) -> str:
        """JupyterHub rendered form template."""
        return options_form(self.user.name)

    def options_from_form(self, formdata) -> dict:
        """Export options from form."""
        return options_from_form(formdata)