"""GLiCID spawner module.""" import asyncio import re from batchspawner import JobStatus, SlurmSpawner from traitlets import Integer, Unicode, default from .form import options_form, options_from_form 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_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)