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 .micromamba import MAMBA_EXE, MAMBA_ROOT_PREFIX
from .progress import ElapseTime, get_progress


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

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

    req_mamba_root_prefix = Unicode(
        MAMBA_ROOT_PREFIX,
        help='Micromamba global root prefix',
    ).tag(config=True)

    req_mamba_exe = Unicode(
        MAMBA_EXE,
        help='Micromamba global exe',
    ).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)

    batch_script = Unicode(
        """#!/bin/bash
#SBATCH --job-name={{job_name}}
#SBATCH --output={{homedir}}/.{{job_name}}.log
#SBATCH --chdir={{workdir}}
#SBATCH --export={{keepvars}}

{% if cluster   -%}#SBATCH --cluster={{cluster}}{%- endif %}
{% if partition -%}#SBATCH --partition={{partition}}{%- endif %}
{% if node      -%}#SBATCH --nodelist={{node}}{%- endif %}
{% if qos       -%}#SBATCH --qos={{qos}}{%- endif %}

{% if runtime   -%}#SBATCH --time={{runtime}}{%- endif %}
{% if nprocs    -%}#SBATCH --cpus-per-task={{nprocs}}{%- endif %}
{% if memory    -%}#SBATCH --mem={{memory}}{%- endif %}
{% if gres      -%}#SBATCH --gres={{gres}}{%- endif %}

# Redirect logs
export JUPYTER_LOG_DIR="{{homedir}}/.jupyter/spawner/logs"
mkdir -p ${JUPYTER_LOG_DIR}

echo "The {{job_name}} logs are located in: ${JUPYTER_LOG_DIR}"

export JUPYTER_JOB_LOG=${JUPYTER_LOG_DIR}/$(date "+%Y-%m-%d")_{{job_name}}_${SLURM_JOB_ID}.log

{   PS4='[$(date "+%Y-%m-%d %T")]\011 ';

    set -xeo pipefail;
    trap 'echo SIGTERM received' TERM;

    {# SLURM config #}
    scontrol write batch_script ${SLURM_JOB_ID} -;

    {# Micromamba config #}
    export MAMBA_EXE={{mamba_exe}};
    export MAMBA_ROOT_PREFIX={{mamba_root_prefix}};
    source $MAMBA_ROOT_PREFIX/etc/profile.d/micromamba.sh;

    {# Activate micromamba env requested by the user #}
    micromamba activate {{ pyenv }};
    export JUPYTER_PATH={{ pyenv }}/share/jupyter;

    {# Prologue #}
    {%- if prologue -%}
    {{prologue}};
    {%- endif -%}

    {# Start Jupyter single-user command #}
    {{cmd}};

    {# Epilogue #}
    {%- if epilogue -%}
    {{epilogue}};
    {%- endif -%}
} > ${JUPYTER_JOB_LOG} 2>&1
""",
        help='Template for SLURM job submission batch script.',
    ).tag(config=True)

    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}') + '?reset'

    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)