spawner/src/glicid_spawner/spawner.py

"""GLiCID spawner module."""

import asyncio
import re

from batchspawner import JobStatus, SlurmSpawner
from traitlets import Bool, Integer, Unicode, default

from .form import options_form, options_from_form
from .progress import ElapseTime, get_progress


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

    disable_user_config = Bool(
        True,
        help='Disable per-user configuration of single-user servers.',
    ).tag(config=True)

    notebook_dir = Unicode(
        '/',
        help='Path to the notebook directory for the single-user server.',
    ).tag(config=True)

    @default('default_url')
    def _default_url_default(self) -> str:
        """The URL the single-user server should start in."""
        return '/lab/tree' + self.user_options.get('workdir', '/home/{username}')

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

    req_job_name = Unicode(
        'jupyterhub_glicid',
        help='SLURM job name',
    ).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)

    progress_rate = Integer(
        5, help='Interval in seconds at which progress is polled for messages'
    ).tag(config=True)

    async def progress(self):
        """Progress bar feedback."""
        elapse_time = 0

        while True:
            if self.state_isrunning():
                job_status = JobStatus.RUNNING
                elapse_time += self.progress_rate
            elif self.state_ispending():
                job_status = JobStatus.PENDING
            else:
                job_status = JobStatus.NOTFOUND

            yield get_progress(job_status, elapse_time)

            if elapse_time >= ElapseTime.CONNECT:
                break

            await asyncio.sleep(self.progress_rate)