127 lines
3.7 KiB
Python
127 lines
3.7 KiB
Python
"""GLiCID spawner module."""
|
|
|
|
import asyncio
|
|
import re
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
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
|
|
from .templates import get_template_src
|
|
|
|
SPAWNER_BIN = Path(sys.exec_prefix) / 'bin'
|
|
|
|
|
|
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(
|
|
get_template_src('slurm_script.jinja'),
|
|
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'
|
|
|
|
def cmd_formatted_for_batch(self):
|
|
"""The command which is substituted inside of the batch script.
|
|
|
|
Here we need the absolute path to the spawner and singleuser commands.
|
|
|
|
"""
|
|
return ' '.join(
|
|
[
|
|
str(SPAWNER_BIN / self.batchspawner_singleuser_cmd),
|
|
str(SPAWNER_BIN / self.cmd[0]),
|
|
*self.get_args(),
|
|
]
|
|
)
|
|
|
|
slurm_job_id_re = Unicode(r'(\d+)(?:;(\w+))?').tag(config=True)
|
|
|
|
def parse_job_id(self, output) -> str:
|
|
"""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 job_id
|
|
|
|
self.log.error(f'GlicidSpawner unable to parse job ID from text: {output}')
|
|
return ''
|
|
|
|
@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)
|